def testReadBatchFromThreeTables(self):
testdata_dir = 'poem/testdata' # Assume $PWD == "google_research/".
table_path = os.path.join(FLAGS.test_srcdir, testdata_dir,
'tfe-2.tfrecords')
parser_fn = tfe_input_layer.create_tfe_parser(
keypoint_names_2d=(keypoint_profiles.LegacyCoco13KeypointProfile2D(
).keypoint_names),
keypoint_names_3d=(keypoint_profiles.LegacyH36m17KeypointProfile3D(
).keypoint_names),
include_keypoint_scores_2d=True,
include_keypoint_scores_3d=False,
num_objects=2)
inputs = tfe_input_layer.read_batch_from_tables(
[table_path, table_path, table_path],
batch_sizes=[4, 2, 3],
drop_remainder=True,
shuffle=True,
parser_fn=parser_fn)
inputs = next(iter(inputs))
self.assertCountEqual(inputs.keys(), [
'image_sizes', 'keypoints_2d', 'keypoint_scores_2d', 'keypoints_3d'
])
self.assertEqual(inputs['image_sizes'].shape, [9, 2, 2])
self.assertEqual(inputs['image_sizes'].shape, [9, 2, 2])
self.assertEqual(inputs['keypoints_2d'].shape, [9, 2, 13, 2])
self.assertEqual(inputs['keypoint_scores_2d'].shape, [9, 2, 13])
self.assertEqual(inputs['keypoints_3d'].shape, [9, 2, 17, 3])
def create_dataset_from_tables(
input_table_patterns,
batch_sizes,
num_instances_per_record,
shuffle=False,
num_epochs=None,
drop_remainder=False,
keypoint_names_3d=None,
keypoint_names_2d=None,
feature_dim=None,
num_classes=None,
num_frames=None,
shuffle_buffer_size=4096,
num_shards=1,
shard_index=None,
common_module=common,
dataset_class=tf.data.TFRecordDataset,
input_example_parser_creator=tfe_input_layer.create_tfe_parser,
seed=None):
"""Reads data from tf.Example table.
Note that this function mainly duplicates `read_batch_from_tfe_tables` in
`v1.pipeline_utils.py` for compatible with tf2.
IMPORTANT: We assume that 2D keypoints from the input have been normalized by
image size. No normalization is expected and no denormalization will be
performed for both 2D and 3D keypoints.
Args:
input_table_patterns: A list of strings for the paths or pattern to input
tables.
batch_sizes: A list of integers for the batch sizes to read from each table.
num_instances_per_record: An integer for the number of instances per
tf.Example record.
shuffle: A boolean for whether to shuffle batch.
num_epochs: An integer for the number of epochs to read. Use `None` to read
indefinitely.
drop_remainder: A boolean for whether to drop remainder batch.
keypoint_names_3d: A list of strings for 3D keypoint names to read
(coordinates). Use None to skip reading 2D keypoints.
keypoint_names_2d: A list of strings for 2D keypoint names to read
(coordinates and scores). Use None to skip reading 2D keypoints.
feature_dim: An integer for size of pre-computed feature vectors. Use None
to skip reading feature vectors.
num_classes: An integer for total number of classification label classes to
read labels for. Use None to skip reading class labels.
num_frames: An integer for the number of frames per object each example has.
Use None to skip adding the frame dimension.
shuffle_buffer_size: An integer for the buffer size used for shuffling. A
large buffer size benefits shuffling quality.
num_shards: An integer for the number of shards to divide the dataset. This
is useful to distributed training. See `tf.data.Dataset.shard` for
details.
shard_index: An integer for the shard index to use. This is useful to
distributed training, and should usually be set to the id of a
synchronized worker. See `tf.data.Dataset.shard` for details. Note this
must be specified if `num_shards` is greater than 1.
common_module: A Python module that defines common constants.
dataset_class: A dataset class to use. Must match input table type.
input_example_parser_creator: A function handle for creating parser
function.
seed: An integer for random seed.
Returns:
A tf.data.Dataset object.
"""
parser_kwargs = {
'num_objects': num_instances_per_record,
}
if keypoint_names_3d:
parser_kwargs.update({
'keypoint_names_3d': keypoint_names_3d,
'include_keypoint_scores_3d': False,
})
if keypoint_names_2d:
parser_kwargs.update({
'keypoint_names_2d': keypoint_names_2d,
'include_keypoint_scores_2d': True,
})
if feature_dim:
parser_kwargs.update({
'feature_dim': feature_dim,
})
if num_classes:
parser_kwargs.update({
'num_classes': num_classes,
})
if num_frames:
parser_kwargs.update({
'sequence_length': num_frames,
})
parser_fn = input_example_parser_creator(
common_module=common_module, **parser_kwargs)
dataset = tfe_input_layer.read_batch_from_tables(
input_table_patterns,
batch_sizes=batch_sizes,
drop_remainder=drop_remainder,
num_epochs=num_epochs,
num_shards=num_shards,
shard_index=shard_index,
shuffle=shuffle,
shuffle_buffer_size=shuffle_buffer_size,
seed=seed,
dataset_class=dataset_class,
parser_fn=parser_fn)
return dataset
def read_batch_from_dataset_tables(input_table_patterns,
batch_sizes,
num_instances_per_record,
shuffle,
num_epochs,
keypoint_names_3d=None,
keypoint_names_2d=None,
min_keypoint_score_2d=-1.0,
shuffle_buffer_size=4096,
num_shards=1,
shard_index=None,
common_module=common,
dataset_class=tf.data.TFRecordDataset,
input_example_parser_creator=None,
seed=None):
"""Reads data from dataset table.
IMPORTANT: We assume that 2D keypoints from the input have been normalized by
image size. This function will reads image sizes from the input and
denormalize the 2D keypoints with them. No normalization is expected and no
denormalization will be performed for 3D keypoints.
Output tensors may include:
keypoints: A tensor for standardized 2D keypoints. Shape = [batch_size,
num_instances_per_record, num_keypoints_2d, 2].
keypoint_scores: A tensor for 2D keypoint scores. Shape = [batch_size,
num_instances_per_record, num_keypoints_2d].
keypoints_3d: A tensor for standardized 3D keypoints. Shape = [batch_size,
num_instances_per_record, num_keypoints_3d, 3].
Args:
input_table_patterns: A list of strings for the paths or pattern to input
tables.
batch_sizes: A list of integers for the batch sizes to read from each table.
num_instances_per_record: An integer for the number of instances per
tf.Example record.
shuffle: A boolean for whether to shuffle batch.
num_epochs: An integer for the number of epochs to read. Use None to read
indefinitely, in which case remainder batch will be dropped.
keypoint_names_3d: A list of strings for 3D keypoint names to read
(coordinates). Use None to skip reading 2D keypoints.
keypoint_names_2d: A list of strings for 2D keypoint names to read
(coordinates and scores). Use None to skip reading 2D keypoints.
min_keypoint_score_2d: A float for the minimum score to consider a 2D
keypoint as invalid.
shuffle_buffer_size: An integer for the buffer size used for shuffling. A
large buffer size benefits shuffling quality.
num_shards: An integer for the number of shards to divide the dataset. This
is useful to distributed training. See `tf.data.Dataset.shard` for
details.
shard_index: An integer for the shard index to use. This is useful to
distributed training, and should usually be set to the id of a
synchronized worker. See `tf.data.Dataset.shard` for details. Note this
must be specified if `num_shards` is greater than 1.
common_module: A Python module that defines common constants.
dataset_class: A dataset class to use. Must match input table type.
input_example_parser_creator: A function handle for creating parser
function. If None, uses the default parser creator.
seed: An integer for random seed.
Returns:
outputs: A dictionary for output tensor inputs.
"""
parser_kwargs = {
'num_objects': num_instances_per_record,
}
if keypoint_names_3d:
parser_kwargs.update({
'keypoint_names_3d': keypoint_names_3d,
'include_keypoint_scores_3d': False,
})
if keypoint_names_2d:
parser_kwargs.update({
'keypoint_names_2d': keypoint_names_2d,
'include_keypoint_scores_2d': True,
})
if input_example_parser_creator is None:
input_example_parser_creator = tfe_input_layer.create_tfe_parser
parser_fn = input_example_parser_creator(common_module=common_module,
**parser_kwargs)
# TODO(lzyuan): consider to refactor read_batch_from_batches into other file.
outputs = tfe_input_layer.read_batch_from_tables(
input_table_patterns,
batch_sizes=batch_sizes,
drop_remainder=num_epochs is None,
num_epochs=num_epochs,
num_shards=num_shards,
shard_index=shard_index,
shuffle=shuffle,
shuffle_buffer_size=shuffle_buffer_size,
dataset_class=dataset_class,
parser_fn=parser_fn,
seed=seed)
outputs = tf.data.make_one_shot_iterator(outputs).get_next()
if keypoint_names_2d:
# Since we assume 2D keypoints from the input have been normalized by image
# size, so we need to denormalize them to restore correctly aspect ratio.
keypoints_2d = keypoint_utils.denormalize_points_by_image_size(
outputs[common_module.KEY_KEYPOINTS_2D],
image_sizes=outputs[common_module.KEY_IMAGE_SIZES])
keypoint_scores_2d = outputs[common_module.KEY_KEYPOINT_SCORES_2D]
if min_keypoint_score_2d < 0.0:
keypoint_masks_2d = tf.ones_like(keypoint_scores_2d,
dtype=tf.float32)
else:
keypoint_masks_2d = tf.cast(tf.math.greater_equal(
keypoint_scores_2d, min_keypoint_score_2d),
dtype=tf.float32)
outputs.update({
common_module.KEY_KEYPOINTS_2D: keypoints_2d,
common_module.KEY_KEYPOINT_MASKS_2D: keypoint_masks_2d
})
return outputs
