Spaces:
Sleeping
Sleeping
File size: 2,313 Bytes
2c806a3 33fa186 b281904 33fa186 2c806a3 20382ea 33fa186 2c806a3 5526a08 33fa186 b281904 33fa186 b281904 33fa186 b281904 33fa186 b281904 2c806a3 20382ea 2c806a3 c2a14be b281904 2c806a3 b281904 2c806a3 b281904 2c806a3 942815c 2c806a3 b281904 2c806a3 b281904 2c806a3 33fa186 2c806a3 b281904 |
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 63 64 65 66 67 68 69 |
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()
|