Added TPU Estimator Prediction

Author: captain-pool

E1_TPU_Sample/image_retraining_tpu.py

@@ -23,11 +23,18 @@
     "horses_or_humans",
     "TFDS Dataset Name. IMAGE Dimension should be >= 224, channel=3")
 flags.DEFINE_string("data_dir", None, "Directory to Save Data to")
+flags.DEFINE_string("infer", None, "Dummy image file to infer")
 
 FLAGS = flags.FLAGS
 NUM_CLASSES = None
 
 
+def resize_and_scale(image, label):
+  image = tf.image.resize(image, size=[224, 224])
+  image = tf.cast(image, tf.float32)
+  image = image / tf.reduce_max(tf.gather(image, 0))
+  return image, label
+
 def input_(mode, batch_size, iterations, **kwargs):
   global NUM_CLASSES
   dataset, info = tfds.load(
@@ -38,13 +45,6 @@ def input_(mode, batch_size, iterations, **kwargs):
       data_dir=kwargs['data_dir']
   )
   NUM_CLASSES = info.features['label'].num_classes
-
-  def resize_and_scale(image, label):
-    image = tf.image.resize(image, size=[224, 224])
-    image = tf.cast(image, tf.float32)
-    image = image / tf.reduce_max(tf.gather(image, 0))
-    return image, label
-
   dataset = dataset.map(resize_and_scale).shuffle(
       1000).repeat(iterations).batch(batch_size, drop_remainder=True)
   return dataset
@@ -135,9 +135,16 @@ def main(_):
       input_fn=lambda params: input_fn(
           mode=tf.estimator.ModeKeys.TRAIN,
           **params),
-      max_steps=1000)
+      max_steps=None, steps=None)
   # TODO(@captain-pool): Implement Evaluation
-
+  if FLAGS.infer:
+    def prepare_input_fn(path):
+      img = tf.image.decode_image(tf.io.read_file(path))
+      return resize_and_scale(img, None)
+
+    predictions = classifer.predict(
+        input_fn=lambda params: prepare_input_fn(FLAGS.infer))
+    print(predictions)
 
 if __name__ == "__main__":
   app.run(main)