Presa-Mano/src/presa_mano/perception/depth_to_rgb.py

"""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