revise clipping for numpy and pytorch; now clips according to the inp…

…ut; also added argument for all functions

equilib/cube2equi/base.py

@@ -46,11 +46,17 @@ class Cube2Equi(object):
     """
 
     def __init__(
-        self, height: int, width: int, cube_format: str, mode: str = "bilinear"
+        self,
+        height: int,
+        width: int,
+        cube_format: str,
+        clip_output: bool = True,
+        mode: str = "bilinear",
     ) -> None:
         self.height = height
         self.width = width
         self.cube_format = cube_format
+        self.clip_output = clip_output
         self.mode = mode
 
     def __call__(self, cubemap: CubeMaps, **kwargs) -> ArrayLike:
@@ -59,6 +65,7 @@ def __call__(self, cubemap: CubeMaps, **kwargs) -> ArrayLike:
             cube_format=self.cube_format,
             width=self.width,
             height=self.height,
+            clip_output=self.clip_output,
             mode=self.mode,
             **kwargs,
         )
@@ -69,6 +76,7 @@ def cube2equi(
     cube_format: str,
     height: int,
     width: int,
+    clip_output: bool = True,
     mode: str = "bilinear",
     **kwargs,
 ) -> ArrayLike:
@@ -124,14 +132,24 @@ def cube2equi(
             cubemap=cubemap, cube_format=cube_format
         )
         out = run_numpy(
-            horizon=horizon, height=height, width=width, mode=mode, **kwargs
+            horizon=horizon,
+            height=height,
+            width=width,
+            clip_output=clip_output,
+            mode=mode,
+            **kwargs,
         )
     elif _type == "torch":
         horizon = convert2horizon_torch(
             cubemap=cubemap, cube_format=cube_format
         )
         out = run_torch(
-            horizon=horizon, height=height, width=width, mode=mode, **kwargs
+            horizon=horizon,
+            height=height,
+            width=width,
+            clip_output=clip_output,
+            mode=mode,
+            **kwargs,
         )
     else:
         raise ValueError("Oops something went wrong here")

equilib/cube2equi/numpy.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-from typing import Any, Callable, Dict, List, Optional, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 
@@ -153,8 +153,7 @@ def create_equi_grid(
     w_face: int,
     batch: int,
     dtype: np.dtype = np.dtype(np.float32),
-) -> np.ndarray:
-
+) -> Tuple[np.ndarray, np.ndarray]:
     w_ratio = (w_out - 1) / w_out
     h_ratio = (h_out - 1) / h_out
     theta = np.linspace(
@@ -224,7 +223,6 @@ def numpy_grid_sample(
 
     # 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 (
@@ -258,6 +256,7 @@ def run(
     height: int,
     width: int,
     mode: str,
+    clip_output: bool = True,
     override_func: Optional[Callable[[], Any]] = None,
 ) -> np.ndarray:
     """Run Cube2Equi
@@ -320,10 +319,11 @@ def run(
             img=horizon, grid=grid, out=out, cube_face_id=tp
         )
 
+    # clip by the input
     out = (
         out.astype(horizon_dtype)
-        if horizon_dtype == np.dtype(np.uint8)
-        else np.clip(out, np.nanmin(out), np.nanmax(out))
+        if horizon_dtype == np.dtype(np.uint8) or not clip_output
+        else np.clip(out, np.min(horizon), np.max(horizon))
     )
 
     return out

equilib/cube2equi/torch.py

@@ -250,6 +250,7 @@ def run(
     height: int,
     width: int,
     mode: str,
+    clip_output: bool = True,
     backend: str = "native",
 ) -> torch.Tensor:
     """Run Cube2Equi
@@ -349,8 +350,8 @@ def run(
 
     out = (
         out.type(horizon_dtype)
-        if horizon_dtype == torch.uint8
-        else torch.clip(out, 0.0, 1.0)
+        if horizon_dtype == torch.uint8 or not clip_output
+        else torch.clip(out, torch.min(horizon), torch.max(horizon))
     )
 
     return out

equilib/equi2cube/base.py

@@ -53,21 +53,25 @@ def __init__(
         w_face: int,
         cube_format: str,
         z_down: bool = False,
+        clip_output: bool = True,
         mode: str = "bilinear",
     ) -> None:
         self.w_face = w_face
         self.cube_format = cube_format
         self.z_down = z_down
+        self.clip_output = clip_output
         self.mode = mode
 
-    def __call__(self, equi: ArrayLike, rots: Rot) -> CubeMaps:
+    def __call__(self, equi: ArrayLike, rots: Rot, **kwargs) -> CubeMaps:
         return equi2cube(
             equi=equi,
             rots=rots,
             w_face=self.w_face,
             cube_format=self.cube_format,
             z_down=self.z_down,
+            clip_output=self.clip_output,
             mode=self.mode,
+            **kwargs,
         )
 
 
@@ -77,6 +81,7 @@ def equi2cube(
     w_face: int,
     cube_format: str,
     z_down: bool = False,
+    clip_output: bool = True,
     mode: str = "bilinear",
     **kwargs,
 ) -> CubeMaps:
@@ -120,6 +125,7 @@ def equi2cube(
             w_face=w_face,
             cube_format=cube_format,
             z_down=z_down,
+            clip_output=clip_output,
             mode=mode,
             **kwargs,
         )
@@ -130,6 +136,7 @@ def equi2cube(
             w_face=w_face,
             cube_format=cube_format,
             z_down=z_down,
+            clip_output=clip_output,
             mode=mode,
             **kwargs,
         )

equilib/equi2cube/numpy.py

@@ -62,7 +62,6 @@ def cube_h2dice(cube_h: np.ndarray) -> np.ndarray:
 
 
 def matmul(m: np.ndarray, R: np.ndarray, method: str = "faster") -> np.ndarray:
-
     if method == "robust":
         # When target image size is smaller, it might be faster with `matmul`
         # but not by much
@@ -89,7 +88,6 @@ def matmul(m: np.ndarray, R: np.ndarray, method: str = "faster") -> np.ndarray:
 def convert_grid(
     xyz: np.ndarray, h_equi: int, w_equi: int, method: str = "robust"
 ) -> np.ndarray:
-
     # convert to rotation
     phi = np.arcsin(xyz[..., 2] / np.linalg.norm(xyz, axis=-1))
     theta = np.arctan2(xyz[..., 1], xyz[..., 0])
@@ -129,6 +127,7 @@ def run(
     cube_format: str,
     z_down: bool,
     mode: str,
+    clip_output: bool = True,
     override_func: Optional[Callable[[], Any]] = None,
 ) -> Union[np.ndarray, List[List[np.ndarray]], List[Dict[str, np.ndarray]]]:
     """Call Equi2Cube
@@ -203,10 +202,11 @@ def run(
     else:
         out = numpy_grid_sample(img=equi, grid=grid, out=out, mode=mode)
 
+    # clip by input
     out = (
         out.astype(equi_dtype)
-        if equi_dtype == np.dtype(np.uint8)
-        else np.clip(out, np.nanmin(out), np.nanmax(out))
+        if equi_dtype == np.dtype(np.uint8) or not clip_output
+        else np.clip(out, np.min(equi), np.max(equi))
     )
 
     # reformat the output

equilib/equi2cube/torch.py

@@ -76,7 +76,6 @@ def cube_h2dice(cube_h: torch.Tensor) -> torch.Tensor:
 
 
 def matmul(m: torch.Tensor, R: torch.Tensor) -> torch.Tensor:
-
     M = torch.matmul(R[:, None, None, ...], m)
     M = M.squeeze(-1)
 
@@ -86,7 +85,6 @@ def matmul(m: torch.Tensor, R: torch.Tensor) -> torch.Tensor:
 def convert_grid(
     xyz: torch.Tensor, h_equi: int, w_equi: int, method: str = "robust"
 ) -> torch.Tensor:
-
     # convert to rotation
     phi = torch.asin(xyz[..., 2] / torch.norm(xyz, dim=-1))
     theta = torch.atan2(xyz[..., 1], xyz[..., 0])
@@ -118,8 +116,11 @@ def run(
     cube_format: str,
     z_down: bool,
     mode: str,
+    clip_output: bool = True,
     backend: str = "native",
-) -> Union[torch.Tensor, List[torch.Tensor], List[Dict[str, torch.Tensor]]]:
+) -> Union[
+    torch.Tensor, List[List[torch.Tensor]], List[Dict[str, torch.Tensor]]
+]:
     """Run Equi2Cube
 
     params:
@@ -232,10 +233,11 @@ def run(
         img=equi, grid=grid, out=out, mode=mode, backend=backend
     )
 
+    # clip by input
     out = (
         out.type(equi_dtype)
-        if equi_dtype == torch.uint8
-        else torch.clip(out, 0.0, 1.0)
+        if equi_dtype == torch.uint8 or not clip_output
+        else torch.clip(out, torch.min(equi), torch.max(equi))
     )
 
     # reformat the output

equilib/equi2equi/base.py

@@ -19,7 +19,7 @@ class Equi2Equi(object):
     """
     params:
     - w_out, h_out (optional int): equi image size
-    - sampling_method (str): defaults to "default"
+    - clip_output (bool) whether to clip values in the range of input
     - mode (str): interpolation mode, defaults to "bilinear"
     - z_down (bool)
 
@@ -35,23 +35,31 @@ def __init__(
         self,
         height: Optional[int] = None,
         width: Optional[int] = None,
+        clip_output: bool = True,
         mode: str = "bilinear",
         z_down: bool = False,
     ) -> None:
         self.height = height
         self.width = width
+        self.clip_output = clip_output
         self.mode = mode
         self.z_down = z_down
 
     def __call__(self, src: ArrayLike, rots: Rot, **kwargs) -> ArrayLike:
         return equi2equi(
-            src=src, rots=rots, mode=self.mode, z_down=self.z_down, **kwargs
+            src=src,
+            rots=rots,
+            clip_output=self.clip_output,
+            mode=self.mode,
+            z_down=self.z_down,
+            **kwargs,
         )
 
 
 def equi2equi(
     src: ArrayLike,
     rots: Rot,
+    clip_output: bool = True,
     mode: str = "bilinear",
     z_down: bool = False,
     height: Optional[int] = None,
@@ -98,6 +106,7 @@ def equi2equi(
             z_down=z_down,
             height=height,
             width=width,
+            clip_output=clip_output,
             **kwargs,
         )
     elif _type == "torch":
@@ -108,6 +117,7 @@ def equi2equi(
             z_down=z_down,
             height=height,
             width=width,
+            clip_output=clip_output,
             **kwargs,
         )
     else:

equilib/equi2equi/numpy.py

@@ -9,7 +9,6 @@
 
 
 def matmul(m: np.ndarray, R: np.ndarray, method: str = "faster") -> np.ndarray:
-
     if method == "robust":
         # When target image size is smaller, it might be faster with `matmul`
         # but not by much
@@ -36,7 +35,6 @@ def matmul(m: np.ndarray, R: np.ndarray, method: str = "faster") -> np.ndarray:
 def convert_grid(
     M: np.ndarray, h_equi: int, w_equi: int, method: str = "robust"
 ) -> np.ndarray:
-
     # convert to rotation
     phi = np.arcsin(M[..., 2] / np.linalg.norm(M, axis=-1))
     theta = np.arctan2(M[..., 1], M[..., 0])
@@ -81,6 +79,7 @@ def run(
     mode: str,
     height: Optional[int] = None,
     width: Optional[int] = None,
+    clip_output: bool = True,
     override_func: Optional[Callable[[], Any]] = None,
 ) -> np.ndarray:
     """Run Equi2Equi
@@ -169,10 +168,11 @@ def run(
     else:
         out = numpy_grid_sample(img=src, grid=grid, out=out, mode=mode)
 
+    # clip by input
     out = (
         out.astype(src_dtype)
-        if src_dtype == np.dtype(np.uint8)
-        else np.clip(out, np.nanmin(out), np.nanmax(out))
+        if src_dtype == np.dtype(np.uint8) or not clip_output
+        else np.clip(out, np.min(src), np.max(src))
     )
 
     return out

equilib/equi2equi/torch.py

@@ -14,7 +14,6 @@
 
 
 def matmul(m: torch.Tensor, R: torch.Tensor) -> torch.Tensor:
-
     M = torch.matmul(R[:, None, None, ...], m)
     M = M.squeeze(-1)
 
@@ -24,7 +23,6 @@ def matmul(m: torch.Tensor, R: torch.Tensor) -> torch.Tensor:
 def convert_grid(
     M: torch.Tensor, h_equi: int, w_equi: int, method: str = "robust"
 ) -> torch.Tensor:
-
     # convert to rotation
     phi = torch.asin(M[..., 2] / torch.norm(M, dim=-1))
     theta = torch.atan2(M[..., 1], M[..., 0])
@@ -61,6 +59,7 @@ def run(
     mode: str,
     height: Optional[int] = None,
     width: Optional[int] = None,
+    clip_output: bool = True,
     backend: str = "native",
 ) -> torch.Tensor:
     """Run Equi2Equi
@@ -191,10 +190,11 @@ def run(
 
     # NOTE: we assume that `out` keeps it's dtype
 
+    # clip by input
     out = (
         out.type(src_dtype)
-        if src_dtype == torch.uint8
-        else torch.clip(out, 0.0, 1.0)
+        if src_dtype == torch.uint8 or not clip_output
+        else torch.clip(out, torch.min(src), torch.max(src))
     )
 
     return out