Skip to content
Snippets Groups Projects
Commit fb637e01 authored by michael.divia's avatar michael.divia
Browse files

And now ?

parent 8472b806
No related branches found
No related tags found
No related merge requests found
import json import json
import argparse import argparse
import os
import numpy as np import numpy as np
from picamera2 import Picamera2 from picamera2 import Picamera2
from picamera2.devices.hailo import Hailo from picamera2.devices.hailo import Hailo
import cv2 import cv2
import os
# --- Argparse --- # --- Argparse ---
parser = argparse.ArgumentParser(description="Pokémon Classifier Inference with Hailo-8") parser = argparse.ArgumentParser(description="Pokémon Classifier Inference with Hailo-8")
parser.add_argument("--model", choices=["1", "2"], required=True, help="1 = ResNet50, 2 = Xception") parser.add_argument("--model", choices=["1", "2"], required=True, help="1 = ResNet50, 2 = Xception")
args = parser.parse_args() args = parser.parse_args()
# --- Paths --- # --- Paths & Input Sizes ---
if args.model == "1": if args.model == "1":
hef_path = "../models/ResNet50/pokedex_ResNet50.hef" hef_path = "../models/ResNet50/pokedex_ResNet50.hef"
json_path = "../models/ResNet50/class_names.json" json_path = "../models/ResNet50/class_names.json"
size=(224, 224) input_shape = (224, 224)
elif args.model == "2": elif args.model == "2":
hef_path = "../models/Xception/pokedex_Xception.hef" hef_path = "../models/Xception/pokedex_Xception.hef"
json_path = "../models/Xception/class_names.json" json_path = "../models/Xception/class_names.json"
size = (256, 256) input_shape = (256, 256)
else: else:
raise ValueError("Invalid model selection") raise ValueError("Invalid model selection")
...@@ -27,15 +27,13 @@ else: ...@@ -27,15 +27,13 @@ else:
with open(json_path, "r") as f: with open(json_path, "r") as f:
class_names = json.load(f) class_names = json.load(f)
# --- Inference --- # --- Hailo Inference with PiCamera2 ---
with Hailo(hef_path) as hailo: with Hailo(hef_path) as hailo:
print(f"-- Hailo model input size: {size}") print(f"-- Model input shape: {input_shape}")
picam2 = Picamera2() picam2 = Picamera2()
# Use *raw* capture configuration with correct size
config = picam2.create_still_configuration( config = picam2.create_still_configuration(
main={"size": size, "format": "RGB888"}, main={"size": input_shape, "format": "RGB888"},
lores=None, lores=None,
display=None display=None
) )
...@@ -46,7 +44,7 @@ with Hailo(hef_path) as hailo: ...@@ -46,7 +44,7 @@ with Hailo(hef_path) as hailo:
frame = picam2.capture_array() frame = picam2.capture_array()
print(f"-- Captured frame shape: {frame.shape}") print(f"-- Captured frame shape: {frame.shape}")
# Optional: display image # Optional: show image
try: try:
cv2.imshow("Captured Image", frame) cv2.imshow("Captured Image", frame)
print("-- Press any key to continue...") print("-- Press any key to continue...")
...@@ -58,8 +56,14 @@ with Hailo(hef_path) as hailo: ...@@ -58,8 +56,14 @@ with Hailo(hef_path) as hailo:
cv2.imwrite(output_path, frame) cv2.imwrite(output_path, frame)
os.system(f"feh --fullscreen {output_path}") os.system(f"feh --fullscreen {output_path}")
# --- Preprocess (normalize + layout + batch) ---
image = frame.astype(np.float32) # Convert to float32
image -= [123.68, 116.779, 103.939] # Subtract ImageNet mean
image = np.transpose(image, (2, 0, 1)) # HWC → CHW
image = np.expand_dims(image, axis=0) # Add batch dimension
print("-- Running inference...") print("-- Running inference...")
inference_results = hailo.run(frame) inference_results = hailo.run(image)
predicted_idx = int(np.argmax(inference_results)) predicted_idx = int(np.argmax(inference_results))
predicted_name = class_names[predicted_idx] predicted_name = class_names[predicted_idx]
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment