def transform_pointcloud_to_another_frame(scene, point_positions,
frame_source_index,
frame_target_index):
"""transform each point in target_frame to source_frame coordinate based on their relative transformation.
Args:
scene: dict of tensor containing scene.
point_positions: the tensor of point positions in inputs.
frame_source_index: source frame index.
frame_target_index: target frame index.
Returns:
A point cloud from frame_target warpped, of [N, 3] shape.
"""
world_rotation_frame0 = scene['frames/pose/R'][frame_source_index]
world_rotation_frame1 = scene['frames/pose/R'][frame_target_index]
world_translation_frame0 = scene['frames/pose/t'][frame_source_index]
world_translation_frame1 = scene['frames/pose/t'][frame_target_index]
point_positions_world = tf.tensordot(
point_positions, world_rotation_frame1,
axes=(1, 1)) + world_translation_frame1
point_positions_frame0 = tf.tensordot(
point_positions_world - world_translation_frame0,
world_rotation_frame0,
axes=(1, 0))
return point_positions_frame0
def compute_motion_labels(scene,
frame0,
frame1,
frame_start_index,
points_key,
box_margin=0.1):
"""Compute motion label for each point.
Args:
scene: dict of tensor containing scene.
frame0: dict of tensor containing points and objects.
frame1: dict of tensor containing points and objects.
frame_start_index: starting frame index.
points_key: A string corresponding to the tensor of point positions in
inputs.
box_margin: A margin value to enlarge box, so that surrounding points are
included.
Returns:
A motion tensor of [N, 3] shape.
"""
point_positions = frame0[points_key]
frame0_object_names = frame0['objects/name']
frame1_object_names = frame1['objects/name']
bool_matrix = tf.math.equal(
tf.expand_dims(frame0_object_names, axis=1),
tf.expand_dims(frame1_object_names, axis=0))
match_indices = tf.where(bool_matrix)
# object box level
box_dimension = tf.gather(
frame0['objects/shape/dimension'], match_indices[:, 0], axis=0)
boxes_length = box_dimension[:, 0:1]
boxes_width = box_dimension[:, 1:2]
boxes_height = box_dimension[:, 2:3]
boxes_rotation_matrix = tf.gather(
frame0['objects/pose/R'], match_indices[:, 0], axis=0)
boxes_center = tf.gather(
frame0['objects/pose/t'], match_indices[:, 0], axis=0)
frame1_box_rotation_matrix = tf.gather(
frame1['objects/pose/R'], match_indices[:, 1], axis=0)
frame1_box_center = tf.gather(
frame1['objects/pose/t'], match_indices[:, 1], axis=0)
# frame level
frame0_rotation = scene['frames/pose/R'][frame_start_index]
frame1_rotation = scene['frames/pose/R'][frame_start_index + 1]
frame0_translation = scene['frames/pose/t'][frame_start_index]
frame1_translation = scene['frames/pose/t'][frame_start_index + 1]
frame1_box_center_global = tf.tensordot(
frame1_box_center, frame1_rotation, axes=(1, 1)) + frame1_translation
frame1_box_center_in_frame0 = tf.tensordot(
frame1_box_center_global - frame0_translation,
frame0_rotation,
axes=(1, 0))
# only find index on boxes that are matched between two frames
points_box_index = box_utils.map_points_to_boxes(
points=point_positions,
boxes_length=boxes_length,
boxes_height=boxes_height,
boxes_width=boxes_width,
boxes_rotation_matrix=boxes_rotation_matrix,
boxes_center=boxes_center,
box_margin=box_margin)
# TODO(huangrui): disappered object box have 0 motion.
# Probably consider set to nan or ignore_label.
# 1. gather points in surviving matched box only,
# and replicate rotation/t to same length;
# 2. get points in box frame, apply new rotation/t per point;
# 3. new location minus old location -> motion vector;
# 4. scatter it to a larger motion_vector with 0 for
# points ouside of matched boxes.
# Need to limit boxes to those matched boxes.
# otherwise the points_box_index will contain useless box.
# index in all point array, of points that are inside the box.
points_inside_box_index = tf.where(points_box_index + 1)[:, 0]
box_index = tf.gather(points_box_index, points_inside_box_index)
points_inside_box = tf.gather(point_positions, points_inside_box_index)
box_rotation_per_point = tf.gather(boxes_rotation_matrix, box_index)
box_center_per_point = tf.gather(boxes_center, box_index)
# Tensor [N, 3, 3] and [N, 3]. note we are transform points reversely.
points_in_box_frame = tf.einsum('ikj,ik->ij', box_rotation_per_point,
points_inside_box - box_center_per_point)
# Transform rotation of box from frame1 coordinate to frame0 coordinate
# note, transpose is implemented via changing summation axis
frame1_box_rotation_matrix_global = tf.transpose(
tf.tensordot(frame1_rotation, frame1_box_rotation_matrix, axes=(1, 1)),
perm=(1, 0, 2))
frame1_box_rotation_matrix_in_frame0 = tf.transpose(
tf.tensordot(
frame0_rotation, frame1_box_rotation_matrix_global, axes=(0, 1)),
perm=(1, 0, 2))
# this is the points_position_after_following_frame1_box's motion.
frame1_box_rotation_in_frame0_per_point = tf.gather(
frame1_box_rotation_matrix_in_frame0, box_index)
frame1_box_center_in_frame0_per_point = tf.gather(frame1_box_center_in_frame0,
box_index)
points_in_box_frame1 = tf.einsum(
'ijk,ik->ij', frame1_box_rotation_in_frame0_per_point,
points_in_box_frame) + frame1_box_center_in_frame0_per_point
motion_vector = points_in_box_frame1 - points_inside_box
scattered_vector = tf.scatter_nd(
indices=tf.expand_dims(points_inside_box_index, axis=1),
updates=motion_vector,
shape=tf.shape(point_positions, out_type=tf.dtypes.int64))
return scattered_vector