fixed body mask

yolo/BodyMask.py

@@ -61,6 +61,9 @@ def combine_masks(masks: List[dict], labels: List[str], is_label=True) -> Image.
     """
     labels_set = set(labels)  # Convert labels list to a set for O(1) lookups
 
+    # Filter out any masks that do not have a label key
+    masks = [mask for mask in masks if "label" in mask]
+
     # Filter and convert mask images based on the specified labels
     mask_images = [
         mask["mask"].convert("L")
@@ -119,14 +122,25 @@ class BodyMask:
         self.body_box = self.get_body_box(
             remove_overlap=remove_overlap, widen=widen_box, elongate=elongate_box
         )
-        if overlay == "box":
-            self.overlay = overlay_mask(
-                self.image, self.body_box, opacity=0.9, color="red"
+        if self.body_mask is not None:
+            self.box = get_bounding_box(self.body_mask)
+            self.body_box = self.get_body_box(
+                remove_overlap=remove_overlap, widen=widen_box, elongate=elongate_box
             )
+            if overlay == "box":
+                self.overlay = overlay_mask(
+                    self.image, self.body_box, opacity=0.9, color="red"
+                )
+            else:
+                self.overlay = overlay_mask(
+                    self.image, self.body_mask, opacity=0.9, color="red"
+                )
         else:
-            self.overlay = overlay_mask(
-                self.image, self.body_mask, opacity=0.9, color="red"
-            )
+            self.box = None
+            self.body_box = None
+            self.overlay = (
+                self.image
+            )  # Just return the original image if no mask is found
 
     def get_image(self, resize_to, resize_to_nearest_eight):
         image = load_image(self.image_path)
@@ -140,7 +154,17 @@ class BodyMask:
 
     def get_body_mask(self):
         body_mask = combine_masks(self.results, self.labels, self.is_label)
-        return dilate_mask(body_mask, self.dilate_factor)
+        if body_mask is not None:
+            return dilate_mask(body_mask, self.dilate_factor)
+        return None
+
+    def get_body_box(self, remove_overlap=True, widen=0, elongate=0):
+        if self.body_mask is None:
+            return None
+        body_box = get_bounding_box_mask(self.body_mask, widen=widen, elongate=elongate)
+        if remove_overlap:
+            body_box = self.remove_overlap(body_box)
+        return body_box
 
     def get_results(self):
         imgsz = max(self.image.size)
@@ -193,12 +217,6 @@ class BodyMask:
             combined_mask = ImageChops.lighter(combined_mask, mask)
         return combined_mask
 
-    def get_body_box(self, remove_overlap=True, widen=0, elongate=0):
-        body_box = get_bounding_box_mask(self.body_mask, widen=widen, elongate=elongate)
-        if remove_overlap:
-            body_box = self.remove_overlap(body_box)
-        return body_box
-
     def remove_overlap(self, body_box):
         """
         Remove mask regions that overlap with unwanted labels