-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathvis_gt_poses.py
189 lines (155 loc) · 6.89 KB
/
vis_gt_poses.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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
# Author: Tomas Hodan ([email protected])
# Center for Machine Perception, Czech Technical University in Prague
# Visualizes the object models at the ground truth poses.
import os
import sys
import numpy as np
import matplotlib.pyplot as plt
import cv2
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from pysixd import inout, misc, renderer
from params.dataset_params import get_dataset_params
dataset = 'hinterstoisser'
#dataset = 'tless'
# dataset = 'tudlight'
# dataset = 'rutgers'
# dataset = 'tejani'
# dataset = 'doumanoglou'
# dataset = 'toyotalight'
# Dataset parameters
dp = get_dataset_params(dataset)
# Select IDs of scenes, images and GT poses to be processed.
# Empty list [] means that all IDs will be used.
scene_ids = []
im_ids = []
gt_ids = []
# Indicates whether to render RGB image
vis_rgb = True
# Indicates whether to resolve visibility in the rendered RGB image (using
# depth renderings). If True, only the part of object surface, which is not
# occluded by any other modeled object, is visible. If False, RGB renderings
# of individual objects are blended together.
vis_rgb_resolve_visib = True
# Indicates whether to render depth image
vis_depth = False
# If to use the original model color
vis_orig_color = False
# Define new object colors (used if vis_orig_colors == False)
colors = inout.load_yaml('../data/colors.yml')
# Path masks for output images
vis_rgb_mpath = '../output/vis_gt_poses_{}/{:02d}/{:04d}.jpg'
vis_depth_mpath = '../output/vis_gt_poses_{}/{:02d}/{:04d}_depth_diff.jpg'
# Whether to consider only the specified subset of images
use_image_subset = True
# Subset of images to be considered
if use_image_subset:
im_ids_sets = inout.load_yaml(dp['test_set_fpath'])
else:
im_ids_sets = None
scene_ids_curr = range(1, dp['scene_count'] + 1)
if scene_ids:
scene_ids_curr = set(scene_ids_curr).intersection(scene_ids)
for scene_id in scene_ids_curr:
misc.ensure_dir(os.path.dirname(vis_rgb_mpath.format(dataset, scene_id, 0)))
# Load scene info and gt poses
scene_info = inout.load_info(dp['scene_info_mpath'].format(scene_id))
scene_gt = inout.load_gt(dp['scene_gt_mpath'].format(scene_id))
# Load models of objects that appear in the current scene
obj_ids = set([gt['obj_id'] for gts in scene_gt.values() for gt in gts])
models = {}
for obj_id in obj_ids:
models[obj_id] = inout.load_ply(dp['model_mpath'].format(obj_id))
# Considered subset of images for the current scene
if im_ids_sets is not None:
im_ids_curr = im_ids_sets[scene_id]
else:
im_ids_curr = sorted(scene_info.keys())
if im_ids:
im_ids_curr = set(im_ids_curr).intersection(im_ids)
# Visualize GT poses in the selected images
for im_id in im_ids_curr:
print('scene: {}, im: {}'.format(scene_id, im_id))
# Load the images
rgb = inout.load_im(dp['test_rgb_mpath'].format(scene_id, im_id))
depth = inout.load_depth(dp['test_depth_mpath'].format(scene_id, im_id))
depth = depth.astype(np.float32) # [mm]
depth *= dp['cam']['depth_scale'] # to [mm]
# Render the objects at the ground truth poses
im_size = (depth.shape[1], depth.shape[0])
ren_rgb = np.zeros(rgb.shape, np.float32)
ren_rgb_info = np.zeros(rgb.shape, np.uint8)
ren_depth = np.zeros(depth.shape, np.float32)
gt_ids_curr = range(len(scene_gt[im_id]))
if gt_ids:
gt_ids_curr = set(gt_ids_curr).intersection(gt_ids)
for gt_id in gt_ids_curr:
gt = scene_gt[im_id][gt_id]
obj_id = gt['obj_id']
if vis_orig_color:
color = (1, 1, 1)
else:
color = tuple(colors[(obj_id - 1) % len(colors)])
color_uint8 = tuple([int(255 * c) for c in color])
model = models[gt['obj_id']]
K = scene_info[im_id]['cam_K']
R = gt['cam_R_m2c']
t = gt['cam_t_m2c']
# Rendering
if vis_rgb:
if vis_orig_color:
m_rgb = renderer.render(model, im_size, K, R, t, mode='rgb')
else:
m_rgb = renderer.render(model, im_size, K, R, t, mode='rgb',
surf_color=color)
if vis_depth or (vis_rgb and vis_rgb_resolve_visib):
m_depth = renderer.render(model, im_size, K, R, t, mode='depth')
# Get mask of the surface parts that are closer than the
# surfaces rendered before
visible_mask = np.logical_or(ren_depth == 0, m_depth < ren_depth)
mask = np.logical_and(m_depth != 0, visible_mask)
ren_depth[mask] = m_depth[mask].astype(ren_depth.dtype)
# Combine the RGB renderings
if vis_rgb:
if vis_rgb_resolve_visib:
ren_rgb[mask] = m_rgb[mask].astype(ren_rgb.dtype)
else:
ren_rgb += m_rgb.astype(ren_rgb.dtype)
# Draw 2D bounding box and info
if True:
obj_mask = np.sum(m_rgb > 0, axis=2)
ys, xs = obj_mask.nonzero()
if len(ys):
# bbox = misc.calc_2d_bbox(xs, ys, im_size)
# im_size = (obj_mask.shape[1], obj_mask.shape[0])
# ren_rgb_bbs = misc.draw_rect(ren_rgb_bbs, bbox, color_uint8)
txt = '{},{}'.format(obj_id, gt_id)
txt_offset = 0
# txt_offset = 5
p_id = np.argmin(ys)
tex_loc = (xs[p_id], ys[p_id] - 5)
# tex_loc = (bbox[0], bbox[1])
cv2.putText(ren_rgb_info, txt, tex_loc,
cv2.FONT_HERSHEY_PLAIN, 0.9, color_uint8, 1)
# Save RGB visualization
if vis_rgb:
vis_im_rgb = 0.5 * rgb.astype(np.float32) +\
0.5 * ren_rgb + \
1.0 * ren_rgb_info
vis_im_rgb[vis_im_rgb > 255] = 255
inout.save_im(vis_rgb_mpath.format(dataset, scene_id, im_id),
vis_im_rgb.astype(np.uint8))
# Save image of depth differences
if vis_depth:
# Calculate the depth difference at pixels where both depth maps
# are valid
valid_mask = (depth > 0) * (ren_depth > 0)
depth_diff = valid_mask * (depth - ren_depth.astype(np.float32))
f, ax = plt.subplots(1, 1)
cax = ax.matshow(depth_diff)
ax.axis('off')
ax.set_title('measured - GT depth [mm]')
f.colorbar(cax, fraction=0.03, pad=0.01)
f.tight_layout(pad=0)
plt.savefig(vis_depth_mpath.format(dataset, scene_id, im_id), pad=0,
bbox_inches='tight')
plt.close()