SpherePad for equi2cube + cache for create_equi_grid and create_normalized_grid

equilib/cube2equi/torch.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import math
+from functools import cache
 
 from typing import Dict, List, Union
 
@@ -169,6 +170,7 @@ def _equirect_facetype(h: int, w: int) -> torch.Tensor:
     return tp.type(int_dtype)
 
 
+@cache
 def create_equi_grid(
     h_out: int,
     w_out: int,

equilib/equi2cube/numpy.py

@@ -130,6 +130,7 @@ def run(
     z_down: bool,
     mode: str,
     override_func: Optional[Callable[[], Any]] = None,
+    pad: float = 0.,
 ) -> Union[np.ndarray, List[List[np.ndarray]], List[Dict[str, np.ndarray]]]:
     """Call Equi2Cube
 
@@ -179,7 +180,7 @@ def run(
     out = np.empty((bs, c, w_face, w_face * 6), dtype=dtype)
 
     # create grid
-    xyz = create_xyz_grid(w_face=w_face, batch=bs, dtype=dtype)
+    xyz = create_xyz_grid(w_face=w_face, pad=pad, batch=bs, dtype=dtype)
     xyz = xyz[..., np.newaxis]
 
     # FIXME: not sure why, but z-axis is facing the opposite

equilib/equi2cube/torch.py

@@ -119,6 +119,7 @@ def run(
     z_down: bool,
     mode: str,
     backend: str = "native",
+    pad: float = 0.,
 ) -> Union[torch.Tensor, List[torch.Tensor], List[Dict[str, torch.Tensor]]]:
     """Run Equi2Cube
 
@@ -201,7 +202,7 @@ def run(
 
     # create grid
     xyz = create_xyz_grid(
-        w_face=w_face, batch=bs, dtype=tmp_dtype, device=tmp_device
+        w_face=w_face, pad=pad, batch=bs, dtype=tmp_dtype, device=tmp_device
     )
     xyz = xyz.unsqueeze(-1)
 

equilib/equi2equi/torch.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
-
-from typing import Dict, List, Optional
+from functools import cache
+from typing import Dict, List, Optional, Tuple, Iterable
 
 import torch
 
@@ -53,6 +53,58 @@ def convert_grid(
 
     return grid
 
+@cache
+def _get_grid(
+    height: int,
+    width: int,
+    bs: int,
+    tmp_dtype: torch.dtype,
+    tmp_device: torch.device,
+    rolls: Tuple[float],
+    pitchs: Tuple[float],
+    yaws: Tuple[float],
+    z_down: float,
+    h_equi: int,
+    w_equi: int,
+) -> torch.Tensor:
+    m = create_normalized_grid(
+        height=height, width=width, batch=bs, dtype=tmp_dtype, device=tmp_device
+    )
+    m = m.unsqueeze(-1)
+
+    rots = [
+        {'roll': roll, 'pitch': pitch, 'yaw': yaw}
+        for (roll, pitch, yaw) in zip(rolls, pitchs, yaws)
+    ]
+
+    # create batched rotation matrices
+    R = create_rotation_matrices(
+        rots=rots, z_down=z_down, dtype=tmp_dtype, device=tmp_device
+    )
+
+    # rotate the grid
+    M = matmul(m, R)
+
+    grid = convert_grid(M=M, h_equi=h_equi, w_equi=w_equi, method="robust")
+
+    return grid
+
+
+def get_grid(
+    rots: Iterable[Dict[str, float]],
+    **kwargs
+) -> torch.Tensor:
+    rolls = tuple([rot['roll'] for rot in rots])
+    pitchs = tuple([rot['pitch'] for rot in rots])
+    yaws = tuple([rot['yaw'] for rot in rots])
+
+    return _get_grid(
+        rolls=rolls,
+        pitchs=pitchs,
+        yaws=yaws,
+        **kwargs
+    )
+
 
 def run(
     src: torch.Tensor,
@@ -158,21 +210,18 @@ def run(
     else:
         tmp_dtype = dtype
 
-    m = create_normalized_grid(
-        height=height, width=width, batch=bs, dtype=tmp_dtype, device=tmp_device
-    )
-    m = m.unsqueeze(-1)
-
-    # create batched rotation matrices
-    R = create_rotation_matrices(
-        rots=rots, z_down=z_down, dtype=tmp_dtype, device=tmp_device
+    grid = get_grid(
+        height=height,
+        width=width,
+        bs=bs,
+        tmp_dtype=tmp_dtype,
+        tmp_device=tmp_device,
+        rots=rots,
+        z_down=z_down,
+        h_equi=h_equi,
+        w_equi=w_equi
     )
 
-    # rotate the grid
-    M = matmul(m, R)
-
-    grid = convert_grid(M=M, h_equi=h_equi, w_equi=w_equi, method="robust")
-
     # FIXME: putting `grid` to device since `pure`'s bilinear interpolation requires it
     # FIXME: better way of forcing `grid` to be the same dtype?
     if src.dtype != grid.dtype:

equilib/numpy_utils/grid.py

@@ -99,6 +99,7 @@ def create_normalized_grid(
 
 def create_xyz_grid(
     w_face: int,
+    pad: float = 0.,
     batch: Optional[int] = None,
     dtype: np.dtype = np.dtype(np.float32),
 ) -> np.ndarray:
@@ -107,7 +108,7 @@ def create_xyz_grid(
     ratio = (w_face - 1) / w_face
 
     out = np.zeros((w_face, w_face * 6, 3), dtype=dtype)
-    rng = np.linspace(-0.5 * ratio, 0.5 * ratio, num=w_face, dtype=dtype)
+    rng = np.linspace(-(0.5 + pad) * ratio, (0.5 + pad) * ratio, num=w_face, dtype=dtype)
 
     # Front face (x = 0.5)
     out[:, 0 * w_face : 1 * w_face, [1, 2]] = np.stack(

equilib/torch_utils/grid.py

@@ -118,6 +118,7 @@ def create_normalized_grid(
 
 def create_xyz_grid(
     w_face: int,
+    pad: float = 0.,
     batch: Optional[int] = None,
     dtype: torch.dtype = torch.float32,
     device: torch.device = torch.device("cpu"),
@@ -128,7 +129,7 @@ def create_xyz_grid(
 
     out = torch.zeros((w_face, w_face * 6, 3), dtype=dtype, device=device)
     rng = torch.linspace(
-        -0.5 * ratio, 0.5 * ratio, w_face, dtype=dtype, device=device
+        -(0.5 + pad) * ratio, (0.5 + pad) * ratio, w_face, dtype=dtype, device=device
     )
 
     # NOTE: https://github.com/pytorch/pytorch/issues/15301