Update app.py

app.py

@@ -1,17 +1,3 @@
-"""
-Ultralytics YOLO26 & YOLOE26 Gradio Demo.
-
-This script creates an interactive Gradio interface showcasing:
-- Ultralytics YOLO26 models across tasks (Detection, Segmentation, Pose, OBB, Classification)
-- Ultralytics YOLOE26 open-vocabulary segmentation with custom text prompts
-
-Docs: https://docs.ultralytics.com/models/yolo26/
-GitHub: https://github.com/ultralytics/ultralytics
-
-Usage:
-    python examples/app.py
-"""
-
 import gradio as gr
 from PIL import Image
 
@@ -27,14 +13,6 @@ if not OBB_IMAGE.exists():
 # Model cache
 model_cache = {}
 
-# Suffixes for filenames (weights naming) and for repo naming
-TASK_FILE_SUFFIX = {
-    "Detection": "",
-    "Segmentation": "-seg",
-    "Classification": "-cls",
-    "Pose": "-pose",
-    "OBB": "-obb",
-}
 TASK_REPO_SUFFIX = {
     "Detection": "",
     "Segmentation": "-seg",
@@ -47,25 +25,26 @@ def _scale_from_label(label: str) -> str:
     # "YOLO26-M" -> "m", "YOLOE-26L" -> "l"
     return label.strip()[-1].lower()
 
-def _get_model(repo_id: str, filename: str) -> YOLO:
-    key = f"{repo_id}:{filename}"
-    if key not in model_cache:
-        path = hf_hub_download(repo_id=repo_id, filename=filename)
-        model_cache[key] = YOLO(path)
-    return model_cache[key]
+def _get_model(repo_id: str) -> YOLO:
+    if repo_id not in model_cache:
+        path = hf_hub_download(repo_id=repo_id, filename="model.pt")
+        model_cache[repo_id] = YOLO(path)
+    return model_cache[repo_id]
 
 def predict_yolo26(image, model_name, task, conf, iou, retina):
-    """Run YOLO26 inference for various tasks."""
-    scale = _scale_from_label(model_name)  # n/s/m/l/x
+    scale = _scale_from_label(model_name)
 
-    # HF repo layout you listed: openvision/yolo26-n, openvision/yolo26-n-seg, ...
+    # openvision/yolo26-n, yolo26-n-seg, yolo26-n-pose, etc.
     repo_id = f"openvision/yolo26-{scale}{TASK_REPO_SUFFIX[task]}"
-    filename = f"yolo26{scale}{TASK_FILE_SUFFIX[task]}.pt"
-
-    model = _get_model(repo_id, filename)
-
-    use_retina = bool(retina) and task == "Segmentation"
-    results = model.predict(source=image, conf=conf, iou=iou, imgsz=640, retina_masks=use_retina)
+    model = _get_model(repo_id)
+
+    results = model.predict(
+        source=image,
+        conf=conf,
+        iou=iou,
+        imgsz=640,
+        retina_masks=bool(retina and task == "Segmentation"),
+    )
 
     if task == "Classification":
         top5 = results[0].probs.top5
@@ -77,111 +56,23 @@ def predict_yolo26(image, model_name, task, conf, iou, retina):
     return Image.fromarray(results[0].plot()[..., ::-1]), None
 
 def predict_yoloe26(image, model_name, classes_text, conf, retina):
-    """Run YOLOE26 open-vocabulary inference with text prompts."""
-    scale = _scale_from_label(model_name)  # n/s/m/l/x
+    scale = _scale_from_label(model_name)
 
-    # HF repo layout you listed: openvision/yoloe26-n-seg (open vocab)
+    # openvision/yoloe26-n-seg (open-vocab)
     repo_id = f"openvision/yoloe26-{scale}-seg"
-    filename = f"yoloe26{scale}-seg.pt"
+    model = _get_model(repo_id)
 
-    model = _get_model(repo_id, filename)
+    names = [c.strip() for c in classes_text.split(",") if c.strip()]
+    if not names:
+        names = ["person", "car", "dog", "cat"]
 
-    names = [c.strip() for c in classes_text.split(",") if c.strip()] or ["person", "car", "dog", "cat"]
     model.set_classes(names, model.get_text_pe(names))
 
-    results = model.predict(source=image, conf=conf, imgsz=640, retina_masks=bool(retina))
-    return Image.fromarray(results[0].plot()[..., ::-1])
-
-theme = gr.themes.Base().set(
-    button_primary_background_fill="#111F68",
-    button_primary_background_fill_hover="#042AFF",
-)
-
-# Build interface
-with gr.Blocks(title="Ultralytics YOLO26 & YOLOE26 Demo") as demo:
-    gr.Markdown(
-        "# 🚀 Ultralytics YOLO26 & YOLOE26 Demo\n"
-        "Showcasing YOLO26 tasks and YOLOE26 open-vocabulary segmentation. "
-        "[GitHub](https://github.com/ultralytics/ultralytics) | [Docs](https://docs.ultralytics.com/models/yolo26/)"
+    results = model.predict(
+        source=image,
+        conf=conf,
+        imgsz=640,
+        retina_masks=bool(retina),
     )
 
-    with gr.Tabs():
-        with gr.Tab("YOLO26 Tasks"):
-            gr.Markdown("### Ultralytics YOLO26: 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", "YOLO26-S", "YOLO26-M", "YOLO26-L", "YOLO26-X"],
-                            label="Model",
-                        )
-                        y26_task = gr.Dropdown(list(TASK_FILE_SUFFIX.keys()), 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(label="Retina Masks", value=True, 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-M", "Detection", 0.25, 0.45, True],
-                    [str(ASSETS / "bus.jpg"), "YOLO26-M", "Segmentation", 0.25, 0.45, True],
-                    [str(ASSETS / "zidane.jpg"), "YOLO26-M", "Pose", 0.25, 0.45, True],
-                    [str(OBB_IMAGE), "YOLO26-M", "OBB", 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("### Ultralytics YOLOE26: Open-Vocabulary Segmentation - Detect any object by text description")
-            with gr.Row():
-                with gr.Column():
-                    ye_image = gr.Image(type="pil", label="Upload Image", value=str(ASSETS / "bus.jpg"))
-                    with gr.Row():
-                        ye_model = gr.Dropdown(
-                            ["YOLOE-26N", "YOLOE-26S", "YOLOE-26M", "YOLOE-26L", "YOLOE-26X"],
-                            value="YOLOE-26L",
-                            label="Model",
-                        )
-                        ye_classes = gr.Textbox(value="person, bus, car", label="Classes", placeholder="person, dog, cat...")
-                    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")
-
-            gr.Examples(
-                examples=[
-                    [str(ASSETS / "bus.jpg"), "YOLOE-26L", "person, bus, car", 0.2, True],
-                    [str(ASSETS / "zidane.jpg"), "YOLOE-26L", "person, football, grass", 0.2, True],
-                ],
-                inputs=[ye_image, ye_model, ye_classes, ye_conf, ye_retina],
-                outputs=ye_output,
-                fn=predict_yoloe26,
-                cache_examples=True,
-            )
-
-            ye_btn.click(predict_yoloe26, [ye_image, ye_model, ye_classes, ye_conf, ye_retina], ye_output)
-
-if __name__ == "__main__":
-    demo.launch(theme=theme, allowed_paths=[str(ASSETS), str(ASSETS.parent)])
+    return Image.fromarray(results[0].plot()[..., ::-1])