Spaces:
Runtime error
Runtime error
File size: 1,650 Bytes
4f23172 4ece717 44a02e2 4f23172 44a02e2 4f23172 44a02e2 4f23172 44a02e2 4ece717 44a02e2 4ece717 44a02e2 4ece717 44a02e2 4ece717 44a02e2 4ece717 44a02e2 4ece717 44a02e2 4f23172 4ece717 44a02e2 4f23172 4ece717 |
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 |
import gradio as gr
import torch
import cv2
import numpy as np
from PIL import Image
# Load the YOLOv8 model
model = torch.hub.load('ultralytics/yolov5', 'custom', path='best.pt', force_reload=True) # Adjust to yolov8 if needed
def process_image(image):
results = model(image)
return results.render()[0] # Returns an image with boxes drawn
def process_video(video):
cap = cv2.VideoCapture(video)
frames = []
while(cap.isOpened()):
ret, frame = cap.read()
if not ret:
break
results = model(frame)
frames.append(results.render()[0])
cap.release()
# Convert frames back to a video format
height, width, layers = frames[0].shape
video_out = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'mp4v'), 30, (width, height))
for frame in frames:
video_out.write(frame)
video_out.release()
return 'output.mp4'
# Create Gradio interface
image_input = gr.inputs.Image(type="numpy", label="Upload an image")
video_input = gr.inputs.Video(type="mp4", label="Upload a video")
image_output = gr.outputs.Image(type="numpy", label="Detected image")
video_output = gr.outputs.Video(type="mp4", label="Detected video")
iface = gr.Interface(fn={'image': process_image, 'video': process_video},
inputs=[image_input, video_input],
outputs=[image_output, video_output],
live=True,
title="YOLOv8 Object Detection",
description="Upload an image or video to detect objects using YOLOv8.")
if __name__ == "__main__":
iface.launch()
|