Update app.py

app.py

@@ -8,11 +8,11 @@ from ultralytics import YOLO
 
 # Define models
 MODEL_OPTIONS = {
-    #"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"
+    #"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
@@ -50,7 +50,13 @@ 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
+    
+    # Handle masks if they exist
+    masks = None
+    if results.masks is not None:
+        masks = results.masks.data.cpu().numpy()
+        # Convert masks to boolean type
+        masks = masks.astype(bool)
     
     # Create Detections object
     detections = sv.Detections(
@@ -69,7 +75,8 @@ def detect_and_annotate(
 
     # Annotate image with masks and labels
     annotated_image = image.copy()
-    annotated_image = MASK_ANNOTATOR.annotate(scene=annotated_image, detections=detections)
+    if masks is not None:
+        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