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
