-
Notifications
You must be signed in to change notification settings - Fork 361
/
demo_segmentation.py
118 lines (94 loc) · 3.72 KB
/
demo_segmentation.py
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
117
118
"""
Play random actions in an environment and render a video that demonstrates segmentation.
"""
import argparse
import colorsys
import json
import random
import imageio
import matplotlib.cm as cm
import numpy as np
from PIL import Image
import robosuite as suite
from robosuite.controllers import load_controller_config
def randomize_colors(N, bright=True):
"""
Modified from https://github.com/matterport/Mask_RCNN/blob/master/mrcnn/visualize.py#L59
Generate random colors.
To get visually distinct colors, generate them in HSV space then
convert to RGB.
"""
brightness = 1.0 if bright else 0.5
hsv = [(1.0 * i / N, 1, brightness) for i in range(N)]
colors = np.array(list(map(lambda c: colorsys.hsv_to_rgb(*c), hsv)))
rstate = np.random.RandomState(seed=20)
np.random.shuffle(colors)
return colors
def segmentation_to_rgb(seg_im, random_colors=False):
"""
Helper function to visualize segmentations as RGB frames.
NOTE: assumes that geom IDs go up to 255 at most - if not,
multiple geoms might be assigned to the same color.
"""
# ensure all values lie within [0, 255]
seg_im = np.mod(seg_im, 256)
if random_colors:
colors = randomize_colors(N=256, bright=True)
return (255.0 * colors[seg_im]).astype(np.uint8)
else:
# deterministic shuffling of values to map each geom ID to a random int in [0, 255]
rstate = np.random.RandomState(seed=8)
inds = np.arange(256)
rstate.shuffle(inds)
# use @inds to map each geom ID to a color
return (255.0 * cm.rainbow(inds[seg_im], 3)).astype(np.uint8)[..., :3]
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--video-path", type=str, default="/tmp/video.mp4", help="Path to video file")
parser.add_argument("--random-colors", action="store_true", help="Radnomize segmentation colors")
parser.add_argument("--segmentation-level", type=str, default="element", help="instance, class, or element")
args = parser.parse_args()
# Create dict to hold options that will be passed to env creation call
options = {}
# Choose environment and add it to options
options["env_name"] = "TwoArmHandover"
options["robots"] = ["Panda", "Panda"]
# Choose controller
controller_name = "OSC_POSE"
# Choose camera
camera = "frontview"
# Choose segmentation type
segmentation_level = args.segmentation_level # Options are {instance, class, element}
# Load the desired controller
options["controller_configs"] = load_controller_config(default_controller=controller_name)
# initialize the task
env = suite.make(
**options,
has_renderer=False,
has_offscreen_renderer=True,
ignore_done=True,
use_camera_obs=True,
control_freq=20,
camera_names=camera,
camera_segmentations=segmentation_level,
camera_heights=512,
camera_widths=512,
)
env.reset()
video_writer = imageio.get_writer(args.video_path, fps=20)
# Get action limits
low, high = env.action_spec
# do visualization
for i in range(100):
action = 0.5 * np.random.uniform(low, high)
obs, reward, done, _ = env.step(action)
video_img = obs[f"{camera}_segmentation_{segmentation_level}"].squeeze(-1)[::-1]
np.savetxt("/tmp/seg_{}.txt".format(i), video_img, fmt="%.2f")
video_img = segmentation_to_rgb(video_img, args.random_colors)
video_writer.append_data(video_img)
image = Image.fromarray(video_img)
image.save("/tmp/seg_{}.png".format(i))
if i % 5 == 0:
print("Step #{} / 100".format(i))
video_writer.close()
print("Video saved to {}".format(args.video_path))