grid-cortex-client 0.5.1__tar.gz → 0.6.0__tar.gz

{grid_cortex_client-0.5.1 → grid_cortex_client-0.6.0}/PKG-INFO

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

{grid_cortex_client-0.5.1 → grid_cortex_client-0.6.0}/src/grid_cortex_client/model_type.py

@@ -13,36 +13,54 @@ from enum import Enum
 
 
 class ModelType(Enum):
-    CONTACT_GRASPNET = "contact-graspnet"
+    C_RADIO = "c-radio"
     """
-    Generate 6-DoF grasp poses from a depth image + camera intrinsics
-    (+ optional segmentation mask) using Contact-GraspNet (PyTorch port).
+    Compute an L2-normalized embedding for a single RGB image.
 
     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.
+        image_input (Union[str, Image.Image, np.ndarray]): RGB image as file path, URL, PIL Image, or numpy array.
+        timeout (float | None): Optional timeout in seconds for the HTTP request.
 
     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.
+        np.ndarray: Embedding as numpy array of shape ``(D,)``, dtype ``float32``,
+            L2-normalized so dot products give cosine similarity directly.
 
     Raises:
-        ValueError: If depth or intrinsics cannot be loaded.
+        ValueError: If image 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
+        >>> from PIL import Image
+        >>> client = CortexClient()
+        >>> image = np.array(Image.open("cat.jpg"))
+        >>> emb = client.run(ModelType.C_RADIO, image_input=image)
+        >>> print(emb.shape, float(np.linalg.norm(emb)))  # (D,) 1.0
+    """
+    CONTACT_GRASPNET = "contact-graspnet"
+    """
+    Generate grasps using ContactGraspNet.
+
+    Args:
+        depth_image (Union[str, Image.Image, np.ndarray, None]): Depth image (required).
+        camera_intrinsics (Union[str, np.ndarray, None]): 3×3 intrinsics matrix (required).
+        seg_image (Union[str, Image.Image, np.ndarray, None]): Optional segmentation mask (uint8 H×W).
+        aux_args (Dict[str, Any] | None): 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 (float | None): 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()
@@ -302,14 +320,18 @@ class ModelType(Enum):
     """
     METRIC3D = "metric3d"
     """
-    Estimate metric depth and surface normals from a single RGB image.
+    Estimate metric depth and/or surface normals from a single RGB image.
 
     Args:
         image_input (Union[str, Image.Image, np.ndarray]): RGB image as file path, URL, PIL Image, or numpy array.
         timeout (float | None): Optional timeout in seconds for the HTTP request.
+        outputs (Optional[Union[str, Sequence[str]]]): 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 keys:
+        Dict with the requested subset of:
             - ``"depth"``: (H, W) float32 array, metric depth in metres.
             - ``"normals"``: (H, W, 3) float32 array, unit surface normals.
 

grid_cortex_client-0.6.0/src/grid_cortex_client/models/c_radio.py

@@ -0,0 +1,129 @@
+"""C-RADIO wrapper.
+
+Client-side helper for the dense visual embedding endpoint powered by
+NVIDIA C-RADIOv4-SO400M.
+"""
+
+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 CRadio(BaseModel):
+    """Dense visual embedding (C-RADIOv4-SO400M).
+
+    Preferred usage (through :class:`CortexClient`)
+    ---------------------------------------------
+    ```pycon
+    >>> emb = CortexClient().run(ModelType.C_RADIO, image_input=img)
+    >>> emb.shape
+    (2304,)
+    >>> float(np.linalg.norm(emb))
+    1.0
+    ```
+
+    Args (for :py:meth:`run`)
+    ------------------------
+    image_input: RGB image (path, URL, ``PIL.Image`` or ``np.ndarray``).
+
+    Returns
+    -------
+    np.ndarray
+        L2-normalized embedding vector (``float32``, 1-D). Cosine similarity
+        between two outputs is a plain ``a @ b``.
+    """
+
+    name: str = "c-radio"
+    model_id: str = "c-radio"
+
+    def __init__(self, *, transport=None) -> None:
+        """Create a C-RADIO wrapper.
+
+        Args:
+            transport: Optional shared HTTP transport injected by
+                :class:`~grid_cortex_client.CortexClient`.
+        """
+        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 JPEG image bytes."""
+        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,
+    ) -> np.ndarray:
+        """Decode backend response into an L2-normalized embedding array.
+
+        Args:
+            response_data: Backend dict (or list-wrapped dict) containing an
+                ``"output"`` key holding ``.npy`` bytes for a 1-D float vector.
+
+        Returns:
+            np.ndarray of shape ``(D,)``, dtype ``float32``, L2-normalized.
+
+        Raises:
+            ValueError: If the payload is missing or cannot be decoded.
+        """
+        if isinstance(response_data, list):
+            response_dict = response_data[0]
+        else:
+            response_dict = response_data
+
+        emb_bytes = response_dict.get("output")
+        if not isinstance(emb_bytes, (bytes, bytearray, memoryview)):
+            raise ValueError("'output' field missing or not bytes")
+
+        arr = np.load(BytesIO(emb_bytes), allow_pickle=False)
+        return arr.astype(np.float32)
+
+    def run(
+        self,
+        image_input: Union[str, Image.Image, np.ndarray],
+        timeout: float | None = None,
+    ) -> np.ndarray:
+        """Compute an L2-normalized embedding for 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:
+            Embedding as numpy array of shape ``(D,)``, dtype ``float32``,
+            L2-normalized so dot products give cosine similarity directly.
+
+        Raises:
+            ValueError: If image 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
+            >>> from PIL import Image
+            >>> client = CortexClient()
+            >>> image = np.array(Image.open("cat.jpg"))
+            >>> emb = client.run(ModelType.C_RADIO, image_input=image)
+            >>> print(emb.shape, float(np.linalg.norm(emb)))  # (D,) 1.0
+        """
+        return super().run(
+            image_input=image_input,
+            timeout=timeout,
+        )

grid_cortex_client-0.6.0/tests/test_c_radio.py

@@ -0,0 +1,124 @@
+"""Test C-RADIO model."""
+
+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, run_concurrent_requests, is_gpu_mem_close
+
+
+class TestCRadioHTTP(BaseModelTestHTTP):
+    """HTTP-based tests for C-RADIO dense visual embedding model."""
+
+    @property
+    def batch_size(self) -> int:
+        """C-RADIO model batch size (matches `max_ongoing_requests` in config.yaml)."""
+        return 8
+
+    @property
+    def gpu_mem(self) -> float:
+        """C-RADIO model GPU memory budget (MB, matches config.yaml `gpu_mem`)."""
+        return 2000.0
+
+    @property
+    def max_input_size(self) -> tuple:
+        """C-RADIO model maximum input size (width, height)."""
+        return (720, 1280)
+
+    def test_basic_inference(self, client: CortexClient, rgb_image_np: np.ndarray):
+        """Test basic embedding extraction with valid input."""
+        emb = client.run(ModelType.C_RADIO, image_input=rgb_image_np)
+
+        assert emb is not None
+        assert isinstance(emb, np.ndarray)
+        assert emb.ndim == 1, f"emb.ndim {emb.ndim} != 1"
+        assert emb.shape[0] == 2304, f"emb dim {emb.shape[0]} != 2304 (C-RADIOv4-SO400M summary)"
+        assert emb.dtype == np.float32, f"emb.dtype {emb.dtype} != float32"
+        assert np.isfinite(emb).all(), "embedding contains NaN/Inf"
+        # L2-normalized → magnitude should be 1.0 within float32 tolerance.
+        norm = float(np.linalg.norm(emb))
+        assert abs(norm - 1.0) < 1e-3, f"embedding not L2-normalized (||emb|| = {norm})"
+
+    def test_image_input_list_rejected(self, client: CortexClient, rgb_image_np: np.ndarray):
+        """List image_input is rejected — at the client (TypeError) or server (422)."""
+        with pytest.raises((TypeError, CortexAPIError)):
+            client.run(ModelType.C_RADIO, image_input=[rgb_image_np, rgb_image_np])
+
+    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)
+        emb = client.run(ModelType.C_RADIO, image_input=large_image)
+
+        assert emb is not None
+        assert isinstance(emb, np.ndarray)
+        assert emb.ndim == 1
+
+        oversized_image = np.random.randint(0, 255, (4096, 4096, 3), dtype=np.uint8)
+        with pytest.raises(CortexAPIError) as exc_info:
+            client.run(ModelType.C_RADIO, image_input=oversized_image)
+        assert exc_info.value.status_code == 413
+
+        large_image = np.random.randint(0, 255, (*self.max_input_size, 3), dtype=np.uint8)
+        emb = client.run(ModelType.C_RADIO, image_input=large_image)
+
+        assert emb is not None
+        assert isinstance(emb, np.ndarray)
+        assert emb.ndim == 1
+
+    def test_inference_latency(self, client: CortexClient, rgb_image_np: np.ndarray):
+        """Test that inference latency is within acceptable range."""
+        start_time = time.time()
+        emb = client.run(ModelType.C_RADIO, image_input=rgb_image_np)
+        elapsed_time = time.time() - start_time
+
+        assert emb is not None
+        assert elapsed_time < 10.0, f"Inference took {elapsed_time:.2f}s, expected < 10s"
+
+    @pytest.mark.asyncio
+    async def test_concurrent_requests(self, async_client: AsyncCortexClient, rgb_image_np: np.ndarray):
+        """Test that multiple concurrent requests are handled correctly."""
+
+        async def make_request():
+            return await async_client.run(ModelType.C_RADIO, image_input=rgb_image_np)
+
+        results = await run_concurrent_requests(make_request, num_requests=self.batch_size)
+
+        assert len(results) == self.batch_size
+        assert all(not isinstance(r, Exception) for r in results)
+        assert all(isinstance(r, np.ndarray) for r in results)
+        assert all(r.ndim == 1 for r in results)
+
+    @pytest.mark.asyncio
+    async def test_gpu_memory_usage(self, async_client: AsyncCortexClient, rgb_image_np: np.ndarray):
+        """Test that GPU memory usage is within expected limits during inference."""
+        large_image = np.random.randint(0, 255, (*self.max_input_size, 3), dtype=np.uint8)
+
+        async def make_request():
+            return await async_client.run(ModelType.C_RADIO, image_input=large_image, raw_response=True)
+
+        results = await run_concurrent_requests(make_request, num_requests=self.batch_size)
+
+        assert len(results) == self.batch_size
+        assert all(not isinstance(r, Exception) for r in results)
+        assert all(is_gpu_mem_close(self.gpu_mem, result["gpu_stats"]["mem_usage_mb"]) for result in results), (
+            f"GPU memory peak too high: {[result['gpu_stats']['mem_usage_mb'] for result in results]} MB"
+        )
+
+    @pytest.mark.asyncio
+    async def test_cortex_hub(self, rgb_image_np: np.ndarray):
+        """Test C-RADIO inference via CortexHubClient (WebSocket)."""
+        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:
+            emb = await hub.run(ModelType.C_RADIO, image_input=rgb_image_np)
+
+        assert emb is not None
+        assert isinstance(emb, np.ndarray)
+        assert emb.ndim == 1