Spaces:
Sleeping
Sleeping
| import cv2 | |
| import torch | |
| import gradio as gr | |
| import numpy as np | |
| from ultralytics import YOLO | |
| # Load YOLOv8 model and set device (GPU if available) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model = YOLO('./data/best.pt') # Path to your model | |
| model.to(device) | |
| # 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) | |
| # Resize to reduce computation (optional) | |
| new_width, new_height = 640, 480 # Resize to 640x480 resolution | |
| frame_width, frame_height = new_width, new_height | |
| # Create a VideoWriter object to write processed frames to an output file | |
| output_video_path = "processed_output.mp4" | |
| fourcc = cv2.VideoWriter_fourcc(*'mp4v') # Codec for .mp4 format | |
| output_video = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height)) | |
| frame_skip = 10 # Skip 10 frames between each processed frame | |
| frame_count = 0 | |
| while True: | |
| # Read a frame from the video | |
| ret, frame = input_video.read() | |
| if not ret: | |
| break # End of video | |
| frame_count += 1 | |
| if frame_count % frame_skip != 0: | |
| continue # Skip frames | |
| # Resize the frame to reduce computational load | |
| frame = cv2.resize(frame, (new_width, new_height)) | |
| # Perform inference on the frame | |
| results = model(frame) # Automatically uses GPU if available | |
| # The results object contains annotations for the frame | |
| annotated_frame = results[0].plot() # Plot the frame with bounding boxes | |
| # Write the annotated frame to the output video | |
| output_video.write(annotated_frame) | |
| # Release resources | |
| input_video.release() | |
| output_video.release() | |
| # Return the processed video file path | |
| return output_video_path | |
| # Create a Gradio interface for video upload | |
| iface = gr.Interface(fn=process_video, | |
| inputs=gr.Video(label="Upload Video"), # Updated line | |
| outputs=gr.Video(label="Processed 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() | |