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:
1 parent
2fe7a7a7fa
commit
f1f9d8f638
2 files changed
+53
-4
No files matched your search
@@ -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")
|
||||
|
||||
Reference in new issue
Block a user