convert_to_graspqp_format.py

# Copyright (c) 2026 ETH Zurich, René Zurbrügg
# SPDX-License-Identifier: MIT

"""Convert an arbitrary grasp data source into the GraspQP dataset layout.

This is a reference implementation of the format documented in
``docs/DATASET.md``. Import :func:`write_graspqp_object` in your own conversion
script, or run this file directly with ``--self-test`` to generate a tiny
synthetic dataset and verify it round-trips through ``GraspDataset``.

Usage
-----
    # Verify the format end-to-end on a synthetic object (needs graspqp + torch):
    python scripts/convert_to_graspqp_format.py --self-test --out /tmp/gqp_demo

    # In your own code:
    from scripts.convert_to_graspqp_format import write_graspqp_object
    write_graspqp_object(asset_dir, object_name, mesh, trans, quats_wxyz,
                         joint_angles, joint_names)
"""

from __future__ import annotations

import argparse
import os
from typing import List

import torch
import trimesh


def write_graspqp_object(
    asset_dir: str,
    object_name: str,
    mesh: trimesh.Trimesh,
    translations: torch.Tensor,   # (G, 3), meters
    quats_wxyz: torch.Tensor,     # (G, 4), unit quaternions in wxyz order
    joint_angles: torch.Tensor,   # (G, N), radians; columns ordered as joint_names
    joint_names: List[str],       # length N, must match the hand model
    hand_type: str = "allegro",
    energy_type: str = "graspqp",
    grasp_type: str = "default",
    n_contacts: int = 12,
    param_type: str = "parameters",
    checkpoint_name: str = "succ_grasps__mined.pt",
) -> str:
    """Write one object into the GraspQP dataset layout.

    Returns the path to the grasp checkpoint that was written.

    See ``docs/DATASET.md`` for the full schema. In short, this writes::

        <asset_dir>/<object_name>/remeshed.obj
        <asset_dir>/<object_name>/grasp_predictions/<hand_type>/
            <n_contacts>_contacts/<energy_type>/<grasp_type>/<checkpoint_name>
    """
    G, N = joint_angles.shape
    if translations.shape != (G, 3):
        raise ValueError(f"translations must be (G, 3), got {tuple(translations.shape)}")
    if quats_wxyz.shape != (G, 4):
        raise ValueError(f"quats_wxyz must be (G, 4), got {tuple(quats_wxyz.shape)}")
    if len(joint_names) != N:
        raise ValueError(
            f"joint_names has {len(joint_names)} entries but joint_angles has {N} columns"
        )

    obj_dir = os.path.join(asset_dir, object_name)
    os.makedirs(obj_dir, exist_ok=True)

    # 1) Mesh — point cloud + normals are sampled from this at train time.
    mesh.export(os.path.join(obj_dir, "remeshed.obj"))

    # 2) Grasp checkpoint: {param_type: {root_pose, <joint>: ...}}.
    root_pose = torch.cat([translations.float(), quats_wxyz.float()], dim=-1)  # (G, 7)
    params = {"root_pose": root_pose}
    for j, name in enumerate(joint_names):
        params[name] = joint_angles[:, j].float().contiguous()

    ckpt_dir = os.path.join(
        obj_dir, "grasp_predictions", hand_type,
        f"{n_contacts}_contacts", energy_type, grasp_type,
    )
    os.makedirs(ckpt_dir, exist_ok=True)
    ckpt_path = os.path.join(ckpt_dir, checkpoint_name)
    torch.save({param_type: params}, ckpt_path)
    return ckpt_path


def _self_test(out_dir: str, hand_type: str = "allegro", num_objects: int = 2,
               grasps_per_object: int = 8) -> None:
    """Generate a synthetic dataset and reload it through GraspDataset."""
    import roma
    from omegaconf import OmegaConf
    from graspqp.hands import get_hand_model
    from grasp_diffuser.data import GraspDataset

    hand_model = get_hand_model(hand_type, "cpu")
    joint_names = list(hand_model.joints_names)
    lower = hand_model.joints_lower
    upper = hand_model.joints_upper
    n_joints = len(joint_names)
    print(f"[self-test] hand={hand_type} n_joints={n_joints}")

    for i in range(num_objects):
        mesh = trimesh.creation.icosphere(subdivisions=2, radius=0.05)
        G = grasps_per_object
        translations = torch.randn(G, 3) * 0.05
        quats_wxyz = roma.quat_xyzw_to_wxyz(roma.random_unitquat(G))
        # Sample joint angles inside the hand's limits.
        joint_angles = torch.rand(G, n_joints) * (upper - lower) + lower

        write_graspqp_object(
            out_dir, f"object_{i:03d}", mesh,
            translations, quats_wxyz, joint_angles, joint_names,
            hand_type=hand_type,
        )
    print(f"[self-test] wrote {num_objects} objects to {out_dir}")

    cfg = OmegaConf.create({
        "name": "Ours", "num_points": 2048, "use_color": False, "use_normal": True,
        "random_rotate": False, "random_translate": False, "with_box": False,
        "energy_type": "graspqp", "hand_type": hand_type, "grasp_type": "default",
        "param_type": "parameters", "iteration_steps": "0", "noise": [],
        "asset_dir": out_dir,
    })
    ds = GraspDataset(cfg, phase="test", hand_model=hand_model)
    sample = ds[0]
    expected_state = 9 + n_joints
    assert sample["x"].shape[-1] == expected_state, (
        f"state length {sample['x'].shape[-1]} != expected {expected_state}"
    )
    assert sample["pos"].shape == (2048, 3), sample["pos"].shape
    print(
        f"[self-test] OK — objects={len(ds.object_names)} samples={len(ds)} "
        f"x={tuple(sample['x'].shape)} pos={tuple(sample['pos'].shape)}"
    )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--self-test", action="store_true",
                        help="Generate a synthetic dataset and reload it.")
    parser.add_argument("--out", default="/tmp/graspqp_demo",
                        help="Output asset directory for --self-test.")
    parser.add_argument("--hand-type", default="allegro")
    args = parser.parse_args()

    if args.self_test:
        _self_test(args.out, hand_type=args.hand_type)
    else:
        parser.print_help()


if __name__ == "__main__":
    main()