File size: 3,416 Bytes
74d7b02
 
 
d4461b5
74d7b02
b48c806
74d7b02
 
 
 
 
 
 
 
 
 
 
 
 
b48c806
74d7b02
 
 
 
 
 
 
 
b48c806
74d7b02
 
 
 
8f78001
b48c806
 
 
 
74d7b02
 
b48c806
74d7b02
 
 
 
 
 
 
b48c806
74d7b02
 
 
b48c806
74d7b02
 
 
be786c9
b48c806
74d7b02
 
be786c9
b48c806
74d7b02
d4461b5
b48c806
be786c9
b48c806
be786c9
b48c806
 
 
 
 
 
74d7b02
b48c806
 
 
 
74d7b02
 
b48c806
74d7b02
b48c806
74d7b02
 
 
b48c806
d4461b5
b48c806
74d7b02
cd645f2
74d7b02
 
 
be786c9
b48c806
74d7b02
 
 
 
b48c806
74d7b02
b48c806
74d7b02
be786c9
b48c806
 
42f5968
 
 
 
 
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
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import gradio as gr
import sahi
import torch
from ultralyticsplus import YOLO, render_model_output

# Download images for the demo
sahi.utils.file.download_from_url(
    "https://raw.githubusercontent.com/kadirnar/dethub/main/data/images/highway.jpg",
    "highway.jpg",
)
sahi.utils.file.download_from_url(
    "https://raw.githubusercontent.com/obss/sahi/main/tests/data/small-vehicles1.jpeg",
    "small-vehicles1.jpeg",
)
sahi.utils.file.download_from_url(
    "https://raw.githubusercontent.com/ultralytics/yolov5/master/data/images/zidane.jpg",
    "zidane.jpg",
)

# Define available YOLOv8 segmentation models
model_names = [
    "yolov8n-seg.pt",
    "yolov8s-seg.pt",
    "yolov8m-seg.pt",
    "yolov8l-seg.pt",
    "yolov8x-seg.pt",
]

# Load the initial YOLOv8 model
current_model_name = "yolov8m-seg.pt"
model = YOLO(current_model_name)

def yolov8_inference(
    image: gr.Image = None,
    model_name: gr.Dropdown = None,
    image_size: gr.Slider = 640,
    conf_threshold: gr.Slider = 0.25,
    iou_threshold: gr.Slider = 0.45,
):
    """
    YOLOv8 inference function
    Args:
        image: Input image
        model_name: Name of the model
        image_size: Image size
        conf_threshold: Confidence threshold
        iou_threshold: IOU threshold
    Returns:
        Rendered image and mask coordinates with labels
    """
    global model
    global current_model_name
    # Switch model if a different one is selected
    if model_name != current_model_name:
        model = YOLO(model_name)
        current_model_name = model_name
    
    # Set model confidence and IOU thresholds
    model.overrides["conf"] = conf_threshold
    model.overrides["iou"] = iou_threshold
    
    # Perform inference with the YOLO model
    results = model.predict(image, imgsz=image_size, return_outputs=True)
    
    masks = []
    for result in results:
        masks.append([result.masks, result.labels])
    
    renders = []
    for image_results in results:
        render = render_model_output(
            model=model, image=image, model_output=image_results
        )
        renders.append(render)

    # Return mask coordinates and labels
    return masks

# Gradio app inputs and outputs
inputs = [
    gr.Image(type="filepath", label="Input Image"),
    gr.Dropdown(model_names, value=current_model_name, label="Model type"),
    gr.Slider(minimum=320, maximum=1280, value=640, step=32, label="Image Size"),
    gr.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold"),
    gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.05, label="IOU Threshold"),
]

outputs = gr.Textbox(label="Mask Coordinates and Labels")

# Example inputs for the Gradio app
examples = [
    ["zidane.jpg", "yolov8m-seg.pt", 640, 0.6, 0.45],
    ["highway.jpg", "yolov8m-seg.pt", 640, 0.25, 0.45],
    ["small-vehicles1.jpeg", "yolov8m-seg.pt", 640, 0.25, 0.45],
]

# Create the Gradio app interface
demo_app = gr.Interface(
    fn=yolov8_inference,
    inputs=inputs,
    outputs=outputs,
    title="Ultralytics YOLOv8 Segmentation Demo",
    examples=examples,
    cache_examples=True,
)

# Launch the Gradio app
demo_app.launch(
    debug=True,         # Show detailed errors in case of issues
    server_name="0.0.0.0",  # Host on all IPs
    server_port=7860,       # Custom port for accessing the app
    share=True              # To make the app accessible from a URL
)