1 to rule them all

python/pokedex_rpi.py

@@ -15,9 +15,11 @@ args = parser.parse_args()
 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)
 else:
     raise ValueError("Invalid model selection")
 
@@ -27,14 +29,13 @@ with open(json_path, "r") as f:
 
 # --- Inference ---
 with Hailo(hef_path) as hailo:
-    model_h, model_w, _ = hailo.get_input_shape()
-    print(f"-- Hailo model input size: {model_w}x{model_h}")
+    print(f"-- Hailo model input size: {size}")
 
     picam2 = Picamera2()
 
     # Use *raw* capture configuration with correct size
     config = picam2.create_still_configuration(
-        main={"size": (model_w, model_h), "format": "RGB888"},
+        main={"size": size, "format": "RGB888"},
         lores=None,
         display=None
     )
@@ -45,10 +46,6 @@ with Hailo(hef_path) as hailo:
     frame = picam2.capture_array()
     print(f"-- Captured frame shape: {frame.shape}")
 
-    if frame.shape[0] != model_h or frame.shape[1] != model_w:
-        print("-- Frame shape doesn't match model input size. Aborting.")
-        exit(1)
-
     # Optional: display image
     try:
         cv2.imshow("Captured Image", frame)