grid-cortex-client 0.4.0__tar.gz → 0.5.0__tar.gz

{grid_cortex_client-0.4.0 → grid_cortex_client-0.5.0}/CLAUDE.md

@@ -10,7 +10,7 @@ uv run --package grid-cortex-client pytest
 
 - Build tool: hatchling + hatch-vcs
 - Version: set via `BUILD_VERSION` env (auto-bumped on main from squash-commit type — see `knowledge/golden-rules/cortex/client-versioning.md`)
-- PyPI history: pre-PR-#44 versions up to `0.2.118` were published by the old pipeline; the new auto-pipeline starts at `0.3.0` and uses `cortex-client/v*` git tags as the version source of truth (first auto-publish lands via PR #45's `feat(cortex):` merge once its review feedback is addressed)
+- PyPI history: pre-PR-#44 versions up to `0.2.118` were published by the old pipeline; the auto-pipeline (live since PR #45) starts at `0.3.0` and uses `cortex-client/v*` git tags as the version source of truth
 
 ## Key Details
 

{grid_cortex_client-0.4.0 → grid_cortex_client-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: grid-cortex-client
-Version: 0.4.0
+Version: 0.5.0
 Summary: Python client for Grid Cortex
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.8

{grid_cortex_client-0.4.0 → grid_cortex_client-0.5.0}/src/grid_cortex_client/model_type.py

@@ -13,6 +13,44 @@ from enum import Enum
 
 
 class ModelType(Enum):
+    CONTACT_GRASPNET = "contact-graspnet"
+    """
+    Generate 6-DoF grasp poses from a depth image + camera intrinsics
+    (+ optional segmentation mask) using Contact-GraspNet (PyTorch port).
+
+    Args:
+        depth_image (Union[str, Image.Image, np.ndarray]): Depth image (path/URL/PIL/ndarray; float32 H×W).
+        camera_intrinsics (Union[str, np.ndarray]): 3×3 intrinsics matrix.
+        seg_image (Union[str, Image.Image, np.ndarray, None]): Optional segmentation mask (uint8 H×W).
+            When supplied + aux_args["filter_grasps"] = True, grasps are restricted to the masked region.
+        aux_args (Dict[str, Any]): Optional dict. Keys:
+            - "top_k" (int, default 64): cap on returned grasps after score-sort.
+            - "z_range" ([float, float], default [0.2, 1.8]): z-axis crop in metres.
+            - "local_regions" (bool, default False): crop 3D regions around segments.
+            - "filter_grasps" (bool, default False): restrict grasps to segmap; requires seg_image.
+            - "forward_passes" (int, default 1): how many forward passes to ensemble.
+        timeout (float | None): Optional HTTP timeout.
+
+    Returns:
+        Dict with:
+            - grasps: Array of 4×4 grasp poses in camera frame, shape (N, 4, 4), N <= top_k.
+            - confidence: Array of confidence scores in [0, 1], shape (N,).
+            - latency_ms: Optional server-reported latency in milliseconds.
+
+    Raises:
+        ValueError: If depth or intrinsics cannot be loaded.
+        RuntimeError: If no HTTP transport is configured.
+        Exception: If the HTTP request fails.
+
+    Example:
+        >>> from grid_cortex_client import CortexClient, ModelType
+        >>> import numpy as np
+        >>> client = CortexClient()
+        >>> depth = np.load("depth.npy")
+        >>> K = np.array([[615,0,320],[0,615,240],[0,0,1]], dtype=np.float64)
+        >>> out = client.run(ModelType.CONTACT_GRASPNET, depth_image=depth, camera_intrinsics=K, aux_args={"top_k": 32})
+        >>> print(out["grasps"].shape)  # (<=32, 4, 4)
+    """
     DA3METRIC = "da3metric"
     """
     Estimate metric depth from single RGB image.

grid_cortex_client-0.5.0/src/grid_cortex_client/models/contact_graspnet.py

@@ -0,0 +1,155 @@
+# grid_cortex_client/src/grid_cortex_client/models/contact_graspnet.py
+"""ContactGraspNet wrapper.
+
+6-DoF grasp pose generation from a depth image + camera intrinsics, with an
+optional segmentation mask to restrict the working region.
+
+Successor to the legacy `graspnet` (TF) model; runs on the PyTorch port via
+the cortex `contact_graspnet` Ray Serve deployment. Wire format mirrors
+`graspgen` so swapping models in client code is just changing the ModelType.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Union
+
+import numpy as np
+from PIL import Image
+
+from ..preprocessing import array_to_npy_bytes, load_image
+from .base_model import BaseModel
+
+
+class ContactGraspNet(BaseModel):
+    """6-DoF grasp generation from depth + intrinsics (+ optional seg).
+
+    Preferred usage
+    ---------------
+    ```pycon
+    >>> grasps, conf = CortexClient().run(
+    ...     "contact_graspnet",
+    ...     depth_image=depth, camera_intrinsics=K,
+    ...     seg_image=seg, aux_args={"top_k": 64}
+    ... )
+    ```
+    """
+
+    name: str = "contact-graspnet"
+    model_id: str = "contact-graspnet"
+
+    def preprocess(
+        self,
+        *,
+        depth_image: Union[str, Image.Image, np.ndarray, None] = None,
+        camera_intrinsics: Union[str, np.ndarray, None] = None,
+        seg_image: Union[str, Image.Image, np.ndarray, None] = None,
+        aux_args: Dict[str, Any] | None = None,
+    ) -> Dict[str, Any]:
+        """Encode depth + intrinsics (+ optional seg) for the server.
+
+        Args:
+            depth_image: Depth image (path/URL/PIL/ndarray; float32 H×W).
+            camera_intrinsics: 3×3 intrinsics matrix (path/ndarray).
+            seg_image: Optional segmentation mask (path/PIL/ndarray; uint8 H×W).
+                When supplied + ``aux_args["filter_grasps"] = True``, grasps are
+                restricted to the masked region (see model README for details).
+            aux_args: Optional dict with ``top_k`` (default 64), ``z_range``
+                (default [0.2, 1.8]), ``local_regions``, ``filter_grasps``,
+                ``forward_passes``.
+        """
+        if depth_image is None or camera_intrinsics is None:
+            raise ValueError("depth_image and camera_intrinsics are required.")
+
+        # Accept .npy paths in addition to image formats. PIL.Image.open chokes
+        # on raw NumPy files; np.load is the right call for them. .png/.exr
+        # depth maps still flow through load_image as before.
+        if isinstance(depth_image, str):
+            if depth_image.endswith(".npy"):
+                depth_array = np.load(depth_image)
+            else:
+                depth_array = np.array(load_image(depth_image))
+        elif isinstance(depth_image, Image.Image):
+            depth_array = np.array(load_image(depth_image))
+        else:
+            depth_array = np.asarray(depth_image)
+
+        if isinstance(camera_intrinsics, str):
+            intrinsics = np.load(camera_intrinsics)
+        else:
+            intrinsics = np.asarray(camera_intrinsics)
+
+        payload: Dict[str, Any] = {
+            "depth_image": array_to_npy_bytes(depth_array.astype(np.float32)),
+            "camera_intrinsics": array_to_npy_bytes(intrinsics.astype(np.float64)),
+            "aux_args": aux_args or {},
+        }
+
+        if seg_image is not None:
+            if isinstance(seg_image, (str, Image.Image)):
+                seg_array = np.array(load_image(seg_image))
+            else:
+                seg_array = np.asarray(seg_image)
+            payload["seg_image"] = array_to_npy_bytes(seg_array.astype(np.uint8))
+
+        return payload
+
+    def postprocess(self, response_data: Dict[str, Any], **_: Any) -> Dict[str, Any]:
+        """Decode grasps and confidence from response."""
+        raw_grasps = response_data["output"]
+        # Keep the (N, 4, 4) shape invariant when no grasps were found —
+        # np.array([]) → (0,) which breaks any caller indexing grasps[:, :3, 3].
+        if not raw_grasps:
+            grasps = np.zeros((0, 4, 4), dtype=np.float32)
+        else:
+            grasps = np.array(raw_grasps)
+        conf = np.array(response_data["confidence"])
+        return {
+            "grasps": grasps,
+            "confidence": conf,
+            "latency_ms": response_data.get("latency_ms"),
+        }
+
+    def run(
+        self,
+        *,
+        depth_image: Union[str, Image.Image, np.ndarray, None] = None,
+        camera_intrinsics: Union[str, np.ndarray, None] = None,
+        seg_image: Union[str, Image.Image, np.ndarray, None] = None,
+        aux_args: Dict[str, Any] | None = None,
+        timeout: float | None = None,
+    ) -> Dict[str, Any]:
+        """Generate grasps using ContactGraspNet.
+
+        Args:
+            depth_image: Depth image (required).
+            camera_intrinsics: 3×3 intrinsics matrix (required).
+            seg_image: Optional segmentation mask (uint8 H×W).
+            aux_args: Optional dict. Keys: ``top_k`` (int, default 64),
+                ``z_range`` ([float, float], default [0.2, 1.8]),
+                ``local_regions`` (bool, default False),
+                ``filter_grasps`` (bool, default False; requires ``seg_image``),
+                ``forward_passes`` (int, default 1).
+            timeout: Optional HTTP timeout.
+
+        Returns:
+            Dict with:
+                - ``grasps``: (N, 4, 4) array of grasp poses in camera frame
+                - ``confidence``: (N,) array of confidence scores in [0, 1]
+                - ``latency_ms``: server-reported latency
+
+        Examples:
+            >>> from grid_cortex_client import CortexClient, ModelType
+            >>> import numpy as np
+            >>> client = CortexClient()
+            >>> depth = np.load("depth.npy")
+            >>> K = np.array([[615,0,320],[0,615,240],[0,0,1]], dtype=np.float64)
+            >>> out = client.run(ModelType.CONTACT_GRASPNET, depth_image=depth, camera_intrinsics=K, aux_args={"top_k": 32})
+            >>> print(out["grasps"].shape)  # (<=32, 4, 4)
+        """
+        return super().run(
+            depth_image=depth_image,
+            camera_intrinsics=camera_intrinsics,
+            seg_image=seg_image,
+            aux_args=aux_args,
+            timeout=timeout,
+        )

grid_cortex_client-0.5.0/src/grid_cortex_client/models/metric3d.py

@@ -0,0 +1,129 @@
+"""Metric3D wrapper.
+
+Client-side helper for monocular metric depth estimation and surface normal
+prediction powered by Metric3D ViT-Large.
+"""
+
+from __future__ import annotations
+
+from io import BytesIO
+from typing import Any, Dict, Optional, Sequence, Union
+
+import numpy as np
+from PIL import Image
+
+from ..preprocessing import load_image, image_to_bytes
+from .base_model import BaseModel
+
+VALID_OUTPUTS = ("depth", "normals")
+
+
+class Metric3D(BaseModel):
+    """Monocular metric depth + surface normals (Metric3D ViT-Large).
+
+    Preferred usage (through :class:`CortexClient`)
+    ------------------------------------------------
+    ```pycon
+    >>> result = CortexClient().run("metric3d", image_input=img)
+    >>> result["depth"].shape   # (H, W) float32 metres
+    >>> result["normals"].shape # (H, W, 3) float32 unit normals
+
+    >>> # Request a single map to cut response size (~75% smaller for depth-only)
+    >>> result = CortexClient().run("metric3d", image_input=img, outputs=["depth"])
+    >>> "normals" in result
+    False
+    ```
+    """
+
+    name: str = "metric3d"
+    model_id: str = "metric3d"
+
+    def __init__(self, *, transport=None) -> None:
+        super().__init__(transport=transport)
+
+    # ------------------------------------------------------------------
+    # BaseModel interface
+    # ------------------------------------------------------------------
+
+    def preprocess(
+        self,
+        *,
+        image_input: Union[str, Image.Image, np.ndarray],
+        outputs: Optional[Union[str, Sequence[str]]] = None,
+    ) -> Dict[str, Any]:
+        """Return msgpack-ready payload with raw image bytes under key ``image_input``.
+
+        Args:
+            image_input: RGB image as file path, URL, PIL Image, or numpy array.
+            outputs: Optional selection of maps to return — ``"depth"``,
+                ``"normals"``, or a non-empty subset of both. ``None`` (default)
+                returns both.
+
+        Raises:
+            ValueError: If ``outputs`` is empty or contains an unknown name.
+        """
+        pil_img = load_image(image_input)
+        payload: Dict[str, Any] = {"image_input": image_to_bytes(pil_img, encoding_format="JPEG")}
+
+        if outputs is not None:
+            selected = [outputs] if isinstance(outputs, str) else [str(o) for o in outputs]
+            if not selected or any(o not in VALID_OUTPUTS for o in selected):
+                raise ValueError(f"'outputs' must be a non-empty subset of {VALID_OUTPUTS}, got {outputs!r}")
+            payload["outputs"] = selected
+
+        return payload
+
+    def postprocess(
+        self,
+        response_data: Dict[str, Any],
+        **_: Any,
+    ) -> Dict[str, np.ndarray]:
+        """Decode the maps present in the server response.
+
+        Returns:
+            Dict with the returned subset of ``"depth"`` (H, W) float32 and
+            ``"normals"`` (H, W, 3) float32.
+        """
+        if isinstance(response_data, list):
+            response_data = response_data[0]
+
+        result: Dict[str, np.ndarray] = {}
+        for key in VALID_OUTPUTS:
+            buf = response_data.get(key)
+            if buf is None:
+                continue
+            if not isinstance(buf, (bytes, bytearray, memoryview)):
+                raise ValueError(f"'{key}' field is not bytes")
+            result[key] = np.load(BytesIO(buf), allow_pickle=False).astype(np.float32)
+
+        if not result:
+            raise ValueError(f"response contains none of {VALID_OUTPUTS}")
+        return result
+
+    def run(
+        self,
+        image_input: Union[str, Image.Image, np.ndarray],
+        timeout: float | None = None,
+        *,
+        outputs: Optional[Union[str, Sequence[str]]] = None,
+    ) -> Dict[str, np.ndarray]:
+        """Estimate metric depth and/or surface normals from a single RGB image.
+
+        Args:
+            image_input: RGB image as file path, URL, PIL Image, or numpy array.
+            timeout: Optional timeout in seconds for the HTTP request.
+            outputs: Keyword-only. Optional selection of maps to return —
+                ``"depth"``, ``"normals"``, or a non-empty subset of both.
+                ``None`` (default) returns both. Requesting only what you need
+                shrinks the response payload (normals are 3x the size of depth).
+
+        Returns:
+            Dict with the requested subset of:
+                - ``"depth"``: (H, W) float32 array, metric depth in metres.
+                - ``"normals"``: (H, W, 3) float32 array, unit surface normals.
+        """
+        return super().run(
+            image_input=image_input,
+            outputs=outputs,
+            timeout=timeout,
+        )

grid_cortex_client-0.5.0/tests/test_contact_graspnet.py

@@ -0,0 +1,282 @@
+"""Test ContactGraspNet model (6-DoF grasp generation, PyTorch port)."""
+
+import os
+import time
+
+import numpy as np
+import pytest
+
+from grid_cortex_client import AsyncCortexClient, CortexClient, CortexHubClient, ModelType
+from grid_cortex_client.client import CortexAPIError
+
+from .conftest import BaseModelTestHTTP, is_gpu_mem_close, run_concurrent_requests
+
+
+class TestContactGraspNetHTTP(BaseModelTestHTTP):
+    """HTTP-based tests for ContactGraspNet.
+
+    Reuses the depth/seg/intrinsics fixtures from the graspgen suite — both
+    grasping models take the same input shape, so a single eval set covers
+    both. ContactGraspNet has no point_cloud direct-input path (the model
+    extracts the cloud itself from depth + K with z-range filtering).
+    """
+
+    @property
+    def batch_size(self) -> int:
+        return 1
+
+    @property
+    def gpu_mem(self) -> float:
+        return 3200.0
+
+    @property
+    def max_input_size(self) -> tuple:
+        return (720, 1280)
+
+    @pytest.fixture(scope="class")
+    def depth_image(self, depth_image_np: np.ndarray) -> np.ndarray:
+        return depth_image_np
+
+    @pytest.fixture(scope="class")
+    def seg_image(self, seg_image_np: np.ndarray) -> np.ndarray:
+        return seg_image_np
+
+    @pytest.fixture(scope="class")
+    def camera_intrinsics(self, camera_intrinsics_np: np.ndarray) -> np.ndarray:
+        return camera_intrinsics_np
+
+    @pytest.fixture(scope="class")
+    def aux_args(self) -> dict:
+        return {
+            "top_k": 32,
+            "z_range": [0.2, 1.8],
+            "local_regions": False,
+            "filter_grasps": False,
+            "forward_passes": 1,
+        }
+
+    def test_basic_inference(
+        self,
+        client: CortexClient,
+        depth_image: np.ndarray,
+        camera_intrinsics: np.ndarray,
+        aux_args: dict,
+    ):
+        result = client.run(
+            ModelType.CONTACT_GRASPNET,
+            depth_image=depth_image,
+            camera_intrinsics=camera_intrinsics,
+            aux_args=aux_args,
+        )
+        assert result is not None
+        assert "grasps" in result
+        assert "confidence" in result
+        assert isinstance(result["grasps"], np.ndarray)
+        assert isinstance(result["confidence"], np.ndarray)
+
+    def test_with_segmentation(
+        self,
+        client: CortexClient,
+        depth_image: np.ndarray,
+        seg_image: np.ndarray,
+        camera_intrinsics: np.ndarray,
+    ):
+        aux = {"top_k": 16, "filter_grasps": True, "local_regions": True}
+        result = client.run(
+            ModelType.CONTACT_GRASPNET,
+            depth_image=depth_image,
+            seg_image=seg_image,
+            camera_intrinsics=camera_intrinsics,
+            aux_args=aux,
+        )
+        assert result is not None
+        assert "grasps" in result
+
+    def test_raw_response_with_gpu_stats(
+        self,
+        client: CortexClient,
+        depth_image: np.ndarray,
+        camera_intrinsics: np.ndarray,
+        aux_args: dict,
+    ):
+        raw = client.run(
+            ModelType.CONTACT_GRASPNET,
+            depth_image=depth_image,
+            camera_intrinsics=camera_intrinsics,
+            aux_args=aux_args,
+            raw_response=True,
+        )
+        assert "output" in raw
+        assert "confidence" in raw
+        assert "latency_ms" in raw
+        assert "gpu_stats" in raw
+        assert "mem_usage_mb" in raw["gpu_stats"]
+
+    def test_grasp_output_structure(
+        self,
+        client: CortexClient,
+        depth_image: np.ndarray,
+        camera_intrinsics: np.ndarray,
+        aux_args: dict,
+    ):
+        result = client.run(
+            ModelType.CONTACT_GRASPNET,
+            depth_image=depth_image,
+            camera_intrinsics=camera_intrinsics,
+            aux_args=aux_args,
+        )
+        grasps = result["grasps"]
+        confidence = result["confidence"]
+        if len(grasps) > 0:
+            assert grasps.ndim == 3
+            assert grasps.shape[1:] == (4, 4)
+            assert confidence.ndim == 1
+            assert len(grasps) == len(confidence)
+            assert len(grasps) <= aux_args["top_k"]
+
+    def test_maximum_payload_size(
+        self,
+        client: CortexClient,
+        camera_intrinsics: np.ndarray,
+        aux_args: dict,
+    ):
+        large_depth = np.random.rand(*self.max_input_size).astype(np.float32) * 2.0
+        result = client.run(
+            ModelType.CONTACT_GRASPNET,
+            depth_image=large_depth,
+            camera_intrinsics=camera_intrinsics,
+            aux_args=aux_args,
+        )
+        assert result is not None
+
+    def test_exceeds_maximum_payload_size(
+        self,
+        client: CortexClient,
+        camera_intrinsics: np.ndarray,
+        aux_args: dict,
+    ):
+        # Handler guard at depth.size > 4_000_000 -> 413.
+        huge_depth = np.random.rand(2001, 2001).astype(np.float32) * 2.0
+        with pytest.raises(CortexAPIError) as exc_info:
+            client.run(
+                ModelType.CONTACT_GRASPNET,
+                depth_image=huge_depth,
+                camera_intrinsics=camera_intrinsics,
+                aux_args=aux_args,
+            )
+        assert exc_info.value.status_code == 413
+
+    def test_inference_latency(
+        self,
+        client: CortexClient,
+        depth_image: np.ndarray,
+        camera_intrinsics: np.ndarray,
+        aux_args: dict,
+    ):
+        start = time.time()
+        result = client.run(
+            ModelType.CONTACT_GRASPNET,
+            depth_image=depth_image,
+            camera_intrinsics=camera_intrinsics,
+            aux_args=aux_args,
+        )
+        elapsed = time.time() - start
+        assert result is not None
+        assert elapsed < 50.0, f"Inference took {elapsed:.2f}s, expected < 50s"
+
+    def test_missing_camera_intrinsics(
+        self,
+        client: CortexClient,
+        depth_image: np.ndarray,
+        aux_args: dict,
+    ):
+        with pytest.raises(ValueError):
+            client.run(
+                ModelType.CONTACT_GRASPNET,
+                depth_image=depth_image,
+                aux_args=aux_args,
+            )
+
+    def test_missing_depth_image(
+        self,
+        client: CortexClient,
+        camera_intrinsics: np.ndarray,
+        aux_args: dict,
+    ):
+        with pytest.raises(ValueError):
+            client.run(
+                ModelType.CONTACT_GRASPNET,
+                camera_intrinsics=camera_intrinsics,
+                aux_args=aux_args,
+            )
+
+    @pytest.mark.asyncio
+    async def test_concurrent_requests(
+        self,
+        async_client: AsyncCortexClient,
+        depth_image: np.ndarray,
+        camera_intrinsics: np.ndarray,
+        aux_args: dict,
+    ):
+        async def make_request():
+            return await async_client.run(
+                ModelType.CONTACT_GRASPNET,
+                depth_image=depth_image,
+                camera_intrinsics=camera_intrinsics,
+                aux_args=aux_args,
+            )
+
+        # Match graspgen's pragmatic cap-of-1 — keeps the test from timing out
+        # at full input size while still exercising the async path.
+        num_requests = min(self.batch_size, 1)
+        results = await run_concurrent_requests(make_request, num_requests=num_requests)
+        assert len(results) == num_requests
+        assert all(not isinstance(r, Exception) for r in results)
+        assert all("grasps" in r for r in results)
+
+    @pytest.mark.asyncio
+    async def test_gpu_memory_usage(
+        self,
+        async_client: AsyncCortexClient,
+        depth_image: np.ndarray,
+        camera_intrinsics: np.ndarray,
+        aux_args: dict,
+    ):
+        async def make_request():
+            return await async_client.run(
+                ModelType.CONTACT_GRASPNET,
+                depth_image=depth_image,
+                camera_intrinsics=camera_intrinsics,
+                aux_args=aux_args,
+                raw_response=True,
+            )
+
+        num_requests = min(self.batch_size, 2)
+        results = await run_concurrent_requests(make_request, num_requests=num_requests)
+        assert len(results) == num_requests
+        assert all(not isinstance(r, Exception) for r in results)
+        for result in results:
+            print(f"Specified GPU mem: {self.gpu_mem} MB, Current GPU mem: {result['gpu_stats']['mem_usage_mb']} MB")
+        assert all(is_gpu_mem_close(self.gpu_mem, result["gpu_stats"]["mem_usage_mb"]) for result in results), (
+            f"GPU memory usage too high: {[r['gpu_stats']['mem_usage_mb'] for r in results]} MB"
+        )
+
+    @pytest.mark.asyncio
+    async def test_cortex_hub(
+        self,
+        depth_image: np.ndarray,
+        camera_intrinsics: np.ndarray,
+        aux_args: dict,
+    ):
+        api_key = os.getenv("GRID_CORTEX_API_KEY") or os.getenv("CORTEX_API_KEY")
+        if not api_key:
+            pytest.skip("GRID_CORTEX_API_KEY or CORTEX_API_KEY not set")
+
+        async with CortexHubClient(api_key=api_key) as hub:
+            result = await hub.run(
+                ModelType.CONTACT_GRASPNET,
+                depth_image=depth_image,
+                camera_intrinsics=camera_intrinsics,
+                aux_args=aux_args,
+            )
+        assert result is not None

{grid_cortex_client-0.4.0 → grid_cortex_client-0.5.0}/tests/test_metric3d.py

@@ -55,6 +55,24 @@ class TestMetric3DHTTP(BaseModelTestHTTP):
         with pytest.raises((TypeError, CortexAPIError)):
             client.run(ModelType.METRIC3D, image_input=[rgb_image_np, rgb_image_np])
 
+    def test_output_selection(self, client: CortexClient, rgb_image_np: np.ndarray):
+        """outputs= returns only the requested maps; invalid names are rejected client-side."""
+        depth_only = client.run(ModelType.METRIC3D, image_input=rgb_image_np, outputs=["depth"])
+        assert set(depth_only) == {"depth"}
+        assert depth_only["depth"].shape == rgb_image_np.shape[:2]
+
+        normals_only = client.run(ModelType.METRIC3D, image_input=rgb_image_np, outputs="normals")
+        assert set(normals_only) == {"normals"}
+        assert normals_only["normals"].shape == (*rgb_image_np.shape[:2], 3)
+
+        both = client.run(ModelType.METRIC3D, image_input=rgb_image_np, outputs=["depth", "normals"])
+        assert set(both) == {"depth", "normals"}
+
+        with pytest.raises(ValueError):
+            client.run(ModelType.METRIC3D, image_input=rgb_image_np, outputs=["depth", "rgb"])
+        with pytest.raises(ValueError):
+            client.run(ModelType.METRIC3D, image_input=rgb_image_np, outputs=[])
+
     def test_maximum_payload_size(self, client: CortexClient):
         """Test that inference works with maximum allowed payload size."""
         large_image = np.random.randint(0, 255, (*self.max_input_size, 3), dtype=np.uint8)
@@ -122,7 +140,9 @@ class TestMetric3DHTTP(BaseModelTestHTTP):
 
         async with CortexHubClient(api_key=api_key) as hub:
             result = await hub.run(ModelType.METRIC3D, image_input=rgb_image_np)
+            depth_only = await hub.run(ModelType.METRIC3D, image_input=rgb_image_np, outputs=["depth"])
         assert result is not None
         assert isinstance(result, dict)
         assert "depth" in result
         assert "normals" in result
+        assert set(depth_only) == {"depth"}

grid_cortex_client-0.4.0/src/grid_cortex_client/models/metric3d.py

@@ -1,95 +0,0 @@
-"""Metric3D wrapper.
-
-Client-side helper for monocular metric depth estimation and surface normal
-prediction powered by Metric3D ViT-Large.
-"""
-
-from __future__ import annotations
-
-from io import BytesIO
-from typing import Any, Dict, Union
-
-import numpy as np
-from PIL import Image
-
-from ..preprocessing import load_image, image_to_bytes
-from .base_model import BaseModel
-
-
-class Metric3D(BaseModel):
-    """Monocular metric depth + surface normals (Metric3D ViT-Large).
-
-    Preferred usage (through :class:`CortexClient`)
-    ------------------------------------------------
-    ```pycon
-    >>> result = CortexClient().run("metric3d", image_input=img)
-    >>> result["depth"].shape   # (H, W) float32 metres
-    >>> result["normals"].shape # (H, W, 3) float32 unit normals
-    ```
-    """
-
-    name: str = "metric3d"
-    model_id: str = "metric3d"
-
-    def __init__(self, *, transport=None) -> None:
-        super().__init__(transport=transport)
-
-    # ------------------------------------------------------------------
-    # BaseModel interface
-    # ------------------------------------------------------------------
-
-    def preprocess(
-        self,
-        *,
-        image_input: Union[str, Image.Image, np.ndarray],
-    ) -> Dict[str, Any]:
-        """Return msgpack-ready payload with raw image bytes under key ``image_input``."""
-        pil_img = load_image(image_input)
-        return {"image_input": image_to_bytes(pil_img, encoding_format="JPEG")}
-
-    def postprocess(
-        self,
-        response_data: Dict[str, Any],
-        **_: Any,
-    ) -> Dict[str, np.ndarray]:
-        """Decode server response into depth and normals arrays.
-
-        Returns:
-            Dict with ``"depth"`` (H, W) float32 and ``"normals"`` (H, W, 3) float32.
-        """
-        if isinstance(response_data, list):
-            response_data = response_data[0]
-
-        depth_bytes = response_data.get("depth")
-        normals_bytes = response_data.get("normals")
-
-        if not isinstance(depth_bytes, (bytes, bytearray, memoryview)):
-            raise ValueError("'depth' field missing or not bytes")
-        if not isinstance(normals_bytes, (bytes, bytearray, memoryview)):
-            raise ValueError("'normals' field missing or not bytes")
-
-        depth = np.load(BytesIO(depth_bytes), allow_pickle=False).astype(np.float32)
-        normals = np.load(BytesIO(normals_bytes), allow_pickle=False).astype(np.float32)
-
-        return {"depth": depth, "normals": normals}
-
-    def run(
-        self,
-        image_input: Union[str, Image.Image, np.ndarray],
-        timeout: float | None = None,
-    ) -> Dict[str, np.ndarray]:
-        """Estimate metric depth and surface normals from a single RGB image.
-
-        Args:
-            image_input: RGB image as file path, URL, PIL Image, or numpy array.
-            timeout: Optional timeout in seconds for the HTTP request.
-
-        Returns:
-            Dict with keys:
-                - ``"depth"``: (H, W) float32 array, metric depth in metres.
-                - ``"normals"``: (H, W, 3) float32 array, unit surface normals.
-        """
-        return super().run(
-            image_input=image_input,
-            timeout=timeout,
-        )