Spaces:

HorizonRobotics
/

EmbodiedGen-Texture-Gen

Running on Zero

EmbodiedGen-Texture-Gen / embodied_gen /scripts /affordance_annot /gen_grasp.py

update

509565d 3 months ago

28 kB

# Project EmbodiedGen
#
# Copyright (c) 2025 Horizon Robotics. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied. See the License for the specific language governing
# permissions and limitations under the License.

import os
import sys
import time
from dataclasses import dataclass, field
from typing import Literal

from embodied_gen.utils.general import filter_warnings

filter_warnings()

import numpy as np
import trimesh
import trimesh.transformations as tra
import tyro
import yaml

sys.path.append(".")
from embodied_gen.utils.geometry import (
    sample_surface_points,
)
from embodied_gen.utils.io_utils import (
    URDFFile,
    load_json,
    load_mesh,
    write_json,
)
from embodied_gen.utils.log import logger
from embodied_gen.utils.vis_utils import show_grasps

__all__ = [
    "GraspGenHyperConfig",
    "GraspGenXHyperConfig",
    "GraspPoseConfig",
    "GraspGenerator",
    "run_grasp",
    "entrypoint",
]

FINGER_CENTER_CACHE: dict[str, np.ndarray] = {
    "franka_panda": np.array([0.0, 0.0, 0.1034], dtype=np.float64),
}

GRIPPER_GEOMETRY_CACHE: dict[str, dict[str, float]] = {
    "franka_panda": {
        "width": 0.10537486,
        "depth": 0.10527314,
    },
}


@dataclass
class GraspGenHyperConfig:
    gripper_config: str = (
        "~/.cache/huggingface/hub/GraspGenModels/checkpoints/graspgen_{}.yml"
    )
    grasp_threshold: float = -1.0
    num_grasps: int = 4000
    return_topk: bool = True
    topk_num_grasps: int = 4000
    mesh_scale: tuple[float, float, float] = (1.0, 1.0, 1.0)
    num_sample_points: int = 2000
    min_grasps: int = 40
    max_tries: int = 6
    remove_outliers: bool = False


@dataclass
class GraspGenXHyperConfig:
    checkpoint_root: str = "~/.cache/huggingface/hub/GraspGenXModels/release"
    gen_pth: str | None = None
    dis_pth: str | None = None
    assets_dir: str = (
        "thirdparty/gripper_descriptions/gripper_descriptions/assets"
    )
    grasp_threshold: float = -1.0
    num_grasps: int = 4000
    return_topk: bool = True
    topk_num_grasps: int = 4000
    mesh_scale: tuple[float, float, float] = (1.0, 1.0, 1.0)
    num_sample_points: int = 2000
    min_grasps: int = 40
    max_tries: int = 6
    remove_outliers: bool = False


@dataclass
class GraspPoseConfig:
    urdf_paths: list[str] = field(default_factory=list)
    model_name: str = "GraspGen"
    gripper_name: str = "franka_panda"
    mesh_type: Literal["collision", "visual"] = "collision"
    update_urdf: bool = True
    overwrite: bool = True
    gripper_mesh_path: str = (
        f"embodied_gen/scripts/affordance_annot/{gripper_name}/coll_mesh.obj"
    )
    grasps_per_part: int = 20
    visualize_grasps: bool = False
    visualizer_port: int = 8080
    graspgen: GraspGenHyperConfig | None = None


class GraspGenerator:
    def __init__(self, cfg: GraspPoseConfig):
        self.cfg = cfg
        self.grasp_visualizer = None
        self.validate_config(cfg)
        self.modules = self.load_model()
        self.pipeline = self.build_sampler()
        self.finger_center = FINGER_CENTER_CACHE.get(
            self.cfg.gripper_name,
            np.array([0.0, 0.0, 0.0], dtype=np.float64),
        )
        (
            self.local_contact_points,
            self.local_closing_directions,
        ) = self.build_local_contact_samples()

    def validate_config(self, cfg: GraspPoseConfig) -> None:
        cfg.model_name = cfg.model_name.lower()
        if cfg.model_name not in {"graspgen", "graspgenx"}:
            raise ValueError(
                "unsupported model '{}', available models: GraspGen, GraspGenX".format(
                    cfg.model_name
                )
            )

        cfg.graspgen = (
            GraspGenHyperConfig()
            if cfg.model_name == "graspgen"
            else GraspGenXHyperConfig()
        )
        if cfg.model_name == "graspgen":
            cfg.graspgen.gripper_config = cfg.graspgen.gripper_config.format(
                cfg.gripper_name
            )

    def load_model(self):
        if self.cfg.model_name == "graspgen":
            sys.path.append("thirdparty/GraspGen")
            from grasp_gen.grasp_server import GraspGenSampler, load_grasp_cfg

            return {
                "GraspGenSampler": GraspGenSampler,
                "model_cfg": load_grasp_cfg,
            }
        if self.cfg.model_name == "graspgenx":
            sys.path.append("thirdparty/GraspGenX")
            from graspgenx.grasp_server import GraspGenXSampler
            from graspgenx.utils.checkpoint_io import load_model_cfg

            return {
                "GraspGenSampler": GraspGenXSampler,
                "model_cfg": load_model_cfg,
            }

    @staticmethod
    def _extract_checkpoint_filenames(yml_path: str) -> list[str]:
        """Recursively collect all `checkpoint: <name>.pth` entries.
        Walks a GraspGen config yaml so we know every weight file it
        requires.
        """
        with open(yml_path) as f:
            cfg = yaml.safe_load(f)

        names: set[str] = set()

        def _walk(node) -> None:
            if isinstance(node, dict):
                for key, value in node.items():
                    if (
                        key == "checkpoint"
                        and isinstance(value, str)
                        and value.strip()
                    ):
                        names.add(value.strip())
                    else:
                        _walk(value)
            elif isinstance(node, list):
                for item in node:
                    _walk(item)

        _walk(cfg)
        return sorted(names)

    @staticmethod
    def _dir_has_pth(directory: str, filename: str | None) -> bool:
        if filename:
            return os.path.exists(os.path.join(directory, filename))
        return os.path.isdir(directory) and any(
            f.endswith(".pth") for f in os.listdir(directory)
        )

    def _ensure_checkpoints(
        self,
        is_ready: callable,
        repo_id: str,
        allow_patterns: str,
        download_dir: str,
        desc: str,
    ) -> None:
        """Download `desc` weights from `repo_id` unless already ready.
        Re-checks afterwards so a stale/partial local cache (e.g. config
        present but a referenced .pth missing) is retried instead of
        silently failing deep inside third-party code.
        """
        if is_ready():
            return

        logger.info(
            f"Downloading {desc} weights from Hugging Face. "
            "First-time use may take several minutes; please wait patiently."
        )
        from huggingface_hub import snapshot_download

        snapshot_download(
            repo_id=repo_id,
            allow_patterns=allow_patterns,
            local_dir=os.path.expanduser(download_dir),
        )

        if not is_ready():
            raise FileNotFoundError(
                f"{desc} checkpoint(s) still missing after downloading from "
                f"Hugging Face repo '{repo_id}'. The local cache may be from "
                "a previous partial download; please remove it and retry, "
                "or download manually."
            )

    def build_sampler(self):
        if self.cfg.model_name == "graspgen":
            gripper_config = os.path.expanduser(
                self.cfg.graspgen.gripper_config
            )
            checkpoints_dir = os.path.dirname(gripper_config)

            def _graspgen_ready() -> bool:
                if not os.path.exists(gripper_config):
                    return False
                required = self._extract_checkpoint_filenames(gripper_config)
                return all(
                    os.path.exists(os.path.join(checkpoints_dir, name))
                    for name in required
                )

            self._ensure_checkpoints(
                _graspgen_ready,
                "adithyamurali/GraspGenModels",
                "checkpoints/*",
                "~/.cache/huggingface/hub/GraspGenModels",
                "GraspGen",
            )
            grasp_cfg = self.modules["model_cfg"](gripper_config)
            return self.modules["GraspGenSampler"](grasp_cfg)
        if self.cfg.model_name == "graspgenx":
            checkpoint_root = os.path.expanduser(
                self.cfg.graspgen.checkpoint_root
            )
            gen_dir = os.path.join(checkpoint_root, "gen")
            dis_dir = os.path.join(checkpoint_root, "dis")

            def _graspgenx_ready() -> bool:
                return (
                    os.path.exists(os.path.join(gen_dir, "config.yaml"))
                    and self._dir_has_pth(gen_dir, self.cfg.graspgen.gen_pth)
                    and self._dir_has_pth(dis_dir, self.cfg.graspgen.dis_pth)
                )

            self._ensure_checkpoints(
                _graspgenx_ready,
                "adithyamurali/GraspGenXModel",
                "release/*",
                "~/.cache/huggingface/hub/GraspGenXModels",
                "GraspGenX",
            )
            model_cfg = self.modules["model_cfg"](
                os.path.join(checkpoint_root, "gen"),
                os.path.join(checkpoint_root, "dis"),
                self.cfg.graspgen.gen_pth,
                self.cfg.graspgen.dis_pth,
            )
            return self.modules["GraspGenSampler"](
                model_cfg,
                self.cfg.gripper_name,
                assets_dir=self.cfg.graspgen.assets_dir,
            )

    def update_urdf(self, urdf_path: str | None, grasp_path: str) -> None:
        if urdf_path is None or not self.cfg.update_urdf:
            return

        URDFFile(urdf_path).write(
            {
                "custom_data/affordance/affordance_annot": os.path.relpath(
                    grasp_path,
                    os.path.dirname(urdf_path),
                ),
            }
        )

    def get_part_seg_mesh_transform(self, urdf: URDFFile) -> dict:
        return {
            "origin_xyz": urdf.read(
                ".//custom_data/affordance/visual_seg/origin",
                attr="xyz",
                default="0 0 0",
            ),
            "origin_rpy": urdf.read(
                ".//custom_data/affordance/visual_seg/origin",
                attr="rpy",
                default="0 0 0",
            ),
            "scale": urdf.read(
                ".//custom_data/affordance/visual_seg/geometry/mesh",
                attr="scale",
                default="1.0 1.0 1.0",
            ),
        }

    def add_grasps_to_affordance(
        self,
        affordance_payload: dict,
        grasps: np.ndarray,
        confidences: np.ndarray,
        part_grasp_ids: dict[int, list[int]] | None,
    ) -> None:
        for affordance in affordance_payload.get("affordances", []):
            if (
                not isinstance(affordance, dict)
                or affordance.get("id") is None
            ):
                continue
            affordance.pop("grasps", None)
            affordance.pop("confidences", None)
            grasp_ids = part_grasp_ids.get(int(affordance["id"]), [])
            grasp_group = {}
            for grasp_idx in grasp_ids:
                grasp = grasps[grasp_idx]
                quaternion = tra.quaternion_from_matrix(grasp[:3, :3])
                grasp_group[f"grasp_{grasp_idx}"] = {
                    "confidence": float(confidences[grasp_idx]),
                    "position": grasp[:3, 3].tolist(),
                    "orientation": {
                        "w": float(quaternion[0]),
                        "xyz": quaternion[1:].tolist(),
                    },
                }
            affordance["grasp_group"] = grasp_group

    def infer_grasps(
        self,
        mesh: trimesh.Trimesh,
        mesh_path: str,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray] | None:
        points = sample_surface_points(
            mesh, self.cfg.graspgen.num_sample_points
        )
        center_transform = tra.translation_matrix(-points.mean(axis=0))
        centered_points = tra.transform_points(
            points, center_transform
        ).astype(np.float32)
        topk_limit = (
            self.cfg.graspgen.topk_num_grasps
            if self.cfg.graspgen.return_topk
            else -1
        )
        grasps, confidences = self.modules["GraspGenSampler"].run_inference(
            centered_points,
            self.pipeline,
            grasp_threshold=self.cfg.graspgen.grasp_threshold,
            num_grasps=self.cfg.graspgen.num_grasps,
            topk_num_grasps=topk_limit,
            min_grasps=self.cfg.graspgen.min_grasps,
            max_tries=self.cfg.graspgen.max_tries,
            remove_outliers=self.cfg.graspgen.remove_outliers,
        )

        confidences = confidences.detach().cpu().numpy()
        grasps = grasps.detach().cpu().numpy()
        if len(grasps) == 0:
            logger.error(f"No grasp poses generated for mesh: {mesh_path}")
            return None

        grasps = np.asarray(
            [tra.inverse_matrix(center_transform) @ grasp for grasp in grasps]
        )
        return grasps, confidences

    def get_grasp_centers(
        self,
        grasps: np.ndarray,
    ) -> np.ndarray:
        local_centers = np.repeat(
            np.asarray(self.finger_center, dtype=np.float64)[None, :],
            len(grasps),
            axis=0,
        )
        grasp_center_points = (
            np.einsum(
                "nij,nj->ni",
                grasps[:, :3, :3],
                local_centers,
            )
            + grasps[:, :3, 3]
        )
        return grasp_center_points

    def build_local_contact_samples(self) -> tuple[np.ndarray, np.ndarray]:
        geometry = GRIPPER_GEOMETRY_CACHE.get(self.cfg.gripper_name)
        if geometry is None:
            return (
                np.asarray(self.finger_center, dtype=np.float64)[None, :],
                np.array([[1.0, 0.0, 0.0]], dtype=np.float64),
            )

        half_width = float(geometry["width"]) / 2.0
        depth = float(geometry["depth"])
        z_values = np.linspace(depth - 0.015, depth, 6, dtype=np.float64)
        points = []
        directions = []
        for z_value in z_values:
            points.extend(
                [
                    [-half_width, 0.0, z_value],
                    [half_width, 0.0, z_value],
                ]
            )
            directions.extend(
                [
                    [1.0, 0.0, 0.0],
                    [-1.0, 0.0, 0.0],
                ]
            )
        return (
            np.asarray(points, dtype=np.float64),
            np.asarray(directions, dtype=np.float64),
        )

    def get_grasp_contact_rays(
        self,
        grasps: np.ndarray,
    ) -> tuple[np.ndarray, np.ndarray]:
        local_contact_points = np.repeat(
            self.local_contact_points[None, :, :],
            len(grasps),
            axis=0,
        )
        local_closing_directions = np.repeat(
            self.local_closing_directions[None, :, :],
            len(grasps),
            axis=0,
        )
        ray_origins = (
            np.einsum(
                "nij,nkj->nki",
                grasps[:, :3, :3],
                local_contact_points,
            )
            + grasps[:, None, :3, 3]
        )
        ray_directions = np.einsum(
            "nij,nkj->nki",
            grasps[:, :3, :3],
            local_closing_directions,
        )
        ray_norms = np.linalg.norm(ray_directions, axis=-1, keepdims=True)
        ray_directions = np.divide(
            ray_directions,
            ray_norms,
            out=np.zeros_like(ray_directions),
            where=ray_norms > 0.0,
        )
        return ray_origins, ray_directions

    def bind_contact_rays_to_faces(
        self,
        part_mesh: trimesh.Trimesh,
        ray_origins: np.ndarray,
        ray_directions: np.ndarray,
    ) -> np.ndarray:
        ray_origins = np.asarray(ray_origins, dtype=np.float64)
        ray_directions = np.asarray(ray_directions, dtype=np.float64)

        if (
            not isinstance(part_mesh, trimesh.Trimesh)
            or len(part_mesh.faces) == 0
        ):
            return np.full(len(ray_origins), -1, dtype=np.int64)

        try:
            face_ids = part_mesh.ray.intersects_first(
                ray_origins,
                ray_directions,
            )
            return np.asarray(face_ids, dtype=np.int64)
        except Exception as exc:
            logger.warning(
                "Failed to raycast grasp closing directions to part mesh: "
                f"{exc}"
            )
            return np.full(len(ray_origins), -1, dtype=np.int64)

    def classify_grasps_by_contact(
        self,
        grasps: np.ndarray,
        part_mesh: trimesh.Trimesh,
        face_ids: np.ndarray | None,
    ) -> tuple[np.ndarray | None, np.ndarray | None]:
        if face_ids is None:
            return None, None

        face_ids = np.asarray(face_ids, dtype=np.int64)
        if len(face_ids) != len(part_mesh.faces):
            raise ValueError(
                "face_ids length must match segmentation mesh faces, got "
                f"{len(face_ids)} and {len(part_mesh.faces)}"
            )

        ray_origins, ray_directions = self.get_grasp_contact_rays(grasps)
        num_grasps, num_contact_points, _ = ray_origins.shape
        logger.info(
            "Binding {} grasp closing rays to {} part faces.".format(
                num_grasps * num_contact_points,
                len(part_mesh.faces),
            )
        )
        nearest_face_ids = self.bind_contact_rays_to_faces(
            part_mesh,
            ray_origins.reshape(-1, 3),
            ray_directions.reshape(-1, 3),
        ).reshape(num_grasps, num_contact_points)
        contact_part_ids = self.parts_for_faces(
            nearest_face_ids.reshape(-1),
            face_ids,
        ).reshape(num_grasps, num_contact_points)

        grasp_part_ids = np.full(num_grasps, -1, dtype=np.int64)
        contact_scores = np.zeros(num_grasps, dtype=np.float64)
        for grasp_idx, part_ids in enumerate(contact_part_ids):
            valid_part_ids = part_ids[part_ids >= 0]
            if len(valid_part_ids) == 0:
                continue
            unique_ids, counts = np.unique(
                valid_part_ids,
                return_counts=True,
            )
            best_idx = int(np.argmax(counts))
            grasp_part_ids[grasp_idx] = int(unique_ids[best_idx])
            contact_scores[grasp_idx] = float(counts[best_idx]) / float(
                num_contact_points
            )

        return grasp_part_ids, contact_scores

    def parts_for_faces(
        self,
        nearest_face_ids: np.ndarray,
        face_ids: np.ndarray | None,
    ) -> np.ndarray | None:
        if face_ids is None:
            return None

        nearest_face_ids = np.asarray(nearest_face_ids, dtype=np.int64)
        grasp_part_ids = np.full(len(nearest_face_ids), -1, dtype=np.int64)
        valid_mask = (nearest_face_ids >= 0) & (
            nearest_face_ids < len(face_ids)
        )
        grasp_part_ids[valid_mask] = face_ids[nearest_face_ids[valid_mask]]
        return grasp_part_ids

    def pick_top_grasps(
        self,
        grasps: np.ndarray,
        confidences: np.ndarray,
        part_ids: np.ndarray | None,
        graspable_part_ids: set[int] | None = None,
        part_scores: np.ndarray | None = None,
    ) -> tuple[
        np.ndarray, np.ndarray, np.ndarray, dict[int, list[int]] | None
    ]:
        if part_ids is None:
            return (
                grasps,
                confidences,
                np.arange(len(grasps), dtype=np.int64),
                None,
            )

        part_ids = np.asarray(part_ids, dtype=np.int64)
        selected_indices = []
        selected_part_ids = []
        for part_id in sorted(
            int(pid) for pid in np.unique(part_ids[part_ids >= 0])
        ):
            if (
                graspable_part_ids is not None
                and part_id not in graspable_part_ids
            ):
                continue
            part_indices = np.flatnonzero(part_ids == part_id)
            if part_scores is None:
                rank = np.argsort(
                    -confidences[part_indices],
                    kind="mergesort",
                )
            else:
                rank = np.lexsort(
                    (
                        -confidences[part_indices],
                        -part_scores[part_indices],
                    )
                )
            ranked_indices = part_indices[rank]
            if len(ranked_indices) > self.cfg.grasps_per_part:
                selected = ranked_indices[: self.cfg.grasps_per_part]
            else:
                selected = ranked_indices
            selected_indices.extend(int(idx) for idx in selected)
            selected_part_ids.extend([part_id] * len(selected))

        selected_indices = np.asarray(selected_indices, dtype=np.int64)
        if len(selected_indices) == 0:
            return (
                grasps[:0],
                confidences[:0],
                selected_indices,
                {},
            )

        part_grasp_ids: dict[int, list[int]] = {}
        for new_idx, part_id in enumerate(selected_part_ids):
            part_grasp_ids.setdefault(part_id, []).append(new_idx)

        return (
            grasps[selected_indices],
            confidences[selected_indices],
            selected_indices,
            part_grasp_ids,
        )

    def select_grasps(
        self,
        grasps: np.ndarray,
        confidences: np.ndarray,
        part_mesh: trimesh.Trimesh,
        face_ids: np.ndarray,
        affordance_payload: dict,
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, dict[int, list[int]]]:
        original_count = len(grasps)

        # Raycast from fingertip samples along the gripper closing direction.
        grasp_center_points = self.get_grasp_centers(grasps)
        grasp_part_ids, contact_scores = self.classify_grasps_by_contact(
            grasps,
            part_mesh,
            face_ids,
        )

        affordances = affordance_payload.get("affordances", [])
        graspable_part_ids = {
            int(affordance["id"])
            for affordance in affordances
            if isinstance(affordance, dict)
            and affordance.get("id") is not None
            and affordance.get("graspable") is True
        }

        # Select top grasps per part
        grasps, confidences, selected_indices, part_grasp_ids = (
            self.pick_top_grasps(
                grasps,
                confidences,
                grasp_part_ids,
                graspable_part_ids,
                contact_scores,
            )
        )
        grasp_center_points = grasp_center_points[selected_indices]
        if len(grasps) == 0:
            logger.info(
                "Selected per-part grasps: {} -> 0 grasps across 0 parts".format(
                    original_count,
                )
            )
            return grasps, confidences, grasp_center_points, part_grasp_ids

        logger.info(
            "Selected per-part grasps: {} -> {} grasps across {} parts, confidences range from {:.4f} to {:.4f}".format(
                original_count,
                len(grasps),
                len(part_grasp_ids),
                np.min(confidences),
                np.max(confidences),
            )
        )
        return grasps, confidences, grasp_center_points, part_grasp_ids

    def process(
        self,
        urdf_path: str,
    ) -> bool:
        try:
            return self._process_impl(urdf_path)
        except Exception as exc:
            logger.error(
                "Grasp generation failed for URDF {}: {}".format(
                    urdf_path,
                    exc,
                )
            )
            return False

    def _process_impl(
        self,
        urdf_path: str,
    ) -> bool:
        urdf = URDFFile(urdf_path)
        mesh_path = urdf.get_mesh_path(self.cfg.mesh_type)
        mesh_transform = urdf.get_mesh_transform(self.cfg.mesh_type)
        seg_mesh_path = urdf.get_mesh_part_seg_path()
        seg_mesh_transform = self.get_part_seg_mesh_transform(urdf)
        affordance_path = urdf.get_affordance_annot_path()
        affordance_payload = load_json(affordance_path)

        if not self.cfg.overwrite:
            has_existing_grasp_group = any(
                isinstance(affordance, dict) and "grasp_group" in affordance
                for affordance in affordance_payload.get("affordances", [])
            )
            if has_existing_grasp_group:
                logger.info(
                    f"Skip existing grasp poses in affordance JSON: {affordance_path}"
                )
                return True

        logger.info(f"Generating grasps for mesh: {mesh_path}")
        mesh = load_mesh(
            mesh_path,
            apply_origin=True,
            apply_scale=True,
            **mesh_transform,
        )

        seg_mesh, face_ids = load_mesh(
            seg_mesh_path,
            apply_origin=True,
            apply_scale=True,
            **seg_mesh_transform,
        )

        grasp_result = self.infer_grasps(mesh, mesh_path)
        if grasp_result is None:
            return False
        grasps, confidences = grasp_result
        grasps, confidences, grasp_center_points, part_grasp_ids = (
            self.select_grasps(
                grasps,
                confidences,
                seg_mesh,
                face_ids,
                affordance_payload,
            )
        )

        self.add_grasps_to_affordance(
            affordance_payload,
            grasps,
            confidences,
            part_grasp_ids,
        )
        write_json(affordance_payload, affordance_path)
        logger.info(f"Saved predicted grasp poses to {affordance_path}")

        self.update_urdf(urdf_path, affordance_path)
        logger.info(
            "Grasp generation succeeded: {} -> {} ({} grasps)".format(
                mesh_path,
                affordance_path,
                len(grasps),
            )
        )
        if self.cfg.visualize_grasps:
            gripper_mesh = load_mesh(self.cfg.gripper_mesh_path)
            self.grasp_visualizer = show_grasps(
                mesh,
                grasps,
                confidences,
                self.cfg.gripper_name,
                self.cfg.model_name,
                grasp_center_points=grasp_center_points,
                visualizer=self.grasp_visualizer,
                visualizer_port=self.cfg.visualizer_port,
                gripper_info=self.pipeline.get_gripper_info()
                if self.cfg.model_name == "graspgenx"
                else None,
                gripper_mesh=gripper_mesh,
            )
        return True


def run_grasp(cfg: GraspPoseConfig) -> None:
    generator = GraspGenerator(cfg)
    for urdf_path in cfg.urdf_paths:
        generator.process(urdf_path)

    if cfg.visualize_grasps and generator.grasp_visualizer is not None:
        logger.info(
            "Keeping grasp visualization alive at http://localhost:{} "
            "(press Ctrl+C to exit).".format(cfg.visualizer_port)
        )
        try:
            while True:
                time.sleep(1)
        except KeyboardInterrupt:
            logger.info("Stopping grasp visualization.")


def entrypoint(*args, **kwargs) -> None:
    cfg = tyro.cli(GraspPoseConfig)
    run_grasp(cfg)


if __name__ == "__main__":
    entrypoint()

	# Project EmbodiedGen
	#
	# Copyright (c) 2025 Horizon Robotics. All Rights Reserved.
	#
	# Licensed under the Apache License, Version 2.0 (the "License");
	# you may not use this file except in compliance with the License.
	# You may obtain a copy of the License at
	#
	# http://www.apache.org/licenses/LICENSE-2.0
	#
	# Unless required by applicable law or agreed to in writing, software
	# distributed under the License is distributed on an "AS IS" BASIS,
	# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
	# implied. See the License for the specific language governing
	# permissions and limitations under the License.

	import os
	import sys
	import time
	from dataclasses import dataclass, field
	from typing import Literal

	from embodied_gen.utils.general import filter_warnings

	filter_warnings()

	import numpy as np
	import trimesh
	import trimesh.transformations as tra
	import tyro
	import yaml

	sys.path.append(".")
	from embodied_gen.utils.geometry import (
	sample_surface_points,
	)
	from embodied_gen.utils.io_utils import (
	URDFFile,
	load_json,
	load_mesh,
	write_json,
	)
	from embodied_gen.utils.log import logger
	from embodied_gen.utils.vis_utils import show_grasps

	__all__ = [
	"GraspGenHyperConfig",
	"GraspGenXHyperConfig",
	"GraspPoseConfig",
	"GraspGenerator",
	"run_grasp",
	"entrypoint",
	]

	FINGER_CENTER_CACHE: dict[str, np.ndarray] = {
	"franka_panda": np.array([0.0, 0.0, 0.1034], dtype=np.float64),
	}

	GRIPPER_GEOMETRY_CACHE: dict[str, dict[str, float]] = {
	"franka_panda": {
	"width": 0.10537486,
	"depth": 0.10527314,
	},
	}


	@dataclass
	class GraspGenHyperConfig:
	gripper_config: str = (
	"~/.cache/huggingface/hub/GraspGenModels/checkpoints/graspgen_{}.yml"
	)
	grasp_threshold: float = -1.0
	num_grasps: int = 4000
	return_topk: bool = True
	topk_num_grasps: int = 4000
	mesh_scale: tuple[float, float, float] = (1.0, 1.0, 1.0)
	num_sample_points: int = 2000
	min_grasps: int = 40
	max_tries: int = 6
	remove_outliers: bool = False


	@dataclass
	class GraspGenXHyperConfig:
	checkpoint_root: str = "~/.cache/huggingface/hub/GraspGenXModels/release"
	gen_pth: str \| None = None
	dis_pth: str \| None = None
	assets_dir: str = (
	"thirdparty/gripper_descriptions/gripper_descriptions/assets"
	)
	grasp_threshold: float = -1.0
	num_grasps: int = 4000
	return_topk: bool = True
	topk_num_grasps: int = 4000
	mesh_scale: tuple[float, float, float] = (1.0, 1.0, 1.0)
	num_sample_points: int = 2000
	min_grasps: int = 40
	max_tries: int = 6
	remove_outliers: bool = False


	@dataclass
	class GraspPoseConfig:
	urdf_paths: list[str] = field(default_factory=list)
	model_name: str = "GraspGen"
	gripper_name: str = "franka_panda"
	mesh_type: Literal["collision", "visual"] = "collision"
	update_urdf: bool = True
	overwrite: bool = True
	gripper_mesh_path: str = (
	f"embodied_gen/scripts/affordance_annot/{gripper_name}/coll_mesh.obj"
	)
	grasps_per_part: int = 20
	visualize_grasps: bool = False
	visualizer_port: int = 8080
	graspgen: GraspGenHyperConfig \| None = None


	class GraspGenerator:
	def __init__(self, cfg: GraspPoseConfig):
	self.cfg = cfg
	self.grasp_visualizer = None
	self.validate_config(cfg)
	self.modules = self.load_model()
	self.pipeline = self.build_sampler()
	self.finger_center = FINGER_CENTER_CACHE.get(
	self.cfg.gripper_name,
	np.array([0.0, 0.0, 0.0], dtype=np.float64),
	)
	(
	self.local_contact_points,
	self.local_closing_directions,
	) = self.build_local_contact_samples()

	def validate_config(self, cfg: GraspPoseConfig) -> None:
	cfg.model_name = cfg.model_name.lower()
	if cfg.model_name not in {"graspgen", "graspgenx"}:
	raise ValueError(
	"unsupported model '{}', available models: GraspGen, GraspGenX".format(
	cfg.model_name
	)
	)

	cfg.graspgen = (
	GraspGenHyperConfig()
	if cfg.model_name == "graspgen"
	else GraspGenXHyperConfig()
	)
	if cfg.model_name == "graspgen":
	cfg.graspgen.gripper_config = cfg.graspgen.gripper_config.format(
	cfg.gripper_name
	)

	def load_model(self):
	if self.cfg.model_name == "graspgen":
	sys.path.append("thirdparty/GraspGen")
	from grasp_gen.grasp_server import GraspGenSampler, load_grasp_cfg

	return {
	"GraspGenSampler": GraspGenSampler,
	"model_cfg": load_grasp_cfg,
	}
	if self.cfg.model_name == "graspgenx":
	sys.path.append("thirdparty/GraspGenX")
	from graspgenx.grasp_server import GraspGenXSampler
	from graspgenx.utils.checkpoint_io import load_model_cfg

	return {
	"GraspGenSampler": GraspGenXSampler,
	"model_cfg": load_model_cfg,
	}

	@staticmethod
	def _extract_checkpoint_filenames(yml_path: str) -> list[str]:
	"""Recursively collect all `checkpoint: <name>.pth` entries.

	Walks a GraspGen config yaml so we know every weight file it
	requires.
	"""
	with open(yml_path) as f:
	cfg = yaml.safe_load(f)

	names: set[str] = set()

	def _walk(node) -> None:
	if isinstance(node, dict):
	for key, value in node.items():
	if (
	key == "checkpoint"
	and isinstance(value, str)
	and value.strip()
	):
	names.add(value.strip())
	else:
	_walk(value)
	elif isinstance(node, list):
	for item in node:
	_walk(item)

	_walk(cfg)
	return sorted(names)

	@staticmethod
	def _dir_has_pth(directory: str, filename: str \| None) -> bool:
	if filename:
	return os.path.exists(os.path.join(directory, filename))
	return os.path.isdir(directory) and any(
	f.endswith(".pth") for f in os.listdir(directory)
	)

	def _ensure_checkpoints(
	self,
	is_ready: callable,
	repo_id: str,
	allow_patterns: str,
	download_dir: str,
	desc: str,
	) -> None:
	"""Download `desc` weights from `repo_id` unless already ready.

	Re-checks afterwards so a stale/partial local cache (e.g. config
	present but a referenced .pth missing) is retried instead of
	silently failing deep inside third-party code.
	"""
	if is_ready():
	return

	logger.info(
	f"Downloading {desc} weights from Hugging Face. "
	"First-time use may take several minutes; please wait patiently."
	)
	from huggingface_hub import snapshot_download

	snapshot_download(
	repo_id=repo_id,
	allow_patterns=allow_patterns,
	local_dir=os.path.expanduser(download_dir),
	)

	if not is_ready():
	raise FileNotFoundError(
	f"{desc} checkpoint(s) still missing after downloading from "
	f"Hugging Face repo '{repo_id}'. The local cache may be from "
	"a previous partial download; please remove it and retry, "
	"or download manually."
	)

	def build_sampler(self):
	if self.cfg.model_name == "graspgen":
	gripper_config = os.path.expanduser(
	self.cfg.graspgen.gripper_config
	)
	checkpoints_dir = os.path.dirname(gripper_config)

	def _graspgen_ready() -> bool:
	if not os.path.exists(gripper_config):
	return False
	required = self._extract_checkpoint_filenames(gripper_config)
	return all(
	os.path.exists(os.path.join(checkpoints_dir, name))
	for name in required
	)

	self._ensure_checkpoints(
	_graspgen_ready,
	"adithyamurali/GraspGenModels",
	"checkpoints/*",
	"~/.cache/huggingface/hub/GraspGenModels",
	"GraspGen",
	)
	grasp_cfg = self.modules["model_cfg"](gripper_config)
	return self.modules["GraspGenSampler"](grasp_cfg)
	if self.cfg.model_name == "graspgenx":
	checkpoint_root = os.path.expanduser(
	self.cfg.graspgen.checkpoint_root
	)
	gen_dir = os.path.join(checkpoint_root, "gen")
	dis_dir = os.path.join(checkpoint_root, "dis")

	def _graspgenx_ready() -> bool:
	return (
	os.path.exists(os.path.join(gen_dir, "config.yaml"))
	and self._dir_has_pth(gen_dir, self.cfg.graspgen.gen_pth)
	and self._dir_has_pth(dis_dir, self.cfg.graspgen.dis_pth)
	)

	self._ensure_checkpoints(
	_graspgenx_ready,
	"adithyamurali/GraspGenXModel",
	"release/*",
	"~/.cache/huggingface/hub/GraspGenXModels",
	"GraspGenX",
	)
	model_cfg = self.modules["model_cfg"](
	os.path.join(checkpoint_root, "gen"),
	os.path.join(checkpoint_root, "dis"),
	self.cfg.graspgen.gen_pth,
	self.cfg.graspgen.dis_pth,
	)
	return self.modules["GraspGenSampler"](
	model_cfg,
	self.cfg.gripper_name,
	assets_dir=self.cfg.graspgen.assets_dir,
	)

	def update_urdf(self, urdf_path: str \| None, grasp_path: str) -> None:
	if urdf_path is None or not self.cfg.update_urdf:
	return

	URDFFile(urdf_path).write(
	{
	"custom_data/affordance/affordance_annot": os.path.relpath(
	grasp_path,
	os.path.dirname(urdf_path),
	),
	}
	)

	def get_part_seg_mesh_transform(self, urdf: URDFFile) -> dict:
	return {
	"origin_xyz": urdf.read(
	".//custom_data/affordance/visual_seg/origin",
	attr="xyz",
	default="0 0 0",
	),
	"origin_rpy": urdf.read(
	".//custom_data/affordance/visual_seg/origin",
	attr="rpy",
	default="0 0 0",
	),
	"scale": urdf.read(
	".//custom_data/affordance/visual_seg/geometry/mesh",
	attr="scale",
	default="1.0 1.0 1.0",
	),
	}

	def add_grasps_to_affordance(
	self,
	affordance_payload: dict,
	grasps: np.ndarray,
	confidences: np.ndarray,
	part_grasp_ids: dict[int, list[int]] \| None,
	) -> None:
	for affordance in affordance_payload.get("affordances", []):
	if (
	not isinstance(affordance, dict)
	or affordance.get("id") is None
	):
	continue
	affordance.pop("grasps", None)
	affordance.pop("confidences", None)
	grasp_ids = part_grasp_ids.get(int(affordance["id"]), [])
	grasp_group = {}
	for grasp_idx in grasp_ids:
	grasp = grasps[grasp_idx]
	quaternion = tra.quaternion_from_matrix(grasp[:3, :3])
	grasp_group[f"grasp_{grasp_idx}"] = {
	"confidence": float(confidences[grasp_idx]),
	"position": grasp[:3, 3].tolist(),
	"orientation": {
	"w": float(quaternion[0]),
	"xyz": quaternion[1:].tolist(),
	},
	}
	affordance["grasp_group"] = grasp_group

	def infer_grasps(
	self,
	mesh: trimesh.Trimesh,
	mesh_path: str,
	) -> tuple[np.ndarray, np.ndarray, np.ndarray] \| None:
	points = sample_surface_points(
	mesh, self.cfg.graspgen.num_sample_points
	)
	center_transform = tra.translation_matrix(-points.mean(axis=0))
	centered_points = tra.transform_points(
	points, center_transform
	).astype(np.float32)
	topk_limit = (
	self.cfg.graspgen.topk_num_grasps
	if self.cfg.graspgen.return_topk
	else -1
	)
	grasps, confidences = self.modules["GraspGenSampler"].run_inference(
	centered_points,
	self.pipeline,
	grasp_threshold=self.cfg.graspgen.grasp_threshold,
	num_grasps=self.cfg.graspgen.num_grasps,
	topk_num_grasps=topk_limit,
	min_grasps=self.cfg.graspgen.min_grasps,
	max_tries=self.cfg.graspgen.max_tries,
	remove_outliers=self.cfg.graspgen.remove_outliers,
	)

	confidences = confidences.detach().cpu().numpy()
	grasps = grasps.detach().cpu().numpy()
	if len(grasps) == 0:
	logger.error(f"No grasp poses generated for mesh: {mesh_path}")
	return None

	grasps = np.asarray(
	[tra.inverse_matrix(center_transform) @ grasp for grasp in grasps]
	)
	return grasps, confidences

	def get_grasp_centers(
	self,
	grasps: np.ndarray,
	) -> np.ndarray:
	local_centers = np.repeat(
	np.asarray(self.finger_center, dtype=np.float64)[None, :],
	len(grasps),
	axis=0,
	)
	grasp_center_points = (
	np.einsum(
	"nij,nj->ni",
	grasps[:, :3, :3],
	local_centers,
	)
	+ grasps[:, :3, 3]
	)
	return grasp_center_points

	def build_local_contact_samples(self) -> tuple[np.ndarray, np.ndarray]:
	geometry = GRIPPER_GEOMETRY_CACHE.get(self.cfg.gripper_name)
	if geometry is None:
	return (
	np.asarray(self.finger_center, dtype=np.float64)[None, :],
	np.array([[1.0, 0.0, 0.0]], dtype=np.float64),
	)

	half_width = float(geometry["width"]) / 2.0
	depth = float(geometry["depth"])
	z_values = np.linspace(depth - 0.015, depth, 6, dtype=np.float64)
	points = []
	directions = []
	for z_value in z_values:
	points.extend(
	[
	[-half_width, 0.0, z_value],
	[half_width, 0.0, z_value],
	]
	)
	directions.extend(
	[
	[1.0, 0.0, 0.0],
	[-1.0, 0.0, 0.0],
	]
	)
	return (
	np.asarray(points, dtype=np.float64),
	np.asarray(directions, dtype=np.float64),
	)

	def get_grasp_contact_rays(
	self,
	grasps: np.ndarray,
	) -> tuple[np.ndarray, np.ndarray]:
	local_contact_points = np.repeat(
	self.local_contact_points[None, :, :],
	len(grasps),
	axis=0,
	)
	local_closing_directions = np.repeat(
	self.local_closing_directions[None, :, :],
	len(grasps),
	axis=0,
	)
	ray_origins = (
	np.einsum(
	"nij,nkj->nki",
	grasps[:, :3, :3],
	local_contact_points,
	)
	+ grasps[:, None, :3, 3]
	)
	ray_directions = np.einsum(
	"nij,nkj->nki",
	grasps[:, :3, :3],
	local_closing_directions,
	)
	ray_norms = np.linalg.norm(ray_directions, axis=-1, keepdims=True)
	ray_directions = np.divide(
	ray_directions,
	ray_norms,
	out=np.zeros_like(ray_directions),
	where=ray_norms > 0.0,
	)
	return ray_origins, ray_directions

	def bind_contact_rays_to_faces(
	self,
	part_mesh: trimesh.Trimesh,
	ray_origins: np.ndarray,
	ray_directions: np.ndarray,
	) -> np.ndarray:
	ray_origins = np.asarray(ray_origins, dtype=np.float64)
	ray_directions = np.asarray(ray_directions, dtype=np.float64)

	if (
	not isinstance(part_mesh, trimesh.Trimesh)
	or len(part_mesh.faces) == 0
	):
	return np.full(len(ray_origins), -1, dtype=np.int64)

	try:
	face_ids = part_mesh.ray.intersects_first(
	ray_origins,
	ray_directions,
	)
	return np.asarray(face_ids, dtype=np.int64)
	except Exception as exc:
	logger.warning(
	"Failed to raycast grasp closing directions to part mesh: "
	f"{exc}"
	)
	return np.full(len(ray_origins), -1, dtype=np.int64)

	def classify_grasps_by_contact(
	self,
	grasps: np.ndarray,
	part_mesh: trimesh.Trimesh,
	face_ids: np.ndarray \| None,
	) -> tuple[np.ndarray \| None, np.ndarray \| None]:
	if face_ids is None:
	return None, None

	face_ids = np.asarray(face_ids, dtype=np.int64)
	if len(face_ids) != len(part_mesh.faces):
	raise ValueError(
	"face_ids length must match segmentation mesh faces, got "
	f"{len(face_ids)} and {len(part_mesh.faces)}"
	)

	ray_origins, ray_directions = self.get_grasp_contact_rays(grasps)
	num_grasps, num_contact_points, _ = ray_origins.shape
	logger.info(
	"Binding {} grasp closing rays to {} part faces.".format(
	num_grasps * num_contact_points,
	len(part_mesh.faces),
	)
	)
	nearest_face_ids = self.bind_contact_rays_to_faces(
	part_mesh,
	ray_origins.reshape(-1, 3),
	ray_directions.reshape(-1, 3),
	).reshape(num_grasps, num_contact_points)
	contact_part_ids = self.parts_for_faces(
	nearest_face_ids.reshape(-1),
	face_ids,
	).reshape(num_grasps, num_contact_points)

	grasp_part_ids = np.full(num_grasps, -1, dtype=np.int64)
	contact_scores = np.zeros(num_grasps, dtype=np.float64)
	for grasp_idx, part_ids in enumerate(contact_part_ids):
	valid_part_ids = part_ids[part_ids >= 0]
	if len(valid_part_ids) == 0:
	continue
	unique_ids, counts = np.unique(
	valid_part_ids,
	return_counts=True,
	)
	best_idx = int(np.argmax(counts))
	grasp_part_ids[grasp_idx] = int(unique_ids[best_idx])
	contact_scores[grasp_idx] = float(counts[best_idx]) / float(
	num_contact_points
	)

	return grasp_part_ids, contact_scores

	def parts_for_faces(
	self,
	nearest_face_ids: np.ndarray,
	face_ids: np.ndarray \| None,
	) -> np.ndarray \| None:
	if face_ids is None:
	return None

	nearest_face_ids = np.asarray(nearest_face_ids, dtype=np.int64)
	grasp_part_ids = np.full(len(nearest_face_ids), -1, dtype=np.int64)
	valid_mask = (nearest_face_ids >= 0) & (
	nearest_face_ids < len(face_ids)
	)
	grasp_part_ids[valid_mask] = face_ids[nearest_face_ids[valid_mask]]
	return grasp_part_ids

	def pick_top_grasps(
	self,
	grasps: np.ndarray,
	confidences: np.ndarray,
	part_ids: np.ndarray \| None,
	graspable_part_ids: set[int] \| None = None,
	part_scores: np.ndarray \| None = None,
	) -> tuple[
	np.ndarray, np.ndarray, np.ndarray, dict[int, list[int]] \| None
	]:
	if part_ids is None:
	return (
	grasps,
	confidences,
	np.arange(len(grasps), dtype=np.int64),
	None,
	)

	part_ids = np.asarray(part_ids, dtype=np.int64)
	selected_indices = []
	selected_part_ids = []
	for part_id in sorted(
	int(pid) for pid in np.unique(part_ids[part_ids >= 0])
	):
	if (
	graspable_part_ids is not None
	and part_id not in graspable_part_ids
	):
	continue
	part_indices = np.flatnonzero(part_ids == part_id)
	if part_scores is None:
	rank = np.argsort(
	-confidences[part_indices],
	kind="mergesort",
	)
	else:
	rank = np.lexsort(
	(
	-confidences[part_indices],
	-part_scores[part_indices],
	)
	)
	ranked_indices = part_indices[rank]
	if len(ranked_indices) > self.cfg.grasps_per_part:
	selected = ranked_indices[: self.cfg.grasps_per_part]
	else:
	selected = ranked_indices
	selected_indices.extend(int(idx) for idx in selected)
	selected_part_ids.extend([part_id] * len(selected))

	selected_indices = np.asarray(selected_indices, dtype=np.int64)
	if len(selected_indices) == 0:
	return (
	grasps[:0],
	confidences[:0],
	selected_indices,
	{},
	)

	part_grasp_ids: dict[int, list[int]] = {}
	for new_idx, part_id in enumerate(selected_part_ids):
	part_grasp_ids.setdefault(part_id, []).append(new_idx)

	return (
	grasps[selected_indices],
	confidences[selected_indices],
	selected_indices,
	part_grasp_ids,
	)

	def select_grasps(
	self,
	grasps: np.ndarray,
	confidences: np.ndarray,
	part_mesh: trimesh.Trimesh,
	face_ids: np.ndarray,
	affordance_payload: dict,
	) -> tuple[np.ndarray, np.ndarray, np.ndarray, dict[int, list[int]]]:
	original_count = len(grasps)

	# Raycast from fingertip samples along the gripper closing direction.
	grasp_center_points = self.get_grasp_centers(grasps)
	grasp_part_ids, contact_scores = self.classify_grasps_by_contact(
	grasps,
	part_mesh,
	face_ids,
	)

	affordances = affordance_payload.get("affordances", [])
	graspable_part_ids = {
	int(affordance["id"])
	for affordance in affordances
	if isinstance(affordance, dict)
	and affordance.get("id") is not None
	and affordance.get("graspable") is True
	}

	# Select top grasps per part
	grasps, confidences, selected_indices, part_grasp_ids = (
	self.pick_top_grasps(
	grasps,
	confidences,
	grasp_part_ids,
	graspable_part_ids,
	contact_scores,
	)
	)
	grasp_center_points = grasp_center_points[selected_indices]
	if len(grasps) == 0:
	logger.info(
	"Selected per-part grasps: {} -> 0 grasps across 0 parts".format(
	original_count,
	)
	)
	return grasps, confidences, grasp_center_points, part_grasp_ids

	logger.info(
	"Selected per-part grasps: {} -> {} grasps across {} parts, confidences range from {:.4f} to {:.4f}".format(
	original_count,
	len(grasps),
	len(part_grasp_ids),
	np.min(confidences),
	np.max(confidences),
	)
	)
	return grasps, confidences, grasp_center_points, part_grasp_ids

	def process(
	self,
	urdf_path: str,
	) -> bool:
	try:
	return self._process_impl(urdf_path)
	except Exception as exc:
	logger.error(
	"Grasp generation failed for URDF {}: {}".format(
	urdf_path,
	exc,
	)
	)
	return False

	def _process_impl(
	self,
	urdf_path: str,
	) -> bool:
	urdf = URDFFile(urdf_path)
	mesh_path = urdf.get_mesh_path(self.cfg.mesh_type)
	mesh_transform = urdf.get_mesh_transform(self.cfg.mesh_type)
	seg_mesh_path = urdf.get_mesh_part_seg_path()
	seg_mesh_transform = self.get_part_seg_mesh_transform(urdf)
	affordance_path = urdf.get_affordance_annot_path()
	affordance_payload = load_json(affordance_path)

	if not self.cfg.overwrite:
	has_existing_grasp_group = any(
	isinstance(affordance, dict) and "grasp_group" in affordance
	for affordance in affordance_payload.get("affordances", [])
	)
	if has_existing_grasp_group:
	logger.info(
	f"Skip existing grasp poses in affordance JSON: {affordance_path}"
	)
	return True

	logger.info(f"Generating grasps for mesh: {mesh_path}")
	mesh = load_mesh(
	mesh_path,
	apply_origin=True,
	apply_scale=True,
	**mesh_transform,
	)

	seg_mesh, face_ids = load_mesh(
	seg_mesh_path,
	apply_origin=True,
	apply_scale=True,
	**seg_mesh_transform,
	)

	grasp_result = self.infer_grasps(mesh, mesh_path)
	if grasp_result is None:
	return False
	grasps, confidences = grasp_result
	grasps, confidences, grasp_center_points, part_grasp_ids = (
	self.select_grasps(
	grasps,
	confidences,
	seg_mesh,
	face_ids,
	affordance_payload,
	)
	)

	self.add_grasps_to_affordance(
	affordance_payload,
	grasps,
	confidences,
	part_grasp_ids,
	)
	write_json(affordance_payload, affordance_path)
	logger.info(f"Saved predicted grasp poses to {affordance_path}")

	self.update_urdf(urdf_path, affordance_path)
	logger.info(
	"Grasp generation succeeded: {} -> {} ({} grasps)".format(
	mesh_path,
	affordance_path,
	len(grasps),
	)
	)
	if self.cfg.visualize_grasps:
	gripper_mesh = load_mesh(self.cfg.gripper_mesh_path)
	self.grasp_visualizer = show_grasps(
	mesh,
	grasps,
	confidences,
	self.cfg.gripper_name,
	self.cfg.model_name,
	grasp_center_points=grasp_center_points,
	visualizer=self.grasp_visualizer,
	visualizer_port=self.cfg.visualizer_port,
	gripper_info=self.pipeline.get_gripper_info()
	if self.cfg.model_name == "graspgenx"
	else None,
	gripper_mesh=gripper_mesh,
	)
	return True


	def run_grasp(cfg: GraspPoseConfig) -> None:
	generator = GraspGenerator(cfg)
	for urdf_path in cfg.urdf_paths:
	generator.process(urdf_path)

	if cfg.visualize_grasps and generator.grasp_visualizer is not None:
	logger.info(
	"Keeping grasp visualization alive at http://localhost:{} "
	"(press Ctrl+C to exit).".format(cfg.visualizer_port)
	)
	try:
	while True:
	time.sleep(1)
	except KeyboardInterrupt:
	logger.info("Stopping grasp visualization.")


	def entrypoint(args, *kwargs) -> None:
	cfg = tyro.cli(GraspPoseConfig)
	run_grasp(cfg)


	if __name__ == "__main__":
	entrypoint()