Spaces:
Sleeping
Sleeping
import cv2 | |
import torch | |
import gradio as gr | |
import numpy as np | |
from ultralytics import YOLO | |
# Load YOLOv8 model | |
model = YOLO('./data/best.pt') # Path to your model | |
# Define the function that processes the uploaded video | |
def process_video(video): | |
# video is now the file path string, not a file object | |
input_video = cv2.VideoCapture(video) # Directly pass the path to cv2.VideoCapture | |
# Get frame width, height, and fps from input video | |
frame_width = int(input_video.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
frame_height = int(input_video.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
fps = input_video.get(cv2.CAP_PROP_FPS) | |
# Create an empty list to store processed frames | |
processed_frames = [] | |
while True: | |
# Read a frame from the video | |
ret, frame = input_video.read() | |
if not ret: | |
break # End of video | |
# Perform inference on the frame | |
results = model(frame) | |
# The results object contains annotations for the frame | |
annotated_frame = results[0].plot() # Plot the frame with bounding boxes | |
# Convert the annotated frame to RGB format | |
annotated_frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB) | |
# Append the frame to the list | |
processed_frames.append(annotated_frame_rgb) | |
# Release resources | |
input_video.release() | |
# Return the processed frames as an output video in Gradio | |
return processed_frames | |
# Create a Gradio interface for video upload | |
iface = gr.Interface(fn=process_video, | |
inputs=gr.Video(label="Upload Video"), # Updated line | |
outputs=gr.Video(), # This will display the output video directly | |
title="YOLOv8 Object Detection on Video", | |
description="Upload a video for object detection using YOLOv8") | |
# Launch the interface | |
iface.launch() | |