Update app.py

app.py

@@ -9,13 +9,11 @@ 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 speed
+# 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)
@@ -25,112 +23,84 @@ else:
     torch.save(model.state_dict(), model_path)
 
 # Optimize model for speed
-model.conf = 0.3  # Confidence threshold
-model.iou = 0.3   # IoU threshold for Non-Maximum Suppression (NMS)
+model.conf = 0.3  # Lower confidence threshold
+model.iou = 0.3   # Non-Maximum Suppression IoU threshold
 model.classes = None  # Detect all classes
-model.eval()
 
 if device.type == "cuda":
-    print("Using FP16 precision for inference (high speed, lower accuracy)")
-    model.half()  # Enable FP16 for faster inference
+    model.half()  # Use FP16 for faster inference
+else:
+    torch.set_num_threads(os.cpu_count())
 
-torch.set_num_threads(os.cpu_count())  # Optimize CPU threading
+model.eval()
 
 # Pre-generate colors for bounding boxes
 np.random.seed(42)
 colors = np.random.uniform(0, 255, size=(len(model.names), 3))
 
-# FPS tracking
-total_inference_time = 0
-inference_count = 0
-
-def detect_objects(image):
-    global total_inference_time, inference_count
-
-    if image is None:
-        return None
-
-    start_time = time.time()
-
-    # Convert image to BGR format for OpenCV
-    image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-
-    # Get image dimensions
-    h, w, _ = image.shape
-
-    with torch.inference_mode():  # Faster than torch.no_grad()
-        results = model(image_bgr)
-
-    inference_time = time.time() - start_time
-    total_inference_time += inference_time
-    inference_count += 1
-    avg_inference_time = total_inference_time / inference_count
-
-    detections = results.xyxy[0].cpu().numpy()  # Use xyxy format
+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)")
 
-    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()
-
-        # Keep bounding boxes within image bounds
-        x1, y1, x2, y2 = max(0, x1), max(0, y1), min(w, x2), min(h, y2)
-
-        # 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
-        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
-
-        # Label background
-        cv2.rectangle(output_image, (x1, y1 - th - 10), (x1 + tw + 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.Row():
+        video_input = gr.Video(label="Upload Video")
+        process_button = gr.Button("Process Video")
 
-with gr.Blocks(title="YOLOv5 Object Detection (High Quality, High FPS)") as demo:
-    gr.Markdown("""
-    # YOLOv5 Object Detection - High Quality & High FPS  
-    Detects objects with full-resolution output and ultra-fast performance.
-    """)
+    video_output = gr.Video(label="Processed Video")
 
-    with gr.Row():
-        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])
+    process_button.click(fn=process_video, inputs=video_input, outputs=video_output)
 
 demo.launch()