File size: 730 Bytes
35bfc71
 
 
 
 
 
 
 
 
 
d3b23a0
 
 
35bfc71
d3b23a0
35bfc71
 
 
 
1310d63
35bfc71
 
 
1310d63
35bfc71
 
 
 
 
 
 
 
 
d3b23a0
35bfc71
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
import gradio as gr

import yolov5

models = {
    "yolov5s": yolov5.load("akbojda/yolov5s-aquarium"),
    "yolov5m": yolov5.load("akbojda/yolov5m-aquarium"),
}

def predict(img, model_type):
    model = models[model_type]
    results = model(img, size=640)
    detection_img = results.render()[0]

    return detection_img

# Interface
inputs = [
    "image",
    gr.Dropdown(["yolov5s", "yolov5m"], label="Model", value="yolov5s"),
]

outputs = [
    "image",
]

examples = [
    ["examples/ex1.jpg", None],
    ["examples/ex2.jpg", None],
    ["examples/ex3.jpg", None],
    ["examples/ex4.jpg", None],
]

iface = gr.Interface(fn=predict, inputs=inputs, outputs=outputs, examples=examples, cache_examples=False)
iface.launch()