"updated app Bodymask and utils"

app.py

@@ -86,6 +86,8 @@ def display_image_with_masks(image, results, cols=4):
 
 def perform_segmentation(input_image):
     bm = BodyMask(input_image, model_id=model_id, resize_to=640)
+    if bm.body_mask is None:
+        return input_image  # Return the original image if no mask is found
     results = bm.results
     buf = display_image_with_masks(input_image, results)
 

yolo/BodyMask.py

@@ -89,7 +89,6 @@ body_labels = ["hair", "face", "arm", "hand", "leg", "foot", "outfit"]
 
 
 class BodyMask:
-
     def __init__(
         self,
         image_path,
@@ -118,29 +117,11 @@ class BodyMask:
         self.results = self.get_results()
         self.dilate_factor = dilate_factor
         self.body_mask = self.get_body_mask()
-        self.box = get_bounding_box(self.body_mask)
+        self.box = self.get_bounding_box()
         self.body_box = self.get_body_box(
             remove_overlap=remove_overlap, widen=widen_box, elongate=elongate_box
         )
-        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.box = None
-            self.body_box = None
-            self.overlay = (
-                self.image
-            )  # Just return the original image if no mask is found
+        self.overlay = self.create_overlay(overlay)
 
     def get_image(self, resize_to, resize_to_nearest_eight):
         image = load_image(self.image_path)
@@ -148,61 +129,59 @@ class BodyMask:
             image = resize_preserve_aspect_ratio(image, resize_to)
         if resize_to_nearest_eight:
             image = resize_image_to_nearest_eight(image)
-        else:
-            image = image
         return image
 
+    def get_results(self):
+        imgsz = max(self.image.size)
+        results = self.model(
+            self.image, retina_masks=True, imgsz=imgsz, verbose=self.verbose
+        )[0]
+        masks, boxes, scores, phrases = unload(results, self.model_labels)
+        results = format_results(
+            masks, boxes, scores, phrases, self.model_labels, person_masks_only=False
+        )
+        masks_to_filter = ["hair"]
+        results = filter_highest_score(results, ["hair", "face", "phone"])
+        return results
+
     def get_body_mask(self):
         body_mask = combine_masks(self.results, self.labels, self.is_label)
         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):
+    def get_bounding_box(self):
         if self.body_mask is None:
             return None
+        return get_bounding_box(self.body_mask)
+
+    def get_body_box(self, remove_overlap=True, widen=0, elongate=0):
+        if self.body_mask is None or self.box is None:
+            return None
         body_box = get_bounding_box_mask(self.body_mask, widen=widen, elongate=elongate)
-        if remove_overlap:
+        if remove_overlap and body_box is not None:
             body_box = self.remove_overlap(body_box)
         return body_box
 
-    def get_results(self):
-        imgsz = max(self.image.size)
-        results = self.model(
-            self.image, retina_masks=True, imgsz=imgsz, verbose=self.verbose
-        )[0]
-        self.masks, self.boxes, self.scores, self.phrases = unload(
-            results, self.model_labels
-        )
-        results = format_results(
-            self.masks,
-            self.boxes,
-            self.scores,
-            self.phrases,
-            self.model_labels,
-            person_masks_only=False,
-        )
-
-        # filter out lower score results
-        masks_to_filter = ["hair"]
-        results = filter_highest_score(results, ["hair", "face", "phone"])
-        return results
-
-    def display_results(self):
-        if len(self.masks) < 4:
-            cols = len(self.masks)
-        else:
-            cols = 4
-        display_image_with_masks(self.image, self.results, cols=cols)
+    def create_overlay(self, overlay_type):
+        if self.body_box is not None and overlay_type == "box":
+            return overlay_mask(self.image, self.body_box, opacity=0.9, color="red")
+        elif self.body_mask is not None:
+            return overlay_mask(self.image, self.body_mask, opacity=0.9, color="red")
+        return self.image
 
-    def get_mask(self, mask_label):
-        assert mask_label in self.phrases, "Mask label not found in results"
-        return [f for f in self.results if f.get("label") == mask_label]
+    def remove_overlap(self, body_box):
+        if body_box is None:
+            return None
+        box_array = np.array(body_box)
+        mask = self.combine_masks(mask_labels=self.labels, is_label=True)
+        if mask is None:
+            return body_box
+        mask_array = np.array(mask)
+        box_array[mask_array == 255] = 0
+        return Image.fromarray(box_array)
 
     def combine_masks(self, mask_labels: List, no_labels=None, is_label=True):
-        """
-        Combine the masks included in the labels list or all of the masks not in the list
-        """
         if not is_label:
             mask_labels = [
                 phrase for phrase in self.phrases if phrase not in mask_labels
@@ -217,50 +196,15 @@ class BodyMask:
             combined_mask = ImageChops.lighter(combined_mask, mask)
         return combined_mask
 
-    def remove_overlap(self, body_box):
-        """
-        Remove mask regions that overlap with unwanted labels
-        """
-        # convert mask to numpy array
-        box_array = np.array(body_box)
-
-        # combine the masks for those labels
-        mask = self.combine_masks(mask_labels=self.labels, is_label=True)
-
-        # convert mask to numpy array
-        mask_array = np.array(mask)
-
-        # where the mask array is white set the box array to black
-        box_array[mask_array == 255] = 0
-
-        # convert the box array to an image
-        mask_image = Image.fromarray(box_array)
-        return mask_image
-
-
-if __name__ == "__main__":
-    url = "https://sjc1.vultrobjects.com/photo-storage/images/525d1f68-314c-455b-a8b6-f5dc3fa044e4.jpeg"
-    image_name = url.split("/")[-1]
-    labels = ["face", "hair", "phone", "hand"]
-    image = load_image(url)
-    image_size = image.size
-    # Get the original size of the image
-    original_size = image.size
-
-    # Create body mask
-    body_mask = BodyMask(
-        image,
-        overlay="box",
-        labels=labels,
-        widen_box=50,
-        elongate_box=10,
-        dilate_factor=0,
-        resize_to=640,
-        is_label=False,
-        remove_overlap=True,
-        verbose=False,
-    )
+    def display_results(self):
+        if not self.results:
+            print("No results to display.")
+            return
+        cols = min(len(self.results), 4)
+        display_image_with_masks(self.image, self.results, cols=cols)
 
-    # Resize the image back to the original size
-    image = body_mask.image.resize(original_size)
-    body_mask.body_box.save(image_name)
+    def get_mask(self, mask_label):
+        if mask_label not in self.phrases:
+            print(f"Mask label '{mask_label}' not found in results.")
+            return None
+        return [f for f in self.results if f.get("label") == mask_label]

yolo/utils.py

@@ -178,16 +178,23 @@ def display_image_with_masks(image, results, cols=4, return_images=False):
 
 
 def get_bounding_box(mask):
-    """
-    Given a segmentation mask, return the bounding box for the mask object.
-    """
-    # Find indices where the mask is non-zero
-    coords = np.argwhere(mask)
-    # Get the minimum and maximum x and y coordinates
-    x_min, y_min = np.min(coords, axis=0)
-    x_max, y_max = np.max(coords, axis=0)
-    # Return the bounding box coordinates
-    return (y_min, x_min, y_max, x_max)
+    if mask is None or not isinstance(mask, np.ndarray):
+        return None
+
+    # Check if the mask is empty
+    if mask.size == 0 or np.all(mask == 0):
+        return None
+
+    # Find the bounding box
+    rows = np.any(mask, axis=1)
+    cols = np.any(mask, axis=0)
+    if not np.any(rows) or not np.any(cols):
+        return None
+
+    rmin, rmax = np.where(rows)[0][[0, -1]]
+    cmin, cmax = np.where(cols)[0][[0, -1]]
+
+    return (int(cmin), int(rmin), int(cmax), int(rmax))
 
 
 def get_bounding_box_mask(segmentation_mask, widen=0, elongate=0):