feat(perception): vectorised IR→RGB depth alignment
Forward-project IR depth pixels into the RGB pixel grid with a z-buffer. Output uint16 mm at RGB resolution, ready to feed SAM21. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
1 parent
f898d9b0da
commit
5354c51d60
2 files changed
+186
No files matched your search
@@ -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)
|
||||||
@@ -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
|
||||||
Reference in new issue
Block a user