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

{grid_cortex_client-0.3.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.3.0 → grid_cortex_client-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: grid-cortex-client
-Version: 0.3.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.3.0 → grid_cortex_client-0.5.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.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.
@@ -181,6 +219,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.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/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.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,
+        )