app.py

import gradio as gr
from PIL import Image

from ultralytics import ASSETS, YOLO
from ultralytics.utils.downloads import safe_download
from huggingface_hub import hf_hub_download

# Download OBB test image if not exists
OBB_IMAGE = ASSETS.parent / "boats.jpg"
if not OBB_IMAGE.exists():
    safe_download("https://ultralytics.com/images/boats.jpg", dir=ASSETS.parent)

TASK_TO_REPO_TEMPLATE = {
    "Detection": "openvision/yolo26-{scale}",
    "Segmentation": "openvision/yolo26-{scale}-seg",
    "Classification": "openvision/yolo26-{scale}-cls",
    "Pose": "openvision/yolo26-{scale}-pose",
    "OBB": "openvision/yolo26-{scale}-obb",
}

YOLOE_REPO_TEMPLATE = "openvision/yoloe26-{scale}-seg"

weights_cache = {}
model_cache = {}


def _scale_from_ui_name(model_name: str) -> str:
    return model_name.split("-")[-1].strip().lower()


def _get_weights(repo_id: str) -> str:
    """Download (if needed) and cache model.pt path."""
    cache_key = f"{repo_id}::model.pt"
    if cache_key not in weights_cache:
        weights_cache[cache_key] = hf_hub_download(repo_id=repo_id, filename="model.pt")
    return weights_cache[cache_key]


def _get_model(repo_id: str) -> YOLO:
    """Download (if needed) and cache YOLO model (safe for YOLO26 tasks)."""
    cache_key = f"{repo_id}::model.pt"
    if cache_key not in model_cache:
        weights_path = _get_weights(repo_id)
        model_cache[cache_key] = YOLO(weights_path)
    return model_cache[cache_key]


def predict_yolo26(image, model_name, task, conf, iou, retina):
    scale = _scale_from_ui_name(model_name)
    repo_id = TASK_TO_REPO_TEMPLATE[task].format(scale=scale)
    model = _get_model(repo_id)

    use_retina = bool(retina) and task == "Segmentation"
    results = model.predict(source=image, conf=conf, iou=iou, imgsz=640, retina_masks=use_retina)

    if task == "Classification":
        top5 = results[0].probs.top5
        return None, {results[0].names[i]: float(results[0].probs.top5conf[j]) for j, i in enumerate(top5)}

    return Image.fromarray(results[0].plot()[..., ::-1]), None


def _parse_classes(classes_text: str):
    if classes_text is None:
        return []
    names = [c.strip() for c in classes_text.split(",") if c.strip()]
    # de-dup while preserving order
    seen = set()
    out = []
    for n in names:
        if n.lower() not in seen:
            out.append(n)
            seen.add(n.lower())
    return out


def predict_yoloe26(image, model_name, classes_text, conf, retina):
    names = _parse_classes(classes_text)
    if not names:
        raise gr.Error("Enter at least 1 class (comma-separated). Example: 'cat, dog, bicycle'")

    scale = _scale_from_ui_name(model_name)
    repo_id = YOLOE_REPO_TEMPLATE.format(scale=scale)

    weights_path = _get_weights(repo_id)
    model = YOLO(weights_path)

    model.set_classes(names, model.get_text_pe(names))
    res = model.predict(source=image, conf=conf, imgsz=640, retina_masks=bool(retina))[0]
    return Image.fromarray(res.plot()[..., ::-1])


theme = gr.themes.Base().set(
    button_primary_background_fill="#111F68",
    button_primary_background_fill_hover="#042AFF",
)

with gr.Blocks(title="Ultralytics YOLO26 & YOLOE26 Demo", theme=theme) as demo:
    gr.Markdown(
        "# 🚀 Ultralytics YOLO26 & YOLOE26 Demo\n"
        "YOLO26 tasks + YOLOE26 open-vocabulary segmentation."
    )

    with gr.Tabs():
        with gr.Tab("YOLO26 Tasks"):
            gr.Markdown("### Detection, Segmentation, Pose, OBB, Classification")
            with gr.Row():
                with gr.Column():
                    y26_image = gr.Image(type="pil", label="Upload Image")
                    with gr.Row():
                        y26_model = gr.Dropdown(["YOLO26-N"], value="YOLO26-N", label="Model")
                        y26_task = gr.Dropdown(list(TASK_TO_REPO_TEMPLATE.keys()), value="Detection", label="Task")
                    with gr.Accordion("Advanced Settings", open=False):
                        y26_conf = gr.Slider(0, 1, value=0.25, label="Confidence Threshold")
                        y26_iou = gr.Slider(0, 1, value=0.45, label="IoU Threshold")
                        y26_retina = gr.Checkbox(value=True, label="Retina Masks", info="Higher quality masks, slower inference")
                    y26_btn = gr.Button("Run Inference", variant="primary")

                with gr.Column():
                    y26_output = gr.Image(type="pil", label="Result")
                    y26_label = gr.Label(label="Classification Results", visible=False)

            y26_task.change(
                lambda t: (gr.update(visible=t != "Classification"), gr.update(visible=t == "Classification")),
                y26_task,
                [y26_output, y26_label],
            )

            gr.Examples(
                examples=[
                    [str(ASSETS / "bus.jpg"), "YOLO26-N", "Detection", 0.25, 0.45, True],
                    [str(ASSETS / "bus.jpg"), "YOLO26-N", "Segmentation", 0.25, 0.45, True],
                    [str(ASSETS / "zidane.jpg"), "YOLO26-N", "Pose", 0.25, 0.45, True],
                    [str(OBB_IMAGE), "YOLO26-N", "OBB", 0.25, 0.45, True],
                    [str(ASSETS / "bus.jpg"), "YOLO26-N", "Classification", 0.25, 0.45, True],
                ],
                inputs=[y26_image, y26_model, y26_task, y26_conf, y26_iou, y26_retina],
                outputs=[y26_output, y26_label],
                fn=predict_yolo26,
                cache_examples=True,
            )

            y26_btn.click(
                predict_yolo26,
                [y26_image, y26_model, y26_task, y26_conf, y26_iou, y26_retina],
                [y26_output, y26_label],
            )

        with gr.Tab("YOLOE26 Open-Vocabulary"):
            gr.Markdown("### Open-Vocabulary Segmentation (text prompts)")

            with gr.Row():
                with gr.Column():
                    ye_image = gr.Image(type="pil", label="Upload Image")
                    ye_model = gr.Dropdown(["YOLOE26-N"], value="YOLOE26-N", label="Model")
                    ye_classes = gr.Textbox(
                        label="Classes (comma-separated)",
                        placeholder="e.g. cat, dog, bicycle",
                        value="person, bus, car",
                    )
                    with gr.Accordion("Advanced Settings", open=False):
                        ye_conf = gr.Slider(0, 1, value=0.2, label="Confidence Threshold")
                        ye_retina = gr.Checkbox(value=True, label="Retina Masks", info="Higher quality masks, slower inference")
                    ye_btn = gr.Button("Run Inference", variant="primary")

                with gr.Column():
                    ye_output = gr.Image(type="pil", label="Result")

            ye_prompt_state = gr.State(ye_classes.value)

            ye_classes.change(lambda s: s, ye_classes, ye_prompt_state)

            gr.Examples(
                examples=[
                    [str(ASSETS / "bus.jpg"), "YOLOE26-N", "person, bus, car", 0.2, True],
                    [str(ASSETS / "zidane.jpg"), "YOLOE26-N", "person, football, grass", 0.2, True],
                    [str(ASSETS / "bus.jpg"), "YOLOE26-N", "bicycle, traffic light, road", 0.2, True],
                ],
                inputs=[ye_image, ye_model, ye_classes, ye_conf, ye_retina],
                outputs=ye_output,
                fn=predict_yoloe26,
            )

            ye_btn.click(
                predict_yoloe26,
                [ye_image, ye_model, ye_prompt_state, ye_conf, ye_retina],
                ye_output,
            )

if __name__ == "__main__":
    demo.launch(allowed_paths=[str(ASSETS), str(ASSETS.parent)])