Standartise everything

Python/convert_onnx.py → python/convert_onnx.py

@@ -2,41 +2,36 @@ import tensorflow as tf
 import tf2onnx
 import argparse
 
-# WHAT ?
+# --- WHAT ? ---
 parser = argparse.ArgumentParser(description="WHAT ?!")
-parser.add_argument("--model", choices=["1", "2"], default="1",
-                    help="1 = ResNet50, 2 = Xception")
+parser.add_argument("--model", choices=["1", "2"], required=True, help="1 = ResNet50, 2 = Xception")
 args = parser.parse_args()
 
-# Load Sequential model
+# Paths
 if args.model == "1":
-    seq_model = tf.keras.models.load_model("../models/ResNet50/pokemon_resnet50.h5", compile=False)
+    h5_path = "../models/ResNet50/pokedex_ResNet50.h5"
+    onnx_path = "../models/ResNet50/pokedex_ResNet50.onnx"
 elif args.model == "2":
-    seq_model = tf.keras.models.load_model("../models/ResNet50/pokemon_xception.h5", compile=False)
+    h5_path = "../models/Xception/pokedex_Xception.h5"
+    onnx_path = "../models/ResNet50/pokedex_Xception.onnx"
 
-# Create input layer with same shape
+# --- Load Sequential model ---
+seq_model = tf.keras.models.load_model(h5_path, compile=False)
+
+# --- Create input layer with same shape ---
 inputs = tf.keras.Input(shape=(224, 224, 3), name="input")
 
-# Call the Sequential model as a function
+# --- Call the Sequential model as a function ---
 outputs = seq_model(inputs)
 
-# Wrap in Functional model
+# --- Wrap in Functional model ---
 model = tf.keras.Model(inputs=inputs, outputs=outputs)
 
-# Convert to ONNX
+# --- Convert to ONNX ---
 spec = (tf.TensorSpec((1, 224, 224, 3), tf.float32, name="input"),)
-
-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"
+    output_path=onnx_path
 )
\ No newline at end of file

python/hailo_compile.py

+import os
+import subprocess
+import argparse
+import onnx
+
+parser = argparse.ArgumentParser(description="Convert, parse, optimize and compile model for Hailo.")
+parser.add_argument("--model", choices=["1", "2"], required=True, help="1 = ResNet50, 2 = Xception")
+parser.add_argument("--calib", default="--use-random-calib-set", help="Calibration data path or '--use-random-calib-set'")
+args = parser.parse_args()
+
+def get_onnx_io_names(onnx_path):
+    model = onnx.load(onnx_path)
+    input_name = model.graph.input[0].name
+    output_name = model.graph.output[0].name
+    return input_name, output_name
+
+# Paths
+if args.model == "1":
+    base_model_path = "../models/ResNet50"
+    model = "pokedex_ResNet50"
+elif args.model == "2":
+    base_model_path = "../models/Xception"
+    model = "pokedex_Xception"
+    
+onnx_path = os.path.join(base_model_path, f"pokedex_{model}.onnx")
+har_path = os.path.join(base_model_path, f"{model}.har")
+optimized_har_path = os.path.join(base_model_path, f"{model}_optimized.har")
+hef_path = os.path.join(base_model_path, f"{model}.hef")
+
+# Node names
+start_node, end_node = get_onnx_io_names(onnx_path)
+print(f"-- Using start_node: {start_node}, end_node: {end_node}")
+
+# Step 1: Parse
+print(f"-- Parsing {onnx_path}...")
+subprocess.run([
+    "hailo", "parser", "onnx", onnx_path,
+    "--start-node-names", start_node,
+    "--end-nodes-names", end_node,
+    "--hw-arch", "hailo8l"
+])
+
+# Step 2: Optimize
+print(f"-- Optimizing to {optimized_har_path}...")
+optimize_cmd = [
+    "hailo", "optimize", har_path,
+    "--hw-arch", "hailo8l",
+    "--use-random-calib-set"
+]
+subprocess.run(optimize_cmd)
+
+# Step 3: Compile
+print(f"-- Compiling to {hef_path}...")
+subprocess.run([
+    "hailo", "compiler", optimized_har_path,
+    "--hw-arch", "hailo8l",
+    "--performance",
+])
+print("-- Done.")

Python/pokedex_ResNet50.py → python/pokedex_ResNet50.py

@@ -117,7 +117,7 @@ model.fit(
 )
 
 # --- Save the model ---
-model_h5_path = os.path.join(model_output_path, "pokemon_resnet50.h5")
+model_h5_path = os.path.join(model_output_path, "pokedex_ResNet50.h5")
 model.save(model_h5_path)
 print(f"Model saved to {model_h5_path}")
 

Python/pokedex_Xception.py → python/pokedex_Xception.py

@@ -135,7 +135,7 @@ model.fit(
 )
 
 # --- Save the model ---
-model_h5_path = os.path.join(model_output_path, "pokemon_xception.h5")
+model_h5_path = os.path.join(model_output_path, "pokedex_Xception.h5")
 model.save(model_h5_path)
 print(f"Model saved to {model_h5_path}")
 

Python/predict_pokemon.py → python/pokedex_rpi.py

@@ -2,13 +2,25 @@ import cv2
 import numpy as np
 import json
 from hailo_platform.pyhailort import HailoRT
+import argparse
+
+# --- WHAT ? ---
+parser = argparse.ArgumentParser(description="WHAT ?!")
+parser.add_argument("--model", choices=["1", "2"], required=True, help="1 = ResNet50, 2 = Xception")
+args = parser.parse_args()
+
+# Paths
+if args.model == "1":
+    hef_path = "../models/ResNet50/pokedex_ResNet50.hef"
+    json_path = "../models/ResNet50/class_names.json"
+elif args.model == "2":
+    hef_path = "../models/Xception/pokedex_Xception.hef"
+    json_path = "../models/Xception/class_names.json"
 
 # Load class names
-with open("../models/ResNet50/class_names.json", "r") as f:
+with open(json_path, "r") as f:
     class_names = json.load(f)
 
-# --- Load HEF and configure device ---
-hef_path = "../models/ResNet50/resnet50.hef"
 device = HailoRT.Device()
 hef = HailoRT.Hef(hef_path)
 configured_network_group = device.create_hef_group(hef)

Python/test_ResNet50.py → python/pokedex_test.py

@@ -5,19 +5,33 @@ import numpy as np
 import os
 import random
 import json
+import argparse
 
-# === Load class names from JSON ===
-with open("../models/ResNet50/class_names.json", "r") as f:
+# --- WHAT ? ---
+parser = argparse.ArgumentParser(description="WHAT ?!")
+parser.add_argument("--model", choices=["1", "2"], required=True, help="1 = ResNet50, 2 = Xception")
+args = parser.parse_args()
+
+# Paths
+if args.model == "1":
+    h5_path = "../models/ResNet50/pokedex_ResNet50.h5"
+    json_path = "../models/ResNet50/class_names.json"
+elif args.model == "2":
+    h5_path = "../models/Xception/pokedex_Xception.h5"
+    json_path = "../models/Xception/class_names.json"
+
+# --- Load class names from JSON ---
+with open(json_path, "r") as f:
     class_names = json.load(f)
     class_names = [class_names[i] for i in range(len(class_names))]  # convert to list
 
-# === Load trained model ===
-model = keras.models.load_model("../models/ResNet50/pokemon_resnet50.h5")
+# --- Load trained model ---
+model = keras.models.load_model(h5_path)
 
-# === Paths ===
+# --- Paths ---
 base_path = "../Combined_Dataset"
 
-# === Prepare 2x2 Plot ===
+# --- Prepare 2x2 Plot ---
 plt.figure(figsize=(10, 10))
 
 for i in range(4):
@@ -30,13 +44,13 @@ for i in range(4):
     ])
     img_path = os.path.join(class_folder, random_image)
 
-    # === Load & Preprocess Image ===
+    # --- Load & Preprocess Image ---
     img = keras.utils.load_img(img_path, target_size=(224, 224))  # resize to match model input
     img_array = keras.utils.img_to_array(img)
     img_array = img_array / 255.0  # normalize if your model expects it
     img_array = tf.expand_dims(img_array, 0)
 
-    # === Predict ===
+    # --- Predict ---
     predictions = model.predict(img_array, verbose=0)
     probabilities = tf.nn.softmax(predictions[0])
     predicted_class_index = np.argmax(probabilities)
@@ -46,7 +60,7 @@ for i in range(4):
     # Compare with actual
     is_correct = predicted_label == random_class
 
-    # === Plot ===
+    # --- Plot ---
     ax = plt.subplot(2, 2, i + 1)
     plt.imshow(img)
     plt.axis("off")