Merge pull request #235 from Visual-Behavior/fix-kumler

Fix kumler
Visual-Behavior · Oct 27, 2022 · db7f770 · db7f770
2 parents 9b6e834 + 673a7ab
commit db7f770
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 11 deletions.
diff --git a/aloscene/depth.py b/aloscene/depth.py
@@ -233,7 +233,7 @@ def as_points3d(
         intrinsic = camera_intrinsic if camera_intrinsic is not None else self.cam_intrinsic
         projection = projection if projection is not None else self.projection
         distortion = distortion if distortion is not None else self.distortion
-        assert projection in ["pinhole", "equidistant"], "Only pinhole and equidistant are supported."
+        assert projection in ["pinhole", "equidistant", "kumler_bauer"], "Only pinhole, equidistant and kumler_bauer are supported."
 
         # if self is not planar depth, we must convert to planar depth before projecting to 3d points
         if self.is_planar:
@@ -285,7 +285,13 @@ def as_points3d(
             for _ in range(len(target_shape[:-1])):
                 theta = theta.unsqueeze(0)
             r = torch.tan(theta)
-            focal_length = focal_length * theta * distortion / r.abs()
+
+            if projection == "equidistant":
+                focal_length = focal_length * theta * distortion / r.abs()
+            elif projection == "kumler_bauer":
+                focal_length = (
+                    distortion[0] * torch.sin(distortion[1] * theta) * focal_length / (distortion[2] * r.abs())
+                )
 
             # find points behind camera
             behind = theta > (np.pi / 2)
@@ -389,7 +395,11 @@ def as_euclidean(
         if not self.is_planar:
             print("This tensor is already a euclidian depth tensor so no transform is performed")
             return self.clone()
-        assert projection in ["pinhole", "equidistant"], "Only pinhole and equidistant projection are supported"
+        assert projection in [
+            "pinhole",
+            "equidistant",
+            "kumler_bauer",
+        ], "Only pinhole, equidistant and kumler_bauer projection are supported"
 
         planar = self
         camera_intrinsic = camera_intrinsic if camera_intrinsic is not None else self.cam_intrinsic
@@ -430,7 +440,11 @@ def as_planar(
         if self.is_planar:
             print("This tensor is already a planar depth tensor so no transform is done.")
             return self.clone()
-        assert projection in ["pinhole", "equidistant"], "Only pinhole and equidistant projection are supported"
+        assert projection in [
+            "pinhole",
+            "equidistant",
+            "kumler_bauer",
+        ], "Only pinhole, equidistant and kumler_bauer projection are supported"
 
         euclidean = self
         camera_intrinsic = camera_intrinsic if camera_intrinsic is not None else self.cam_intrinsic

diff --git a/aloscene/tensors/spatial_augmented_tensor.py b/aloscene/tensors/spatial_augmented_tensor.py
@@ -72,8 +72,8 @@ def __init__(self, x, *args, **kwargs):
     def dummy(cls, size: tuple, names: tuple):
         dummy_class = cls(torch.ones(size))
         mask_size = list(size)
-        if 'C' in names:
-            mask_size[names.index('C')] = 1
+        if "C" in names:
+            mask_size[names.index("C")] = 1
         mask_size = tuple(mask_size)
         dummy_class.append_mask(aloscene.Mask(torch.ones(mask_size), names=names))
         return dummy_class

diff --git a/aloscene/utils/depth_utils.py b/aloscene/utils/depth_utils.py
@@ -7,10 +7,7 @@
 
 
 def coords2rtheta(
-    K,
-    size: Tuple[int, int],
-    distortion: Union[float, Tuple[float, float]],
-    projection: str = "pinhole"
+    K, size: Tuple[int, int], distortion: Union[float, Tuple[float, float]], projection: str = "pinhole"
 ):
     """Compute r_d and theta from image coordinates.
 
@@ -47,7 +44,9 @@ def coords2rtheta(
         theta = r_d / (focal * distortion)
     elif projection == "kumler_bauer":
         # distortion [k1, k2, focal_meter]
-        assert isinstance(distortion, Sequence) and len(distortion) == 2, "Kumler-Bauer projection needs two distortion coefficients (alpha, beta)"
+        assert (
+            isinstance(distortion, Sequence) and len(distortion) == 3
+        ), "Kumler-Bauer projection needs two distortion coefficients (alpha, beta)"
         fm = distortion[2]
         theta = torch.arcsin(fm / distortion[0] * r_d / focal) / distortion[1]
     else: