Add routine to check TFRecord files.

Change-Id: I42e50880bfc04b1eea25b3f8ecb6be482f1200af

Author: cfezequiel

tfrutil/accessor.py

@@ -34,12 +34,13 @@ def __init__(self, pandas_obj):
     self._df = pandas_obj
 
   # pylint: disable=too-many-arguments
-  def to_tfr(self,
-             output_path: str,
-             runner: str = "DirectRunner",
-             job_label: str = "to-tfr",
-             compression: Union[str, None] = "gzip",
-             num_shards: int = 0) -> str:
+  def to_tfr(
+      self,
+      output_dir: str,
+      runner: str = "DirectRunner",
+      job_label: str = "to-tfr",
+      compression: Union[str, None] = "gzip",
+      num_shards: int = 0) -> str:
     """TFRUtil Pandas Accessor.
 
     TFRUtil provides an easy interface to create image-based tensorflow records
@@ -48,16 +49,17 @@ def to_tfr(self,
     Usage:
       import tfrutil
 
-      df.tfrutil.to_tfr(runner="local",
-                        output_path="gcs://foo/bar/train",
-                        compression="gzip",
-                        num_shards=10,
-                        image_col="image",
-                        label_col="label)
+      df.tfrutil.to_tfr(
+          output_dir="gcs://foo/bar/train",
+          runner="DirectRunner",
+          compression="gzip",
+          num_shards=10,
+          image_col="image",
+          label_col="label)
 
     Args:
       runner: Beam runner. Can be "local" or "DataFlow"
-      output_path: Local directory or GCS Location to save TFRecords to.
+      output_dir: Local directory or GCS Location to save TFRecords to.
       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
@@ -70,9 +72,9 @@ def to_tfr(self,
     """
     print("Starting DataFlow Transform. This may take a while. Please wait.")
     client.create_tfrecords(self._df,
+                            output_dir=output_dir,
                             runner=runner,
-                            output_path=output_path,
                             job_label=job_label,
                             compression=compression,
                             num_shards=num_shards)
-    print("TFRecords created. Output stored in {}".format(output_path))
+    print("TFRecords created. Output stored in {}".format(output_dir))

tfrutil/accessor_test.py

@@ -32,9 +32,9 @@ def setUp(self):
 
   def test_accessor(self):
     """Tests pandas accessor."""
-    self.assertIsNone(self.test_df.tensorflow.to_tfr(runner="DirectRunner",
-                                                     output_path="/tmp/train"))
 
+    self.assertIsNone(self.test_df.tensorflow.to_tfr(runner="DirectRunner",
+                                                     output_dir="/tmp/train"))
 
 if __name__ == "__main__":
   unittest.main()

tfrutil/beam_image.py

@@ -24,27 +24,51 @@
 from PIL import Image
 
 
-def load(image_uri):
-  """Loads an image."""
+BASE64_ALTCHARS = b'-_'
 
-  try:
-    with tf.io.gfile.GFile(image_uri, "rb") as f:
-      return Image.open(f)
-  except tf.python.framework.errors_impl.NotFoundError:
-    raise OSError("File {} was not found.".format(image_uri))
 
+def mode_to_channel(mode: str) -> int:
+  """Returns number of channels depending on PIL image `mode`s."""
 
-def encode(image):
+  return 1 if 'L' in mode else 3
+
+
+def channel_to_mode(channels: int) -> str:
+  """Returns PIL image mode depending on number of `channels`."""
+
+  return 'L' if channels == 1 else 'RGB'
+
+
+def encode(image: Image):
   """Returns base64-encoded image data.
 
   Args:
     image: PIL image.
   """
 
-  return base64.b64encode(image.tobytes(), altchars=b"-_")
+  return base64.b64encode(image.tobytes(), altchars=BASE64_ALTCHARS)
+
+
+def decode(b64_bytes, width, height, channels) -> Image:
+  """Decodes an image from base64-encoded data."""
+
+  image_bytes = base64.b64decode(b64_bytes, altchars=BASE64_ALTCHARS)
+  mode = channel_to_mode(channels)
+  return Image.frombytes(mode, (width, height), image_bytes)
+
+
+def load(image_uri):
+  """Loads an image."""
+
+  try:
+    with tf.io.gfile.GFile(image_uri, "rb") as f:
+      return Image.open(f)
+  except tf.python.framework.errors_impl.NotFoundError:
+    raise OSError("File {} was not found.".format(image_uri))
+
 
+# pylint: disable=abstract-method
 
-#pylint: disable=abstract-method
 class ExtractImagesDoFn(beam.DoFn):
   """Adds image to PCollection."""
 
@@ -53,7 +77,7 @@ def __init__(self, image_key: str):
     super().__init__()
     self.image_key = image_key
 
-  #pylint: disable=unused-argument
+  # pylint: disable=unused-argument
   def process(
       self,
       element: Dict,
@@ -70,8 +94,8 @@ def process(
       image_uri = element[self.image_key]
       image = load(image_uri)
       d["image"] = encode(image)
-      d["image_height"], d["image_width"] = image.size
-      d["image_channels"] = 1 if "L" in image.mode else 3
+      d["image_width"], d["image_height"] = image.size
+      d["image_channels"] = mode_to_channel(image.mode)
 
     #pylint: disable=broad-except
     except Exception as e:

tfrutil/beam_image_test.py

@@ -21,7 +21,9 @@
 
 import apache_beam as beam
 from apache_beam.testing import util
+import numpy as np
 import PIL
+from PIL import Image
 import tensorflow_transform as tft
 
 from tfrutil import beam_image
@@ -35,26 +37,51 @@ class BeamImageTests(unittest.TestCase):
   def setUp(self):
     self.pipeline = test_utils.get_test_pipeline()
     self.df = test_utils.get_test_df()
+    self.image_file = "tfrutil/test_data/images/cat/cat-640x853-1.jpg"
 
   def test_load(self):
     """Tests the image loading function."""
-    img = beam_image.load("tfrutil/test_data/images/cat/cat-640x853-1.jpg")
+    img = beam_image.load(self.image_file)
     self.assertIsInstance(img, PIL.JpegImagePlugin.JpegImageFile)
 
   def test_file_not_found_load(self):
     """Test loading an image that doesn"t exist."""
     with self.assertRaises(OSError):
       _ = beam_image.load("tfrutil/test_data/images/cat/food.jpg")
 
+  def test_mode_to_channel(self):
+    """Tests `mode_to_channel`."""
+
+    actual = [beam_image.mode_to_channel(mode)
+              for mode in ('L', 'RGB', 'whatever')]
+    self.assertEqual(actual, [1, 3, 3])
+
+  def test_channel_to_mode(self):
+    """Tests `channel_to_mode`."""
+
+    actual = [beam_image.channel_to_mode(channel) for channel in (1, 2, 3)]
+    self.assertEqual(actual, ['L', 'RGB', 'RGB'])
+
   def test_base64_encode(self):
     """Tests encode function."""
-    img = beam_image.load("tfrutil/test_data/images/cat/cat-640x853-1.jpg")
+    img = beam_image.load(self.image_file)
     enc = beam_image.encode(img)
-    decode = base64.b64decode(enc, altchars=b"-_")
+    decode = base64.b64decode(enc, altchars=beam_image.BASE64_ALTCHARS)
     self.assertEqual(img.tobytes(), decode)
 
+  def test_base64_decode(self):
+    """Tests `decode` function."""
+
+    image = Image.open(self.image_file)
+    b64_bytes = base64.b64encode(
+        image.tobytes(), altchars=beam_image.BASE64_ALTCHARS)
+    width, height = image.size
+    actual = beam_image.decode(b64_bytes, width, height, 3)
+    np.testing.assert_array_equal(np.asarray(actual), np.asarray(image))
+
   def test_extract_image_dofn(self):
     """Tests ExtractImageDoFn."""
+
     with self.pipeline as p:
 
       converter = tft.coders.CsvCoder(constants.IMAGE_CSV_COLUMNS,
@@ -76,9 +103,13 @@ def key_matcher(expected_keys):
         def _equal(actual):
           """ _equal raises a BeamAssertException when an element in the
               PCollection doesn't contain the image extraction keys."""
+          expected_keys_ = set(expected_keys)
           for element in actual:
-            if set(element.keys()) != set(expected_keys):
-              raise util.BeamAssertException("PCollection key match failed.")
+            actual_keys = set(element.keys())
+            if actual_keys != expected_keys_:
+              raise util.BeamAssertException(
+                  "PCollection key match failed. Actual ({}) vs. expected ({})"
+                  .format(actual_keys, expected_keys_))
         return _equal
 
       expected_keys = ["image_uri", "label", "split", "image",

tfrutil/beam_pipeline.py

@@ -18,7 +18,7 @@
 
 This file implements the full beam pipeline for TFRUtil.
 """
-import datetime
+
 import functools
 import logging
 import os
@@ -30,6 +30,7 @@
 from tensorflow_transform import beam as tft_beam
 
 from tfrutil import beam_image
+from tfrutil import common
 from tfrutil import constants
 
 
@@ -44,7 +45,7 @@ def _get_job_name(job_label: str = None) -> str:
       insure uniqueness.
   """
 
-  job_name = "tfrutil-" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
+  job_name = "tfrutil-" + common.get_timestamp()
   if job_label:
     job_label = job_label.replace("_", "-")
     job_name += "-" + job_label
@@ -89,7 +90,7 @@ def _partition_fn(
     index = constants.DISCARD_INDEX
   return index
 
-def _get_write_to_tfrecord(output_path: str,
+def _get_write_to_tfrecord(output_dir: str,
                            prefix: str,
                            compress: bool = True,
                            num_shards: int = 0) \
@@ -99,13 +100,13 @@ def _get_write_to_tfrecord(output_path: str,
   This configures a Beam sink to output TFRecord files.
 
   Args:
-    output_path: Directory to output TFRecord files.
+    output_dir: Directory to output TFRecord files.
     prefix: TFRecord file prefix.
     compress: If True, GZip compress TFRecord files.
     num_shards: Number of file shards to split the TFRecord data.
   """
 
-  path = os.path.join(output_path, prefix)
+  path = os.path.join(output_dir, prefix)
   suffix = '.tfrecord'
   if compress:
     compression_type = 'gzip'
@@ -135,28 +136,29 @@ def _preprocessing_fn(inputs, integer_label: bool = False):
 
 # pylint: disable=too-many-arguments
 # pylint: disable=too-many-locals
-def run_pipeline(df: pd.DataFrame,
-                 job_label: str,
-                 runner: str,
-                 output_path: str,
-                 compression: str,
-                 num_shards: int,
-                 integer_label: bool = False):
+def run_pipeline(
+    df: pd.DataFrame,
+    job_label: str,
+    runner: str,
+    output_dir: str,
+    compression: str,
+    num_shards: int,
+    integer_label: bool = False):
   """Runs TFRUtil Beam Pipeline.
 
   Args:
     df: Pandas Dataframe
     job_label: User description for the beam job.
     runner: Beam Runner: (e.g. DataFlowRunner, DirectRunner).
-    output_path: GCS or Local Path for output.
+    output_dir: GCS or Local Path for output.
     compression: gzip or None.
     num_shards: Number of shards.
 
   Note: These inputs must be validated upstream (by client.create_tfrecord())
   """
 
   job_name = _get_job_name(job_label)
-  job_dir = _get_job_dir(output_path, job_name)
+  job_dir = _get_job_dir(output_dir, job_name)
   popts = {}  # TODO(mikebernico): consider how/if to pass pipeline options.
   options = _get_pipeline_options(job_name, job_dir, **popts)
 
@@ -215,7 +217,7 @@ def run_pipeline(df: pd.DataFrame,
 
       # Sinks for TFRecords and metadata.
       tfr_writer = functools.partial(_get_write_to_tfrecord,
-                                     output_path=job_dir,
+                                     output_dir=job_dir,
                                      compress=compression,
                                      num_shards=num_shards)
 
@@ -239,7 +241,6 @@ def run_pipeline(df: pd.DataFrame,
           | 'DiscardDataWriter' >> beam.io.WriteToText(
               os.path.join(job_dir, "discarded-data")))
 
-
       # Output transform function and metadata
       _ = (transform_fn | 'WriteTransformFn' >> tft_beam.WriteTransformFn(
           job_dir))

tfrutil/beam_pipeline_test.py

@@ -55,7 +55,7 @@ def test_processing_fn_with_string_label(self, mock_transform):
   def test_write_to_tfrecord(self):
     """Test _write_to_tfrecord() fn."""
     tfr_writer = beam_pipeline._get_write_to_tfrecord(
-        output_path='tmp',
+        output_dir='tmp',
         prefix='foo',
         compress=True,
         num_shards=2)