From 7daacca2317f3e61dd559e5722d7c54539f76c18 Mon Sep 17 00:00:00 2001 From: Luuk Jensen Date: Mon, 17 Jun 2019 15:39:03 +0200 Subject: [PATCH 1/3] Add function to export to tflite --- aocr/__main__.py | 2 +- aocr/util/export.py | 19 +++++++++++++++++++ setup.py | 2 +- 3 files changed, 21 insertions(+), 2 deletions(-) diff --git a/aocr/__main__.py b/aocr/__main__.py index b55ab140..53f19748 100644 --- a/aocr/__main__.py +++ b/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))) diff --git a/aocr/util/export.py b/aocr/util/export.py index 57307cc7..60f8e4aa 100644 --- a/aocr/util/export.py +++ b/aocr/util/export.py @@ -51,3 +51,22 @@ 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') + output = { + 'output': graph.get_tensor_by_name('prediction:0'), + 'probability': graph.get_tensor_by_name('probability:0') + } + + 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') \ No newline at end of file diff --git a/setup.py b/setup.py index 3474f180..78d70302 100644 --- a/setup.py +++ b/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') From 6f2ba4274fecb6d669afeae8120add789e7d8f8a Mon Sep 17 00:00:00 2001 From: Luuk Jensen Date: Tue, 18 Jun 2019 10:21:59 +0200 Subject: [PATCH 2/3] Add shape to input tensor for export to tf lite Fixed output tensors for export to tf lite Added width, height and channels parameters to the Exporter.save() method to create a tensor shape. --- aocr/__main__.py | 2 +- aocr/util/export.py | 13 +++++++------ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/aocr/__main__.py b/aocr/__main__.py index 53f19748..e2283242 100644 --- a/aocr/__main__.py +++ b/aocr/__main__.py @@ -273,7 +273,7 @@ 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 diff --git a/aocr/util/export.py b/aocr/util/export.py index 60f8e4aa..4ecfd1bc 100644 --- a/aocr/util/export.py +++ b/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.") @@ -60,12 +60,13 @@ def save(self, path, model_format): graph = self.model.sess.graph input = graph.get_tensor_by_name('input_image_as_bytes:0') - output = { - 'output': graph.get_tensor_by_name('prediction:0'), - 'probability': graph.get_tensor_by_name('probability: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]) + converter = tf.lite.TFLiteConverter.from_session(self.model.sess, [input], output) tflite_model = converter.convert() open("saved_model.tflite", "wb").write(tflite_model) From 4e98596ca5730756a9e3d0925c21bd435097edf2 Mon Sep 17 00:00:00 2001 From: stickler-ci Date: Tue, 18 Jun 2019 08:27:03 +0000 Subject: [PATCH 3/3] Fixing style errors. --- aocr/__main__.py | 3 ++- aocr/util/export.py | 6 +++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/aocr/__main__.py b/aocr/__main__.py index e2283242..9ad1c213 100644 --- a/aocr/__main__.py +++ b/aocr/__main__.py @@ -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, parameters.max_width, parameters.max_height, parameters.channels) + exporter.save(parameters.export_path, parameters.format, + parameters.max_width, parameters.max_height, parameters.channels) return else: raise NotImplementedError diff --git a/aocr/util/export.py b/aocr/util/export.py index 4ecfd1bc..f7b48be9 100644 --- a/aocr/util/export.py +++ b/aocr/util/export.py @@ -60,14 +60,14 @@ def save(self, path, model_format, max_width, max_height, channels): graph = self.model.sess.graph input = graph.get_tensor_by_name('input_image_as_bytes:0') - shape=[1,max_height,max_width,channels] + 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] + 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') \ No newline at end of file + logging.info("Exported as %s", path + '/saved_model.tflite')