Update app.py for YOLO11 segmentation models

app.py

@@ -8,11 +8,11 @@ from ultralytics import YOLO
 
 # 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"
+    "YOLOv11-Nano": "medieval-yolo11n-seg.pt",
+    "YOLOv11-Small": "medieval-yolo11s-seg.pt",
+    "YOLOv11-Medium": "medieval-yolo11m-seg.pt",
+    "YOLOv11-Large": "medieval-yolo11l-seg.pt",
+    "YOLOv11-XLarge": "medieval-yolo11x-seg.pt"
 }
 
 # Dictionary to store loaded models
@@ -21,14 +21,14 @@ 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",
+        repo_id="johnlockejrr/medieval-manuscript-yolov11-seg",
         filename=model_file
     )
     models[name] = YOLO(model_path)
 
 # Create annotators
 LABEL_ANNOTATOR = sv.LabelAnnotator(text_color=sv.Color.BLACK)
-BOX_ANNOTATOR = sv.BoxAnnotator()
+MASK_ANNOTATOR = sv.MaskAnnotator()
 
 def detect_and_annotate(
     image: np.ndarray,
@@ -50,12 +50,14 @@ def detect_and_annotate(
     boxes = results.boxes.xyxy.cpu().numpy()
     confidence = results.boxes.conf.cpu().numpy()
     class_ids = results.boxes.cls.cpu().numpy().astype(int)
+    masks = results.masks.data.cpu().numpy() if results.masks is not None else None
     
     # Create Detections object
     detections = sv.Detections(
         xyxy=boxes,
         confidence=confidence,
-        class_id=class_ids
+        class_id=class_ids,
+        mask=masks
     )
     
     # Create labels with confidence scores
@@ -65,9 +67,9 @@ def detect_and_annotate(
         in zip(class_ids, confidence)
     ]
 
-    # Annotate image
+    # Annotate image with masks and labels
     annotated_image = image.copy()
-    annotated_image = BOX_ANNOTATOR.annotate(scene=annotated_image, detections=detections)
+    annotated_image = MASK_ANNOTATOR.annotate(scene=annotated_image, detections=detections)
     annotated_image = LABEL_ANNOTATOR.annotate(scene=annotated_image, detections=detections, labels=labels)
     
     return annotated_image
@@ -142,4 +144,4 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(debug=True, show_error=True)
+    demo.launch(debug=True, show_error=True)