Update app.py

app.py

@@ -9,11 +9,13 @@ from pathlib import Path
 # Create cache directory for models
 os.makedirs("models", exist_ok=True)
 
+# Select device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
-# Use YOLOv5 Nano for better speed
+# Use YOLOv5 Nano for speed
 model_path = Path("models/yolov5n.pt")
+
 if model_path.exists():
     print(f"Loading model from cache: {model_path}")
     model = torch.hub.load("ultralytics/yolov5", "custom", path=str(model_path), source="local").to(device)
@@ -23,31 +25,36 @@ else:
     torch.save(model.state_dict(), model_path)
 
 # Optimize model for speed
-model.conf = 0.3  # Lower confidence threshold
-model.iou = 0.3   # Non-Maximum Suppression IoU threshold
+model.conf = 0.3  # Confidence threshold
+model.iou = 0.3   # IoU threshold for Non-Maximum Suppression (NMS)
 model.classes = None  # Detect all classes
+model.eval()
 
 if device.type == "cuda":
-    model.half()  # Use FP16 for faster inference
-else:
-    torch.set_num_threads(os.cpu_count())
+    print("Using FP16 precision for inference (high speed, lower accuracy)")
+    model.half()  # Enable FP16 for faster inference
 
-model.eval()
+torch.set_num_threads(os.cpu_count())  # Optimize CPU threading
 
 # Pre-generate colors for bounding boxes
 np.random.seed(42)
 colors = np.random.uniform(0, 255, size=(len(model.names), 3))
 
-# Track FPS
+# FPS tracking
 total_inference_time = 0
 inference_count = 0
 
 def preprocess_image(image):
-    """ Prepares image for YOLOv5 detection. """
-    input_size = 640
-    image = cv2.resize(image, (input_size, input_size))
-    image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  # Convert to BGR for OpenCV
-    return image
+    """Prepares image for YOLOv5 detection while maintaining aspect ratio."""
+    h, w, _ = image.shape
+    scale = 640 / max(h, w)
+    new_w, new_h = int(w * scale), int(h * scale)
+
+    resized_image = cv2.resize(image, (new_w, new_h))
+    padded_image = np.full((640, 640, 3), 114, dtype=np.uint8)  # Gray padding
+    padded_image[:new_h, :new_w] = resized_image
+
+    return cv2.cvtColor(padded_image, cv2.COLOR_RGB2BGR)  # Convert to BGR for OpenCV
 
 def detect_objects(image):
     global total_inference_time, inference_count
@@ -68,7 +75,7 @@ def detect_objects(image):
     inference_count += 1
     avg_inference_time = total_inference_time / inference_count
 
-    detections = results.pred[0].cpu().numpy()
+    detections = results.xyxy[0].cpu().numpy()  # Use xyxy format
 
     output_image = image.copy()
 
@@ -77,6 +84,9 @@ def detect_objects(image):
         class_id = int(cls)
         color = colors[class_id].tolist()
 
+        # Keep bounding boxes within image bounds
+        x1, y1, x2, y2 = max(0, x1), max(0, y1), min(640, x2), min(640, y2)
+
         # Draw bounding box
         cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 3, lineType=cv2.LINE_AA)
 
@@ -108,7 +118,7 @@ os.makedirs("examples", exist_ok=True)
 
 with gr.Blocks(title="Optimized YOLOv5 Object Detection") as demo:
     gr.Markdown("""
-    # Optimized YOLOv5 Object Detection
+    # Optimized YOLOv5 Object Detection  
     Detects objects using YOLOv5 with enhanced visualization and FPS tracking.
     """)