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)