Update app.py

app.py

@@ -2,59 +2,62 @@ import gradio as gr
 from ultralytics import YOLO
 import tempfile
 import cv2
-import os
-import torch
-
 
+# Load YOLOv8 model once (you can use 'yolov8n.pt' for better speed if needed)
 model = YOLO("best.pt")
 
-# Image prediction
+# Inference on image
 def predict_image(image):
-    results = model.predict(image, imgsz=640, conf=0.3, verbose=False)
+    results = model.predict(image, imgsz=480, conf=0.5, verbose=False)
     return results[0].plot()
 
-# Video prediction
+# Inference on video file
 def predict_video(video_path):
     cap = cv2.VideoCapture(video_path)
-    fps = cap.get(cv2.CAP_PROP_FPS) or 25
-    width, height = 640, 360
-
-    fd, output_path = tempfile.mkstemp(suffix=".mp4")
-    os.close(fd)
-
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
-
-    for result in model.track(source=video_path, stream=True, imgsz=640, conf=0.3, verbose=False):
-        if result.orig_img is None:
-            continue
-        frame = cv2.resize(result.orig_img, (width, height))
-        annotated = result.plot()
+    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
+
+    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=480, conf=0.5, verbose=False)
+        annotated = results[0].plot()
         out.write(annotated)
 
     cap.release()
     out.release()
+    return temp_output.name
 
-    # Check if file is valid
-    if not os.path.exists(output_path) or os.path.getsize(output_path) == 0:
-        raise RuntimeError("Output video file is missing or corrupted.")
-
-    return output_path
-
-# Gradio UI
+# Gradio Interface
 with gr.Blocks() as demo:
-    gr.Markdown("# YOLOv8 Object Detection\nFast & Accurate for Images and Videos")
+    gr.Markdown("# 🚀 YOLOv8 Object Detection")
+    gr.Markdown("Detect objects in images, videos, or live webcam feed using a YOLOv8 model.")
 
-    with gr.Tab("Image"):
+    # Image detection tab
+    with gr.Tab("🖼️ Image Detection"):
         img_input = gr.Image(type="pil")
-        img_output = gr.Image()
+        img_output = gr.Image(label="Detected Objects")
         img_btn = gr.Button("Run Detection")
         img_btn.click(predict_image, inputs=img_input, outputs=img_output)
 
-    with gr.Tab("Video"):
+    # Video detection tab
+    with gr.Tab("🎥 Video Detection"):
         vid_input = gr.Video()
         vid_output = gr.Video()
         vid_btn = gr.Button("Run Detection on Video")
         vid_btn.click(predict_video, inputs=vid_input, outputs=vid_output)
 
+    # Live webcam (browser-based) tab
+    with gr.Tab("📷 Live Webcam"):
+        gr.Markdown("This uses your browser webcam (works on Hugging Face Spaces).")
+        webcam_input = gr.Image(source="webcam", streaming=True, type="pil")
+        webcam_output = gr.Image()
+        webcam_input.change(fn=predict_image, inputs=webcam_input, outputs=webcam_output)
+
 demo.launch()