"""Back-project an RGB pixel to 3D using the depth sensor, with proper IR↔RGB registration. Why this exists --------------- Nuwa-HP60C is RGB-D with two physically distinct cameras (RGB + IR) that have different intrinsics and a small baseline: fx_rgb=593.6 fx_ir=445.8 cx_rgb=348.1 cx_ir=316.6 t_ir_to_rgb ≈ (-12.05, +0.08, -0.03) mm Nuwa's `/api/measure?u=&v=` interprets `(u,v)` as **IR/depth-image** pixel coords and back-projects with IR intrinsics. So if we have a pixel `(u,v)` in the **RGB** image (e.g. a YOLO bbox center, an ArUco corner) and we want the actual 3D point in the RGB camera frame, we need to handle the registration ourselves. Algorithm (single fixed-point iteration is enough at typical robot ranges): 1. Initial guess depth z0 from sampling depth[v_rgb, u_rgb] (depth is in IR coords — at small baselines it's already close). 2. Hypothesise the 3D point on the RGB ray at distance z0. 3. Project that 3D point into the IR image → (u_ir, v_ir). 4. Sample depth at (u_ir, v_ir) → z1. 5. Use (u_ir, v_ir, z1) to back-project with IR intrinsics → P_ir. 6. Transform P_ir → P_rgb via the IR→RGB extrinsics. Done. Two iterations of step 2-4 reduce residual error to sub-mm even with mediocre initial guesses; we do two unconditionally because they're cheap. """ from __future__ import annotations from dataclasses import dataclass from typing import Any import cv2 import numpy as np from ..clients import NuwaClient @dataclass(frozen=True) class CameraRegistration: """Cached intrinsics+extrinsics for a Nuwa-HP60C-style RGB-D camera.""" K_rgb: np.ndarray # 3x3 K_ir: np.ndarray # 3x3 R_ir_to_rgb: np.ndarray # 3x3 t_ir_to_rgb: np.ndarray # (3,) in **meters** rgb_width: int rgb_height: int ir_width: int ir_height: int @classmethod def from_nuwa(cls, intr: dict[str, Any]) -> "CameraRegistration": rgb = intr.get("rgb") or {} ir = intr.get("ir") or {} ext = intr.get("extrinsics_ir_to_rgb") or {} if not rgb or not ir: raise ValueError(f"Nuwa intrinsics must contain 'rgb' and 'ir' blocks: {intr!r}") K_rgb = np.array([[rgb["fx"], 0, rgb["cx"]], [0, rgb["fy"], rgb["cy"]], [0, 0, 1]], dtype=np.float64) K_ir = np.array([[ir["fx"], 0, ir["cx"]], [0, ir["fy"], ir["cy"]], [0, 0, 1]], dtype=np.float64) R = np.asarray(ext.get("rotation") or [1, 0, 0, 0, 1, 0, 0, 0, 1], dtype=np.float64) if R.size == 9: R = R.reshape(3, 3) elif R.shape != (3, 3): raise ValueError(f"extrinsics rotation must be 9 elements or 3x3, got {R.shape}") # Translation in Nuwa is in mm; we convert to meters once. t_mm = np.asarray(ext.get("translation") or [0, 0, 0], dtype=np.float64).flatten() t = t_mm / 1000.0 return cls( K_rgb=K_rgb, K_ir=K_ir, R_ir_to_rgb=R, t_ir_to_rgb=t, rgb_width=int(rgb.get("width") or 640), rgb_height=int(rgb.get("height") or 480), ir_width=int(ir.get("width") or 640), ir_height=int(ir.get("height") or 480), ) def fetch_depth_mm(nuwa: NuwaClient) -> np.ndarray: """Fetch the current depth frame as an int32 array (mm), shape (H, W). Zero values mean 'no depth' — caller should mask them out.""" png = nuwa.depth_png() arr = np.frombuffer(png, dtype=np.uint8) img = cv2.imdecode(arr, cv2.IMREAD_UNCHANGED) if img is None: raise RuntimeError("could not decode Nuwa depth PNG") if img.dtype != np.uint16: # Some configurations return 8-bit (color preview); we need raw mm. raise RuntimeError(f"expected 16-bit depth PNG, got dtype={img.dtype}") return img.astype(np.int32) def _sample_depth(depth_mm: np.ndarray, u: float, v: float, radius: int = 4) -> int: """Nearest valid (non-zero) depth value within `radius` of (u,v) on the depth image. Returns 0 if nothing found.""" h, w = depth_mm.shape u0, v0 = int(round(u)), int(round(v)) if u0 < 0 or v0 < 0 or u0 >= w or v0 >= h: return 0 z = int(depth_mm[v0, u0]) if z > 0: return z # Spiral / box search outward to `radius`. for r in range(1, radius + 1): u_lo, u_hi = max(0, u0 - r), min(w, u0 + r + 1) v_lo, v_hi = max(0, v0 - r), min(h, v0 + r + 1) patch = depth_mm[v_lo:v_hi, u_lo:u_hi] valid = patch[patch > 0] if valid.size: return int(valid.min()) # nearest along ray = smallest depth return 0 def back_project_rgb_pixel( u_rgb: float, v_rgb: float, depth_mm: np.ndarray, reg: CameraRegistration, *, sample_radius: int = 4, n_iters: int = 2, depth_correction: dict[str, Any] | None = None, ) -> dict[str, Any]: """Back-project an RGB pixel using the registered depth image. If `depth_correction` is supplied (dict with `enabled`, `scale`, `offset_mm`), each raw depth sample is corrected before back-projection. See `depth_correction.apply_correction`. Returns a dict: { "ok": bool, "rgb_pixel": [u, v], "ir_pixel": [u_ir, v_ir], "depth_mm_raw": int, # what the sensor reported "depth_mm": int, # corrected (== raw if correction off) "point_rgb_m": [X, Y, Z], # in RGB camera frame, OpenCV convention "point_ir_m": [X, Y, Z], # in IR camera frame } """ from .depth_correction import apply_correction out: dict[str, Any] = {"ok": False, "rgb_pixel": [float(u_rgb), float(v_rgb)]} K_rgb_inv = np.linalg.inv(reg.K_rgb) ray_rgb = K_rgb_inv @ np.array([u_rgb, v_rgb, 1.0]) z_raw_first = _sample_depth(depth_mm, u_rgb, v_rgb, radius=sample_radius) if z_raw_first <= 0: out["error"] = f"no valid depth near RGB pixel ({u_rgb:.0f},{v_rgb:.0f})" return out z_guess = apply_correction(z_raw_first, depth_correction) P_rgb = (z_guess / 1000.0) * ray_rgb u_ir = v_ir = None z_raw = z_raw_first z_corr = z_guess for _ in range(max(1, n_iters)): P_ir = reg.R_ir_to_rgb.T @ (P_rgb - reg.t_ir_to_rgb) if P_ir[2] <= 0: out["error"] = "RGB ray hits behind IR camera" return out u_ir = reg.K_ir[0, 0] * P_ir[0] / P_ir[2] + reg.K_ir[0, 2] v_ir = reg.K_ir[1, 1] * P_ir[1] / P_ir[2] + reg.K_ir[1, 2] z_raw = _sample_depth(depth_mm, u_ir, v_ir, radius=sample_radius) if z_raw <= 0: break z_corr = apply_correction(z_raw, depth_correction) K_ir_inv = np.linalg.inv(reg.K_ir) P_ir_new = (z_corr / 1000.0) * (K_ir_inv @ np.array([u_ir, v_ir, 1.0])) P_rgb = reg.R_ir_to_rgb @ P_ir_new + reg.t_ir_to_rgb if z_raw <= 0: out["error"] = f"no valid depth near IR pixel ({u_ir:.0f},{v_ir:.0f})" return out out.update({ "ok": True, "ir_pixel": [float(u_ir), float(v_ir)], "depth_mm_raw": int(z_raw), "depth_mm": int(z_corr), "point_rgb_m": P_rgb.tolist(), "point_ir_m": (reg.R_ir_to_rgb.T @ (P_rgb - reg.t_ir_to_rgb)).tolist(), }) return out