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