def main(argv):
beam_utils.BeamInit()
assert FLAGS.input_file_pattern
assert FLAGS.output_filebase
# Construct pipeline options from argv.
options = beam.options.pipeline_options.PipelineOptions(argv[1:])
reader = beam_utils.GetReader('tfrecord',
FLAGS.input_file_pattern,
value_coder=beam.coders.ProtoCoder(
dataset_pb2.Frame))
writer = beam_utils.GetWriter('tfrecord',
file_pattern=FLAGS.output_filebase,
value_coder=beam.coders.ProtoCoder(
tf.train.Example))
emitter_fn = beam_utils.GetEmitterFn('tfrecord')
with beam_utils.GetPipelineRoot(options=options) as root:
_ = (root
| 'Read' >> reader
| 'ConvertToTFExample' >> beam.ParDo(
waymo_proto_to_tfe.WaymoOpenDatasetConverter(emitter_fn))
| 'Write' >> writer)
def testGetEmitterFn(self):
_ = beam_utils.GetEmitterFn('tfrecord')
with self.assertRaises(ValueError):
_ = beam_utils.GetEmitterFn('unknown')