Merge pull request #16 from Leoooo333/master-1

Add a functional argument to decide whether skip clip before final result.
haruishi43 · Aug 16, 2023 · 120d055 · 120d055
2 parents e7bf731 + 7dd2916
commit 120d055
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/equilib/equi2pers/numpy.py b/equilib/equi2pers/numpy.py
@@ -131,6 +131,7 @@ def run(
     z_down: bool,
     mode: str,
     override_func: Optional[Callable[[], Any]] = None,
+    clip_output: bool = True,
 ) -> np.ndarray:
     """Run Equi2Pers
 
@@ -224,7 +225,7 @@ def run(
 
     out = (
         out.astype(equi_dtype)
-        if equi_dtype == np.dtype(np.uint8)
+        if equi_dtype == np.dtype(np.uint8) or not clip_output
         else np.clip(out, 0.0, 1.0)
     )
 

diff --git a/equilib/equi2pers/torch.py b/equilib/equi2pers/torch.py
@@ -117,6 +117,7 @@ def run(
     z_down: bool,
     mode: str,
     backend: str = "native",
+    clip_output: bool = True,
 ) -> torch.Tensor:
     """Run Equi2Pers
 
@@ -242,7 +243,7 @@ def run(
 
     out = (
         out.type(equi_dtype)
-        if equi_dtype == torch.uint8
+        if equi_dtype == torch.uint8 or not clip_output
         else torch.clip(out, 0.0, 1.0)
     )