Spaces:
Sleeping
Sleeping
| import torch | |
| import cv2 | |
| import numpy as np | |
| import gradio as gr | |
| from ultralytics import YOLO | |
| # Check if CUDA (GPU support) is available | |
| # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| # Load the YOLOv8 model | |
| model = YOLO('yolov8x-seg.pt') | |
| def process_video(input_video_path): | |
| cap = cv2.VideoCapture(input_video_path) | |
| if not cap.isOpened(): | |
| print("Error: Couldn't open the video file.") | |
| return | |
| frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fps = cap.get(cv2.CAP_PROP_FPS) | |
| fourcc = cv2.VideoWriter_fourcc(*'mp4v') | |
| out = cv2.VideoWriter("output_video.mp4", fourcc, fps, (frame_width, frame_height)) | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: | |
| break | |
| # Resize frame to match the expected input shape of the model | |
| resized_frame = cv2.resize(frame, (640, 640)) | |
| # Convert resized frame to torch tensor and move it to GPU | |
| frame_tensor = torch.from_numpy(resized_frame).permute(2, 0, 1).unsqueeze(0).float() / 255.0 | |
| threshold = 0.2 | |
| frame_copy = frame.copy() | |
| results = model(frame_tensor)[0] | |
| for result in results.boxes.data.tolist(): | |
| x1, y1, x2, y2, score, class_id = result | |
| if score > threshold: | |
| cv2.rectangle(frame_copy, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 4) | |
| cv2.putText(frame_copy, results.names[int(class_id)].upper(), (int(x1), int(y1 - 10)), | |
| cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 255, 0), 3, cv2.LINE_AA) | |
| cv2.putText(frame_copy, str(score), (int(x1), int(y2 + 10)), | |
| cv2.FONT_HERSHEY_SIMPLEX, 1.3, (0, 0, 255), 3, cv2.LINE_AA) | |
| out.write(frame_copy) | |
| cap.release() | |
| out.release() | |
| cv2.destroyAllWindows() | |
| return "output_video.mp4" | |
| # Define the input and output interfaces for Gradio | |
| inputs_video = gr.Video(label="Input Video") | |
| outputs_video = gr.Video(label="Output Video") | |
| # Create the Gradio interface | |
| demo = gr.Interface( | |
| fn=process_video, | |
| inputs=inputs_video, | |
| outputs=outputs_video, | |
| title="Animal detector using YOLOv8 NANO for Videos (GPU)", | |
| ) | |
| # Launch the interface | |
| demo.launch() | |