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