Update app.py

app.py

@@ -8,30 +8,34 @@ import onnxruntime
 from pathlib import Path
 from ultralytics import YOLO
 
-# Load YOLOv5 model for ONNX export
-model = YOLO("yolov5n.pt")  
-
-# Export to ONNX format
-model.export(format="onnx", dynamic=True)
-
-os.makedirs("models", exist_ok=True)
+# Load YOLOv5 model without AutoShape
+model = torch.hub.load("ultralytics/yolov5", "yolov5n", source="local")
 
+# Set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-print(f"Using device: {device}")
+model.to(device)
+model.eval()
 
-model_path = Path("models/yolov5n.onnx")
+# Fuse layers for optimization
+model.fuse()
 
-if not model_path.exists():
-    print("Downloading YOLOv5n model and converting to ONNX...")
-    model = torch.hub.load("ultralytics/yolov5", "yolov5n", pretrained=True).to(device)
-    model.eval()
-    
-    # Exporting to ONNX
-    model.export(format="onnx", dynamic=True)
-    os.rename("yolov5n.onnx", model_path)
-    del model  # Free memory
+# Export to ONNX format
+os.makedirs("models", exist_ok=True)
+model_path = Path("models/yolov5n.onnx")
 
-# Loading ONNX model for ultra-fast inference
+torch.onnx.export(
+    model,
+    torch.zeros(1, 3, 640, 640).to(device),  # Input tensor
+    str(model_path),
+    export_params=True,
+    opset_version=11,
+    do_constant_folding=True,
+    input_names=["images"],
+    output_names=["output"],
+    dynamic_axes={"images": {0: "batch_size"}, "output": {0: "batch_size"}}
+)
+
+# Load ONNX model for inference
 session = onnxruntime.InferenceSession(str(model_path), providers=['CUDAExecutionProvider'])
 
 # Generate random colors for each class
@@ -50,7 +54,7 @@ def detect_objects(image):
     
     # Preprocess image
     original_shape = image.shape
-    input_shape = (416, 416)
+    input_shape = (640, 640)
     image_resized = cv2.resize(image, input_shape)
     image_norm = image_resized.astype(np.float32) / 255.0
     image_transposed = np.transpose(image_norm, (2, 0, 1))