diff --git a/src/presa_mano/api/grasp_routes.py b/src/presa_mano/api/grasp_routes.py index 2318fe2..3500bfc 100644 --- a/src/presa_mano/api/grasp_routes.py +++ b/src/presa_mano/api/grasp_routes.py @@ -2,8 +2,10 @@ from __future__ import annotations +import logging from typing import Any, Optional +import cv2 import numpy as np import time from dataclasses import asdict @@ -11,7 +13,7 @@ from fastapi import APIRouter, HTTPException from pydantic import BaseModel, Field from ..calibration import CalibrationStore -from ..clients import G1Unavailable +from ..clients import G1Unavailable, Sam21Error, Sam21Unavailable from ..config import get_config from ..geometry import transform_points from ..grasp import GraspPlanner, compute_axes @@ -24,13 +26,27 @@ from ..grasp.visual_servo import ( set_work_plane, ) from ..perception import ObjectDetector, filter_outliers, nuwa_to_cv_frame +from ..perception.depth_align import align_depth_to_rgb +from ..perception.depth_to_rgb import CameraRegistration from .perception_routes import _bbox_to_roi from ..timing import Stopwatch from . import _deps +log = logging.getLogger(__name__) + router = APIRouter() +class _Sam21Fallback(Exception): + """Internal sentinel: the SAM21 path could not produce a usable cloud. + Never escapes _object_points_in_base — caught and turned into a log line + + legacy bbox path.""" + + +_SAM21_BUG_CODES = frozenset({"SHAPE_MISMATCH", "BBOX_OUT_OF_BOUNDS", + "IMAGE_DECODE_ERROR", "INVALID_INPUT"}) + + class PlanReq(BaseModel): conf: Optional[float] = None class_name: Optional[str] = None @@ -66,14 +82,23 @@ def _object_points_in_base(conf: float | None, class_name: str | None, sw: Stopw dets.sort(key=lambda d: d["confidence"], reverse=True) top = dets[0] bbox = _bbox_to_roi(top["bbox"]) + meta: dict[str, Any] = {"detection": top, "T_base_camera": T_base_cam.tolist()} + + if cfg.get("grasp", {}).get("use_sam21", True): + try: + pts_cv = _points_via_sam21(nuwa, bbox, sw) + with sw.step("to_base_frame"): + pts_base = transform_points(T_base_cam, pts_cv) + meta["source"] = "sam21" + meta["n_points"] = int(len(pts_base)) + return pts_base, meta + except _Sam21Fallback as e: + log.warning("SAM21 path failed, falling back to bbox: %s", e) + meta["sam21_fallback_reason"] = str(e) pts_base = _points_via_bbox(nuwa, T_base_cam, bbox, sw) - meta = { - "detection": top, - "T_base_camera": T_base_cam.tolist(), - "n_points": int(len(pts_base)), - "source": "bbox_roi", - } + meta["source"] = "bbox_roi" + meta["n_points"] = int(len(pts_base)) return pts_base, meta @@ -104,6 +129,75 @@ def _points_via_bbox(nuwa, T_base_cam: np.ndarray, bbox: tuple[int, int, int, in return pts_base +def _points_via_sam21(nuwa, bbox: tuple[int, int, int, int], sw: Stopwatch) -> np.ndarray: + """SAM21 path: bbox → segmented point cloud in RGB-camera/OpenCV frame. + + Returns points in **meters**. Raises _Sam21Fallback on any failure that + should trigger the bbox path; lets unexpected exceptions propagate as + _Sam21Fallback with repr(e). + """ + try: + with sw.step("nuwa_grab"): + rgb_jpeg = nuwa.rgb_jpeg() + depth_png_in = nuwa.depth_png() + intr_full = nuwa.intrinsics() + rgb_intr = intr_full.get("rgb") + if not rgb_intr: + raise _Sam21Fallback("nuwa intrinsics missing 'rgb' block") + + # 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. Verified by inspection. + bgr = cv2.imdecode(np.frombuffer(rgb_jpeg, np.uint8), cv2.IMREAD_COLOR) + if bgr is None: + raise _Sam21Fallback("could not decode Nuwa RGB JPEG") + ok, rgb_png_buf = cv2.imencode(".png", bgr) + if not ok: + 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: + raise _Sam21Fallback("could not decode Nuwa depth PNG") + if depth_ir.dtype != np.uint16: + raise _Sam21Fallback(f"depth PNG dtype expected uint16, got {depth_ir.dtype}") + 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: + raise _Sam21Fallback("could not encode aligned depth as PNG") + depth_png_out = bytes(depth_png_buf) + except _Sam21Fallback: + raise + except Exception as e: + raise _Sam21Fallback(f"prep failed: {e!r}") from e + + with sw.step("sam21"): + try: + pts_cv = _deps.sam21().segment_pointcloud( + rgb_png=rgb_png, + depth_png=depth_png_out, + intrinsics=rgb_intr, + bbox=bbox, + depth_scale_mm=1.0, + ) + except Sam21Unavailable as 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) + 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 + + if pts_cv.shape[0] < 10: + raise _Sam21Fallback(f"too few points: {pts_cv.shape[0]}") + return pts_cv + + @router.post("/plan") def plan(req: PlanReq) -> dict[str, Any]: sw = Stopwatch()