graspfactory/grasp_visualizer.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
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
import trimesh
import pyglet.app

def create_gripper_mesh_panda(color=[0, 0, 255], tube_radius=2, sections=6, width=80, finger_length=66):
    ''' Create a simple gripper mesh for Franka Panda robot'''

    half_width = width / 2
    cb2_1 = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[[0, 0, finger_length], [half_width, 0, finger_length]],
    )
    cb2_2 = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[[-half_width, 0, finger_length], [0, 0, finger_length]],
    )
    cb_1 = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[[0, 0, finger_length*2], [0, 0, finger_length]],
    )
    cf_l = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[[half_width, 0, finger_length], [half_width, 0, 0]],
    )
    cf_r = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[[-half_width, 0, finger_length], [-half_width, 0, 0]],
    )
    tmp = trimesh.util.concatenate([cb_1, cb2_1, cb2_2, cf_r, cf_l])
    tmp.visual.face_colors = color
    return tmp

def create_gripper_mesh_robotiq(color=[0, 0, 255], tube_radius=2, sections=6, width=85, base_height=100, finger_length=75):
    ''' Create a simple gripper mesh for Robotiq 2F-85 gripper'''

    half_width = width / 2
    cb2_1 = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[[0, 0, base_height], [half_width, 0, base_height]],
    )
    cb2_2 = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[[-half_width, 0, base_height], [0, 0, base_height]],
    )
    cb_1 = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[[0, 0, base_height + finger_length], [0, 0, base_height]],
    )
    cf_l = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[[half_width, 0, base_height], [half_width, 0, 0]],
    )
    cf_r = trimesh.creation.cylinder(
        radius=tube_radius,
        sections=sections,
        segment=[[-half_width, 0, base_height], [-half_width, 0, 0]],
    )
    tmp = trimesh.util.concatenate([cb_1, cb2_1, cb2_2, cf_r, cf_l])
    tmp.visual.face_colors = color
    return tmp


def get_gripper_mesh(grasps, grasp_widths=None, gripper_type="robotiq"):
    ''' Generate gripper meshes for a set of grasps'''
    grips = []
    for k in range(grasps.shape[0]):
        H = grasps[k,...]
        c_vis = [0, 255, 0]
        grasp_width = grasp_widths[k] if grasp_widths is not None else 0.08
        if gripper_type == "robotiq":
            grips.append(create_gripper_mesh_robotiq(color=c_vis, width=grasp_width).apply_transform(H))
        else:
            grips.append(create_gripper_mesh_panda(color=c_vis, width=grasp_width).apply_transform(H))
    return grips


def visualize_grasps(grasps, mesh=None, grasp_widths=None, gripper_type="robotiq"):
    ''' Visualize a set of grasps on a mesh'''
    assert gripper_type in ["robotiq", "panda"]
    event_loop_constructor = pyglet.app.EventLoop
    event_loop_instance = pyglet.app.event_loop
    pyglet.app.EventLoop = pyglet.app.base.EventLoop
    pyglet.app.event_loop = pyglet.app.EventLoop()

    grasp_widths = grasp_widths.squeeze()
    grips = get_gripper_mesh(grasps, grasp_widths=grasp_widths, gripper_type=gripper_type)
    scene = trimesh.Scene([mesh]+ grips)
    scene.show(smooth=False)

    pyglet.app.EventLoop = event_loop_constructor
    pyglet.app.event_loop = event_loop_instance