diff --git a/src/presa_mano/api/grasp_routes.py b/src/presa_mano/api/grasp_routes.py index ad2085f..efa07a1 100644 --- a/src/presa_mano/api/grasp_routes.py +++ b/src/presa_mano/api/grasp_routes.py @@ -232,9 +232,13 @@ def find_model(req: FindModelReq) -> dict[str, Any]: top = dets[0] bbox = _bbox_to_roi(top["bbox"]) - # 2. SAM21 (no plane filter) + # 2. SAM21, with depth pre-filter so the segmentation cannot leak onto + # the table around the object: pixels whose back-projection sits below + # the plane have their depth zeroed BEFORE SAM21 sees them, so they + # never end up in the cloud regardless of the mask outline. try: - seg = segment_with_sam21(nuwa, sam21, bbox, sw) + seg = segment_with_sam21(nuwa, sam21, bbox, sw, + pre_filter_plane=get_work_plane()) except Sam21Fallback as e: raise HTTPException(status_code=422, detail=f"SAM21 fail: {e}") diff --git a/src/presa_mano/perception/sam21_bridge.py b/src/presa_mano/perception/sam21_bridge.py index af26f45..90b9626 100644 --- a/src/presa_mano/perception/sam21_bridge.py +++ b/src/presa_mano/perception/sam21_bridge.py @@ -18,11 +18,44 @@ from ..clients import NuwaError, Sam21Error, Sam21Unavailable from ..timing import Stopwatch from .depth_align import align_depth_to_rgb from .depth_to_rgb import CameraRegistration -from .work_plane import apply_below_plane_filter, get_work_plane +from .work_plane import WorkPlane, apply_below_plane_filter, get_work_plane log = logging.getLogger(__name__) +def zero_below_plane_depth( + depth_aligned_mm: np.ndarray, + intr_rgb: dict, + wp: WorkPlane, +) -> tuple[np.ndarray, int]: + """Zero-out depth pixels whose 3D back-projection sits below the plane. + + Per-pixel back-projection: P = K_rgb^-1 * [u,v,1] * depth(u,v) (mm). + Signed distance to the plane (in m): n . (P/1000) + offset. + Pixel is dropped iff `signed < -offset_mm/1000` — same threshold as + apply_below_plane_filter so the SAM21-then-partition path and the + depth-prefilter-then-SAM21 path agree. + + Returns (filtered_depth_mm, n_zeroed). + """ + H, W = depth_aligned_mm.shape + fx = float(intr_rgb["fx"]); fy = float(intr_rgb["fy"]) + cx = float(intr_rgb["cx"]); cy = float(intr_rgb["cy"]) + z_mm = depth_aligned_mm.astype(np.float64) + u = np.arange(W, dtype=np.float64)[None, :] + v = np.arange(H, dtype=np.float64)[:, None] + x_mm = (u - cx) * z_mm / fx + y_mm = (v - cy) * z_mm / fy + n = np.asarray(wp.normal, dtype=np.float64) + signed_m = (n[0] * x_mm + n[1] * y_mm + n[2] * z_mm) / 1000.0 + float(wp.offset) + threshold_m = -float(wp.offset_mm) / 1000.0 + invalid = (z_mm > 0) & (signed_m < threshold_m) + n_zeroed = int(invalid.sum()) + out = depth_aligned_mm.copy() + out[invalid] = 0 + return out, n_zeroed + + class Sam21Fallback(Exception): """The SAM21 path could not produce a usable cloud. Callers should fall back to the non-SAM21 path (bbox ROI cloud).""" @@ -41,9 +74,16 @@ class Sam21Segmentation: def segment_with_sam21(nuwa, sam21, bbox: tuple[int, int, int, int], - sw: Stopwatch) -> Sam21Segmentation: + sw: Stopwatch, + pre_filter_plane: WorkPlane | None = None, + ) -> Sam21Segmentation: """SAM21 puro: nuwa_grab + depth_align + segment, SENZA plane filter. + If `pre_filter_plane` is provided, depth pixels whose back-projection + sits below the plane are zeroed BEFORE SAM21 is called. This prevents + the segmentation from leaking onto the table around the object: SAM21 + cannot back-project pixels with depth=0 into the cloud. + Raises Sam21Fallback per failure semantiche (mask vuota, no depth, SAM21 unreachable, decode fail). Re-raises HTTPException(502) per failure Nuwa. """ @@ -95,6 +135,11 @@ def segment_with_sam21(nuwa, sam21, bbox: tuple[int, int, int, int], } reg = CameraRegistration.from_nuwa(intr_full) depth_aligned = align_depth_to_rgb(depth_ir, reg) + if pre_filter_plane is not None: + depth_aligned, n_zeroed = zero_below_plane_depth( + depth_aligned, intr_full["rgb"], pre_filter_plane) + log.info("pre_filter_plane: zeroed %d/%d depth pixels below plane", + n_zeroed, depth_aligned.size) ok, depth_png_buf = cv2.imencode(".png", depth_aligned) if not ok: log.warning("SAM21 path skipped: could not encode aligned depth as PNG")