Update app.py

app.py

@@ -12,67 +12,124 @@ os.makedirs("models", exist_ok=True)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 print(f"Using device: {device}")
 
-# Load YOLOv5n model
+# Use YOLOv5 Nano for better speed
 model_path = Path("models/yolov5n.pt")
-if not model_path.exists():
-    torch.hub.download_url_to_file("https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5n.pt", "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)
+else:
+    print("Downloading YOLOv5n model and caching...")
+    model = torch.hub.load("ultralytics/yolov5", "yolov5n", pretrained=True).to(device)
+    torch.save(model.state_dict(), model_path)
 
-model = torch.hub.load("ultralytics/yolov5", "custom", path=str(model_path)).to(device)
+# Optimize model for speed
+model.conf = 0.3  # Lower confidence threshold
+model.iou = 0.3   # Non-Maximum Suppression IoU threshold
+model.classes = None  # Detect all classes
 
-# Model optimizations
-model.conf = 0.5
-model.iou = 0.45
 if device.type == "cuda":
-    model.half()
+    model.half()  # Use FP16 for faster inference
 else:
-    model.float()
-    torch.set_num_threads(2)
+    torch.set_num_threads(os.cpu_count())
 
 model.eval()
 
-colors = np.random.rand(80, 3) * 255  # COCO classes
+# Pre-generate colors for bounding boxes
+np.random.seed(42)
+colors = np.random.uniform(0, 255, size=(len(model.names), 3))
+
+# Track FPS
+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
 
 def detect_objects(image):
+    global total_inference_time, inference_count
+
     if image is None:
         return None
-    
-    start = time.perf_counter()
-    
-    # Preprocess
-    im = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    im = cv2.resize(im, (320, 320))
-    tensor = torch.from_numpy(im).to(device)
-    tensor = tensor.half() if device.type == "cuda" else tensor.float()
-    tensor = tensor.permute(2, 0, 1).unsqueeze(0) / 255
-    
-    # Inference
-    with torch.no_grad():
-        pred = model(tensor)[0]
-    
-    # Post-process
-    pred = pred.float() if device.type == "cpu" else pred.half()
-    pred = non_max_suppression(pred, model.conf, model.iou)[0]
-    
-    # Visualization
-    output = image.copy()
-    if pred is not None:
-        pred[:, :4] = scale_coords(tensor.shape[2:], pred[:, :4], image.shape).round()
-        for *xyxy, conf, cls in pred:
-            x1, y1, x2, y2 = map(int, xyxy)
-            cv2.rectangle(output, (x1, y1), (x2, y2), colors[int(cls)].tolist(), 2)
-    
-    # FPS counter
-    fps = 1 / (time.perf_counter() - start)
-    cv2.putText(output, f"FPS: {fps:.1f}", (10, 30), 
-                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)
-    
-    return output
-
-with gr.Blocks() as demo:
-    gr.Markdown("# Real-Time Object Detection")
+
+    start_time = time.time()
+
+    # Preprocess image
+    image = preprocess_image(image)
+
+    with torch.inference_mode():  # Faster than torch.no_grad()
+        results = model(image, size=640)
+
+    inference_time = time.time() - start_time
+    total_inference_time += inference_time
+    inference_count += 1
+    avg_inference_time = total_inference_time / inference_count
+
+    detections = results.pred[0].cpu().numpy()
+
+    output_image = image.copy()
+
+    for *xyxy, conf, cls in detections:
+        x1, y1, x2, y2 = map(int, xyxy)
+        class_id = int(cls)
+        color = colors[class_id].tolist()
+
+        # Draw bounding box
+        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 3, lineType=cv2.LINE_AA)
+
+        label = f"{model.names[class_id]} {conf:.2f}"
+        font_scale, font_thickness = 0.9, 2
+        (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
+
+        # Label background
+        cv2.rectangle(output_image, (x1, y1 - h - 10), (x1 + w + 10, y1), color, -1)
+        cv2.putText(output_image, label, (x1 + 5, y1 - 5),
+                    cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), font_thickness, lineType=cv2.LINE_AA)
+
+    fps = 1 / inference_time
+
+    # Display FPS
+    overlay = output_image.copy()
+    cv2.rectangle(overlay, (10, 10), (300, 80), (0, 0, 0), -1)
+    output_image = cv2.addWeighted(overlay, 0.6, output_image, 0.4, 0)
+    cv2.putText(output_image, f"FPS: {fps:.2f}", (20, 40),
+                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, lineType=cv2.LINE_AA)
+    cv2.putText(output_image, f"Avg FPS: {1/avg_inference_time:.2f}", (20, 70),
+                cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, lineType=cv2.LINE_AA)
+
+    return output_image
+
+# Gradio UI
+example_images = ["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
+os.makedirs("examples", exist_ok=True)
+
+with gr.Blocks(title="Optimized YOLOv5 Object Detection") as demo:
+    gr.Markdown("""
+    # Optimized YOLOv5 Object Detection
+    Detects objects using YOLOv5 with enhanced visualization and FPS tracking.
+    """)
+
     with gr.Row():
-        input_img = gr.Image(label="Input", streaming=True)  # Modified webcam handling
-        output_img = gr.Image(label="Output")
-    input_img.change(detect_objects, input_img, output_img)
+        with gr.Column(scale=1):
+            input_image = gr.Image(label="Input Image", type="numpy")
+            submit_button = gr.Button("Submit", variant="primary")
+            clear_button = gr.Button("Clear")
+
+        with gr.Column(scale=1):
+            output_image = gr.Image(label="Detected Objects", type="numpy")
+
+    gr.Examples(
+        examples=example_images,
+        inputs=input_image,
+        outputs=output_image,
+        fn=detect_objects,
+        cache_examples=True
+    )
+
+    submit_button.click(fn=detect_objects, inputs=input_image, outputs=output_image)
+    clear_button.click(lambda: (None, None), None, [input_image, output_image])
 
-demo.launch()
+demo.launch()