Update app.py

app.py

@@ -1,13 +1,57 @@
-import gradio as gr
+# Code inspired from ultralyrics example with gradio
+import gradio as gradio
 import PIL.Image as Image
+import os
+import shutil
 
-from ultralytics import ASSETS, YOLO
+from ultralytics import YOLO
+from huggingface_hub import hf_hub_download
 
-model = YOLO("./yolov8m_seg-speech-bubble.pt")
+# Directory where downloaded model will be stored
+MODEL_DIR = "cached_models"
+os.makedirs(MODEL_DIR, exist_ok=True)
 
+# List of models available in the gradio ui
+AVAILABLE_MODELS = {
+    "YOLOv8m Speech Bubble (kitsumed)": {
+        "repo_id": "kitsumed/yolov8m_seg-speech-bubble",
+        # Filename, include sub-directory if model not at root (models/v1/model.pt)
+        "filename": "model.pt"
+    },
+    # Add more models here
+}
+
+# Cache for currently loaded model
+current_model = None
+current_model_name = None
+
+
+def load_model(model_name):
+    global current_model, current_model_name
+
+    if model_name == current_model_name:
+        return current_model
+
+    info = AVAILABLE_MODELS[model_name]
+    local_model_path = os.path.join(MODEL_DIR, f"{model_name}_{info['filename']}")
+
+    # If the model is yet existing in the cache, download and cache it
+    if not os.path.exists(local_model_path):
+        print(f"Downloading model '{model_name}'...")
+        downloaded_path = hf_hub_download(
+            repo_id=info["repo_id"],
+            filename=info["filename"]
+        )
+        shutil.move(downloaded_path, local_model_path)
+
+    current_model = YOLO(local_model_path)
+    current_model_name = model_name
+    return current_model
+
+
+def predict_image(img, conf_threshold, iou_threshold, model_name):
+    model = load_model(model_name)
 
-def predict_image(img, conf_threshold, iou_threshold):
-    """Predicts objects in an image using a YOLO11 model with adjustable confidence and IOU thresholds."""
     results = model.predict(
         source=img,
         conf=conf_threshold,
@@ -24,17 +68,18 @@ def predict_image(img, conf_threshold, iou_threshold):
     return im
 
 
-iface = gr.Interface(
+iface = gradio.Interface(
     fn=predict_image,
     inputs=[
-        gr.Image(type="pil", label="Upload Image"),
-        gr.Slider(minimum=0, maximum=1, value=0.25, label="Confidence threshold"),
-        gr.Slider(minimum=0, maximum=1, value=0.45, label="IoU threshold"),
+        gradio.Image(type="pil", label="Upload Image"),
+        gradio.Slider(minimum=0, maximum=1, value=0.20, label="Confidence threshold"),
+        gradio.Slider(minimum=0, maximum=1, value=0.40, label="IoU threshold"),
+        gradio.Dropdown(choices=list(AVAILABLE_MODELS.keys()), label="Select Model", value=list(AVAILABLE_MODELS.keys())[0])
     ],
-    outputs=gr.Image(type="pil", label="Result"),
-    title="Ultralytics Gradio",
-    description="Upload images for inference.",
+    outputs=gradio.Image(type="pil", label="Result"),
+    title="Try out kitsumed YOLO models",
+    description="Select a model from kitsumed on Hugging Face and upload an image to perform predictions.",
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()