Add undistorted input to InstanceGroup numpy and FrameGroup numpy

sleap/io/cameras.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 from typing import Any, Dict, Iterator, List, Optional, Set, Tuple, Union
 
+import cv2
 import cattr
 import numpy as np
 import toml
@@ -614,6 +615,7 @@ def numpy(
         pred_as_nan: bool = False,
         invisible_as_nan=True,
         cams_to_include: Optional[List[Camcorder]] = None,
+        undistort: bool = False,
     ) -> np.ndarray:
         """Return instances as a numpy array of shape (n_views, n_nodes, 2).
 
@@ -631,6 +633,8 @@ def numpy(
             cams_to_include: List of `Camcorder`s to include in the numpy array. If
                 None, then all `Camcorder`s in the `CameraCluster` are included. Default
                 is None.
+            undistort: If True, then undistort the points using cv2.undistortPoints.
+                Default is False.
 
         Returns:
             Numpy array of shape (n_views, n_nodes, 2).
@@ -655,6 +659,16 @@ def numpy(
             instance_numpy: np.ndarray = instance.get_points_array(
                 invisible_as_nan=invisible_as_nan
             )  # N x 2
+
+            if undistort:
+                instance_numpy_shape = instance_numpy.shape
+                instance_numpy = instance_numpy.reshape(-1, 2)
+                instance_numpy = cv2.undistortPoints(
+                    instance_numpy.astype("float64"),
+                    cameraMatrix=cam.camera.matrix,
+                    distCoeffs=cam.camera.dist,
+                ).reshape(instance_numpy_shape)
+
             instance_numpys.append(instance_numpy)
 
         return np.stack(instance_numpys, axis=0)  # M x N x 2
@@ -1844,6 +1858,7 @@ def numpy(
         instance_groups: Optional[List[InstanceGroup]] = None,
         pred_as_nan: bool = False,
         invisible_as_nan: bool = True,
+        undistort: bool = False,
     ) -> np.ndarray:
         """Numpy array of all `InstanceGroup`s in `FrameGroup.cams_to_include`.
 
@@ -1854,6 +1869,8 @@ def numpy(
                 self.dummy_instance. Default is False.
             invisible_as_nan: If True, then replaces invisible points with nan. Default
                 is True.
+            undistort: If True, then undistort the points. Default is False.
+
         Returns:
             Numpy array of shape (M, T, N, 2) where M is the number of views (determined
             by self.cams_to_include), T is the number of `InstanceGroup`s, N is the
@@ -1878,6 +1895,7 @@ def numpy(
                 pred_as_nan=pred_as_nan,
                 invisible_as_nan=invisible_as_nan,
                 cams_to_include=self.cams_to_include,
+                undistort=undistort,
             )  # M=include x N x 2
             instance_group_numpys.append(instance_group_numpy)
 

tests/io/test_cameras.py

@@ -693,28 +693,10 @@ def test_instance_group(
     frame_group = session.frame_groups[frame_idx]
     instance_group = frame_group.instance_groups[0]
 
-    # Test `numpy` method
-    instance_group_numpy = instance_group.numpy()
-    n_views, n_nodes, n_coords = instance_group_numpy.shape
-    assert n_views == len(instance_group.camera_cluster.cameras)
-    assert n_nodes == len(instance_group.dummy_instance.skeleton.nodes)
-    assert n_coords == 2
-    # Different instance groups should have different coordinates
-    for inst_idx, _ in enumerate(instance_group.instances[:-1]):
-        assert not np.allclose(
-            instance_group_numpy[:, inst_idx],
-            instance_group_numpy[:, inst_idx + 1],
-            equal_nan=True,
-        )
-    # Different views should have different coordinates
-    for view_idx, _ in enumerate(instance_group.camera_cluster.cameras[:-1]):
-        assert not np.allclose(
-            instance_group_numpy[view_idx],
-            instance_group_numpy[view_idx + 1],
-            equal_nan=True,
-        )
-
     # Test `update_points` method
+    n_views = len(session.cams_to_include)
+    n_nodes = len(instance_group.dummy_instance.skeleton.nodes)
+    n_coords = 2
     assert not np.all(instance_group.numpy(invisible_as_nan=False) == 72317)
     # Remove some Instances to "expose" underlying PredictedInstances
     for inst in instance_group.instances[:2]:
@@ -752,6 +734,46 @@ def test_instance_group(
     assert cam not in instance_group.cameras
 
 
+def test_instance_group_numpy(multiview_min_session_frame_groups: Labels):
+    """Test `InstanceGroup.numpy` method."""
+    labels = multiview_min_session_frame_groups
+    session = labels.sessions[0]
+    frame_group = session.frame_groups[0]
+    instance_group = frame_group.instance_groups[0]
+
+    instance_group_numpy = instance_group.numpy()
+    n_views, n_nodes, n_coords = instance_group_numpy.shape
+    assert n_views == len(instance_group.camera_cluster.cameras)
+    assert n_nodes == len(instance_group.dummy_instance.skeleton.nodes)
+    assert n_coords == 2
+
+    # Different instance groups should have different coordinates
+    for inst_idx, _ in enumerate(instance_group.instances[:-1]):
+        assert not np.allclose(
+            instance_group_numpy[:, inst_idx],
+            instance_group_numpy[:, inst_idx + 1],
+            equal_nan=True,
+        )
+
+    # Different views should have different coordinates
+    for view_idx, _ in enumerate(instance_group.camera_cluster.cameras[:-1]):
+        assert not np.allclose(
+            instance_group_numpy[view_idx],
+            instance_group_numpy[view_idx + 1],
+            equal_nan=True,
+        )
+
+    # Test for undisorted points
+    instance_group_numpy_0 = instance_group.numpy(undistort=False)
+    instance_group_numpy_undistorted = instance_group.numpy(undistort=True)
+    assert np.allclose(
+        instance_group_numpy_0, instance_group_numpy, atol=1e-3, equal_nan=True
+    )
+    assert not np.allclose(
+        instance_group_numpy, instance_group_numpy_undistorted, equal_nan=True
+    )
+
+
 def test_frame_group(
     multiview_min_session_labels: Labels, multiview_min_session_frame_groups: Labels
 ):
@@ -857,18 +879,6 @@ def test_frame_group(
     with pytest.raises(ValueError):
         frame_group.cams_to_include = session.linked_cameras
 
-    # Test `numpy` method
-    frame_group_np = frame_group.numpy()
-    n_views, n_inst_groups, n_nodes, n_coords = frame_group_np.shape
-    assert n_views == len(frame_group.cams_to_include)
-    assert n_inst_groups == len(frame_group.instance_groups)
-    assert n_nodes == len(labels.skeleton.nodes)
-    assert n_coords == 2
-    # Different instance groups should have different coordinates
-    assert not np.allclose(frame_group_np[:, 0], frame_group_np[:, 1], equal_nan=True)
-    # Different views should have different coordinates
-    assert not np.allclose(frame_group_np[0], frame_group_np[1], equal_nan=True)
-
     # Test `get_instance_group`
     instance_group = frame_group.instance_groups[0]
     camera = session.cameras[0]
@@ -1049,6 +1059,29 @@ def test_frame_group(
     assert np.all(frame_group_numpy[value_mask] == value)  # Updated to value
 
 
+def test_frame_group_numpy(multiview_min_session_frame_groups: Labels):
+    """Test `FrameGroup.numpy` method."""
+    labels = multiview_min_session_frame_groups
+    session = labels.sessions[0]
+    frame_group = session.frame_groups[0]
+
+    # Test `numpy` method
+    frame_group_np = frame_group.numpy()
+    n_views, n_inst_groups, n_nodes, n_coords = frame_group_np.shape
+    assert n_views == len(frame_group.cams_to_include)
+    assert n_inst_groups == len(frame_group.instance_groups)
+    assert n_nodes == len(labels.skeleton.nodes)
+    assert n_coords == 2
+    # Different instance groups should have different coordinates
+    assert not np.allclose(frame_group_np[:, 0], frame_group_np[:, 1], equal_nan=True)
+    # Different views should have different coordinates
+    assert not np.allclose(frame_group_np[0], frame_group_np[1], equal_nan=True)
+    frame_group_np_0 = frame_group.numpy(undistort=False)
+    frame_group_np_undistorted = frame_group.numpy(undistort=True)
+    assert np.allclose(frame_group_np_0, frame_group_np, atol=1e-3, equal_nan=True)
+    assert not np.allclose(frame_group_np, frame_group_np_undistorted, equal_nan=True)
+
+
 def test_cameras_are_not_sorted():
     """Test that cameras are not sorted in `RecordingSession`.