feat(find_model): zero below-plane depth before SAM21 segmentation

Previously SAM21 ran on the raw aligned depth and frequently leaked the
mask onto the table around the object: ~88% of segmented points (2136
of 2440 in a sample run) were table pixels filtered away by the post
partition_by_plane step. Open3DModel got a noisy cluster mostly off
the object.

New pipeline: with a calibrated work plane, depth pixels whose 3D
back-projection sits below the plane are zeroed BEFORE SAM21 sees the
depth PNG. SAM21 cannot back-project zero-depth pixels, so only true
above-plane points end up in the cloud regardless of mask outline.

Same sample run after the change: 513 SAM21 -> 513 accepted, 0 excluded.

The post-partition still runs (now a no-op on a calibrated plane), so
the find_model response shape is unchanged.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
Adriano Dal PastroandClaude Opus 4.7 committed 2026-05-13 11:50:01 +02:00
1 parent 2fe7a7a7fa
commit f1f9d8f638
2 files changed
+53 -4

No files matched your search

+6 -2
View File
@@ -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}")
+47 -2
View File
@@ -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")