def estimate_current_pose(self, previous_pose, current_rgb, current_depth, debug=False):
render_rgb, render_depth = self.renderer.render(previous_pose.transpose())
# todo implement this part on gpu...
rgbA, depthA = normalize_scale(render_rgb, render_depth, previous_pose.inverse(), self.camera, self.image_size,
self.object_width)
rgbB, depthB = normalize_scale(current_rgb, current_depth, previous_pose.inverse(), self.camera, self.image_size,
self.object_width)
depthA = normalize_depth(depthA, previous_pose)
depthB = normalize_depth(depthB, previous_pose)
if debug:
show_frames(rgbA, depthA, rgbB, depthB)
rgbA, depthA = normalize_channels(rgbA, depthA, self.mean[:4], self.std[:4])
rgbB, depthB = normalize_channels(rgbB, depthB, self.mean[4:], self.std[4:])
self.input_buffer[0, 0:3, :, :] = rgbA
self.input_buffer[0, 3, :, :] = depthA
self.input_buffer[0, 4:7, :, :] = rgbB
self.input_buffer[0, 7, :, :] = depthB
self.prior_buffer[0] = np.array(previous_pose.to_parameters(isQuaternion=True))
prediction = self.tracker_model.test([self.input_buffer, self.prior_buffer]).asNumpyTensor()
prediction = unnormalize_label(prediction, self.translation_range, self.rotation_range)
if debug:
print("Prediction : {}".format(prediction))
prediction = Transform.from_parameters(*prediction[0], is_degree=True)
current_pose = combine_view_transform(previous_pose, prediction)
self.debug_rgb = render_rgb
return current_pose
OBJECT_WIDTH,
scale=(1000, -1000, -1000))
rgbA, depthA = normalize_scale(rgbA, depthA, bb,
real_dataset.camera, IMAGE_SIZE)
rgbB, depthB = normalize_scale(rotated_rgb, rotated_depth, bb,
real_dataset.camera, IMAGE_SIZE)
index = output_dataset.add_pose(rgbA, depthA, previous_pose)
output_dataset.add_pair(rgbB, depthB, random_transform, index)
iteration = i * SAMPLE_QUANTITY + j
sys.stdout.write(
"Progress: %d%% \r" %
(int(iteration /
(SAMPLE_QUANTITY * real_dataset.size()) * 100)))
sys.stdout.flush()
if iteration % 500 == 0:
output_dataset.dump_images_on_disk()
if iteration % 5000 == 0:
output_dataset.save_json_files(metadata)
if args.verbose:
show_frames(rgbA, depthA, rgbB, depthB)
cv2.imshow("testB", rgbB[:, :, ::-1])
k = cv2.waitKey(1)
if k == ESCAPE_KEY:
break
output_dataset.dump_images_on_disk()
output_dataset.save_json_files(metadata)
def estimate_current_pose(self,
previous_pose,
current_rgb,
current_depth,
debug=False,
debug_time=False):
if debug_time:
start_time = time.time()
bb = compute_2Dboundingbox(previous_pose,
self.camera,
self.object_width,
scale=(1000, 1000, -1000))
bb2 = compute_2Dboundingbox(previous_pose,
self.camera,
self.object_width,
scale=(1000, -1000, -1000))
if debug_time:
print("Compute BB : {}".format(time.time() - start_time))
start_time = time.time()
rgbA, depthA = self.compute_render(previous_pose, bb)
if debug_time:
print("Render : {}".format(time.time() - start_time))
start_time = time.time()
rgbB, depthB = normalize_scale(current_rgb, current_depth, bb2,
self.camera, self.image_size)
debug_info = (rgbA, bb2, np.hstack((rgbA, rgbB)))
rgbA = rgbA.astype(np.float)
rgbB = rgbB.astype(np.float)
depthA = depthA.astype(np.float)
depthB = depthB.astype(np.float)
depthA = normalize_depth(depthA, previous_pose)
depthB = normalize_depth(depthB, previous_pose)
if debug:
show_frames(rgbA, depthA, rgbB, depthB)
rgbA, depthA = normalize_channels(rgbA, depthA, self.mean[:4],
self.std[:4])
rgbB, depthB = normalize_channels(rgbB, depthB, self.mean[4:],
self.std[4:])
self.input_buffer[0, 0:3, :, :] = rgbA
self.input_buffer[0, 3, :, :] = depthA
self.input_buffer[0, 4:7, :, :] = rgbB
self.input_buffer[0, 7, :, :] = depthB
self.prior_buffer[0] = np.array(
previous_pose.to_parameters(isQuaternion=True))
if debug_time:
print("Normalize : {}".format(time.time() - start_time))
start_time = time.time()
prediction = self.tracker_model.test(
[self.input_buffer, self.prior_buffer]).asNumpyTensor()
if debug_time:
print("Network time : {}".format(time.time() - start_time))
prediction = unnormalize_label(prediction, self.translation_range,
self.rotation_range)
if debug:
print("Prediction : {}".format(prediction))
prediction = Transform.from_parameters(*prediction[0], is_degree=True)
current_pose = combine_view_transform(previous_pose, prediction)
return current_pose, debug_info