[WIP] Export to Tensorflow Lite for Edge TPU

aocr/__main__.py

@@ -182,7 +182,7 @@ def process_args(args, defaults):
                                      % (defaults.EXPORT_PATH)))
     parser_export.add_argument('--format', dest="format",
                                type=str, default=defaults.EXPORT_FORMAT,
-                               choices=['frozengraph', 'savedmodel'],
+                               choices=['frozengraph', 'savedmodel', 'lite'],
                                help=('export format'
                                      ' (default: %s)'
                                      % (defaults.EXPORT_FORMAT)))
@@ -273,7 +273,8 @@ def main(args=None):
                 logging.info('Result: OK. %s %s', '{:.2f}'.format(probability), text)
         elif parameters.phase == 'export':
             exporter = Exporter(model)
-            exporter.save(parameters.export_path, parameters.format)
+            exporter.save(parameters.export_path, parameters.format,
+                          parameters.max_width, parameters.max_height, parameters.channels)
             return
         else:
             raise NotImplementedError

aocr/util/export.py

@@ -10,7 +10,7 @@ class Exporter(object):
     def __init__(self, model):
         self.model = model
 
-    def save(self, path, model_format):
+    def save(self, path, model_format, max_width, max_height, channels):
         if model_format == "savedmodel":
             logging.info("Creating a SavedModel.")
 
@@ -51,3 +51,23 @@ def save(self, path, model_format):
                 outfile.write(output_graph_def.SerializeToString())
 
             logging.info("Exported as %s", path + '/frozen_graph.pb')
+        elif model_format == "lite":
+            logging.info("Creating Tensorflow Lite graph.")
+
+            if not os.path.exists(path):
+                os.makedirs(path)
+
+            graph = self.model.sess.graph
+
+            input = graph.get_tensor_by_name('input_image_as_bytes:0')
+            shape = [1, max_height, max_width, channels]
+            input.set_shape(shape)
+            prediction_output = graph.get_tensor_by_name('prediction:0')
+            probability_output = graph.get_tensor_by_name('probability:0')
+            output = [prediction_output, probability_output]
+
+            converter = tf.lite.TFLiteConverter.from_session(self.model.sess, [input], output)
+            tflite_model = converter.convert()
+            open("saved_model.tflite", "wb").write(tflite_model)
+
+            logging.info("Exported as %s", path + '/saved_model.tflite')

setup.py

@@ -2,7 +2,7 @@
 from setuptools import setup
 
 REQUIRED_PACKAGES = ['distance', 'numpy', 'six', 'pillow']
-VERSION = '0.7.6'
+VERSION = '0.7.7'
 try:
     import pypandoc
     README = pypandoc.convert('README.md', 'rst')