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" size=(224,224) elif args.model == "2": hef_path = "../models/Xception/pokedex_Xception.hef" json_path = "../models/Xception/class_names.json" size=(256,256) # Load class names with open(json_path, "r") as f: class_names = json.load(f) device = HailoRT.Device() hef = HailoRT.Hef(hef_path) configured_network_group = device.create_hef_group(hef) input_vstream_info = configured_network_group.get_input_vstream_infos()[0] output_vstream_info = configured_network_group.get_output_vstream_infos()[0] # --- Open webcam --- cap = cv2.VideoCapture(0) if not cap.isOpened(): print("-- Unable to open webcam") exit() print("-- Taking picture...") ret, frame = cap.read() cap.release() if not ret: print("-- Failed to capture image") exit() # --- Preprocess image --- image = cv2.resize(frame, size) image = image.astype(np.float32) / 255.0 # Normalize to [0, 1] image = np.expand_dims(image, axis=0) # Add batch dimension image = np.transpose(image, (0, 3, 1, 2)) # NHWC ? NCHW if required (check your model) # --- Inference --- with HailoRT.VirtualStreams(input_vstream_info, output_vstream_info, configured_network_group) as (input_vstreams, output_vstreams): input_vstreams[0].send(image) output_data = output_vstreams[0].recv() # --- Postprocess --- predicted_idx = int(np.argmax(output_data)) predicted_name = class_names[predicted_idx] print(f"-- Predicted Pokémon: {predicted_name}")