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