Update app.py

app.py

@@ -2,17 +2,12 @@ from typing import Tuple, Dict
 import gradio as gr
 import supervision as sv
 import numpy as np
-from PIL import Image
 from huggingface_hub import hf_hub_download
 from ultralytics import YOLO
 
 # Define models
 MODEL_OPTIONS = {
-    #"YOLOv11-Nano": "medieval-yolov11n.pt",
     "YOLOv11-Small": "medieval-yolo11s-seg.pt"
-    #"YOLOv11-Medium": "medieval-yolov11m.pt",
-    #"YOLOv11-Large": "medieval-yolov11l.pt",
-    #"YOLOv11-XLarge": "medieval-yolov11x.pt"
 }
 
 # Dictionary to store loaded models
@@ -55,8 +50,17 @@ def detect_and_annotate(
     masks = None
     if results.masks is not None:
         masks = results.masks.data.cpu().numpy()
-        # Convert masks to boolean type
-        masks = masks.astype(bool)
+        # Resize masks to match original image dimensions
+        h, w = image.shape[:2]
+        masks = [
+            cv2.resize(mask.astype(float), (w, h), 
+            interpolation=cv2.INTER_LINEAR
+        ).astype(bool)
+            for mask in masks
+        ]
+        masks = np.array(masks)
+        # Transpose to (H, W, num_masks)
+        masks = np.transpose(masks, (1, 2, 0))
     
     # Create Detections object
     detections = sv.Detections(
@@ -73,7 +77,7 @@ def detect_and_annotate(
         in zip(class_ids, confidence)
     ]
 
-    # Annotate image with masks and labels
+    # Annotate image
     annotated_image = image.copy()
     if masks is not None:
         annotated_image = MASK_ANNOTATOR.annotate(scene=annotated_image, detections=detections)
@@ -83,14 +87,11 @@ def detect_and_annotate(
 
 # Create Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("# Medieval Manuscript Detection with YOLO")
+    gr.Markdown("# Medieval Manuscript Segmentation with YOLO")
     
     with gr.Row():
         with gr.Column():
-            input_image = gr.Image(
-                label="Input Image",
-                type='numpy'
-            )
+            input_image = gr.Image(label="Input Image", type='numpy')
             with gr.Accordion("Detection Settings", open=True):
                 model_selector = gr.Dropdown(
                     choices=list(MODEL_OPTIONS.keys()),
@@ -119,17 +120,9 @@ with gr.Blocks() as demo:
                 detect_btn = gr.Button("Detect", variant="primary")
                 
         with gr.Column():
-            output_image = gr.Image(
-                label="Detection Result",
-                type='numpy'
-            )
+            output_image = gr.Image(label="Segmentation Result", type='numpy')
 
-    def process_image(
-        image: np.ndarray,
-        model_name: str,
-        conf_threshold: float,
-        iou_threshold: float
-    ) -> Tuple[np.ndarray, np.ndarray]:
+    def process_image(image, model_name, conf_threshold, iou_threshold):
         if image is None:
             return None, None
         annotated_image = detect_and_annotate(image, model_name, conf_threshold, iou_threshold)
@@ -138,17 +131,12 @@ with gr.Blocks() as demo:
     def clear():
         return None, None
 
-    # Connect buttons to functions
     detect_btn.click(
         process_image,
         inputs=[input_image, model_selector, conf_threshold, iou_threshold],
         outputs=[input_image, output_image]
     )
-    clear_btn.click(
-        clear,
-        inputs=None,
-        outputs=[input_image, output_image]
-    )
+    clear_btn.click(clear, inputs=None, outputs=[input_image, output_image])
 
 if __name__ == "__main__":
     demo.launch(debug=True, show_error=True)