REF: SAM2 AMG and the corresponding test case.

SegmentAnything2AssistApp.py

@@ -257,25 +257,27 @@ def generate_auto_mask(
     if VERBOSE:
         print("SegmentAnything2AssistApp::generate_auto_mask::Called.")
 
-    __auto_masks, masks, bboxes = segment_anything2assist.generate_automatic_masks(
-        image,
-        points_per_side,
-        points_per_batch,
-        pred_iou_thresh,
-        stability_score_thresh,
-        stability_score_offset,
-        mask_threshold,
-        box_nms_thresh,
-        crop_n_layers,
-        crop_nms_thresh,
-        crop_overlay_ratio,
-        crop_n_points_downscale_factor,
-        min_mask_region_area,
-        use_m2m,
-        multimask_output,
+    masks, bboxes, predicted_iou, stability_score = (
+        segment_anything2assist.generate_automatic_masks(
+            image,
+            points_per_side,
+            points_per_batch,
+            pred_iou_thresh,
+            stability_score_thresh,
+            stability_score_offset,
+            mask_threshold,
+            box_nms_thresh,
+            crop_n_layers,
+            crop_nms_thresh,
+            crop_overlay_ratio,
+            crop_n_points_downscale_factor,
+            min_mask_region_area,
+            use_m2m,
+            multimask_output,
+        )
     )
 
-    if len(__auto_masks) == 0:
+    if len(masks) == 0:
         gradio.Warning(
             "No masks generated, please tweak the advanced parameters.", duration=5
         )
@@ -294,7 +296,7 @@ def generate_auto_mask(
             ),
         )
     else:
-        choices = [str(i) for i in range(len(__auto_masks))]
+        choices = [str(i) for i in range(len(masks))]
 
         returning_image = __generate_auto_mask(
             image, ["0"], output_mode, False, masks, bboxes

src/SegmentAnything2Assist/SegmentAnything2Assist.py

@@ -98,7 +98,7 @@ class SegmentAnything2Assist:
             )
 
         if download:
-            self.download_model()
+            self.__download_model()
 
         if self.is_model_available():
             self.sam2 = sam2.build_sam.build_sam2(
@@ -121,14 +121,14 @@ class SegmentAnything2Assist:
             print(f"SegmentAnything2Assist::is_model_available::{ret}")
         return ret
 
-    def load_model(self) -> bool:
+    def __load_model(self) -> bool:
         if self.is_model_available():
             self.sam2 = sam2.build_sam(checkpoint=self.model_path)
             return True
 
         return False
 
-    def download_model(self, force: bool = False) -> bool:
+    def __download_model(self, force: bool = False) -> bool:
         if not force and self.is_model_available():
             print(f"{self.model_path} already exists. Skipping download.")
             return False
@@ -162,7 +162,17 @@ class SegmentAnything2Assist:
         min_mask_region_area=0,
         use_m2m=False,
         multimask_output=True,
-    ):
+    ) -> typing.Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray, numpy.ndarray]:
+        """
+        Generates automatic masks from the given image.
+
+        Returns:
+            typing.Tuple: Four numpy arrays where:
+                - segmentation_masks: Numpy array shape (N, H, W, C) where N is the number of masks, H is the height of the image, W is the width of the image, and C is the number of channels. Each N is a binary mask of the image of shape (H, W, C).
+                - bbox_masks: Numpy array of shape (N, 4) where N is the number of masks and 4 is the bounding box coordinates. Each mask is a bounding box of shape (x, y, w, h).
+                - predicted_iou: Numpy array of shape (N,) where N is the number of masks. Each value is the predicted IOU of the mask.
+                - stability_score: Numpy array of shape (N,) where N is the number of masks. Each value is the stability score of the mask.
+        """
         if self.sam2 is None:
             print(
                 "SegmentAnything2Assist::generate_automatic_masks::SAM2 is not loaded."
@@ -196,8 +206,15 @@ class SegmentAnything2Assist:
             cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR) for mask in segmentation_masks
         ]
         bbox_masks = [mask["bbox"] for mask in masks]
-
-        return masks, segmentation_masks, bbox_masks
+        predicted_iou = [mask["predicted_iou"] for mask in masks]
+        stability_score = [mask["stability_score"] for mask in masks]
+
+        return (
+            numpy.array(segmentation_masks, dtype=numpy.uint8),
+            numpy.array(bbox_masks, dtype=numpy.uint32),
+            numpy.array(predicted_iou, dtype=numpy.float32),
+            numpy.array(stability_score, dtype=numpy.float32),
+        )
 
     def generate_masks_from_image(
         self,
@@ -208,7 +225,15 @@ class SegmentAnything2Assist:
         mask_threshold=0.0,
         max_hole_area=0.0,
         max_sprinkle_area=0.0,
-    ):
+    ) -> typing.Tuple[numpy.ndarray, numpy.ndarray]:
+        """
+        Generates masks from the given image.
+
+        Returns:
+            typing.Tuple: Two numpy arrays where:
+                - masks_chw: Numpy array shape (1, H, W) for the mask, H is the height of the image, and W is the width of the image.
+                - mask_iou: Numpy array of shape (1,) for IOU of the mask.
+        """
         generator = sam2.sam2_image_predictor.SAM2ImagePredictor(
             self.sam2,
             mask_threshold=mask_threshold,
@@ -240,8 +265,6 @@ class SegmentAnything2Assist:
         image_with_bounding_boxes = image.copy()
         all_masks = None
 
-        cv2.imwrite(".tmp/mask_2.png", masks[3])
-
         for _ in auto_list:
             mask = masks[_]
             mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
@@ -252,8 +275,6 @@ class SegmentAnything2Assist:
             else:
                 all_masks = cv2.bitwise_or(all_masks, mask)
 
-            cv2.imwrite(".tmp/mask_3.png", masks[3])
-
             random_color = numpy.random.randint(0, 255, size=3)
             image_with_bounding_boxes = cv2.rectangle(
                 image_with_bounding_boxes,

test/test_module.py

@@ -2,6 +2,8 @@ import unittest
 import src.SegmentAnything2Assist.SegmentAnything2Assist as SegmentAnything2Assist
 import cv2
 
+import numpy
+
 
 class TestSegmentAnything2Assist(unittest.TestCase):
     def setUp(self) -> None:
@@ -39,21 +41,46 @@ class TestSegmentAnything2Assist(unittest.TestCase):
                     device="cpu",
                 )
 
-    def test_generate_automatic_mask(self):
+    def _generate_automatic_mask(self):
         image = cv2.imread("test/assets/liberty.jpg")
 
         sam_model = SegmentAnything2Assist.SegmentAnything2Assist(
             sam_model_name="sam2_hiera_tiny", download=True, device="cpu"
         )
 
-        masks, segmentation_masks, bboxes = sam_model.generate_automatic_masks(image)
+        segmentation_masks, bboxes, predicted_iou, stability_score = (
+            sam_model.generate_automatic_masks(image)
+        )
 
-        print(type(masks[0]))
-        print(type(segmentation_masks[0]))
-        print(type(bboxes[0]))
+        self.assertEqual(len(segmentation_masks.shape), 4)
+        self.assertEqual(segmentation_masks[0].shape, image.shape)
+        self.assertEqual(segmentation_masks.shape[3], 3)
+        self.assertEqual(type(segmentation_masks[0][0][0][0]), numpy.uint8)
+        self.assertEqual(len(bboxes.shape), 2)
+        self.assertEqual(bboxes[0].shape, (4,))
+        self.assertEqual(type(bboxes[0][0]), numpy.uint32)
+        self.assertEqual(len(predicted_iou.shape), 1)
+        self.assertEqual(type(predicted_iou[0]), numpy.float32)
+        self.assertEqual(len(stability_score.shape), 1)
+        self.assertEqual(type(stability_score[0]), numpy.float32)
 
-        self.assertEqual(len(masks), len(segmentation_masks))
-        self.assertEqual(len(masks), len(bboxes))
+        for segmentation_mask in segmentation_masks:
+            self.assertEqual(segmentation_mask.shape, image.shape)
 
-        # for mask, segmentation_mask, bbox in zip(masks, segmentation_masks, bboxes):
-        self.assertEqual(segmentation_masks[0].shape, image.shape)
+    def test_generate_masks_from_image(self):
+        image = cv2.imread("test/assets/liberty.jpg")
+
+        sam_model = SegmentAnything2Assist.SegmentAnything2Assist(
+            sam_model_name="sam2_hiera_tiny", download=True, device="cpu"
+        )
+
+        mask_chw, mask_iou = sam_model.generate_masks_from_image(
+            image, None, None, None
+        )
+
+        self.assertEqual(len(mask_chw.shape), 3)
+        self.assertEqual(mask_chw[0].shape, image.shape)
+        self.assertEqual(mask_chw.shape[0], 1)
+
+        self.assertEqual(len(mask_iou.shape), 1)
+        self.assertEqual(mask_iou.shape[0], 1)