def test_create_tfrecords_dataflow_runner(self, mock_beam): """Tests `create_tfrecords` Dataflow case.""" mock_beam.build_pipeline().run().job_id.return_value = 'foo_id' df2 = self.test_df.copy() df2[constants.IMAGE_URI_KEY] = 'gs://' + df2[constants.IMAGE_URI_KEY] outdir = '/tmp/dataflow_runner' expected = { 'job_id': 'foo_id', 'dataflow_url': 'https://console.cloud.google.com/dataflow/jobs/' + 'us-central1/foo_id?project=foo' } os.makedirs(outdir, exist_ok=True) r = client.create_tfrecords(df2, runner='DataflowRunner', output_dir=outdir, region=self.test_region, project=self.test_project, tfrecorder_wheel=self.test_wheel) self.assertEqual(r, expected)
def to_tfr(self, output_dir: str, schema_map: Dict[str, schema.SchemaMap] = schema.image_csv_schema, runner: str = 'DirectRunner', project: Optional[str] = None, region: Optional[str] = None, tfrecorder_wheel: Optional[str] = None, dataflow_options: Union[Dict[str, Any], None] = None, job_label: str = 'to-tfr', compression: Optional[str] = 'gzip', num_shards: int = 0) -> Dict[str, Any]: """TFRecorder Pandas Accessor. TFRecorder provides an easy interface to create image-based tensorflow records from a dataframe containing GCS locations of the images and labels. Usage: import tfrecorder df.tfrecorder.to_tfr( output_dir='gcs://foo/bar/train', runner='DirectRunner', compression='gzip', num_shards=10) Args: schema_map: A dict mapping column names to supported types. output_dir: Local directory or GCS Location to save TFRecords to. Note: GCS required for DataflowRunner runner: Beam runner. Can be DirectRunner or DataflowRunner. project: GCP project name (Required if DataflowRunner). region: GCP region name (Required if DataflowRunner). tfrecorder_wheel: Path to the tfrecorder wheel Dataflow will run. (create with 'python setup.py sdist' or 'pip download tfrecorder --no-deps') dataflow_options: Optional dictionary containing Dataflow options. job_label: User supplied description for the beam job name. compression: Can be 'gzip' or None for no compression. num_shards: Number of shards to divide the TFRecords into. Default is 0 = no sharding. Returns: job_results: A dictionary of job results. """ display.display( display.HTML('<b>Logging output to /tmp/{} </b>'.format( constants.LOGFILE))) r = client.create_tfrecords(self._df, output_dir=output_dir, schema_map=schema_map, runner=runner, project=project, region=region, tfrecorder_wheel=tfrecorder_wheel, dataflow_options=dataflow_options, job_label=job_label, compression=compression, num_shards=num_shards) return r
def test_create_tfrecords_direct_runner(self, mock_beam): """Tests `create_tfrecords` Direct case.""" mock_beam.build_pipeline().run().wait_until_finished.return_value = { 'rows': 6 } r = client.create_tfrecords(self.test_df, runner='DirectRunner', output_dir='/tmp/direct_runner') self.assertTrue('metrics' in r)
def to_tfr(self, output_dir: str, runner: str = 'DirectRunner', project: Optional[str] = None, region: Optional[str] = None, dataflow_options: Union[Dict[str, Any], None] = None, job_label: str = 'to-tfr', compression: Optional[str] = 'gzip', num_shards: int = 0) -> Dict[str, Any]: """TFRecorder Pandas Accessor. TFRecorder provides an easy interface to create image-based tensorflow records from a dataframe containing GCS locations of the images and labels. Usage: import tfrecorder df.tfrecorder.to_tfr( output_dir='gcs://foo/bar/train', runner='DirectRunner', compression='gzip', num_shards=10) Args: output_dir: Local directory or GCS Location to save TFRecords to. runner: Beam runner. Can be DirectRunner or DataFlowRunner. project: GCP project name (Required if DataFlowRunner). region: GCP region name (Required if DataFlowRunner). dataflow_options: Optional dictionary containing DataFlow options. job_label: User supplied description for the beam job name. compression: Can be 'gzip' or None for no compression. num_shards: Number of shards to divide the TFRecords into. Default is 0 = no sharding. Returns: job_results: A dictionary of job results. """ display.display( display.HTML('<b>Logging output to /tmp/{} </b>'.format( constants.LOGFILE))) r = client.create_tfrecords(self._df, output_dir=output_dir, runner=runner, project=project, region=region, dataflow_options=dataflow_options, job_label=job_label, compression=compression, num_shards=num_shards) return r