def test_segment_init(self): # valid data random_valid_data = (255.0 * np.random.rand(IM_HEIGHT, IM_WIDTH)).astype( np.uint8) im = SegmentationImage(random_valid_data) self.assertEqual(im.height, IM_HEIGHT) self.assertEqual(im.width, IM_WIDTH) self.assertEqual(im.channels, 1) self.assertTrue(np.allclose(im.data, random_valid_data)) # invalid channels random_data = np.random.rand(IM_HEIGHT, IM_WIDTH, 3).astype(np.uint8) caught_bad_channels = False try: im = SegmentationImage(random_data) except: caught_bad_channels = True self.assertTrue(caught_bad_channels) # invalid type random_data = np.random.rand(IM_HEIGHT, IM_WIDTH).astype(np.float32) caught_bad_dtype = False try: im = SegmentationImage(random_data) except: caught_bad_dtype = True self.assertTrue(caught_bad_dtype)
def load(save_dir): if not os.path.exists(save_dir): raise ValueError('Directory %s does not exist!' % (save_dir)) color_image_filename = os.path.join(save_dir, 'color.png') depth_image_filename = os.path.join(save_dir, 'depth.npy') camera_intr_filename = os.path.join(save_dir, 'camera.intr') segmask_filename = os.path.join(save_dir, 'segmask.npy') obj_segmask_filename = os.path.join(save_dir, 'obj_segmask.npy') state_filename = os.path.join(save_dir, 'state.pkl') camera_intr = CameraIntrinsics.load(camera_intr_filename) color = ColorImage.open(color_image_filename, frame=camera_intr.frame) depth = DepthImage.open(depth_image_filename, frame=camera_intr.frame) segmask = None if os.path.exists(segmask_filename): segmask = BinaryImage.open(segmask_filename, frame=camera_intr.frame) obj_segmask = None if os.path.exists(obj_segmask_filename): obj_segmask = SegmentationImage.open(obj_segmask_filename, frame=camera_intr.frame) fully_observed = None if os.path.exists(state_filename): fully_observed = pkl.load(open(state_filename, 'rb')) return RgbdImageState(RgbdImage.from_color_and_depth(color, depth), camera_intr, segmask=segmask, obj_segmask=obj_segmask, fully_observed=fully_observed)
# Read out all points in large clusters. cur_i = 0 for j, indices in enumerate(cluster_indices): num_points = len(indices) points = np.zeros([3, num_points]) for i, index in enumerate(indices): points[0, i] = pcl_cloud[index][0] points[1, i] = pcl_cloud[index][1] points[2, i] = pcl_cloud[index][2] segment = PointCloud(points, frame=camera_intr.frame) depth_segment = camera_intr.project_to_image(segment) obj_segmask_data[depth_segment.data > 0] = j + 1 obj_segmask = SegmentationImage(obj_segmask_data.astype(np.uint8)) obj_segmask = obj_segmask.mask_binary(segmask) # Inpaint. depth_im = depth_im.inpaint(rescale_factor=inpaint_rescale_factor) if "input_images" in policy_config["vis"] and policy_config["vis"][ "input_images"]: vis.figure(size=(10, 10)) num_plot = 3 vis.subplot(1, num_plot, 1) vis.imshow(depth_im) vis.subplot(1, num_plot, 2) vis.imshow(segmask) vis.subplot(1, num_plot, 3) vis.imshow(obj_segmask)