Update app.py

app.py

@@ -1,4 +1,4 @@
-from typing import Tuple
+from typing import Tuple, Dict
 import gradio as gr
 import supervision as sv
 import numpy as np
@@ -6,13 +6,25 @@ from PIL import Image
 from huggingface_hub import hf_hub_download
 from ultralytics import YOLO
 
-# Load the YOLO model from Hugging Face
-model_path = hf_hub_download(
-    repo_id="cultural-heritage/medieval-manuscript-yolov11",
-    filename="medieval-yolov11n.pt"
-)
-# Load the YOLO model from local path
-model = YOLO(model_path)
+# Define models
+MODEL_OPTIONS = {
+    "YOLOv11-Nano": "medieval-yolov11n.pt",
+    "YOLOv11-Small": "medieval-yolov11s.pt",
+    "YOLOv11-Medium": "medieval-yolov11m.pt",
+    "YOLOv11-Large": "medieval-yolov11l.pt",
+    "YOLOv11-XLarge": "medieval-yolov11x.pt"
+}
+
+# Dictionary to store loaded models
+models: Dict[str, YOLO] = {}
+
+# Load all models
+for name, model_file in MODEL_OPTIONS.items():
+    model_path = hf_hub_download(
+        repo_id="biglam/medieval-manuscript-yolov11",
+        filename=model_file
+    )
+    models[name] = YOLO(model_path)
 
 # Create annotators
 LABEL_ANNOTATOR = sv.LabelAnnotator(text_color=sv.Color.BLACK)
@@ -20,9 +32,13 @@ BOX_ANNOTATOR = sv.BoxAnnotator()
 
 def detect_and_annotate(
     image: np.ndarray,
+    model_name: str,
     conf_threshold: float,
     iou_threshold: float
 ) -> np.ndarray:
+    # Get the selected model
+    model = models[model_name]
+    
     # Perform inference
     results = model.predict(
         image,
@@ -67,6 +83,12 @@ with gr.Blocks() as demo:
                 type='numpy'
             )
             with gr.Accordion("Detection Settings", open=True):
+                model_selector = gr.Dropdown(
+                    choices=list(MODEL_OPTIONS.keys()),
+                    value=list(MODEL_OPTIONS.keys())[0],
+                    label="Model",
+                    info="Select YOLO model variant"
+                )
                 with gr.Row():
                     conf_threshold = gr.Slider(
                         label="Confidence Threshold",
@@ -95,12 +117,13 @@ with gr.Blocks() as demo:
 
     def process_image(
         image: np.ndarray,
+        model_name: str,
         conf_threshold: float,
         iou_threshold: float
     ) -> Tuple[np.ndarray, np.ndarray]:
         if image is None:
             return None, None
-        annotated_image = detect_and_annotate(image, conf_threshold, iou_threshold)
+        annotated_image = detect_and_annotate(image, model_name, conf_threshold, iou_threshold)
         return image, annotated_image
 
     def clear():
@@ -109,7 +132,7 @@ with gr.Blocks() as demo:
     # Connect buttons to functions
     detect_btn.click(
         process_image,
-        inputs=[input_image, conf_threshold, iou_threshold],
+        inputs=[input_image, model_selector, conf_threshold, iou_threshold],
         outputs=[input_image, output_image]
     )
     clear_btn.click(