diff --git a/src/presa_mano/clients/sam21.py b/src/presa_mano/clients/sam21.py new file mode 100644 index 0000000..f19e3af --- /dev/null +++ b/src/presa_mano/clients/sam21.py @@ -0,0 +1,94 @@ +"""Client for the SAM21 segmentation service (port 8002).""" +from __future__ import annotations + +import base64 +from typing import Any + +import httpx +import numpy as np + + +class Sam21Unavailable(RuntimeError): + """Transport-level failure (connection refused, timeout, model loading).""" + + +class Sam21Error(RuntimeError): + """Semantic failure surfaced by the service (e.g. empty segmentation).""" + + def __init__(self, code: str, message: str): + super().__init__(f"{code}: {message}") + self.code = code + self.message = message + + +_TRANSIENT_CODES = {"MODEL_NOT_READY"} + + +class Sam21Client: + """Thin wrapper around POST /api/segment_pointcloud. + + Returns points in **meters**, OpenCV RGB-camera frame (X-right, Y-down, + Z-forward) — directly compatible with `T_base_camera`. + """ + + def __init__(self, base_url: str, timeout_s: float = 10.0): + self.base_url = base_url.rstrip("/") + self.timeout_s = timeout_s + + def segment_pointcloud( + self, + rgb_png: bytes, + depth_png: bytes, + intrinsics: dict[str, Any], + bbox: tuple[int, int, int, int], + depth_scale_mm: float = 1.0, + ) -> np.ndarray: + body = { + "rgb_png_base64": base64.b64encode(rgb_png).decode("ascii"), + "depth_png_base64": base64.b64encode(depth_png).decode("ascii"), + "intrinsics": { + "fx": float(intrinsics["fx"]), + "fy": float(intrinsics["fy"]), + "cx": float(intrinsics["cx"]), + "cy": float(intrinsics["cy"]), + "width": int(intrinsics["width"]), + "height": int(intrinsics["height"]), + }, + "bbox": [int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])], + "depth_scale_mm": float(depth_scale_mm), + "return_format": "json", + "return_mask": False, + } + url = f"{self.base_url}/api/segment_pointcloud" + try: + with httpx.Client(timeout=self.timeout_s) as c: + r = c.post(url, json=body) + except (httpx.ConnectError, httpx.TimeoutException, httpx.NetworkError) as e: + raise Sam21Unavailable(f"sam21 unreachable: {e}") from e + + # Try to parse the envelope even on non-2xx, because SAM21 always + # returns its `{success, data, error}` shape. + try: + payload = r.json() + except ValueError as e: + raise Sam21Unavailable(f"sam21 returned non-JSON (status {r.status_code}): {e}") from e + + if not isinstance(payload, dict): + raise Sam21Unavailable(f"sam21 returned unexpected payload type: {type(payload)}") + + if payload.get("success") is True: + data = payload.get("data") or {} + points = data.get("points") or [] + if not points: + return np.empty((0, 3), dtype=np.float64) + arr_mm = np.asarray(points, dtype=np.float64) + if arr_mm.ndim != 2 or arr_mm.shape[1] < 3: + raise Sam21Unavailable(f"unexpected points shape: {arr_mm.shape}") + return arr_mm[:, :3] / 1000.0 + + err = payload.get("error") or {} + code = str(err.get("code") or "UNKNOWN") + msg = str(err.get("message") or f"http {r.status_code}") + if code in _TRANSIENT_CODES: + raise Sam21Unavailable(f"{code}: {msg}") + raise Sam21Error(code, msg) diff --git a/tests/test_sam21_client.py b/tests/test_sam21_client.py new file mode 100644 index 0000000..1ce411d --- /dev/null +++ b/tests/test_sam21_client.py @@ -0,0 +1,131 @@ +"""Unit tests for Sam21Client.""" +from __future__ import annotations + +import json +from unittest.mock import MagicMock, patch + +import httpx +import numpy as np +import pytest + +from presa_mano.clients.sam21 import Sam21Client, Sam21Error, Sam21Unavailable + + +def _intrinsics(): + return {"fx": 600.0, "fy": 600.0, "cx": 320.0, "cy": 240.0, + "width": 640, "height": 480} + + +def _ok_response(points_mm): + body = { + "success": True, + "data": { + "n_points": len(points_mm), + "n_pixels_in_mask": len(points_mm), + "n_pixels_with_valid_depth": len(points_mm), + "elapsed_ms": {"segment": 100.0, "project": 5.0, "total": 105.0}, + "points": [list(p) + [0.0, 0.0, 0.0] for p in points_mm], + }, + "error": None, + } + return httpx.Response(200, json=body) + + +def _err_response(status: int, code: str, message: str): + body = {"success": False, "data": None, + "error": {"code": code, "message": message}} + return httpx.Response(status, json=body) + + +def test_segment_pointcloud_happy_path(): + points_mm = [(100.0, 200.0, 1000.0), (110.0, 200.0, 1010.0)] + response = _ok_response(points_mm) + client = Sam21Client("http://sam21:8002", timeout_s=5.0) + with patch("httpx.Client.post", return_value=response) as mock_post: + out = client.segment_pointcloud( + rgb_png=b"FAKE_PNG", depth_png=b"FAKE_DEPTH", + intrinsics=_intrinsics(), bbox=(10, 20, 100, 200), + ) + assert mock_post.call_count == 1 + posted_url = mock_post.call_args.args[0] + assert posted_url.endswith("/api/segment_pointcloud") + body = mock_post.call_args.kwargs["json"] + assert body["bbox"] == [10, 20, 100, 200] + assert body["depth_scale_mm"] == 1.0 + assert body["return_format"] == "json" + assert body["return_mask"] is False + assert isinstance(body["rgb_png_base64"], str) and len(body["rgb_png_base64"]) > 0 + assert isinstance(body["depth_png_base64"], str) and len(body["depth_png_base64"]) > 0 + # Output is in meters: mm/1000. + assert out.shape == (2, 3) + assert out.dtype == np.float64 + assert np.allclose(out, np.array([[0.1, 0.2, 1.0], [0.11, 0.2, 1.01]])) + + +def test_segment_pointcloud_connect_error_raises_unavailable(): + client = Sam21Client("http://sam21:8002") + with patch("httpx.Client.post", side_effect=httpx.ConnectError("refused")): + with pytest.raises(Sam21Unavailable): + client.segment_pointcloud( + rgb_png=b"x", depth_png=b"y", + intrinsics=_intrinsics(), bbox=(0, 0, 10, 10), + ) + + +def test_segment_pointcloud_timeout_raises_unavailable(): + client = Sam21Client("http://sam21:8002") + with patch("httpx.Client.post", side_effect=httpx.TimeoutException("timeout")): + with pytest.raises(Sam21Unavailable): + client.segment_pointcloud( + rgb_png=b"x", depth_png=b"y", + intrinsics=_intrinsics(), bbox=(0, 0, 10, 10), + ) + + +def test_segment_pointcloud_model_not_ready_raises_unavailable(): + """503 MODEL_NOT_READY is transient → treated as unavailable (fallback).""" + response = _err_response(503, "MODEL_NOT_READY", "loading") + client = Sam21Client("http://sam21:8002") + with patch("httpx.Client.post", return_value=response): + with pytest.raises(Sam21Unavailable): + client.segment_pointcloud( + rgb_png=b"x", depth_png=b"y", + intrinsics=_intrinsics(), bbox=(0, 0, 10, 10), + ) + + +def test_segment_pointcloud_segmentation_empty_raises_error(): + """422 SEGMENTATION_EMPTY is a semantic outcome, not transport.""" + response = _err_response(422, "SEGMENTATION_EMPTY", "no mask") + client = Sam21Client("http://sam21:8002") + with patch("httpx.Client.post", return_value=response): + with pytest.raises(Sam21Error) as info: + client.segment_pointcloud( + rgb_png=b"x", depth_png=b"y", + intrinsics=_intrinsics(), bbox=(0, 0, 10, 10), + ) + assert "SEGMENTATION_EMPTY" in str(info.value) + + +def test_segment_pointcloud_no_valid_depth_raises_error(): + response = _err_response(422, "NO_VALID_DEPTH", "mask found but no depth") + client = Sam21Client("http://sam21:8002") + with patch("httpx.Client.post", return_value=response): + with pytest.raises(Sam21Error) as info: + client.segment_pointcloud( + rgb_png=b"x", depth_png=b"y", + intrinsics=_intrinsics(), bbox=(0, 0, 10, 10), + ) + assert "NO_VALID_DEPTH" in str(info.value) + + +def test_segment_pointcloud_empty_points_returns_empty_array(): + """Service returned success but 0 points — return (0, 3) array.""" + response = _ok_response([]) + client = Sam21Client("http://sam21:8002") + with patch("httpx.Client.post", return_value=response): + out = client.segment_pointcloud( + rgb_png=b"x", depth_png=b"y", + intrinsics=_intrinsics(), bbox=(0, 0, 10, 10), + ) + assert out.shape == (0, 3)