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

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

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

{grid_cortex_client-0.3.0 → grid_cortex_client-0.4.0}/src/grid_cortex_client/cortex_client.py

@@ -62,11 +62,20 @@ class _CortexClientBase:
     _MODEL_ID_TO_HANDLER_CLASS: Dict[str, Type[BaseModel]] = get_registry()
 
     def _get_handler_class(self, model_id: Union[str, ModelType]) -> Optional[Type[BaseModel]]:
-        """Look up the handler class for a model ID."""
-        model_id_str = model_id.value if isinstance(model_id, ModelType) else str(model_id)
-        for keyword, HClass in self._MODEL_ID_TO_HANDLER_CLASS.items():
-            if keyword in model_id_str.lower():
-                return HClass
+        """Look up the handler class for a model ID.
+
+        Prefer an exact match so a registered id that is a substring of another
+        (e.g. ``graspgen`` vs ``graspgenx``) can't shadow it. Fall back to
+        substring matching, longest keyword first, to keep the lenient behavior
+        for free-form ids while still preventing short-id shadowing.
+        """
+        model_id_str = (model_id.value if isinstance(model_id, ModelType) else str(model_id)).lower()
+        exact = self._MODEL_ID_TO_HANDLER_CLASS.get(model_id_str)
+        if exact is not None:
+            return exact
+        for keyword in sorted(self._MODEL_ID_TO_HANDLER_CLASS, key=len, reverse=True):
+            if keyword in model_id_str:
+                return self._MODEL_ID_TO_HANDLER_CLASS[keyword]
         return None
 
     def _preprocess(

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

@@ -181,6 +181,33 @@ class ModelType(Enum):
         >>> grasps, conf = client.run(ModelType.GRASPGEN, depth_image=depth_image, seg_image=seg_image, camera_intrinsics=K, aux_args=aux)
         >>> print(f"Generated {len(grasps)} grasps")
     """
+    GRASPGENX = "graspgenx"
+    """
+    Generate 6-DoF grasps for a named gripper with GraspGenX.
+
+    Args:
+        aux_args (Dict[str, Any]): Must include "gripper_name"; optional
+            "num_grasps", "planner" ("graspmoe"|"diffusion"), "grasp_threshold",
+            "camera_extrinsics".
+        point_cloud (Union[str, np.ndarray, Sequence, None]): (N, 3) object point cloud.
+        depth_image / seg_image / camera_intrinsics: alternative to point_cloud.
+        timeout (float | None): Optional HTTP timeout.
+
+    Returns:
+        Dict[str, Any]: Dict with:
+            - grasps: (N, 4, 4) array of grasp poses
+            - confidence: (N,) array of scores
+            - latency_ms: server-reported latency
+
+    Examples:
+        >>> from grid_cortex_client import CortexClient, ModelType
+        >>> import numpy as np
+        >>> client = CortexClient()
+        >>> pc = np.load("object_pc.npy")  # (N, 3)
+        >>> out = client.run(ModelType.GRASPGENX, point_cloud=pc,
+        ...                   aux_args={"gripper_name": "robotiq_2f_85", "num_grasps": 64})
+        >>> print(out["grasps"].shape)
+    """
     GSAM2 = "gsam2"
     """
     Segment objects in image using text prompt.

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

@@ -0,0 +1,151 @@
+# grid_cortex_client/src/grid_cortex_client/models/graspgenx.py
+"""GraspGenX wrapper.
+
+Cross-embodiment 6-DoF grasp generation (NVlabs/GraspGenX). A single model
+conditioned on a gripper representation, so grasps are requested for a named
+gripper (``gripper_name``) rather than a per-gripper model.
+
+Input: a segmented object point cloud (or depth+seg+intrinsics, converted
+server-side). Output: 6-DoF grasp poses + confidences.
+
+NOTE: weights are under the NVIDIA Open Model License — see the model recipe.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Sequence, Union
+
+import numpy as np
+from PIL import Image
+
+from ..preprocessing import array_to_npy_bytes, load_image
+from .base_model import BaseModel
+
+
+class GraspGenX(BaseModel):
+    """Cross-embodiment grasp generation (GraspGenX).
+
+    Preferred usage
+    ---------------
+    ```pycon
+    >>> out = CortexClient().run(
+    ...     "graspgenx",
+    ...     point_cloud=pc,                      # (N, 3) float32
+    ...     aux_args={"gripper_name": "robotiq_2f_85", "num_grasps": 64},
+    ... )
+    >>> grasps, conf = out["grasps"], out["confidence"]
+    ```
+    """
+
+    name: str = "graspgenx"
+    model_id: str = "graspgenx"
+
+    # ------------------------------------------------------------------
+    # BaseModel implementation
+    # ------------------------------------------------------------------
+
+    def preprocess(
+        self,
+        *,
+        aux_args: Dict[str, Any] | None = None,
+        point_cloud: Union[str, np.ndarray, Sequence[Sequence[float]], None] = None,
+        depth_image: Union[str, Image.Image, np.ndarray, None] = None,
+        seg_image: Union[str, Image.Image, np.ndarray, None] = None,
+        camera_intrinsics: Union[str, np.ndarray, None] = None,
+    ) -> Dict[str, Any]:
+        """Prepare the msgpack payload.
+
+        Args:
+            aux_args: Dict with "gripper_name", "num_grasps", "planner",
+                "grasp_threshold", "camera_extrinsics".
+            point_cloud: (N, 3) point cloud (ndarray/list/path to .npy).
+            depth_image / seg_image / camera_intrinsics: alternative to point_cloud;
+                converted to a point cloud server-side.
+        """
+        if aux_args is None:
+            raise ValueError("'aux_args' is required (must include 'gripper_name').")
+
+        has_pc = point_cloud is not None
+        has_depth_triplet = any(x is not None for x in (depth_image, seg_image, camera_intrinsics))
+        if has_pc and has_depth_triplet:
+            raise ValueError("Provide either point_cloud or depth_image+seg_image+camera_intrinsics, not both.")
+        if not has_pc and not all(x is not None for x in (depth_image, seg_image, camera_intrinsics)):
+            raise ValueError("Provide point_cloud, or all of depth_image + seg_image + camera_intrinsics.")
+
+        if has_pc:
+            pc_array = np.load(point_cloud) if isinstance(point_cloud, str) else np.asarray(point_cloud)
+            if pc_array.ndim != 2 or pc_array.shape[1] != 3:
+                raise ValueError("point_cloud must be an (N, 3) array.")
+            return {"point_cloud": array_to_npy_bytes(pc_array.astype(np.float32)), "aux_args": aux_args}
+
+        depth_array = (
+            np.array(load_image(depth_image))
+            if isinstance(depth_image, (str, Image.Image))
+            else np.asarray(depth_image)
+        )
+        seg_array = (
+            np.array(load_image(seg_image)) if isinstance(seg_image, (str, Image.Image)) else np.asarray(seg_image)
+        )
+        intrinsics = np.load(camera_intrinsics) if isinstance(camera_intrinsics, str) else np.asarray(camera_intrinsics)
+        return {
+            "depth_image": array_to_npy_bytes(depth_array),
+            "seg_image": array_to_npy_bytes(seg_array),
+            "camera_intrinsics": array_to_npy_bytes(intrinsics),
+            "aux_args": aux_args,
+        }
+
+    def postprocess(self, response_data: Dict[str, Any], **_: Any) -> Dict[str, Any]:  # noqa: D401
+        """Decode grasps + confidence from the response."""
+        grasps = np.array(response_data["output"])
+        if grasps.size == 0:
+            # Keep the documented (N, 4, 4) contract on the empty path.
+            grasps = np.empty((0, 4, 4), dtype=np.float32)
+        return {
+            "grasps": grasps,
+            "confidence": np.array(response_data["confidence"]),
+            "latency_ms": response_data.get("latency_ms"),
+        }
+
+    def run(
+        self,
+        *,
+        aux_args: Dict[str, Any],
+        point_cloud: Union[str, np.ndarray, Sequence[Sequence[float]], None] = None,
+        depth_image: Union[str, Image.Image, np.ndarray, None] = None,
+        seg_image: Union[str, Image.Image, np.ndarray, None] = None,
+        camera_intrinsics: Union[str, np.ndarray, None] = None,
+        timeout: float | None = None,
+    ) -> Dict[str, Any]:
+        """Generate 6-DoF grasps for a named gripper with GraspGenX.
+
+        Args:
+            aux_args (Dict[str, Any]): Must include "gripper_name"; optional
+                "num_grasps", "planner" ("graspmoe"|"diffusion"), "grasp_threshold",
+                "camera_extrinsics".
+            point_cloud (Union[str, np.ndarray, Sequence, None]): (N, 3) object point cloud.
+            depth_image / seg_image / camera_intrinsics: alternative to point_cloud.
+            timeout (float | None): Optional HTTP timeout.
+
+        Returns:
+            Dict[str, Any]: Dict with:
+                - grasps: (N, 4, 4) array of grasp poses
+                - confidence: (N,) array of scores
+                - latency_ms: server-reported latency
+
+        Examples:
+            >>> from grid_cortex_client import CortexClient, ModelType
+            >>> import numpy as np
+            >>> client = CortexClient()
+            >>> pc = np.load("object_pc.npy")  # (N, 3)
+            >>> out = client.run(ModelType.GRASPGENX, point_cloud=pc,
+            ...                   aux_args={"gripper_name": "robotiq_2f_85", "num_grasps": 64})
+            >>> print(out["grasps"].shape)
+        """
+        return super().run(
+            aux_args=aux_args,
+            point_cloud=point_cloud,
+            depth_image=depth_image,
+            seg_image=seg_image,
+            camera_intrinsics=camera_intrinsics,
+            timeout=timeout,
+        )

grid_cortex_client-0.4.0/tests/test_graspgenx.py

@@ -0,0 +1,143 @@
+"""Test GraspGenX model (cross-embodiment 6-DoF grasp generation).
+
+Runs against a live Ray Serve + CortexHub via CortexClient.run(...). Uses a
+synthetic object point cloud so the test does not depend on staged eval data.
+"""
+
+import os
+import time
+
+import numpy as np
+import pytest
+
+from grid_cortex_client import AsyncCortexClient, CortexClient, CortexHubClient, ModelType
+
+from .conftest import BaseModelTestHTTP, is_gpu_mem_close, run_concurrent_requests
+
+
+def _synthetic_object_pc(n: int = 2048, seed: int = 0) -> np.ndarray:
+    """A small noisy point cloud standing in for a segmented object."""
+    rng = np.random.default_rng(seed)
+    return ((rng.random((n, 3), dtype=np.float32) - 0.5) * 0.1).astype(np.float32)
+
+
+class TestGraspGenXHTTP(BaseModelTestHTTP):
+    """HTTP-based tests for GraspGenX."""
+
+    @property
+    def batch_size(self) -> int:
+        """max_ongoing_requests from config."""
+        return 1
+
+    @property
+    def gpu_mem(self) -> float:
+        """Declared GPU memory (MB). ~688 MB peak in CI smoke (num_grasps=16); 1000 = headroom."""
+        return 1000.0
+
+    @property
+    def max_input_size(self) -> tuple:
+        """Not image-based; placeholder for the ABC contract."""
+        return (480, 640)
+
+    @pytest.fixture(scope="class")
+    def point_cloud(self) -> np.ndarray:
+        return _synthetic_object_pc()
+
+    @pytest.fixture(scope="class")
+    def aux_args(self) -> dict:
+        return {"gripper_name": "robotiq_2f_85", "num_grasps": 16}
+
+    def test_basic_inference(self, client: CortexClient, point_cloud: np.ndarray, aux_args: dict):
+        """Grasp generation from a direct point cloud."""
+        result = client.run(ModelType.GRASPGENX, point_cloud=point_cloud, 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_grasp_output_structure(self, client: CortexClient, point_cloud: np.ndarray, aux_args: dict):
+        """Grasps are (N, 4, 4) and confidence is (N,) when any are returned."""
+        result = client.run(ModelType.GRASPGENX, point_cloud=point_cloud, aux_args=aux_args)
+        grasps, confidence = result["grasps"], result["confidence"]
+        if len(grasps) > 0:
+            assert grasps.ndim == 3 and grasps.shape[1:] == (4, 4)
+            assert confidence.ndim == 1
+            assert len(grasps) == len(confidence)
+
+    def test_maximum_payload_size(self, client: CortexClient, aux_args: dict):
+        """Inference works on a large object point cloud (well above the basic case).
+
+        Uses 8192 points — a realistic dense segmented object, ~4x the basic test.
+        (The handler's hard guard is MAX_POINTS=50000, but a single object cloud is
+        never that large, and graspmoe's outlier-removal step is superlinear in N.)
+        """
+        large_pc = _synthetic_object_pc(n=8192)
+        result = client.run(ModelType.GRASPGENX, point_cloud=large_pc, aux_args=aux_args)
+        assert result is not None
+        assert "grasps" in result
+        assert "confidence" in result
+
+    def test_raw_response_with_gpu_stats(self, client: CortexClient, point_cloud: np.ndarray, aux_args: dict):
+        """Raw response carries output/confidence/latency/gpu_stats."""
+        raw = client.run(ModelType.GRASPGENX, point_cloud=point_cloud, 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 and "mem_usage_mb" in raw["gpu_stats"]
+
+    def test_missing_aux_args(self, client: CortexClient, point_cloud: np.ndarray):
+        """Missing aux_args raises a client-side ValueError."""
+        with pytest.raises(ValueError):
+            client.run(ModelType.GRASPGENX, point_cloud=point_cloud)
+
+    def test_missing_input(self, client: CortexClient, aux_args: dict):
+        """Neither point_cloud nor depth-triplet raises a client-side ValueError."""
+        with pytest.raises(ValueError):
+            client.run(ModelType.GRASPGENX, aux_args=aux_args)
+
+    def test_inference_latency(self, client: CortexClient, point_cloud: np.ndarray, aux_args: dict):
+        """Latency within an acceptable bound."""
+        start = time.time()
+        result = client.run(ModelType.GRASPGENX, point_cloud=point_cloud, 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"
+
+    @pytest.mark.asyncio
+    async def test_concurrent_requests(self, async_client: AsyncCortexClient, point_cloud: np.ndarray, aux_args: dict):
+        """Concurrent requests are handled."""
+
+        async def make_request():
+            return await async_client.run(ModelType.GRASPGENX, point_cloud=point_cloud, aux_args=aux_args)
+
+        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, point_cloud: np.ndarray, aux_args: dict):
+        """GPU memory within the declared budget."""
+
+        async def make_request():
+            return await async_client.run(
+                ModelType.GRASPGENX, point_cloud=point_cloud, aux_args=aux_args, raw_response=True
+            )
+
+        results = await run_concurrent_requests(make_request, num_requests=min(self.batch_size, 2))
+        assert all(not isinstance(r, Exception) for r in results)
+        assert all(is_gpu_mem_close(self.gpu_mem, r["gpu_stats"]["mem_usage_mb"]) for r 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, point_cloud: np.ndarray, aux_args: dict):
+        """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:
+            result = await hub.run(ModelType.GRASPGENX, point_cloud=point_cloud, aux_args=aux_args)
+        assert result is not None