Merge pull request #161 from SthPhoenix/add_onnx_export

Add export to ONNX

model2onnx.py

@@ -0,0 +1,32 @@
+
+# import the necessary packages
+from tensorflow.keras.models import load_model, save_model
+import argparse
+import tf2onnx
+import onnx
+
+def model2onnx():
+ # construct the argument parser and parse the arguments
+ ap = argparse.ArgumentParser()
+ ap.add_argument("-m", "--model", type=str,
+ default="mask_detector.model",
+ help="path to trained face mask detector model")
+ ap.add_argument("-o", "--output", type=str,
+ default='mask_detector.onnx',
+ help="path to trained face mask detector model")
+ args = vars(ap.parse_args())
+
+
+ # load the face mask detector model from disk
+ print("[INFO] loading face mask detector model...")
+ model = load_model(args["model"])
+ onnx_model, _ = tf2onnx.convert.from_keras(model, opset=13)
+
+ onnx_model.graph.input[0].type.tensor_type.shape.dim[0].dim_param = '?'
+ onnx_model.graph.output[0].type.tensor_type.shape.dim[0].dim_param = '?'
+
+ onnx.save(onnx_model, args['output'])
+
+
+if __name__ == "__main__":
+ model2onnx()

requirements.txt

@@ -9,3 +9,5 @@ scipy==1.6.2
 scikit-learn==0.24.1
 pillow>=8.3.2
 streamlit==0.79.0
+onnx==1.10.1
+tf2onnx==1.9.3