diff --git a/src/presa_mano/perception/depth_align.py b/src/presa_mano/perception/depth_align.py new file mode 100644 index 0000000..7c3fbfd --- /dev/null +++ b/src/presa_mano/perception/depth_align.py @@ -0,0 +1,98 @@ +"""Forward-project an IR depth image into RGB pixel coordinates. + +Vectorised equivalent of librealsense `rs.align(stream.color)`. Each IR pixel +with valid depth is back-projected to 3D in the IR camera frame, transformed to +the RGB camera frame via the rigid extrinsics, and projected with RGB +intrinsics. On collisions the closest point wins (z-buffer). + +Output is a uint16 array in millimeters, shape (rgb_height, rgb_width), with 0 +on pixels that received no IR projection. +""" +from __future__ import annotations + +import numpy as np + +from .depth_to_rgb import CameraRegistration + +_SENTINEL = np.iinfo(np.uint32).max + + +def align_depth_to_rgb( + depth_ir_mm: np.ndarray, + reg: CameraRegistration, +) -> np.ndarray: + """Forward-project IR depth → RGB pixel grid. + + Parameters + ---------- + depth_ir_mm : (H_ir, W_ir) uint16 + Depth in millimeters, IR camera frame. Zero means "no depth". + reg : CameraRegistration + IR and RGB intrinsics + IR→RGB rigid transform (translation in meters). + + Returns + ------- + (reg.rgb_height, reg.rgb_width) uint16 + Depth in millimeters in RGB pixel coordinates; 0 where no IR pixel + projected. + """ + if depth_ir_mm.ndim != 2: + raise ValueError(f"depth_ir_mm must be 2D, got shape {depth_ir_mm.shape}") + if depth_ir_mm.shape != (reg.ir_height, reg.ir_width): + raise ValueError( + f"depth_ir_mm shape {depth_ir_mm.shape} ≠ IR resolution " + f"({reg.ir_height}, {reg.ir_width}) from registration" + ) + + H_rgb, W_rgb = reg.rgb_height, reg.rgb_width + H_ir, W_ir = reg.ir_height, reg.ir_width + + valid = depth_ir_mm > 0 + if not valid.any(): + return np.zeros((H_rgb, W_rgb), dtype=np.uint16) + + v_idx, u_idx = np.where(valid) + z_ir_m = depth_ir_mm[v_idx, u_idx].astype(np.float64) / 1000.0 + + # Back-project IR pixel → 3D point in IR camera frame. + K_ir_inv = np.linalg.inv(reg.K_ir) + pixels_h = np.vstack([u_idx.astype(np.float64), + v_idx.astype(np.float64), + np.ones_like(u_idx, dtype=np.float64)]) + rays_ir = K_ir_inv @ pixels_h + P_ir = rays_ir * z_ir_m + + # IR camera frame → RGB camera frame. + P_rgb = reg.R_ir_to_rgb @ P_ir + reg.t_ir_to_rgb.reshape(3, 1) + + z_rgb = P_rgb[2] + in_front = z_rgb > 1e-6 + P_rgb = P_rgb[:, in_front] + if P_rgb.shape[1] == 0: + return np.zeros((H_rgb, W_rgb), dtype=np.uint16) + + z_rgb = P_rgb[2] + u_rgb = reg.K_rgb[0, 0] * P_rgb[0] / z_rgb + reg.K_rgb[0, 2] + v_rgb = reg.K_rgb[1, 1] * P_rgb[1] / z_rgb + reg.K_rgb[1, 2] + + u_int = np.round(u_rgb).astype(np.int64) + v_int = np.round(v_rgb).astype(np.int64) + z_mm = np.round(z_rgb * 1000.0).astype(np.int64) + + in_bounds = ( + (u_int >= 0) & (u_int < W_rgb) + & (v_int >= 0) & (v_int < H_rgb) + & (z_mm > 0) & (z_mm < 65535) + ) + u_int = u_int[in_bounds] + v_int = v_int[in_bounds] + z_mm = z_mm[in_bounds].astype(np.uint32) + if z_mm.size == 0: + return np.zeros((H_rgb, W_rgb), dtype=np.uint16) + + # Z-buffer: initialise to sentinel, reduce-min, then sentinel→0. + zbuf = np.full(H_rgb * W_rgb, _SENTINEL, dtype=np.uint32) + flat_idx = v_int * W_rgb + u_int + np.minimum.at(zbuf, flat_idx, z_mm) + zbuf[zbuf == _SENTINEL] = 0 + return zbuf.reshape(H_rgb, W_rgb).astype(np.uint16) diff --git a/tests/test_depth_align.py b/tests/test_depth_align.py new file mode 100644 index 0000000..095ac27 --- /dev/null +++ b/tests/test_depth_align.py @@ -0,0 +1,88 @@ +"""Unit tests for align_depth_to_rgb.""" +from __future__ import annotations + +import numpy as np + +from presa_mano.perception.depth_align import align_depth_to_rgb +from presa_mano.perception.depth_to_rgb import CameraRegistration + + +def _identity_reg(W: int = 64, H: int = 48) -> CameraRegistration: + """Same K for RGB and IR, no extrinsic offset.""" + K = np.array([[60.0, 0.0, W / 2.0], + [0.0, 60.0, H / 2.0], + [0.0, 0.0, 1.0]], dtype=np.float64) + return CameraRegistration( + K_rgb=K.copy(), K_ir=K.copy(), + R_ir_to_rgb=np.eye(3, dtype=np.float64), + t_ir_to_rgb=np.zeros(3, dtype=np.float64), + rgb_width=W, rgb_height=H, + ir_width=W, ir_height=H, + ) + + +def test_identity_alignment_preserves_depth(): + """K_ir == K_rgb and no extrinsic → every IR pixel maps to itself.""" + reg = _identity_reg() + depth_ir = np.full((reg.ir_height, reg.ir_width), 1000, dtype=np.uint16) + depth_rgb = align_depth_to_rgb(depth_ir, reg) + assert depth_rgb.shape == (reg.rgb_height, reg.rgb_width) + assert depth_rgb.dtype == np.uint16 + nonzero = depth_rgb[depth_rgb > 0] + # Coverage close to full (rounding can drop a few border pixels). + assert nonzero.size >= 0.95 * depth_rgb.size, f"coverage too low: {nonzero.size}/{depth_rgb.size}" + # All filled pixels equal the input depth (±1 mm rounding). + assert np.all(np.abs(nonzero.astype(int) - 1000) <= 1) + + +def test_zero_depth_pixels_stay_zero(): + """IR pixels with depth=0 must not produce any RGB output.""" + reg = _identity_reg() + depth_ir = np.zeros((reg.ir_height, reg.ir_width), dtype=np.uint16) + depth_rgb = align_depth_to_rgb(depth_ir, reg) + assert depth_rgb.shape == (reg.rgb_height, reg.rgb_width) + assert int(depth_rgb.sum()) == 0 + + +def test_sparse_pixels_preserved(): + """Two isolated IR depths produce exactly two filled RGB pixels (identity reg).""" + reg = _identity_reg(W=64, H=48) + depth_ir = np.zeros((reg.ir_height, reg.ir_width), dtype=np.uint16) + depth_ir[0, 0] = 1000 + depth_ir[reg.ir_height - 1, reg.ir_width - 1] = 1000 + depth_rgb = align_depth_to_rgb(depth_ir, reg) + assert (depth_rgb > 0).sum() == 2 + + +def test_z_buffer_keeps_closer_point(): + """When two IR pixels collide on the same RGB pixel, the smaller depth wins. + + Force a collision with K_ir = 2 × K_rgb (same resolution), so adjacent IR + pixels round to the same RGB pixel.""" + W = H = 16 + K_rgb = np.array([[60.0, 0.0, W / 2.0], + [0.0, 60.0, H / 2.0], + [0.0, 0.0, 1.0]], dtype=np.float64) + K_ir = K_rgb.copy() + K_ir[0, 0] *= 2.0 # double focal length → two IR pixels collapse to one RGB pixel + K_ir[1, 1] *= 2.0 + reg = CameraRegistration( + K_rgb=K_rgb, K_ir=K_ir, + R_ir_to_rgb=np.eye(3, dtype=np.float64), + t_ir_to_rgb=np.zeros(3, dtype=np.float64), + rgb_width=W, rgb_height=H, ir_width=W, ir_height=H, + ) + depth_ir = np.zeros((H, W), dtype=np.uint16) + # Two IR pixels (5, 5) and (6, 5) both project to the same RGB pixel near (4, 5) + # (when projected through K_ir with depth 1m and re-projected through K_rgb). + depth_ir[5, 5] = 1500 + depth_ir[5, 6] = 800 + depth_rgb = align_depth_to_rgb(depth_ir, reg) + nonzero = depth_rgb[depth_rgb > 0] + # At least one RGB pixel must receive the closer (800) value; nothing further + # may be ≥ 1500 if a collision actually occurred. If they didn't collide, we + # see exactly two values [800, 1500]; if they did, we see only [800]. + assert nonzero.size in (1, 2) + assert int(nonzero.min()) == 800 + if nonzero.size == 1: + assert int(nonzero[0]) == 800 # collision happened, closer point won