Merge pull request #12 from haruishi43/fix/equi2ecube_cube2equi

A robust `equi2cube` and `cube2equi` function thanks to Oletus

equilib/cube2equi/numpy.py

@@ -4,7 +4,8 @@
 
 import numpy as np
 
-from equilib.grid_sample import numpy_grid_sample
+# from equilib.grid_sample import numpy_grid_sample
+from equilib.grid_sample.numpy.bilinear import interp2d
 
 __all__ = ["convert2horizon", "run"]
 
@@ -124,26 +125,17 @@ def _equirect_facetype(h: int, w: int) -> np.ndarray:
 
     int_dtype = np.dtype(np.int64)
 
+    w_ratio = (w - 1) / w
+    h_ratio = (h - 1) / h
+
     tp = np.roll(
         np.arange(4).repeat(w // 4)[None, :].repeat(h, 0), 3 * w // 8, 1
     )
 
-    half_pixel_angular_width = np.pi / w
-    pixel_angular_height = np.pi / h
-
     # Prepare ceil mask
     mask = np.zeros((h, w // 4), np.bool)
-    idx = (
-        np.linspace(
-            -np.pi + half_pixel_angular_width,
-            np.pi - half_pixel_angular_width,
-            num=w // 4,
-        )
-        / 4
-    )
-    idx = h // 2 - np.around(
-        np.arctan(np.cos(idx)) * h / (np.pi - pixel_angular_height)
-    )
+    idx = np.linspace(-(np.pi * w_ratio), np.pi * w_ratio, w // 4) / 4
+    idx = h // 2 - np.around(np.arctan(np.cos(idx)) * h / (np.pi * h_ratio))
     idx = idx.astype(int_dtype)
     for i, j in enumerate(idx):
         mask[:j, i] = 1
@@ -162,19 +154,14 @@ def create_equi_grid(
     batch: int,
     dtype: np.dtype = np.dtype(np.float32),
 ) -> np.ndarray:
-    half_pixel_angular_width = np.pi / w_out
-    half_pixel_angular_height = np.pi / h_out / 2
+
+    w_ratio = (w_out - 1) / w_out
+    h_ratio = (h_out - 1) / h_out
     theta = np.linspace(
-        -np.pi + half_pixel_angular_width,
-        np.pi - half_pixel_angular_width,
-        num=w_out,
-        dtype=dtype,
+        -(np.pi * w_ratio), np.pi * w_ratio, num=w_out, dtype=dtype
     )
     phi = np.linspace(
-        np.pi / 2 - half_pixel_angular_height,
-        -np.pi / 2 + half_pixel_angular_height,
-        num=h_out,
-        dtype=dtype,
+        np.pi * h_ratio / 2, -(np.pi * h_ratio) / 2, num=h_out, dtype=dtype
     )
     theta, phi = np.meshgrid(theta, phi)
 
@@ -203,18 +190,67 @@ def create_equi_grid(
             coor_y[mask] = -c * np.cos(theta[mask])
 
     # Final renormalize
-    coor_x = np.clip(coor_x + 0.5, 0, 1) * w_face
-    coor_y = np.clip(coor_y + 0.5, 0, 1) * w_face
+    # coor_x = np.clip(np.clip(coor_x + 0.5, 0, 1) * w_face, 0, w_face - 1)
+    # coor_y = np.clip(np.clip(coor_y + 0.5, 0, 1) * w_face, 0, w_face - 1)
+
+    coor_x = (np.clip(coor_x, -0.5, 0.5) + 0.5) * w_face
+    coor_y = (np.clip(coor_y, -0.5, 0.5) + 0.5) * w_face
 
     # change x axis of the x coordinate map
     for i in range(6):
         mask = tp == i
         coor_x[mask] = coor_x[mask] + w_face * i
 
-    grid = np.stack((coor_y, coor_x), axis=0)
+    grid = np.stack((coor_y, coor_x), axis=0) - 0.5
     grid = np.concatenate([grid[np.newaxis, ...]] * batch)
-    grid = grid - 0.5  # Offset pixel center
-    return grid
+    return grid, tp
+
+
+def numpy_grid_sample(
+    img: np.ndarray, grid: np.ndarray, out: np.ndarray, cube_face_id: np.ndarray
+):
+    b, _, h, w = img.shape
+
+    min_grid = np.floor(grid).astype(np.int64)
+    max_grid = min_grid + 1
+    d_grid = grid - min_grid
+
+    _, _, grid_h, grid_w = grid.shape
+    cube_face_min_grid = min_grid // h
+    cube_face_max_grid = max_grid // h
+
+    min_grid[:, 0, :, :] = np.clip(min_grid[:, 0, :, :], 0, None)
+    max_grid[:, 0, :, :] = np.clip(max_grid[:, 0, :, :], None, h - 1)
+
+    # FIXME: any way to do efficient batch?
+    for i in range(b):
+
+        for y in range(grid_h):
+            for x in range(grid_w):
+                if (
+                    cube_face_max_grid[i, 1, y, x]
+                    != cube_face_min_grid[i, 1, y, x]
+                ):
+                    if cube_face_max_grid[i, 1, y, x] != cube_face_id[y, x]:
+                        max_grid[i, 1, y, x] -= 1
+                    else:
+                        min_grid[i, 1, y, x] += 1
+
+        dy = d_grid[i, 0, ...]
+        dx = d_grid[i, 1, ...]
+        min_ys = min_grid[i, 0, ...]
+        min_xs = min_grid[i, 1, ...]
+        max_ys = max_grid[i, 0, ...]
+        max_xs = max_grid[i, 1, ...]
+
+        p00 = img[i][:, min_ys, min_xs]
+        p10 = img[i][:, max_ys, min_xs]
+        p01 = img[i][:, min_ys, max_xs]
+        p11 = img[i][:, max_ys, max_xs]
+
+        out[i, ...] = interp2d(p00, p10, p01, p11, dy, dx)
+
+    return out
 
 
 def run(
@@ -269,7 +305,7 @@ def run(
     out = np.empty((bs, c, height, width), dtype=dtype)
 
     # create sampling grid
-    grid = create_equi_grid(
+    grid, tp = create_equi_grid(
         h_out=height, w_out=width, w_face=w_face, batch=bs, dtype=dtype
     )
 
@@ -279,7 +315,10 @@ def run(
             img=horizon, grid=grid, out=out, mode=mode
         )
     else:
-        out = numpy_grid_sample(img=horizon, grid=grid, out=out, mode=mode)
+        # out = numpy_grid_sample(img=horizon, grid=grid, out=out, mode=mode)
+        out = numpy_grid_sample(
+            img=horizon, grid=grid, out=out, cube_face_id=tp
+        )
 
     out = (
         out.astype(horizon_dtype)

equilib/cube2equi/torch.py

@@ -137,6 +137,9 @@ def _equirect_facetype(h: int, w: int) -> torch.Tensor:
 
     int_dtype = torch.int64
 
+    w_ratio = (w - 1) / w
+    h_ratio = (h - 1) / h
+
     tp = torch.roll(
         torch.arange(4)  # 1
         .repeat_interleave(w // 4)  # 2 same as np.repeat
@@ -151,8 +154,10 @@ def _equirect_facetype(h: int, w: int) -> torch.Tensor:
 
     # Prepare ceil mask
     mask = torch.zeros((h, w // 4), dtype=torch.bool)
-    idx = torch.linspace(-math.pi, math.pi, w // 4) / 4
-    idx = h // 2 - torch.round(torch.atan(torch.cos(idx)) * h / math.pi)
+    idx = torch.linspace(-(math.pi * w_ratio), math.pi * w_ratio, w // 4) / 4
+    idx = h // 2 - torch.round(
+        torch.atan(torch.cos(idx)) * h / (math.pi * h_ratio)
+    )
     idx = idx.type(int_dtype)
     for i, j in enumerate(idx):
         mask[:j, i] = 1
@@ -172,19 +177,18 @@ def create_equi_grid(
     dtype: torch.dtype = torch.float32,
     device: torch.device = torch.device("cpu"),
 ) -> torch.Tensor:
-
-    half_pixel_angular_width = math.pi / w_out
-    half_pixel_angular_height = math.pi / h_out / 2
+    w_ratio = (w_out - 1) / w_out
+    h_ratio = (h_out - 1) / h_out
     theta = torch.linspace(
-        -math.pi + half_pixel_angular_width,
-        math.pi - half_pixel_angular_width,
+        -(math.pi * w_ratio),
+        math.pi * w_ratio,
         steps=w_out,
         dtype=dtype,
         device=device,
     )
     phi = torch.linspace(
-        math.pi / 2 - half_pixel_angular_height,
-        -math.pi / 2 + half_pixel_angular_height,
+        (math.pi * h_ratio) / 2,
+        -(math.pi * h_ratio) / 2,
         steps=h_out,
         dtype=dtype,
         device=device,
@@ -221,20 +225,23 @@ def create_equi_grid(
             coor_y[mask] = -c * torch.cos(theta[mask])
 
     # Final renormalize
-    coor_x = torch.clamp(coor_x + 0.5, 0, 1) * w_face
-    coor_y = torch.clamp(coor_y + 0.5, 0, 1) * w_face
+    coor_x = torch.clamp(
+        torch.clamp(coor_x + 0.5, 0, 1) * w_face, 0, w_face - 1
+    )
+    coor_y = torch.clamp(
+        torch.clamp(coor_y + 0.5, 0, 1) * w_face, 0, w_face - 1
+    )
 
     # change x axis of the x coordinate map
     for i in range(6):
         mask = tp == i
         coor_x[mask] = coor_x[mask] + w_face * i
 
     # repeat batch
-    coor_x = coor_x.repeat(batch, 1, 1)
-    coor_y = coor_y.repeat(batch, 1, 1)
+    coor_x = coor_x.repeat(batch, 1, 1) - 0.5
+    coor_y = coor_y.repeat(batch, 1, 1) - 0.5
 
     grid = torch.stack((coor_y, coor_x), dim=-3).to(device)
-    grid = grid - 0.5  # offset pixel center
     return grid
 
 

equilib/equi2cube/numpy.py

@@ -98,8 +98,8 @@ def convert_grid(
         # convert to pixel
         # I thought it would be faster if it was done all at once,
         # but it was faster separately
-        ui = (theta - np.pi) * w_equi / (2 * np.pi)
-        uj = (np.pi / 2 - phi) * h_equi / np.pi  # NOTE: fixed here
+        ui = (theta / (2 * np.pi) - 0.5) * w_equi - 0.5
+        uj = (0.5 - phi / np.pi) * h_equi - 0.5
         ui %= w_equi
         uj %= h_equi
     elif method == "faster":
@@ -108,8 +108,8 @@ def convert_grid(
         # the range of phi is -pi/2 ~ pi/2
         # this means that if the input `rots` have rotations that are
         # out of range, it will not work with `faster`
-        ui = (theta - np.pi) * w_equi / (2 * np.pi)
-        uj = (np.pi / 2 - phi) * h_equi / np.pi  # NOTE: fixed here
+        ui = (theta / (2 * np.pi) - 0.5) * w_equi - 0.5
+        uj = (0.5 - phi / np.pi) * h_equi - 0.5
         ui = np.where(ui < 0, ui + w_equi, ui)
         ui = np.where(ui >= w_equi, ui - w_equi, ui)
         uj = np.where(uj < 0, uj + h_equi, uj)
@@ -119,7 +119,6 @@ def convert_grid(
 
     # stack the pixel maps into a grid
     grid = np.stack((uj, ui), axis=-3)
-    grid = grid - 0.5  # offset pixel center
     return grid
 
 

equilib/equi2cube/torch.py

@@ -92,13 +92,13 @@ def convert_grid(
     theta = torch.atan2(xyz[..., 1], xyz[..., 0])
 
     if method == "robust":
-        ui = (theta - pi) * w_equi / (2 * pi)
-        uj = (pi / 2 - phi) * h_equi / pi  # FIXME: fixed here
+        ui = (theta / (2 * pi) - 0.5) * w_equi - 0.5
+        uj = (0.5 - phi / pi) * h_equi - 0.5
         ui %= w_equi
         uj %= h_equi
     elif method == "faster":
-        ui = (theta - pi) * w_equi / (2 * pi)
-        uj = (pi / 2 - phi) * h_equi / pi  # FIXME: fixed here
+        ui = (theta / (2 * pi) - 0.5) * w_equi - 0.5
+        uj = (0.5 - phi / pi) * h_equi - 0.5
         ui = torch.where(ui < 0, ui + w_equi, ui)
         ui = torch.where(ui >= w_equi, ui - w_equi, ui)
         uj = torch.where(uj < 0, uj + h_equi, uj)
@@ -108,7 +108,6 @@ def convert_grid(
 
     # stack the pixel maps into a grid
     grid = torch.stack((uj, ui), dim=-3)
-    grid = grid - 0.5  # offset pixel center
     return grid
 
 

equilib/grid_sample/numpy/grid_sample.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+import warnings
+
 import numpy as np
 
 from .bicubic import bicubic
@@ -33,6 +35,10 @@ def grid_sample(
     elif mode == "bilinear":
         out = bilinear(img, grid, out)
     elif mode == "bicubic":
+        # FIXME: bicubic algorithm is not perfect yet
+        warnings.warn(
+            "Bicubic interpolation is not perfect (especially when upsampling). Use with care!"
+        )
         out = bicubic(img, grid, out)
     else:
         raise ValueError(f"ERR: {mode} is not supported")

equilib/numpy_utils/grid.py

@@ -103,11 +103,11 @@ def create_xyz_grid(
     dtype: np.dtype = np.dtype(np.float32),
 ) -> np.ndarray:
     """xyz coordinates of the faces of the cube"""
+
+    ratio = (w_face - 1) / w_face
+
     out = np.zeros((w_face, w_face * 6, 3), dtype=dtype)
-    pixel_half_width = 0.5 / w_face
-    rng = np.linspace(
-        -0.5 + pixel_half_width, 0.5 - pixel_half_width, num=w_face, dtype=dtype
-    )
+    rng = np.linspace(-0.5 * ratio, 0.5 * 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

@@ -123,8 +123,13 @@ def create_xyz_grid(
     device: torch.device = torch.device("cpu"),
 ) -> torch.Tensor:
     """xyz coordinates of the faces of the cube"""
+
+    ratio = (w_face - 1) / w_face
+
     out = torch.zeros((w_face, w_face * 6, 3), dtype=dtype, device=device)
-    rng = torch.linspace(-0.5, 0.5, w_face, dtype=dtype, device=device)
+    rng = torch.linspace(
+        -0.5 * ratio, 0.5 * ratio, w_face, dtype=dtype, device=device
+    )
 
     # NOTE: https://github.com/pytorch/pytorch/issues/15301
     # Torch meshgrid behaves differently than numpy

tests/test_cube2equi/test_numpy.py

@@ -143,6 +143,7 @@ def bench_baselines(
             save(out_scipy[b], os.path.join(SAVE_ROOT, f"out_scipy_{b}.jpg"))
 
 
+@pytest.mark.skip(reason="shouldn't compare with other sampling methods")
 @pytest.mark.skipif(cv2 is None, reason="cv2 is None; not installed")
 @pytest.mark.skipif(
     map_coordinates is None,

tests/test_cube2equi/test_torch.py

@@ -301,6 +301,7 @@ def bench_gpu(
             save(pure_out[b], os.path.join(SAVE_ROOT, f"out_gpu_pure_{b}.jpg"))
 
 
+@pytest.mark.skip(reason="shouldn't compare with other sampling methods")
 @pytest.mark.parametrize("bs", [1, 4])
 @pytest.mark.parametrize("height", [256])
 @pytest.mark.parametrize("width", [512])
@@ -312,6 +313,7 @@ def test_cube2equi_cpu(
     bench_cpu(bs=bs, height=height, width=width, mode=mode, dtype=dtype)
 
 
+@pytest.mark.skip(reason="shouldn't compare with other sampling methods")
 @pytest.mark.skipif(
     not torch.cuda.is_available(), reason="cuda device is not available"
 )