File size: 1,920 Bytes
4f23172
4ece717
 
44a02e2
 
4f23172
44a02e2
4b1eaff
4f23172
44a02e2
4b1eaff
 
 
 
 
4f23172
4b1eaff
 
 
 
4ece717
4b1eaff
 
44a02e2
4ece717
44a02e2
4ece717
 
 
4b1eaff
 
4ece717
4b1eaff
 
 
 
44a02e2
4ece717
 
44a02e2
 
4b1eaff
 
44a02e2
 
 
 
 
 
 
4ece717
 
4b1eaff
44a02e2
 
4b1eaff
44a02e2
 
 
 
 
 
 
4f23172
4ece717
44a02e2
4f23172
 
4b1eaff
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
57
58
59
60
61
62
63
64
65
66
67
68
69
70
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/yolov8', 'custom', path='best.pt')  # YOLOv8 specific

def process_image(image):
    # Convert PIL image to numpy array if necessary
    if isinstance(image, Image.Image):
        image = np.array(image)

    # Perform detection
    results = model(image)
    
    # Render results
    annotated_image = results.render()[0]
    return Image.fromarray(annotated_image)

def process_video(video_path):
    cap = cv2.VideoCapture(video_path)
    frames = []

    while(cap.isOpened()):
        ret, frame = cap.read()
        if not ret:
            break
        
        # Perform detection
        results = model(frame)
        
        # Render results
        annotated_frame = results.render()[0]
        frames.append(annotated_frame)

    cap.release()
    
    # Convert frames back to a video format
    height, width, layers = frames[0].shape
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    video_out = cv2.VideoWriter('output.mp4', fourcc, 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="pil", label="Upload an image")
video_input = gr.inputs.Video(type="mp4", label="Upload a video")

image_output = gr.outputs.Image(type="pil", 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],
                     title="YOLOv8 Object Detection",
                     description="Upload an image or video to detect objects using YOLOv8.")

if __name__ == "__main__":
    iface.launch()