One fix in Image package and one fix in google.datalab.ml (#359)

* Fix an issue in image classification package: trainer incorrectly depend on DataFlow.
Fix an issue in google.datalab.ml.Summary and google.datalab.FeatureSliceView.

* Remove unused import in image classification package.

* Fix trailing whitespace and indent.

* Fix indent.

Author: qimingj

google/datalab/ml/_feature_slice_view.py

@@ -58,10 +58,8 @@ def plot(self, data):
     """
     import IPython
 
-    if sys.version_info.major > 2:
-      basestring = (str, bytes)  # for python 3 compatibility
-
-    if isinstance(data, basestring):
+    if ((sys.version_info.major > 2 and isinstance(data, str)) or
+       (sys.version_info.major <= 2 and isinstance(data, basestring))):
       data = bq.Query(data)
 
     if isinstance(data, bq.Query):

google/datalab/ml/_summary.py

@@ -95,10 +95,9 @@ def get_events(self, event_names):
           events (i.e. 'loss' under trains_set/, and 'loss' under eval_set/.)
     """
 
-    if sys.version_info.major > 2:
-      basestring = (str, bytes)  # for python 3 compatibility
-
-    event_names = [event_names] if isinstance(event_names, basestring) else event_names
+    if ((sys.version_info.major > 2 and isinstance(event_names, str)) or
+       (sys.version_info.major <= 2 and isinstance(event_names, basestring))):
+      event_names = [event_names]
 
     all_events = self.list_events()
     dirs_to_look = set()
@@ -152,10 +151,10 @@ def plot(self, event_names, x_axis='step'):
       x_axis: whether to use step or time as x axis.
     """
 
-    if sys.version_info.major > 2:
-      basestring = (str, bytes)  # for python 3 compatibility
+    if ((sys.version_info.major > 2 and isinstance(event_names, str)) or
+       (sys.version_info.major <= 2 and isinstance(event_names, basestring))):
+      event_names = [event_names]
 
-    event_names = [event_names] if isinstance(event_names, basestring) else event_names
     events_list = self.get_events(event_names)
     for event_name, dir_event_dict in zip(event_names, events_list):
       for dir, df in dir_event_dict.iteritems():

solutionbox/image_classification/mltoolbox/image/classification/_cloud.py

@@ -16,15 +16,13 @@
 """Cloud implementation for preprocessing, training and prediction for inception model.
 """
 
-import apache_beam as beam
+
 import base64
 import datetime
 import logging
 import os
 import urllib
 
-from . import _predictor
-from . import _preprocess
 from . import _util
 
 
@@ -39,7 +37,9 @@ class Cloud(object):
   def preprocess(train_dataset, output_dir, eval_dataset, checkpoint, pipeline_option):
     """Preprocess data in Cloud with DataFlow."""
 
+    import apache_beam as beam
     import google.datalab.utils
+    from . import _preprocess
 
     if checkpoint is None:
       checkpoint = _util._DEFAULT_CHECKPOINT_GSURL
@@ -154,7 +154,9 @@ def predict(model_id, image_files, resize, show_image):
   def batch_predict(dataset, model_dir, output_csv, output_bq_table, pipeline_option):
     """Batch predict running in cloud."""
 
+    import apache_beam as beam
     import google.datalab.utils
+    from . import _predictor
 
     if output_csv is None and output_bq_table is None:
       raise ValueError('output_csv and output_bq_table cannot both be None.')

solutionbox/image_classification/mltoolbox/image/classification/_local.py

@@ -16,13 +16,11 @@
 """Local implementation for preprocessing, training and prediction for inception model.
 """
 
-import apache_beam as beam
+
 import datetime
 
 
 from . import _model
-from . import _predictor
-from . import _preprocess
 from . import _trainer
 from . import _util
 
@@ -34,7 +32,9 @@ class Local(object):
   def preprocess(train_dataset, output_dir, eval_dataset, checkpoint):
     """Preprocess data locally."""
 
+    import apache_beam as beam
     from google.datalab.utils import LambdaJob
+    from . import _preprocess
 
     if checkpoint is None:
       checkpoint = _util._DEFAULT_CHECKPOINT_GSURL
@@ -70,6 +70,8 @@ def train(input_dir, batch_size, max_steps, output_dir, checkpoint):
   def predict(model_dir, image_files, resize, show_image):
     """Predict using an model in a local or GCS directory."""
 
+    from . import _predictor
+
     images = _util.load_images(image_files, resize=resize)
     labels_and_scores = _predictor.predict(model_dir, images)
     results = zip(image_files, images, labels_and_scores)
@@ -80,7 +82,9 @@ def predict(model_dir, image_files, resize, show_image):
   def batch_predict(dataset, model_dir, output_csv, output_bq_table):
     """Batch predict running locally."""
 
+    import apache_beam as beam
     from google.datalab.utils import LambdaJob
+    from . import _predictor
 
     if output_csv is None and output_bq_table is None:
       raise ValueError('output_csv and output_bq_table cannot both be None.')

solutionbox/image_classification/mltoolbox/image/classification/_predictor.py

@@ -95,7 +95,9 @@ class LoadImagesDoFn(beam.DoFn):
   """A DoFn that reads image from url."""
 
   def process(self, element):
-    with _util.open_local_or_gcs(element['image_url'], 'r') as ff:
+    from tensorflow.python.lib.io import file_io as tf_file_io
+
+    with tf_file_io.FileIO(element['image_url'], 'r') as ff:
       image_bytes = ff.read()
     out_element = {'image_bytes': image_bytes}
     out_element.update(element)

solutionbox/image_classification/mltoolbox/image/classification/_preprocess.py

@@ -87,10 +87,11 @@ class ReadImageAndConvertToJpegDoFn(beam.DoFn):
   """
 
   def process(self, element):
-    uri, label_id = element
+    from tensorflow.python.lib.io import file_io as tf_file_io
 
+    uri, label_id = element
     try:
-      with _util.open_local_or_gcs(uri, mode='r') as f:
+      with tf_file_io.FileIO(uri, 'r') as f:
         img = Image.open(f).convert('RGB')
     # A variety of different calling libraries throw different exceptions here.
     # They all correspond to an unreadable file so we treat them equivalently.

solutionbox/image_classification/mltoolbox/image/classification/_util.py

@@ -16,13 +16,9 @@
 """Reusable utility functions.
 """
 
-from apache_beam.io import gcsio
 import collections
-import glob
 import multiprocessing
 import os
-import logging
-import time
 
 import tensorflow as tf
 from tensorflow.python.lib.io import file_io
@@ -44,41 +40,12 @@ def default_project():
   return Context.default().project_id
 
 
-def open_local_or_gcs(path, mode):
-  """Opens the given path."""
-  if path.startswith('gs://'):
-    try:
-      return gcsio.GcsIO().open(path, mode)
-    except Exception as e:  # pylint: disable=broad-except
-      # Currently we retry exactly once, to work around flaky gcs calls.
-      logging.error('Retrying after exception reading gcs file: %s', e)
-      time.sleep(10)
-      return gcsio.GcsIO().open(path, mode)
-  else:
-    return open(path, mode)
-
-
-def file_exists(path):
-  """Returns whether the file exists."""
-  if path.startswith('gs://'):
-    return gcsio.GcsIO().exists(path)
-  else:
-    return os.path.exists(path)
-
-
-def glob_files(path):
-  if path.startswith('gs://'):
-    return gcsio.GcsIO().glob(path)
-  else:
-    return glob.glob(path)
-
-
 def _get_latest_data_dir(input_dir):
   latest_file = os.path.join(input_dir, 'latest')
-  if not file_exists(latest_file):
+  if not file_io.file_exists(latest_file):
     raise Exception(('Cannot find "latest" file in "%s". ' +
                     'Please use a preprocessing output dir.') % input_dir)
-  with open_local_or_gcs(latest_file, 'r') as f:
+  with file_io.FileIO(latest_file, 'r') as f:
     dir_name = f.read().rstrip()
   return os.path.join(input_dir, dir_name)
 
@@ -88,16 +55,16 @@ def get_train_eval_files(input_dir):
   data_dir = _get_latest_data_dir(input_dir)
   train_pattern = os.path.join(data_dir, 'train*.tfrecord.gz')
   eval_pattern = os.path.join(data_dir, 'eval*.tfrecord.gz')
-  train_files = glob_files(train_pattern)
-  eval_files = glob_files(eval_pattern)
+  train_files = file_io.get_matching_files(train_pattern)
+  eval_files = file_io.get_matching_files(eval_pattern)
   return train_files, eval_files
 
 
 def get_labels(input_dir):
   """Get a list of labels from preprocessed output dir."""
   data_dir = _get_latest_data_dir(input_dir)
   labels_file = os.path.join(data_dir, 'labels')
-  with open_local_or_gcs(labels_file, mode='r') as f:
+  with file_io.FileIO(labels_file, 'r') as f:
     labels = f.read().rstrip().split('\n')
   return labels
 
@@ -254,7 +221,7 @@ def load_images(image_files, resize=True):
 
   images = []
   for image_file in image_files:
-    with open_local_or_gcs(image_file, 'r') as ff:
+    with file_io.FileIO(image_file, 'r') as ff:
       images.append(ff.read())
   if resize is False:
     return images