graspfactory/visualize.py at main · AutodeskRoboticsLab/graspfactory · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73


import numpy as np
import glob
import argparse
from grasp_visualizer import visualize_grasps
import trimesh

def main():
    parser = argparse.ArgumentParser(description='Visualize GraspFactory grasps')
    parser.add_argument('--gripper', '-g',
                        choices=['panda', 'robotiq'],
                        default='robotiq',
                        help='Gripper type to visualize (default: robotiq)')
    parser.add_argument('--num_grasps', '-n',
                        type=int,
                        default=25,
                        help='Number of grasps to visualize per object (default: 25)')
    parser.add_argument('--data_path', '-p',
                        type=str,
                        default="./sample_data/robotiq_2f85/grasps/",
                        help='Path to dataset directory')

    args = parser.parse_args()

    grasps_path = args.data_path

    print(f"Visualizing {args.gripper} grasps from: {grasps_path}")
    print(f"Number of grasps per object: {args.num_grasps}")

    grasp_files = glob.glob(grasps_path + "/*.npz")

    if not grasp_files:
        print(f"No grasp files found in {grasps_path}")
        return

    print(f"Found {len(grasp_files)} grasp files")

    for i, grasp_file in enumerate(grasp_files):
        print(f"Processing file {i+1}/{len(grasp_files)}: {grasp_file}")

        try:
            data = np.load(grasp_file, allow_pickle=True)
            grasps = data["grasps"]
            grasps_widths = data["grasp_widths"]
            success_indices = data["success_indices"]

            if len(success_indices) == 0:
                print(f"No successful grasps found in {grasp_file}, skipping...")
                continue

            successful_grasps = grasps[success_indices]
            successful_grasps_widths = grasps_widths[success_indices]

            # Load mesh
            mesh_path = f"{grasps_path[:-len('grasps')]}/meshes/{data['mesh_uid']}.obj"
            mesh = trimesh.load(mesh_path)

            # Sample random grasps
            num_to_show = min(args.num_grasps, len(successful_grasps))
            random_indices = np.random.choice(len(successful_grasps), size=num_to_show, replace=False)
            selected_grasps = successful_grasps[random_indices]
            selected_widths = successful_grasps_widths[random_indices]

            print(f"Showing {num_to_show} grasps for object {data['mesh_uid']}")
            visualize_grasps(selected_grasps, grasp_widths=selected_widths, mesh=mesh, gripper_type=args.gripper)

        except Exception as e:
            print(f"Error processing {grasp_file}: {e}")
            continue

if __name__ == "__main__":
    main()