trash-detector / app.py
kbarnard's picture
Initial HF space
06b9c08
raw history blame
No virus
668 Bytes
import glob
import gradio as gr
from ultralytics import YOLO
model_path = "trash_mbari_09072023_640imgsz_50epochs_yolov8.pt"
model = YOLO(model_path)
PREDICT_KWARGS = {
"classes": 0,
}
def run(image_path):
results = model.predict(image_path, **PREDICT_KWARGS)
return results[0].plot()[:, :, ::-1] # reverse channels for gradio
title = "Trash Detector"
description = (
""
)
examples = glob.glob("images/*.png")
interface = gr.Interface(
run,
inputs=[gr.components.Image(type="filepath")],
outputs=gr.components.Image(type="numpy"),
title=title,
description=description,
examples=examples,
)
interface.queue().launch()