Update app.py

app.py

@@ -1,73 +1,62 @@
 import gradio as gr
 from ultralytics import YOLO
-import cv2
 import tempfile
-import numpy as np
+import cv2
 
-# Load YOLOv8 model once
-model = YOLO("best.pt")  
+# ✅ Load YOLOv8 model once on GPU (if available)
+model = YOLO("best.pt")
+model.to("cuda" if model.device.type == "cpu" else model.device)
 
-# 📸 Image prediction
+# ✅ Inference on image
 def predict_image(image):
-    results = model.predict(image, imgsz=320, conf=0.4, verbose=False)
+    results = model.predict(image, imgsz=640, conf=0.5, verbose=False)
     return results[0].plot()
 
-# 🎥 Video file prediction
+# ✅ Optimized video detection
 def predict_video(video_path):
+    import os
+
+    # Open input video
     cap = cv2.VideoCapture(video_path)
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) // 2)
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) // 2)
+    fps = cap.get(cv2.CAP_PROP_FPS) or 25
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH) * 0.5)
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT) * 0.5)
 
+    # Create output video file
     temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
-    out = cv2.VideoWriter(temp_output.name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
-
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-        frame = cv2.resize(frame, (width, height))
-        results = model.predict(frame, imgsz=320, conf=0.4, verbose=False)
-        annotated = results[0].plot()
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    out = cv2.VideoWriter(temp_output.name, fourcc, fps, (width, height))
+
+    # Efficient streaming-based inference
+    for result in model.track(source=video_path, stream=True, imgsz=480, conf=0.5, verbose=False):
+        if result.orig_img is None:
+            continue
+        # Resize original image
+        frame = cv2.resize(result.orig_img, (width, height))
+        # Draw annotations
+        annotated = result.plot()
+        # Write to output video
         out.write(annotated)
 
     cap.release()
     out.release()
     return temp_output.name
 
-# 🔴 Live webcam stream detection
-def live_feed():
-    cap = cv2.VideoCapture(0)
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-        frame = cv2.resize(frame, (640, 480))
-        results = model.predict(frame, imgsz=320, conf=0.4, verbose=False)
-        annotated = results[0].plot()
-        yield annotated
-    cap.release()
-
-# 🚀 Gradio App
+# ✅ Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 🔍 YOLOv8 Object Detection")
-    gr.Markdown("Upload an image or video, or use your webcam for real-time detection.")
+    gr.Markdown("# 🚀 YOLOv8 Fast Detection\nOptimized for Speed on Images & Videos")
 
-    with gr.Tab("🖼️ Image Detection"):
+    with gr.Tab("Image"):
         img_input = gr.Image(type="pil")
-        img_output = gr.Image()
-        img_btn = gr.Button("Detect Objects")
+        img_output = gr.Image(label="Detected")
+        img_btn = gr.Button("Run Detection")
         img_btn.click(predict_image, inputs=img_input, outputs=img_output)
 
-    with gr.Tab("🎞️ Video Detection"):
+    with gr.Tab("Video"):
         vid_input = gr.Video()
         vid_output = gr.Video()
-        vid_btn = gr.Button("Detect in Video")
+        vid_btn = gr.Button("Run Detection on Video")
         vid_btn.click(predict_video, inputs=vid_input, outputs=vid_output)
 
-    with gr.Tab("📷 Live Webcam"):
-        gr.Markdown("Note: Webcam streaming may not work on Hugging Face Spaces.")
-        live_img = gr.Image(streaming=True)
-        live_img.stream(live_feed)
-
+# ✅ Launch Gradio
 demo.launch()