def register_example(reg_example, aligner, model):
ses = reg_example.salient_edge_set
# Pre-process mesh
mesh = ses.mesh
edge_mask = ses.edge_mask
edge_pc_obj = PointCloud(generate_canonical_pc(mesh, edge_mask).T,
frame='obj')
# Process Depth Image
ci = reg_example.camera_intrs
depth_im = reg_example.depth_im
point_cloud_im = ci.deproject_to_image(depth_im)
normal_cloud_im = point_cloud_im.normal_cloud_im()
joined = np.dstack((depth_im.data, normal_cloud_im.data))
mask = model.predict(joined[np.newaxis, :, :, :])[0]
mask *= 255.0
mask = BinaryImage(mask.astype(np.uint8))
depth_im_edges = depth_im.mask_binary(mask)
edge_pc_cam = ci.deproject(depth_im_edges)
edge_pc_cam.remove_zero_points()
# Align the two point sets with Super4PCS
T_obj_camera_est = aligner.align(edge_pc_cam, edge_pc_obj)
return T_obj_camera_est
def get_state(self, depth, segmask):
# Read images.
depth_im = DepthImage(depth, frame=self.camera_intr.frame)
color_im = ColorImage(np.zeros([depth_im.height, depth_im.width,
3]).astype(np.uint8),
frame=self.camera_intr.frame)
segmask = BinaryImage(segmask.astype(np.uint8) * 255,
frame=self.camera_intr.frame)
# Inpaint.
depth_im = depth_im.inpaint(rescale_factor=self.inpaint_rescale_factor)
# Create state.
rgbd_im = RgbdImage.from_color_and_depth(color_im, depth_im)
state = RgbdImageState(rgbd_im, self.camera_intr, segmask=segmask)
return state, rgbd_im
