def _pad_or_clip_voxels(voxel_features, voxel_indices, num_valid_voxels, segment_ids, voxels_pad_or_clip_size): """Pads or clips voxels.""" if voxels_pad_or_clip_size: num_valid_voxels = tf.minimum(num_valid_voxels, voxels_pad_or_clip_size) num_channels = voxel_features.get_shape().as_list()[-1] if len(voxel_features.shape.as_list()) == 2: output_shape = [voxels_pad_or_clip_size, num_channels] elif len(voxel_features.shape.as_list()) == 3: num_samples_per_voxel = voxel_features.get_shape().as_list()[1] if num_samples_per_voxel is None: num_samples_per_voxel = tf.shape(voxel_features)[1] output_shape = [ voxels_pad_or_clip_size, num_samples_per_voxel, num_channels ] else: raise ValueError('voxel_features should be either rank 2 or 3.') voxel_features = shape_utils.pad_or_clip_nd(tensor=voxel_features, output_shape=output_shape) voxel_indices = shape_utils.pad_or_clip_nd( tensor=voxel_indices, output_shape=[voxels_pad_or_clip_size, 3]) valid_segment_ids_mask = tf.cast(tf.less(segment_ids, num_valid_voxels), dtype=tf.int32) segment_ids *= valid_segment_ids_mask return voxel_features, voxel_indices, num_valid_voxels, segment_ids
def _pad_or_clip_point_properties(inputs, pad_or_clip_size): """Pads or clips the inputs point properties. If pad_or_clip_size is None, it won't perform any action. Args: inputs: A dictionary containing input tensors. pad_or_clip_size: Number of target points to pad or clip to. If None, it will not perform the padding. """ inputs[standard_fields.InputDataFields.num_valid_points] = tf.shape( inputs[standard_fields.InputDataFields.point_positions])[0] if pad_or_clip_size is not None: inputs[standard_fields.InputDataFields.num_valid_points] = tf.minimum( inputs[standard_fields.InputDataFields.num_valid_points], pad_or_clip_size) for key in sorted(standard_fields.get_input_point_fields()): if key == standard_fields.InputDataFields.num_valid_points: continue if key in inputs: tensor_rank = len(inputs[key].get_shape().as_list()) padding_shape = [pad_or_clip_size] for i in range(1, tensor_rank): padding_shape.append(inputs[key].get_shape().as_list()[i]) inputs[key] = shape_utils.pad_or_clip_nd( tensor=inputs[key], output_shape=padding_shape)
def pad_or_clip(mesh_inputs, view_indices_2d_inputs, pad_or_clip_size): """Pads and clips the points and correspondences.""" if standard_fields.InputDataFields.point_positions not in mesh_inputs: return num_valid_points = tf.shape( mesh_inputs[standard_fields.InputDataFields.point_positions])[0] if pad_or_clip_size: num_valid_points = tf.minimum(num_valid_points, pad_or_clip_size) for key in sorted(mesh_inputs.keys()): num_channels = mesh_inputs[key].get_shape().as_list()[1] mesh_inputs[key] = shape_utils.pad_or_clip_nd( tensor=mesh_inputs[key], output_shape=[pad_or_clip_size, num_channels]) for key in sorted(view_indices_2d_inputs): num_images = view_indices_2d_inputs[key].get_shape().as_list()[0] if num_images is None: num_images = tf.shape(view_indices_2d_inputs[key])[0] view_indices_2d_inputs[key] = shape_utils.pad_or_clip_nd( tensor=(view_indices_2d_inputs[key] + 1), output_shape=[num_images, pad_or_clip_size, 2]) - 1 mesh_inputs[ standard_fields.InputDataFields.num_valid_points] = num_valid_points