Spaces:
Running
Running
| from pathlib import Path | |
| import gradio as gr | |
| from ultralytics import YOLO | |
| # Load models | |
| MODEL_PATH = "weights/mbari_315k_yolov8.pt" | |
| MODEL = YOLO(MODEL_PATH) | |
| PREDICT_KWARGS = { | |
| "conf": 0.15, | |
| } | |
| # Get example images | |
| EXAMPLES_DIR = Path("examples") | |
| EXAMPLES = list(EXAMPLES_DIR.glob("*.png")) if EXAMPLES_DIR.exists() else [] | |
| def detect_objects(image): | |
| results = MODEL.predict(image, **PREDICT_KWARGS) | |
| return results[0].plot() | |
| # Gradio interface with gr.Interface instead of gr.Blocks | |
| demo = gr.Interface( | |
| fn=detect_objects, | |
| inputs=gr.Image(type="numpy"), | |
| outputs=gr.Image(type="numpy"), | |
| title="MBARI 315k", | |
| examples=EXAMPLES if EXAMPLES else None, | |
| cache_examples=True, | |
| ) | |
| demo.queue().launch() | |