feat: add distortion parameters to BaseDataset (#88)

* feat: add distortion parameters to base dataset

dgp/datasets/base_dataset.py

@@ -34,6 +34,28 @@
 from dgp.utils.protobuf import open_pbobject
 
 AVAILABLE_DATUM_TYPES = ("image", "point_cloud")
+AVAILABLE_DISTORTION_PARAMS = (
+    'k1',
+    'k2',
+    'k3',
+    'k4',
+    'k5',
+    'k6',
+    'p1',
+    'p2',
+    'alpha',
+    'beta',
+    'xi',
+    's1',
+    's2',
+    's3',
+    's4',
+    'taux',
+    'tauy',
+    'fov',
+    'fisheye',
+    'w',
+)
 
 
 class SceneContainer:
@@ -164,7 +186,7 @@ def get_sample(self, sample_idx_in_scene):
         assert sample_idx_in_scene >= 0 and sample_idx_in_scene < len(self.datum_index)
         return self.samples[self.datum_index.coords["samples"][sample_idx_in_scene].data]
 
-    @lru_cache(maxsize=None)
+    @lru_cache(maxsize=1024)
     def get_datum_type(self, datum_name):
         """Get datum type based on the datum name"""
         for datum in self.data:
@@ -1062,9 +1084,20 @@ def _get_scene_calibration_table(calibration_files):
             for (name, intrinsic, extrinsic) in zip(calibration.names, calibration.intrinsics, calibration.extrinsics):
                 p_WS = Pose.load(extrinsic)
                 # If the intrinsics are invalid, i.e. fx = fy = 0, then it is
-                # assumed to be a LIDAR or RADAR sensor.
+                # assumed to be a LIDAR sensor.
+
+                # TODO: this needs a refactor for two reasons,
+                # 1. This uses a hardcoded list of distortion parameters, it should instead use the proto defintion
+                # 2. We probably want the camera class to calculate and cache the remaps for undistortion
+
+                # Get a dictionary of distortion parameters
+                distortion = {}
+                for k in AVAILABLE_DISTORTION_PARAMS:
+                    if hasattr(intrinsic, k):
+                        distortion[k] = getattr(intrinsic, k)
+
                 cam = Camera.from_params(
-                    intrinsic.fx, intrinsic.fy, intrinsic.cx, intrinsic.cy, p_WS
+                    intrinsic.fx, intrinsic.fy, intrinsic.cx, intrinsic.cy, p_WS, distortion=distortion
                 ) if intrinsic.fx > 0 and intrinsic.fy > 0 else None
                 calibration_table[(calibration_key, name.lower())] = (p_WS, cam)
         return calibration_table
@@ -1577,12 +1610,15 @@ def get_image_from_datum(self, scene_idx, sample_idx_in_scene, datum_name):
         sample = self.get_sample(scene_idx, sample_idx_in_scene)
 
         if self.calibration_table:
-            camera_intrinsics = self.get_camera_calibration(sample.calibration_key, datum.id.name).K
+            camera = self.get_camera_calibration(sample.calibration_key, datum.id.name)
+            camera_intrinsics = camera.K
+            camera_distortion = camera.D
             pose_VC = self.get_sensor_extrinsics(sample.calibration_key, datum.id.name)
             # Get ego-pose for the image (at the corresponding image timestamp t=Tc)
             pose_WC_Tc = Pose.load(datum.datum.image.pose)
         else:
             camera_intrinsics = None
+            camera_distortion = None
             pose_VC = None
             pose_WC_Tc = Pose()
 
@@ -1594,6 +1630,7 @@ def get_image_from_datum(self, scene_idx, sample_idx_in_scene, datum_name):
             "datum_name": datum.id.name,
             "rgb": image,
             "intrinsics": camera_intrinsics,
+            "distortion": camera_distortion,
             "extrinsics": pose_VC,
             "pose": pose_WC_Tc
         })

dgp/utils/camera.py

@@ -1,5 +1,6 @@
-# Copyright 2021 Toyota Research Institute.  All rights reserved.
+# Copyright 2021-2022 Toyota Research Institute.  All rights reserved.
 """General-purpose class for cameras."""
+import logging
 from functools import lru_cache
 
 import cv2
@@ -45,14 +46,17 @@ def generate_depth_map(camera, Xw, shape):
     return depth
 
 
-def pbobject_from_camera_matrix(K):
+def pbobject_from_camera_matrix(K, distortion=None):
     """Convert camera intrinsic matrix into pb object.
- 
+
     Parameters
     ----------
     K: np.ndarray
         Camera Intrinsic Matrix
- 
+
+    distortion: dict[str, float]
+        Dictionary of distortion params i.e, k1,k2,p1,p2,k3,k4,xi,alpha etc
+
     Returns
     -------
     intrinsics: geometry_pb2.CameraIntrinsics
@@ -65,6 +69,12 @@ def pbobject_from_camera_matrix(K):
         intrinsics.cx = K[0, 2]
         intrinsics.cy = K[1, 2]
         intrinsics.skew = K[0, 1]
+
+    if distortion is not None:
+        for k, v in distortion.items():
+            # TODO: assert the proto contains this value
+            setattr(intrinsics, k, v)
+
     return intrinsics
 
 
@@ -114,13 +124,14 @@ def __init__(self, K, D=None, p_cw=None):
         K: np.ndarray (3x3)
             Camera calibration matrix.
 
-        D: np.ndarray (5,) or (H x W)
-            Distortion parameters or distortion map.
+        D: np.ndarray (5,) or (H x W) or Dict[str,float]
+            Distortion parameters or distortion map or dictionary of distortion values
 
         p_cw: dgp.utils.pose.Pose
             Pose from world to camera frame.
         """
         self.K = K
+        # TODO: refactor this class to support other camera models. This assumes Brown Conrady
         self.D = Distortion() if D is None else D
         self.p_cw = Pose() if p_cw is None else p_cw
         assert isinstance(self.K, np.ndarray)
@@ -140,7 +151,7 @@ def rodrigues(self):
         return rvec
 
     @classmethod
-    def from_params(cls, fx, fy, cx, cy, p_cw=None):
+    def from_params(cls, fx, fy, cx, cy, p_cw=None, distortion=None):
         """Create camera batch from calibration parameters.
 
         Parameters
@@ -160,17 +171,21 @@ def from_params(cls, fx, fy, cx, cy, p_cw=None):
         p_cw: Pose
             Pose from world to camera frame.
 
+        distortion_params: dict[str, float], default: None
+            Optional dictionary of distortion parameters k1,k2,.. etc.
+
         Returns
         ----------
         Camera
             Camera object with relevant intrinsics.
         """
+        # TODO: add skew
         K = np.float32([
             [fx, 0, cx],
             [0, fy, cy],
             [0, 0, 1],
         ])
-        return cls(K=K, p_cw=p_cw)
+        return cls(K=K, D=distortion, p_cw=p_cw)
 
     @property
     def fx(self):
@@ -237,7 +252,22 @@ def project(self, Xw):
         """
         _, C = Xw.shape
         assert C == 3
-        uv, _ = cv2.projectPoints(Xw, self.rodrigues, self.p_cw.tvec, self.K, self.D.coefficients)
+
+        # Since self.D can be a distoriton object or a dictionary, handle the appropriate case and
+        # throw a warning about the model being used. This currenty does not support fisheye.
+        distortion = np.zeros(5, dtype=np.float32)
+        if isinstance(self.D, Distortion):
+            distortion = self.D.coefficients
+        elif isinstance(self.D, dict):
+            logging.warning('Using Brown-Conrady (Opencv default) distortion model for projection.')
+            k1 = self.D.get('k1', 0.0)
+            k2 = self.D.get('k2', 0.0)
+            p1 = self.D.get('p1', 0.0)
+            p2 = self.D.get('p2', 0.0)
+            k3 = self.D.get('k3', 0.0)
+            distortion = np.array([k1, k2, p1, p2, k3])
+
+        uv, _ = cv2.projectPoints(Xw, self.rodrigues, self.p_cw.tvec, self.K, distortion)
         return uv.reshape(-1, 2)
 
     @staticmethod
@@ -279,6 +309,7 @@ def unproject(self, points_in_2d):
         points_in_3d: np.ndarray
             Array of shape (N, 3) of points projected into 3D
         """
+        logging.warning('unproject currently does not consider distortion parameters')
         rays = cv2.undistortPoints(points_in_2d[:, None], self.K, None)
         return cv2.convertPointsToHomogeneous(rays).reshape(-1, 3).astype(np.float32)
 

requirements.txt

@@ -6,7 +6,7 @@ coverage==6.0.2
 diskcache>=4.1.0
 grpcio==1.41.0
 grpcio-tools==1.41.0
-isort==4.3.21
+isort >=5.0,<6.0
 matplotlib>=3.0.3,<4.0
 opencv-python==4.5.3.56
 Pillow>=8.3.2

tests/test_agent_dataset.py

@@ -32,6 +32,7 @@ def test_prediction_agent_dataset_lite(self):
             'datum_name',
             'rgb',
             'intrinsics',
+            'distortion',
             'extrinsics',
             'pose',
             'datum_type',
@@ -80,6 +81,7 @@ def test_prediction_agent_dataset(self):
             'datum_name',
             'rgb',
             'intrinsics',
+            'distortion',
             'extrinsics',
             'pose',
             'datum_type',

tests/test_datasets.py

@@ -28,6 +28,7 @@ def _test_labeled_dataset(dataset):
             'datum_name',
             'pose',
             'intrinsics',
+            'distortion',
             'extrinsics',
             'bounding_box_2d',
             'bounding_box_3d',
@@ -80,6 +81,7 @@ def test_labeled_synchronized_scene_dataset(self):
             'datum_name',
             'pose',
             'intrinsics',
+            'distortion',
             'extrinsics',
             'bounding_box_2d',
             'bounding_box_3d',