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)])