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')