Update app.py

app.py

@@ -9,12 +9,11 @@ from pathlib import Path
 from ultralytics import YOLO
 
 # Load YOLOv5 model for ONNX export
-model = YOLO("yolov5n.pt")  # Use "yolov5x.pt" if you want the larger model
+model = YOLO("yolov5n.pt")  
 
 # Export to ONNX format
 model.export(format="onnx", dynamic=True)
 
-
 os.makedirs("models", exist_ok=True)
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -44,49 +43,61 @@ inference_count = 0
 
 def detect_objects(image):
     global total_inference_time, inference_count
-
     if image is None:
         return None
-
+        
     start_time = time.time()
-
-    image = cv2.resize(image, (416, 416))  
-    image = image.astype(np.float32) / 255.0 
-    image = np.transpose(image, (2, 0, 1))  
-    image = np.expand_dims(image, axis=0)  
-
-    # Run inference
-    inputs = {session.get_inputs()[0].name: image}
-    output = session.run(None, inputs)
-    detections = output[0][0]
-
+    
+    # Preprocess image
+    original_shape = image.shape
+    input_shape = (416, 416)
+    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))
+    image_batch = np.expand_dims(image_transposed, axis=0)
+    
+    # Get input name and run inference
+    input_name = session.get_inputs()[0].name
+    outputs = session.run(None, {input_name: image_batch})
+    
+    # Process detections
+    detections = outputs[0][0]  # First batch, all detections
+    
+    # Calculate timing
     inference_time = time.time() - start_time
     total_inference_time += inference_time
     inference_count += 1
     avg_inference_time = total_inference_time / inference_count
     fps = 1 / inference_time
-
-    # Draw bounding boxes
-    output_image = image[0].transpose(1, 2, 0) * 255 
-    output_image = output_image.astype(np.uint8)
-
+    
+    # Create a copy of the original image for visualization
+    output_image = image.copy()
+    
+    # Scale factor for bounding box coordinates
+    scale_x = original_shape[1] / input_shape[0]
+    scale_y = original_shape[0] / input_shape[1]
+    
+    # Draw bounding boxes and labels
     for det in detections:
-        x1, y1, x2, y2, conf, class_id = map(int, det[:6])
+        x1, y1, x2, y2, conf, class_id = det[:6]
         if conf < 0.3:  # Confidence threshold
             continue
-
-        color = colors[class_id].tolist()
-        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 3)
+            
+        # Convert to original image coordinates
+        x1, x2 = int(x1 * scale_x), int(x2 * scale_x)
+        y1, y2 = int(y1 * scale_y), int(y2 * scale_y)
+        class_id = int(class_id)
+        
+        # Draw rectangle and label
+        color = tuple(map(int, colors[class_id]))
+        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 2)
         label = f"Class {class_id} {conf:.2f}"
-        cv2.putText(output_image, label, (x1, y1 - 10),
-                    cv2.FONT_HERSHEY_SIMPLEX, 0.9, color, 2)
-
+        cv2.putText(output_image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+    
     # Display FPS
-    cv2.putText(output_image, f"FPS: {fps:.2f}", (20, 40),
-                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
-    cv2.putText(output_image, f"Avg FPS: {1/avg_inference_time:.2f}", (20, 70),
-                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
-
+    cv2.putText(output_image, f"FPS: {fps:.2f}", (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+    cv2.putText(output_image, f"Avg FPS: {1/avg_inference_time:.2f}", (20, 70), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+    
     return output_image
 
 # Gradio Interface