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`.

```python
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.

```python
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()`:

```python
def sam21() -> Sam21Client:
    cfg = get_config()
    s = cfg.get("services", {}).get("sam21", {})
    return Sam21Client(
        base_url=s.get("base_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`.

```python
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)

```yaml
services:
  sam21:
    base_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.