refactor(sam21_bridge): extract segment_with_sam21 (no plane filter)
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
1 parent
684bd2de5e
commit
87abe26476
2 files changed
+121
-12
No files matched your search
@@ -7,6 +7,7 @@ clients built per-request via api/_deps.py.
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
|
||||
import cv2
|
||||
import httpx
|
||||
@@ -31,18 +32,20 @@ SAM21_BUG_CODES = frozenset({"SHAPE_MISMATCH", "BBOX_OUT_OF_BOUNDS",
|
||||
"IMAGE_DECODE_ERROR", "INVALID_INPUT"})
|
||||
|
||||
|
||||
def points_via_sam21(nuwa, sam21, bbox: tuple[int, int, int, int],
|
||||
sw: Stopwatch) -> tuple[np.ndarray, dict]:
|
||||
"""SAM21 path: bbox → segmented point cloud in RGB-camera/OpenCV frame.
|
||||
@dataclass
|
||||
class Sam21Segmentation:
|
||||
"""Output di segment_with_sam21 — SAM21 puro senza plane filter."""
|
||||
points_cv_m: np.ndarray # (N,3) OpenCV camera frame, metri
|
||||
rgb_jpeg: bytes # snapshot Nuwa usato per la segmentazione
|
||||
intrinsics_rgb: dict # con width/height iniettati dal decoded RGB
|
||||
|
||||
Returns points in **meters**. Raises Sam21Fallback on failures that
|
||||
should trigger the bbox path. Re-raises HTTPException(502) for Nuwa
|
||||
failures (the bbox path would fail too — better to fail fast). All
|
||||
SAM21-specific logging happens here; callers do not log on fallback.
|
||||
|
||||
After the SAM21 segmentation the work-plane filter is applied (no-op if no
|
||||
plane is calibrated). Returns (points, plane_info) where plane_info has
|
||||
`plane_applied` and `n_points_before_plane`.
|
||||
def segment_with_sam21(nuwa, sam21, bbox: tuple[int, int, int, int],
|
||||
sw: Stopwatch) -> Sam21Segmentation:
|
||||
"""SAM21 puro: nuwa_grab + depth_align + segment, SENZA plane filter.
|
||||
|
||||
Raises Sam21Fallback per failure semantiche (mask vuota, no depth, SAM21
|
||||
unreachable, decode fail). Re-raises HTTPException(502) per failure Nuwa.
|
||||
"""
|
||||
try:
|
||||
with sw.step("nuwa_grab"):
|
||||
@@ -105,7 +108,7 @@ def points_via_sam21(nuwa, sam21, bbox: tuple[int, int, int, int],
|
||||
|
||||
with sw.step("sam21"):
|
||||
try:
|
||||
pts_cv = sam21.segment_pointcloud(
|
||||
pts_cv_m = sam21.segment_pointcloud(
|
||||
rgb_png=rgb_png,
|
||||
depth_png=depth_png_out,
|
||||
intrinsics=intr_full["rgb"],
|
||||
@@ -125,7 +128,20 @@ def points_via_sam21(nuwa, sam21, bbox: tuple[int, int, int, int],
|
||||
log.error("Unexpected SAM21 error: %r", e)
|
||||
raise Sam21Fallback(f"unexpected: {e!r}") from e
|
||||
|
||||
pts_cv, plane_info = apply_below_plane_filter(pts_cv, get_work_plane())
|
||||
return Sam21Segmentation(
|
||||
points_cv_m=pts_cv_m,
|
||||
rgb_jpeg=rgb_jpeg,
|
||||
intrinsics_rgb=intr_full["rgb"],
|
||||
)
|
||||
|
||||
|
||||
def points_via_sam21(nuwa, sam21, bbox: tuple[int, int, int, int],
|
||||
sw: Stopwatch) -> tuple[np.ndarray, dict]:
|
||||
"""Legacy wrapper: segment_with_sam21 + apply_below_plane_filter.
|
||||
Returns (points_cv_m, plane_info). Raises Sam21Fallback if <10 points
|
||||
after plane filter (unchanged semantics)."""
|
||||
seg = segment_with_sam21(nuwa, sam21, bbox, sw)
|
||||
pts_cv, plane_info = apply_below_plane_filter(seg.points_cv_m, get_work_plane())
|
||||
if pts_cv.shape[0] < 10:
|
||||
log.warning("SAM21 cloud has %d points after plane filter (<10), falling back to bbox",
|
||||
pts_cv.shape[0])
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
"""Tests for the segment_with_sam21/points_via_sam21 decomposition.
|
||||
|
||||
We exercise the path that the find_model endpoint will take:
|
||||
segment_with_sam21(nuwa, sam21, bbox, sw) -> Sam21Segmentation
|
||||
and the legacy wrapper still works.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from presa_mano.perception.sam21_bridge import (
|
||||
Sam21Segmentation,
|
||||
segment_with_sam21,
|
||||
points_via_sam21,
|
||||
Sam21Fallback,
|
||||
)
|
||||
from presa_mano.timing import Stopwatch
|
||||
|
||||
|
||||
def _fake_nuwa():
|
||||
"""Returns a MagicMock Nuwa that yields a 4x4 black RGB jpeg and a 4x4
|
||||
uint16 depth PNG (all 1000 mm) + a minimal Nuwa intrinsics dict."""
|
||||
rgb = np.zeros((4, 4, 3), dtype=np.uint8)
|
||||
ok, jpg = cv2.imencode(".jpg", rgb)
|
||||
assert ok
|
||||
depth = np.full((4, 4), 1000, dtype=np.uint16)
|
||||
ok, dpng = cv2.imencode(".png", depth)
|
||||
assert ok
|
||||
intr = {
|
||||
"rgb": {"fx": 600.0, "fy": 600.0, "cx": 2.0, "cy": 2.0},
|
||||
"ir": {"fx": 600.0, "fy": 600.0, "cx": 2.0, "cy": 2.0,
|
||||
"extrinsics": {"R": [[1,0,0],[0,1,0],[0,0,1]],
|
||||
"t": [0.0, 0.0, 0.0]}},
|
||||
}
|
||||
n = MagicMock()
|
||||
n.rgb_jpeg.return_value = bytes(jpg)
|
||||
n.depth_png.return_value = bytes(dpng)
|
||||
n.intrinsics.return_value = intr
|
||||
return n, bytes(jpg)
|
||||
|
||||
|
||||
def _fake_sam21(points_m: np.ndarray):
|
||||
s = MagicMock()
|
||||
s.segment_pointcloud.return_value = points_m
|
||||
return s
|
||||
|
||||
|
||||
def test_segment_with_sam21_returns_segmentation(monkeypatch):
|
||||
nuwa, rgb_jpeg = _fake_nuwa()
|
||||
pts_m = np.array([[0.01, 0.02, 0.5], [0.011, 0.022, 0.49]])
|
||||
sam21 = _fake_sam21(pts_m)
|
||||
# Niente plane filter qui — segment_with_sam21 NON deve applicarlo
|
||||
monkeypatch.setattr(
|
||||
"presa_mano.perception.sam21_bridge.get_work_plane", lambda: None
|
||||
)
|
||||
sw = Stopwatch()
|
||||
seg = segment_with_sam21(nuwa, sam21, bbox=(0, 0, 4, 4), sw=sw)
|
||||
assert isinstance(seg, Sam21Segmentation)
|
||||
np.testing.assert_array_equal(seg.points_cv_m, pts_m)
|
||||
assert seg.rgb_jpeg == rgb_jpeg
|
||||
assert "width" in seg.intrinsics_rgb and "height" in seg.intrinsics_rgb
|
||||
|
||||
|
||||
def test_points_via_sam21_still_applies_plane_filter(monkeypatch):
|
||||
"""L'API legacy non deve cambiare: ritorna (pts, plane_info) e applica
|
||||
apply_below_plane_filter come prima."""
|
||||
nuwa, _rgb = _fake_nuwa()
|
||||
pts_m = np.array([[0.0, 0.0, 0.5]] * 20)
|
||||
sam21 = _fake_sam21(pts_m)
|
||||
monkeypatch.setattr(
|
||||
"presa_mano.perception.sam21_bridge.get_work_plane", lambda: None
|
||||
)
|
||||
sw = Stopwatch()
|
||||
pts, info = points_via_sam21(nuwa, sam21, bbox=(0, 0, 4, 4), sw=sw)
|
||||
assert pts.shape == (20, 3)
|
||||
assert info["plane_applied"] is False
|
||||
assert info["n_points_before_plane"] == 20
|
||||
|
||||
|
||||
def test_points_via_sam21_too_few_points_raises_fallback(monkeypatch):
|
||||
nuwa, _rgb = _fake_nuwa()
|
||||
pts_m = np.array([[0.0, 0.0, 0.5]] * 5) # <10
|
||||
sam21 = _fake_sam21(pts_m)
|
||||
monkeypatch.setattr(
|
||||
"presa_mano.perception.sam21_bridge.get_work_plane", lambda: None
|
||||
)
|
||||
sw = Stopwatch()
|
||||
with pytest.raises(Sam21Fallback):
|
||||
points_via_sam21(nuwa, sam21, bbox=(0, 0, 4, 4), sw=sw)
|
||||
Reference in new issue
Block a user