scikit-robot/examples/trimesh_scene_viewer.py at main · iory/scikit-robot · 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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
#!/usr/bin/env python

import argparse
import time

import numpy as np

import skrobot
from skrobot.utils.video import record_viewer
from skrobot.viewers import VIEWER_HELP
from skrobot.viewers import VIEWER_TYPES


def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        '--interactive',
        action='store_true',
        help='enter interactive shell'
    )
    parser.add_argument(
        '--no-interactive',
        action='store_true',
        help="Run in non-interactive mode (do not wait for user input)"
    )
    parser.add_argument(
        '--viewer', type=str,
        choices=VIEWER_TYPES, default='pyrender',
        help=VIEWER_HELP)
    parser.add_argument(
        '--save-video', type=str, default=None,
        help='Record the animation to this video file (e.g. out.mp4). Works '
             'with any --viewer; use --viewer mitsuba to record headlessly.')
    args = parser.parse_args()

    robot = skrobot.models.Kuka()

    viewer = skrobot.viewers.create_viewer(args.viewer, resolution=(640, 480))

    # base plane
    plane = skrobot.model.Box(
        extents=(2, 2, 0.01), face_colors=(0.75, 0.75, 0.75)
    )
    viewer.add(plane)

    viewer.add(robot)

    viewer.set_camera(angles=[np.deg2rad(45), 0, 0], distance=4)

    viewer.show()

    box = skrobot.model.Box(
        extents=(0.05, 0.05, 0.05), face_colors=(1., 0, 0)
    )
    box.translate((0.5, 0, 0.3))
    viewer.add(box)

    # Record every redraw of the pose sequence below when --save-video is given.
    recorder = record_viewer(viewer, args.save_video, fps=2)

    if args.interactive:
        print('''\
>>> # Usage

>>> robot.reset_manip_pose()
>>> viewer.redraw()
>>> robot.init_pose()
>>> robot.inverse_kinematics(box, rotation_mask='xz')
''')

        import IPython

        IPython.embed()
    else:
        print('==> Waiting 3 seconds')
        time.sleep(3)

        print('==> Moving to reset_manip_pose')
        robot.reset_manip_pose()
        print(robot.angle_vector())
        time.sleep(1)
        viewer.redraw()

        print('==> Waiting 3 seconds')
        time.sleep(3)

        print('==> Moving to init_pose')
        robot.init_pose()
        print(robot.angle_vector())
        time.sleep(1)
        viewer.redraw()

        print('==> Waiting 3 seconds')
        time.sleep(3)

        print('==> IK to box')
        robot.reset_manip_pose()
        robot.inverse_kinematics(box, rotation_mask='xz')
        print(robot.angle_vector())
        time.sleep(1)
        viewer.redraw()

        if recorder is not None:
            print('saving video to {}'.format(recorder.save()))

        if not args.no_interactive:
            viewer.wait_until_close()
        else:
            viewer.close()
        time.sleep(1.0)


if __name__ == '__main__':
    main()