From 9e0298bc753c7a558464ab1bf9de487b0042ee09 Mon Sep 17 00:00:00 2001
From: "michael.divia" <michael.divia@etu.hesge.ch>
Date: Wed, 9 Apr 2025 16:09:39 +0200
Subject: [PATCH] Dynamic ONNX converter
---
Python/convert_onnx.py | 33 ++++++++++++++++++++++++++-------
1 file changed, 26 insertions(+), 7 deletions(-)
diff --git a/Python/convert_onnx.py b/Python/convert_onnx.py
index 29f3050..d59ebff 100644
--- a/Python/convert_onnx.py
+++ b/Python/convert_onnx.py
@@ -1,8 +1,18 @@
import tensorflow as tf
import tf2onnx
+import argparse
+
+# WHAT ?
+parser = argparse.ArgumentParser(description="WHAT ?!")
+parser.add_argument("--model", choices=["1", "2"], default="1",
+ help="1 = ResNet50, 2 = Xception")
+args = parser.parse_args()
# Load Sequential model
-seq_model = tf.keras.models.load_model("../models/resnet50.h5", compile=False)
+if args.model == "1":
+ seq_model = tf.keras.models.load_model("../models/ResNet50/pokemon_resnet50.h5", compile=False)
+elif args.model == "2":
+ seq_model = tf.keras.models.load_model("../models/ResNet50/pokemon_xception.h5", compile=False)
# Create input layer with same shape
inputs = tf.keras.Input(shape=(224, 224, 3), name="input")
@@ -15,9 +25,18 @@ model = tf.keras.Model(inputs=inputs, outputs=outputs)
# Convert to ONNX
spec = (tf.TensorSpec((1, 224, 224, 3), tf.float32, name="input"),)
-onnx_model, _ = tf2onnx.convert.from_keras(
- model,
- input_signature=spec,
- opset=13,
- output_path="../models/resnet50.onnx"
-)
\ No newline at end of file
+
+if args.model == "1":
+ onnx_model, _ = tf2onnx.convert.from_keras(
+ model,
+ input_signature=spec,
+ opset=13,
+ output_path="../models/ResNet50/pokemon_resnet50.onnx"
+ )
+elif args.model == "2":
+ onnx_model, _ = tf2onnx.convert.from_keras(
+ model,
+ input_signature=spec,
+ opset=13,
+ output_path="../models/ResNet50/pokemon_xception.onnx"
+ )
\ No newline at end of file
--
GitLab