feat(project): add points_to_depths function to calculate z-depths

camtools/project.py

@@ -63,6 +63,37 @@ def points_to_pixels(
     return pixels
 
 
+def points_to_depths(
+    points: Float[np.ndarray, "n 3"],
+    T: Float[np.ndarray, "4 4"],
+) -> Float[np.ndarray, "n"]:
+    """
+    Convert 3D points in world coordinates to z-depths in camera coordinates.
+
+    Args:
+        points: (N, 3) array of 3D points in world coordinates.
+        T: (4, 4) camera extrinsic matrix (world-to-camera transformation).
+
+    Returns:
+        (N,) array of z-depths in camera coordinates. Positive values indicate
+        points in front of the camera, negative values indicate points behind
+        the camera.
+
+    Note: The depth is z-depth instead of distance to the camera center.
+    """
+    sanity.assert_T(T)
+    sanity.assert_shape_nx3(points, name="points")
+
+    # (N, 3) -> (N, 4)
+    points_homo = convert.to_homo(points)
+    # Transform to camera coordinates: (N, 4)
+    points_camera = (T @ points_homo.T).T
+    # Extract z-coordinate: (N,)
+    depths = points_camera[:, 2]
+
+    return depths
+
+
 def im_depth_to_point_cloud(
     im_depth: Float[np.ndarray, "h w"],
     K: Float[np.ndarray, "3 3"],

test/test_project.py

@@ -0,0 +1,38 @@
+import numpy as np
+import open3d as o3d
+
+import camtools as ct
+
+
+def test_points_to_depths(visualize=True):
+    # Identity camera pose (looking at +Z axis)
+    T = np.eye(4)
+    fx = fy = 500
+    cx = 320
+    cy = 240
+    K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
+
+    # Create points on a plane z units away from origin
+    num_points = 1000
+    z = 3.0
+    x = np.random.uniform(-4, 4, num_points)
+    y = np.random.uniform(-4, 4, num_points)
+    points = np.column_stack([x, y, np.full(num_points, z)])
+
+    # Test depths are all close to z
+    depths = ct.project.points_to_depths(points, T)
+    assert np.allclose(
+        depths, z
+    ), f"Depths should be {z}, got {depths.min():.2f}-{depths.max():.2f}"
+
+    if visualize:
+        pcd = o3d.geometry.PointCloud()
+        pcd.points = o3d.utility.Vector3dVector(points)
+        pcd.paint_uniform_color([1, 0, 0])  # Red points
+        frustum = ct.camera.create_camera_frustums(
+            Ks=[K],
+            Ts=[T],
+            size=1.0,
+        )
+        axes = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1)
+        o3d.visualization.draw_geometries([pcd, frustum, axes])