Aleksander Bojda
Limit output image height
081f351
raw
history blame
926 Bytes
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 = [
gr.Image(),
gr.Dropdown(["yolov5s", "yolov5m"], 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],
]
css = ".output-image {height: 700px !important}"
iface = gr.Interface(fn=predict,
inputs=inputs,
outputs=outputs,
examples=examples,
cache_examples=False,
css=css)
iface.launch()