CL_Assets/view_urdf.py at main · correlllab/CL_Assets · 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
import os
import argparse
import pinocchio as pin
from pinocchio.visualize import MeshcatVisualizer


def _attach_frame_shape(handle):
    try:
        import meshcat_shapes
    except ImportError as exc:
        raise RuntimeError("--frames requires the meshcat-shapes package") from exc

    meshcat_shapes.frame(
        handle,
        axis_length=0.08,
        axis_thickness=0.003,
        opacity=0.9,
        origin_radius=0.008,
    )


def _resolve_frames(model, frame_names):
    resolved = []
    for name in frame_names:
        try:
            frame_id = model.getFrameId(name)
        except Exception:
            print(f"Skipping unknown frame: {name}")
            continue

        if frame_id < 0 or frame_id >= len(model.frames):
            print(f"Skipping unknown frame: {name}")
            continue

        resolved.append((name, frame_id))
    return resolved


def _viewer_root_name(filename, fallback):
    stem = os.path.splitext(os.path.basename(filename))[0]
    return stem or fallback


def _load_urdf(filename):
    # package dir as parent folder
    model_path = os.path.dirname(os.path.abspath(filename))

    model, visual, collision = pin.buildModelsFromUrdf(
        filename,
        package_dirs=model_path,
    )
    return model, visual, collision


def view_urdf(filename, frame_names=None, overlay_filename=None):
    model, visual, collision = _load_urdf(filename)
    viz = MeshcatVisualizer(model, visual, collision)
    viz.initViewer()
    viz.loadViewerModel(_viewer_root_name(filename, "robot"))

    overlay_viz = None
    overlay_model = None
    if overlay_filename:
        overlay_model, overlay_visual, overlay_collision = _load_urdf(overlay_filename)
        overlay_viz = MeshcatVisualizer(overlay_model, overlay_visual, overlay_collision)
        overlay_viz.initViewer(viewer=viz.viewer)
        overlay_viz.loadViewerModel(_viewer_root_name(overlay_filename, "overlay"))

    frame_nodes = []
    if frame_names:
        data = model.createData()
        for name, frame_id in _resolve_frames(model, frame_names):
            node = viz.viewer[f"frames/{name}"]
            _attach_frame_shape(node)
            frame_nodes.append((node, frame_id))

    while True:
        q = pin.neutral(model)
        # display robot at neutral configuration
        viz.display(q)

        if overlay_viz is not None and overlay_model is not None:
            overlay_q = pin.neutral(overlay_model)
            overlay_viz.display(overlay_q)

        if frame_nodes:
            pin.forwardKinematics(model, data, q)
            pin.updateFramePlacements(model, data)
            for node, frame_id in frame_nodes:
                node.set_transform(data.oMf[frame_id].homogeneous)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Visualize URDF model using Meshcat')
    parser.add_argument('filename', type=str, help='URDF filename')
    parser.add_argument('overlay_filename', nargs='?', default=None, help='Optional second URDF filename to overlay')
    parser.add_argument('--frames', nargs='+', default=[], help='Frame names to show in Meshcat')
    args = parser.parse_args()
    view_urdf(args.filename, args.frames, args.overlay_filename)