import gradio as gr

import yolov5

models = {
    "yolov5s": yolov5.load("akbojda/yolov5s-aquarium"),
    "yolov5m": yolov5.load("akbojda/yolov5m-aquarium"),
    "yolov5l": yolov5.load("akbojda/yolov5l-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 = [
    gr.Image(),
    gr.Dropdown(["yolov5s", "yolov5m", "yolov5l"],
                label="Model",
                value="yolov5s"),
]

outputs = [
    gr.Image(elem_classes="output-image")
]

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

title = "Aquarium object detection"
description = """
Based on [Aquarium Combined dataset](https://universe.roboflow.com/brad-dwyer/aquarium-combined)

Supports following 7 classes: [fish, jellyfish, penguin, puffin, shark, starfish, stingray]
"""

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