Spaces:
Sleeping
Sleeping
| 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() |