Spaces:
Runtime error
Runtime error
File size: 2,595 Bytes
790c62b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 |
import gradio as gr
import cv2
import tempfile
import os
from ultralytics import YOLO
model = YOLO("yolov8x.pt")
vehicle_classes = {
"bicycle": 1,
"car": 2,
"motorcycle": 3,
"bus": 5,
"truck": 7
}
def detect_vehicles_with_logs(video_input):
cap = cv2.VideoCapture(video_input)
if not cap.isOpened():
return None, "", "❌ Error: Could not open video."
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = cap.get(cv2.CAP_PROP_FPS)
with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmpfile:
out_path = tmpfile.name
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
out_vid = cv2.VideoWriter(out_path, fourcc, fps, (width, height))
logs = ""
frame_index = 0
while True:
ret, frame = cap.read()
if not ret:
break
frame_index += 1
results = model(frame, verbose=False)[0]
detected = []
for box in results.boxes:
cls_id = int(box.cls.item())
if cls_id in vehicle_classes.values():
xyxy = box.xyxy[0].cpu().numpy().astype(int)
conf = box.conf.item()
label = model.names[cls_id]
detected.append(f"{label} ({conf:.2f})")
cv2.rectangle(frame, tuple(xyxy[:2]), tuple(xyxy[2:]), (0, 255, 0), 2)
cv2.putText(frame, f"{label} {conf:.2f}", (xyxy[0], xyxy[1] - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
out_vid.write(frame)
logs += f"Frame {frame_index}: {', '.join(detected) if detected else 'No vehicles detected'}\n"
cap.release()
out_vid.release()
if not os.path.exists(out_path):
return None, "", "❌ Error: Processed video not found."
# Return three outputs:
# 1) path for video player to display (as video file)
# 2) path for download button (same video file)
# 3) logs text
return out_path, out_path, logs
demo = gr.Interface(
fn=detect_vehicles_with_logs,
inputs=gr.Video(label="Upload a video"),
outputs=[
gr.Video(label="Processed Video"), # show inline video
gr.File(label="Download Processed Video"), # download button
gr.Textbox(label="Detection Logs", lines=20, interactive=False)
],
title="YOLOv8x Vehicle Detection with Inline Video and Download",
description="Upload a video, detect vehicles, see the processed video inline, download it, and view logs."
)
if __name__ == "__main__":
demo.launch(share=True)
|