depth_image = np.array(
imageio.imread(os.path.join(args.dir, depth_image_name))) / 1000.
depth_image_drake = Image[PixelType.kDepth32F](depth_image.shape[1],
depth_image.shape[0])
depth_image_drake.mutable_data[:, :, 0] = depth_image[:, :]
points = camera.ConvertDepthImageToPointCloud(
depth_image_drake, camera.depth_camera_info())
good_points_mask = np.all(np.isfinite(points), axis=0)
points = points[:, good_points_mask]
points = np.vstack([points, np.ones([1, points.shape[1]])])
# Transform them to camera base frame
points = np.dot(depth_camera_pose, points)
# and then to world frame
# Last body = camera floating base
kinsol = rbt.doKinematics(q0)
tf = rbt.relativeTransform(kinsol, 0, rbt.get_num_bodies() - 1)
points = tf.dot(points)[:3, :]
label_image_name = "%09d_mask.png" % (i)
labels = np.array(
imageio.imread(os.path.join(args.dir, label_image_name))).ravel()
labels = labels[good_points_mask]
normal_image_name = "%09d_normal.png" % (i)
# Empirically-derived transposing and memory reording to get into
# 3xN format...
normals = np.array(
imageio.imread(os.path.join(args.dir, normal_image_name)))
normals = normals.T.reshape(3, width * height,
order='F')[:, good_points_mask]
# Rescale from int8 to float values
depth_image * 1000.)
# Compute normal image from the depth image
structured_point_cloud = camera.ConvertDepthImageToPointCloud(
u_data, camera.depth_camera_info())
structured_point_cloud = structured_point_cloud.T.reshape([h, w, 3])
normal_image = structured_normal_estimate(structured_point_cloud)
# TODO: Precision is bad here (8-bit per channel), but imageio
# doesn't want to export 3-channel 16bit images. OK with the loss
# for now, it's not too horrible for normals...
save_image_uint8("%s/%09d_normal.png" % (args.dir, i),
(normal_image + 1.) * 127)
u_data = output.get_data(
camera.label_image_output_port().get_index()).get_value()
h, w, _ = u_data.data.shape
label_image = np.empty((h, w), dtype=np.int8)
label_image[:, :] = u_data.data[:, :, 0]
# Make labels 0 where there's no label
label_image = np.where(
np.logical_and(label_image > 0,
label_image < rbt.get_num_bodies()), label_image,
label_image * 0)
save_image_uint8("%s/%09d_mask.png" % (args.dir, i), label_image)
# Save a human-readable version for debugging
save_image_uint8(
"%s/%09d_mask_color.png" % (args.dir, i),
plt.cm.jet(label_image.astype(float) / np.max(label_image)) * 255)
print "Completed rendering of viewpoint %d" % i