Presa-Mano/docs/superpowers/specs/2026-05-11-sam21-segmentation-design.md

SAM21 segmentation in the grasp pipeline

Status: approved (brainstorming → ready for implementation plan) Date: 2026-05-11 Scope: presa-mano orchestrator, grasp/TrovaUtensile flow

Goal

Replace the current bbox-only point extraction in the grasp pipeline with a SAM2.1-segmented point cloud, so that the PCA → GraspPlanner stage operates on points that belong to the object itself, not to the bbox background. The bbox remains the prompt; SAM21 turns it into a tight 2D mask, and the depth+RGB intrinsics produce the 3D points directly in the OpenCV RGB-camera frame.

Non-goals

• Modify /api/perception/scan or /api/perception/object_points. They keep the existing bbox→ROI cloud path. SAM21 is grasp-only.
• Fix the pre-existing nuwa_to_cv_frame extra flip on the ROI cloud path (already flagged in CLAUDE.md as suspicious — out of scope here).
• Expose a SAM21 passthrough endpoint from the orchestrator. SAM21 stays internal.
• Surface the 2D mask in the GUI. return_mask=false.
• Add a SAM21 health prefetch. Fallback already handles "not ready".

High-level flow

YOLO bbox
    │
    ▼
config.grasp.use_sam21 == true ?  ── no ──┐
    │ yes                                  │
    ▼                                      │
[Nuwa: rgb_jpeg, depth_png, intrinsics]    │
    │                                      │
    ▼                                      │
[presa-mano: align_depth_to_rgb]           │
    │                                      │
    ▼                                      │
[SAM21 /api/segment_pointcloud]            │
    │                                      │
    ▼                                      │
points in RGB-camera/OpenCV frame ──┐      │
                                    │      │
exception during SAM21 path ────────┤      │
                                    ▼      ▼
                          [legacy bbox path]
                          (pointcloud_roi + filter_outliers + nuwa_to_cv_frame)
                                    │
                                    ▼
                       transform_points(T_base_camera, …)
                                    │
                                    ▼
                       PCA → GraspPlanner → wrist pose

Architecture

Five touchpoints, all inside src/presa_mano/:

clients/sam21.py (new)

Thin httpx wrapper, peer of nuwa.py / marker.py / yolo.py.

class Sam21Unavailable(RuntimeError): ...    # transport/connection
class Sam21Error(RuntimeError): ...          # semantic 4xx/5xx from service

class Sam21Client:
    def __init__(self, base_url: str, timeout_s: float = 10.0): ...
    def segment_pointcloud(
        self,
        rgb_png: bytes,
        depth_png: bytes,
        intrinsics: dict,   # {fx, fy, cx, cy, width, height}
        bbox: tuple[int, int, int, int],
        depth_scale_mm: float = 1.0,
    ) -> np.ndarray:        # (N, 3) float64, OpenCV RGB-camera frame
        ...

Internally:

• base64-encode the two PNGs.
• POST to {base_url}/api/segment_pointcloud with return_format="json", return_mask=False.
• Unwrap the {success, data, error} envelope. On success=false raise Sam21Error(code, message). On httpx.ConnectError/TimeoutException raise Sam21Unavailable. On HTTP 503 with code MODEL_NOT_READY raise Sam21Unavailable("model loading") (it's transient, same fallback behaviour).
• Parse data.points (list of [x,y,z,r,g,b]) into an (N,3) numpy array (colours discarded — grasp planner doesn't use them).

perception/depth_align.py (new)

Vectorised IR→RGB depth alignment, built on the math already in perception/depth_to_rgb.py. back_project_rgb_pixel does the reverse direction per-pixel and iteratively; here we go forward in one pass.

def align_depth_to_rgb(
    depth_ir_mm: np.ndarray,        # (H_ir, W_ir) uint16, mm, IR-frame
    reg: CameraRegistration,        # IR & RGB intrinsics + extrinsics
) -> np.ndarray:                    # (H_rgb, W_rgb) uint16, mm, RGB-pixel-aligned

Algorithm:

1. Build IR pixel grid (u_ir, v_ir) masked by depth_ir_mm > 0.
2. Back-project to 3D in IR camera frame: P_ir = depth * K_ir^-1 * [u, v, 1].
3. Transform to RGB camera frame: P_rgb = R_ir_to_rgb @ P_ir + t_ir_to_rgb.
4. Project to RGB pixel grid: (u_rgb, v_rgb) = K_rgb @ P_rgb / P_rgb.z.
5. Z-buffer: for each RGB pixel that receives ≥1 projected point, keep the minimum P_rgb.z (closest = occluder wins). Initialise the buffer to a large sentinel (e.g. np.iinfo(np.uint32).max), use np.minimum.at for the in-place reduce, then convert sentinel→0 at the end. (Initialising to 0 and using np.minimum.at would be wrong: every projected depth would compare against 0 and the buffer would stay 0.)
6. Output uint16 mm at RGB resolution; unset pixels stay 0.

Inputs come from CameraRegistration.from_nuwa(nuwa.intrinsics()) which already parses both IR and RGB intrinsics and the IR→RGB rigid transform.

Unit test: synthetic IR depth of a plane → align → back-project from RGB pixels with back_project_rgb_pixel → check that the round-trip recovers the plane within sub-pixel/sub-mm tolerance on the bulk of pixels (edge artefacts acceptable).

api/_deps.py (modify)

One factory, mirroring nuwa()/marker()/yolo():

def sam21() -> Sam21Client:
    cfg = get_config()
    s = cfg.get("services", {}).get("sam21", {})
    return Sam21Client(
        base_url=s.get("url", "http://localhost:8002"),
        timeout_s=float(s.get("timeout_s", 10.0)),
    )

Per-request construction, no module-level cache — PUT /api/config hot-reload keeps working.

api/grasp_routes.py::_object_points_in_base (modify)

Insert the SAM21 branch before the legacy branch. Signature unchanged. Return value unchanged shape; meta enriched with source ∈ {"sam21","bbox_roi"} and (on fallback) sam21_fallback_reason.

def _object_points_in_base(conf, class_name, sw):
    cfg = get_config()
    nuwa = _deps.nuwa()
    yolo = _deps.yolo()
    side = cfg["robot"]["wrist_side"]
    T_base_cam = CalibrationStore(_deps.calibration_file_path(side)).T_base_camera()
    if T_base_cam is None:
        raise HTTPException(412, f"calibration missing for side='{side}'")

    detector = ObjectDetector(yolo, nuwa)
    with sw.step("yolo"):
        dets = detector.detect(conf=conf)
    if class_name:
        dets = [d for d in dets if (d.get("class") or "").lower() == class_name.lower()]
    if not dets:
        raise HTTPException(404, "no object detected")
    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)
            pts_base = transform_points(T_base_cam, pts_cv)
            meta["source"] = "sam21"
            meta["n_points"] = 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["source"] = "bbox_roi"
    meta["n_points"] = len(pts_base)
    return pts_base, meta

Two helpers, both private to grasp_routes.py. Sentinel exception _Sam21Fallback is module-local and never escapes — it only signals "the SAM21 branch couldn't produce a usable result, drop to bbox".

_points_via_sam21(nuwa, bbox, sw) -> np.ndarray:

1. sw.step("nuwa_grab"): fetch rgb_jpeg, depth_png, intrinsics.
2. Convert RGB JPEG → PNG bytes (PIL.Image.open(...).save(buf, "PNG")).
3. sw.step("depth_align"): decode depth PNG → numpy → align_depth_to_rgb → encode back to PNG bytes.
4. sw.step("sam21"): _deps.sam21().segment_pointcloud(...) → (N,3).
5. If N < 10, raise _Sam21Fallback("too few points: N") — same minimum the legacy path enforces, applied uniformly so PCA never runs on degenerate clouds.
6. Raise _Sam21Fallback on:
   • Sam21Unavailable (any transport error)
   • Sam21Error with code in {SEGMENTATION_EMPTY, NO_VALID_DEPTH, MODEL_NOT_READY}
   • Any unexpected exception (catch-all with repr(e) in the reason)
7. Log ERROR (not WARNING) on Sam21Error with code in {SHAPE_MISMATCH, BBOX_OUT_OF_BOUNDS, IMAGE_DECODE_ERROR, INVALID_INPUT} — those signal a bug in the payload assembly that we want to notice, even though we still fall back instead of failing the request.

_points_via_bbox(nuwa, T_base_cam, bbox, sw) -> np.ndarray: extracted from the current body of _object_points_in_base, identical behaviour (pointcloud_roi → filter_outliers → nuwa_to_cv_frame → transform_points). Still raises HTTPException(422) on <10 valid points — that is a true failure of the grasp request.

config.yaml (additions)

services:
  sam21:
    url: http://localhost:8002
    timeout_s: 10.0

grasp:
  use_sam21: true

grasp.use_sam21=false skips the SAM21 branch entirely (no exception, no fallback log) — useful in dev when SAM21 isn't running.

Frame convention summary

Source	Native frame	Treated as
Nuwa /api/rgb	RGB-camera (OpenCV)	input to SAM21, RGB intrinsics
Nuwa /api/depth	IR-camera	input to align_depth_to_rgb only; never back-projected directly
Nuwa /api/intrinsics	RGB-camera (OpenCV)	passed to SAM21 as-is
SAM21 output	RGB-camera (OpenCV: X-right, Y-down, Z-forward)	directly T_base_camera-compatible — no nuwa_to_cv_frame flip
T_base_camera	base ← RGB-camera (OpenCV)	unchanged

The SAM21 path skips nuwa_to_cv_frame. The legacy path still calls it, as today. The two paths converge in base frame and feed compute_axes / GraspPlanner identically.

Error handling matrix

Origin	Cause	Behaviour
httpx.ConnectError / TimeoutException	SAM21 container down, GPU OOM, network	log WARNING, fall back to bbox
HTTP 503 MODEL_NOT_READY	SAM21 still loading checkpoint	log WARNING, fall back to bbox
HTTP 422 SEGMENTATION_EMPTY	SAM2 found no object in bbox	log WARNING, fall back to bbox
HTTP 422 NO_VALID_DEPTH	mask OK but 0 pixels with depth	log WARNING, fall back to bbox
HTTP 400 SHAPE_MISMATCH / BBOX_OUT_OF_BOUNDS / IMAGE_DECODE_ERROR / INVALID_INPUT	payload assembly bug	log ERROR + fall back to bbox
Any other exception	unexpected	log ERROR with repr(e) + fall back to bbox
Nuwa fails (intrinsics/depth/rgb)	camera down	propagate 502 (would fail bbox path too)
Calibration missing	T_base_camera is None	propagate 412 (existing behaviour)
<10 valid depth points after fallback	bbox region too empty	propagate 422 (existing behaviour)

Performance budget

GPU mode (per SAM21 README + measured costs):

• nuwa_grab (RGB jpeg + depth png + intrinsics): ~30 ms
• depth_align (vectorised numpy): ~5–15 ms
• sam21 (segment + back-project, GPU): ~300–500 ms
• to_base_frame: <1 ms

Total SAM21 path: ~350–550 ms (vs legacy ~50–100 ms). Acceptable cost for the precision gain; the user-perceived grasp loop is dominated by wrist move time anyway.

CPU mode (no NVIDIA): SAM21 alone ~2–5 s. Still well below the 10 s client timeout, but the user should know. README of SAM21 already documents this.

Testing

• Unit tests/test_depth_align.py: synthetic plane in IR frame → align → per-pixel round-trip via back_project_rgb_pixel reproduces the plane within tolerance on ≥90% of valid pixels.
• Unit tests/test_sam21_client.py: mock httpx responses for happy path, 503 MODEL_NOT_READY, 422 SEGMENTATION_EMPTY, ConnectError. Assert the exception type and message.
• Integration (manual, against live stack):
  1. SAM21 up, object in frame → POST /api/grasp/plan → response has source: "sam21", n_points > 0, timings_ms.sam21 < 1000.
  2. SAM21 stopped (docker compose stop segmentation-sam21) → same call → source: "bbox_roi", sam21_fallback_reason: "unreachable: …".
  3. SAM21 up, bbox over empty desk → source: "bbox_roi", sam21_fallback_reason: "empty mask".

Out-of-scope follow-ups (logged, not done)

• /api/perception/scan and /object_points could later switch to SAM21 with a query param, once we've validated the grasp path in production.
• Surfacing the SAM21 mask in the GUI (overlay on the live RGB) would help debug bad detections — a UI-only change, easy when wanted.
• The pre-existing nuwa_to_cv_frame extra flip on ROI cloud (CLAUDE.md "Frame conventions (gotcha-prone)") deserves a clean-up but in its own PR.