Update app.py

app.py

@@ -3,64 +3,104 @@ import numpy as np
 import gradio as gr
 import cv2
 import time
-# Check device availability
+import os
+from pathlib import Path
+
+# Create cache directory for models
+os.makedirs("models", exist_ok=True)
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-# Load smaller YOLOv5 model
-model = torch.hub.load("ultralytics/yolov5", "yolov5x", pretrained=True).to(device)
-# Optimization configurations
-model.conf = 0.3  # Confidence threshold
-model.iou = 0.3   # NMS IoU threshold
+print(f"Using device: {device}")
+
+# Load YOLOv5 Nano model
+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)
+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)
+
+# 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
+
 if device.type == "cuda":
-    model.half().to(device)  # Use FP16 for performance boost
-model.eval()  # Set model to evaluation mode
-# Assign fixed colors to each class for bounding boxes
+    model.half()  # Use FP16 for faster inference
+else:
+    torch.set_num_threads(os.cpu_count())
+
+model.eval()
+
+# Pre-generate colors for bounding boxes
+np.random.seed(42)
 colors = np.random.uniform(0, 255, size=(len(model.names), 3))
-def detect_objects(image):
-    start_time = time.time()
-
-    # Convert BGR to RGB (if needed, Gradio might already provide RGB)
-    # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-
-    # Perform inference
-    with torch.no_grad():
-        results = model(image, size=640)  # Fixed inference size
-
-    # Process results directly on numpy array
-    output_image = image.copy()
-
-    # Extract detections
-    detections = results.pred[0].cpu().numpy()
-
-    for *xyxy, conf, cls in detections:
-        x1, y1, x2, y2 = map(int, xyxy)
-        class_id = int(cls)
-
-        # Draw bounding box
-        color = colors[class_id].tolist()
-        cv2.rectangle(output_image, (x1, y1), (x2, y2), color, 2)
-
-        # Create label
-        label = f"{model.names[class_id]} {conf:.2f}"
-
-        # Draw label background
-        (w, h), * = cv2.getTextSize(label, cv2.FONT*HERSHEY_SIMPLEX, 0.5, 1)
-        cv2.rectangle(output_image, (x1, y1 - 20), (x1 + w, y1), color, -1)
-
-        # Draw label text
-        cv2.putText(output_image, label, (x1, y1 - 5),
-                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
-    # Calculate FPS
-    fps = 1 / (time.time() - start_time)
-    print(f"FPS: {fps:.2f}")
-    return output_image
-# Gradio interface
-iface = gr.Interface(
-    fn=detect_objects,
-    inputs=gr.Image(type="numpy", label="Upload Image"),
-    outputs=gr.Image(type="numpy", label="Detected Objects"),
-    title="Optimized Object Detection with YOLOv5",
-    description="Faster detection using YOLOv5s with FP16 and optimized processing",
-    allow_flagging="never",
-    examples=["spring_street_after.jpg", "pexels-hikaique-109919.jpg"],
-)
-iface.launch()
+
+def process_video(video_path):
+    cap = cv2.VideoCapture(video_path)
+    
+    if not cap.isOpened():
+        return "Error: Could not open video file."
+
+    frame_width = int(cap.get(3))
+    frame_height = int(cap.get(4))
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    output_path = "output_video.mp4"
+    out = cv2.VideoWriter(output_path, fourcc, fps, (frame_width, frame_height))
+
+    total_frames = 0
+    total_time = 0
+
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break  # Break if no more frames
+        
+        start_time = time.time()
+        
+        # Convert frame for YOLOv5
+        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        results = model(img, size=640)
+        
+        inference_time = time.time() - start_time
+        total_time += inference_time
+        total_frames += 1
+        
+        detections = results.pred[0].cpu().numpy()
+
+        for *xyxy, conf, cls in detections:
+            x1, y1, x2, y2 = map(int, xyxy)
+            class_id = int(cls)
+            color = colors[class_id].tolist()
+            cv2.rectangle(frame, (x1, y1), (x2, y2), color, 3, lineType=cv2.LINE_AA)
+            label = f"{model.names[class_id]} {conf:.2f}"
+            cv2.putText(frame, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), 2)
+
+        # Calculate FPS
+        avg_fps = total_frames / total_time if total_time > 0 else 0
+        cv2.putText(frame, f"FPS: {avg_fps:.2f}", (20, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
+
+        out.write(frame)
+
+    cap.release()
+    out.release()
+    
+    return output_path
+
+# Gradio Interface
+with gr.Blocks(title="Real-Time YOLOv5 Video Detection") as demo:
+    gr.Markdown("# Real-Time YOLOv5 Video Detection (30+ FPS)")
+
+    with gr.Row():
+        video_input = gr.Video(label="Upload Video")
+        process_button = gr.Button("Process Video")
+
+    video_output = gr.Video(label="Processed Video")
+
+    process_button.click(fn=process_video, inputs=video_input, outputs=video_output)
+
+demo.launch()