Spaces:
Sleeping
Sleeping
File size: 2,247 Bytes
391e560 b250ef1 391e560 4fe0657 391e560 1f89f37 391e560 4fe0657 391e560 57e64d1 b250ef1 391e560 4fe0657 391e560 4fe0657 391e560 57e64d1 391e560 |
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 |
import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
import tempfile # For creating temporary files
import os
# Load the model (do this outside the function for efficiency)
try:
model = YOLO("yolov8n.pt") # Or your model path
except Exception as e:
print(f"Error loading model: {e}")
exit() # Exit if the model fails to load
def detect_pedestrians(video):
try:
cap = cv2.VideoCapture(video)
fps = cap.get(cv2.CAP_PROP_FPS) # Get the video's FPS
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fourcc = cv2.VideoWriter_fourcc(*'mp4v') # Codec for MP4 (you can change this)
# Create a temporary file for the processed video
with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_video:
out = cv2.VideoWriter(temp_video.name, fourcc, fps, (width, height))
temp_file_name = temp_video.name #store the file name
while True:
ret, frame = cap.read()
if not ret:
break
results = model(frame)
for result in results:
boxes = result.boxes
for box in boxes:
if result.names[int(box.cls)] == 'person':
x1, y1, x2, y2 = map(int, box.xyxy[0])
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(frame, 'Person', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
out.write(frame) # Write the processed frame to the video
cap.release()
out.release()
return temp_file_name # Return the path to the temporary video file
except Exception as e:
print(f"Error in detection: {e}")
return f"Error: {e}" # Return the error message as a string
iface = gr.Interface(
fn=detect_pedestrians,
inputs=gr.Video(), # No 'source' argument here
outputs=gr.Video(), # Output is now a video
title="Pedestrian Detection",
description="Upload a video to detect pedestrians.",
allow_flagging="never",
)
iface.launch() |