docs: design SAM21 segmentation in grasp pipeline
Spec brainstorming validato: SAM21 si inserisce tra YOLO bbox e PCA con fallback automatico al path bbox legacy. Tre componenti nuovi (clients/sam21.py, perception/depth_align.py, helper _points_via_sam21) e due modifiche minori (config.yaml, api/_deps.py). Out-of-scope: scan/ object_points, mask UI, nuwa_to_cv_frame cleanup. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Adriano Dal PastroandClaude Opus 4.7 committed
1 parent
0c211630cc
commit
adefc93949
1 file changed
+301
@@ -0,0 +1,301 @@
|
|||||||
|
# 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("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:
|
||||||
|
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.
|
||||||
Reference in new issue
Block a user