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

"""Bridge: bbox → SAM21-segmented point cloud in OpenCV RGB-camera frame.

Shared by /api/grasp/* and /api/perception/* endpoints. The helper is
stateless and pure HTTP-call wiring; callers pass in the Nuwa and SAM21
clients built per-request via api/_deps.py.
"""
from __future__ import annotations

import logging
from dataclasses import dataclass

import cv2
import httpx
import numpy as np
from fastapi import HTTPException

from ..clients import NuwaError, Sam21Error, Sam21Unavailable
from ..timing import Stopwatch
from .depth_align import align_depth_to_rgb
from .depth_to_rgb import CameraRegistration
from .work_plane import apply_below_plane_filter, get_work_plane

log = logging.getLogger(__name__)


class Sam21Fallback(Exception):
    """The SAM21 path could not produce a usable cloud. Callers should
    fall back to the non-SAM21 path (bbox ROI cloud)."""


SAM21_BUG_CODES = frozenset({"SHAPE_MISMATCH", "BBOX_OUT_OF_BOUNDS",
                              "IMAGE_DECODE_ERROR", "INVALID_INPUT"})


@dataclass
class Sam21Segmentation:
    """Output di segment_with_sam21 — SAM21 puro senza plane filter."""
    points_cv_m: np.ndarray         # (N,3) OpenCV camera frame, metri
    rgb_jpeg: bytes                 # snapshot Nuwa usato per la segmentazione
    intrinsics_rgb: dict            # con width/height iniettati dal decoded RGB


def segment_with_sam21(nuwa, sam21, bbox: tuple[int, int, int, int],
                       sw: Stopwatch) -> Sam21Segmentation:
    """SAM21 puro: nuwa_grab + depth_align + segment, SENZA plane filter.

    Raises Sam21Fallback per failure semantiche (mask vuota, no depth, SAM21
    unreachable, decode fail). Re-raises HTTPException(502) per failure Nuwa.
    """
    try:
        with sw.step("nuwa_grab"):
            rgb_jpeg = nuwa.rgb_jpeg()
            depth_png_in = nuwa.depth_png()
            intr_full = nuwa.intrinsics()
    except (NuwaError, httpx.RequestError) as e:
        raise HTTPException(status_code=502, detail=f"Nuwa unreachable: {e}")

    if not intr_full.get("rgb"):
        log.warning("SAM21 path skipped: nuwa intrinsics missing 'rgb' block")
        raise Sam21Fallback("nuwa intrinsics missing 'rgb' block")

    try:
        # JPEG → PNG. cv2.imdecode gives BGR; cv2.imencode(".png", bgr) writes
        # PNG bytes with standard R,G,B color order (libpng handles the swap),
        # which SAM21's PIL-based decoder reads as RGB.
        bgr = cv2.imdecode(np.frombuffer(rgb_jpeg, np.uint8), cv2.IMREAD_COLOR)
        if bgr is None:
            log.warning("SAM21 path skipped: could not decode Nuwa RGB JPEG")
            raise Sam21Fallback("could not decode Nuwa RGB JPEG")
        h_rgb, w_rgb = bgr.shape[:2]
        ok, rgb_png_buf = cv2.imencode(".png", bgr)
        if not ok:
            log.warning("SAM21 path skipped: could not re-encode RGB as PNG")
            raise Sam21Fallback("could not re-encode RGB as PNG")
        rgb_png = bytes(rgb_png_buf)

        with sw.step("depth_align"):
            depth_ir = cv2.imdecode(np.frombuffer(depth_png_in, np.uint8),
                                    cv2.IMREAD_UNCHANGED)
            if depth_ir is None:
                log.warning("SAM21 path skipped: could not decode Nuwa depth PNG")
                raise Sam21Fallback("could not decode Nuwa depth PNG")
            if depth_ir.dtype != np.uint16:
                log.warning("SAM21 path skipped: depth PNG dtype %s (expected uint16)",
                            depth_ir.dtype)
                raise Sam21Fallback(f"depth PNG dtype expected uint16, got {depth_ir.dtype}")
            h_ir, w_ir = depth_ir.shape
            # Nuwa /api/intrinsics omits width/height — inject the real ones
            # from the decoded images so CameraRegistration and SAM21 agree on
            # frame sizes.
            intr_full = {
                **intr_full,
                "rgb": {**intr_full["rgb"], "width": w_rgb, "height": h_rgb},
                "ir": {**intr_full.get("ir", {}), "width": w_ir, "height": h_ir},
            }
            reg = CameraRegistration.from_nuwa(intr_full)
            depth_aligned = align_depth_to_rgb(depth_ir, reg)
            ok, depth_png_buf = cv2.imencode(".png", depth_aligned)
            if not ok:
                log.warning("SAM21 path skipped: could not encode aligned depth as PNG")
                raise Sam21Fallback("could not encode aligned depth as PNG")
            depth_png_out = bytes(depth_png_buf)
    except Sam21Fallback:
        raise
    except Exception as e:
        log.error("Unexpected SAM21 prep error: %r", e)
        raise Sam21Fallback(f"prep failed: {e!r}") from e

    with sw.step("sam21"):
        try:
            pts_cv_m = sam21.segment_pointcloud(
                rgb_png=rgb_png,
                depth_png=depth_png_out,
                intrinsics=intr_full["rgb"],
                bbox=bbox,
                depth_scale_mm=1.0,
            )
        except Sam21Unavailable as e:
            log.warning("SAM21 unreachable, falling back to bbox: %s", e)
            raise Sam21Fallback(f"unreachable: {e}") from e
        except Sam21Error as e:
            if e.code in SAM21_BUG_CODES:
                log.error("SAM21 rejected payload (%s): %s", e.code, e.message)
            else:
                log.warning("SAM21 returned %s: %s — falling back to bbox", e.code, e.message)
            raise Sam21Fallback(f"{e.code}: {e.message}") from e
        except Exception as e:
            log.error("Unexpected SAM21 error: %r", e)
            raise Sam21Fallback(f"unexpected: {e!r}") from e

    return Sam21Segmentation(
        points_cv_m=pts_cv_m,
        rgb_jpeg=rgb_jpeg,
        intrinsics_rgb=intr_full["rgb"],
    )


def points_via_sam21(nuwa, sam21, bbox: tuple[int, int, int, int],
                     sw: Stopwatch) -> tuple[np.ndarray, dict]:
    """Legacy wrapper: segment_with_sam21 + apply_below_plane_filter.
    Returns (points_cv_m, plane_info). Raises Sam21Fallback if <10 points
    after plane filter (unchanged semantics)."""
    seg = segment_with_sam21(nuwa, sam21, bbox, sw)
    pts_cv, plane_info = apply_below_plane_filter(seg.points_cv_m, get_work_plane())
    if pts_cv.shape[0] < 10:
        log.warning("SAM21 cloud has %d points after plane filter (<10), falling back to bbox",
                    pts_cv.shape[0])
        raise Sam21Fallback(
            f"too few points after plane filter: {pts_cv.shape[0]}"
            f" (before plane: {plane_info['n_points_before_plane']})"
        )
    return pts_cv, plane_info