UPD: Added a YOLOv10 assisted Segmentation Mode.

SegmentAnything2AssistApp.py

@@ -50,11 +50,11 @@ DEBUG = False
 
 
 segment_anything2assist = SegmentAnything2Assist.SegmentAnything2Assist(
-    model_name="sam2_hiera_tiny", device=torch.device("cuda")
+    sam_model_name="sam2_hiera_tiny", device=torch.device("cuda")
 )
 
 
-def __change_base_model(model_name, device):
+def change_base_model(model_name, device):
     global segment_anything2assist
     gradio.Info(f"Changing model to {model_name} on {device}", duration=3)
     try:
@@ -113,7 +113,7 @@ def __post_process_annotator_inputs(value):
 
 
 @spaces.GPU(duration=60)
-def __generate_mask(
+def generate_image_mask(
     value,
     mask_threshold,
     max_hole_area,
@@ -128,7 +128,7 @@ def __generate_mask(
     )
 
     if VERBOSE:
-        print("SegmentAnything2AssistApp::__generate_mask::Called.")
+        print("SegmentAnything2AssistApp::generate_image_mask::Called.")
     mask_chw, mask_iou = segment_anything2assist.generate_masks_from_image(
         value["image"],
         image_point_coords,
@@ -140,14 +140,16 @@ def __generate_mask(
     )
 
     if VERBOSE:
-        print("SegmentAnything2AssistApp::__generate_mask::Masks generated.")
+        print("SegmentAnything2AssistApp::generate_image_mask::Masks generated.")
 
     __current_mask, __current_segment = segment_anything2assist.apply_mask_to_image(
         value["image"], mask_chw[0]
     )
 
     if VERBOSE:
-        print("SegmentAnything2AssistApp::__generate_mask::Masks and Segments created.")
+        print(
+            "SegmentAnything2AssistApp::generate_image_mask::Masks and Segments created."
+        )
 
     __image_box = gradio.DataFrame(value=[[]])
     __image_point_coords = gradio.DataFrame(value=[[]])
@@ -201,9 +203,9 @@ def __generate_mask(
         )
 
 
-def __change_output_mode(image_input, radio, __current_mask, __current_segment):
+def on_image_output_mode_change(image_input, radio, __current_mask, __current_segment):
     if VERBOSE:
-        print("SegmentAnything2AssistApp::__generate_mask::Called.")
+        print("SegmentAnything2AssistApp::generate_image_mask::Called.")
     if __current_mask is None or __current_segment is None:
         gradio.Warning("Configuration was changed, generate the mask again", duration=5)
         return gradio_imageslider.ImageSlider(render=True)
@@ -216,9 +218,7 @@ def __change_output_mode(image_input, radio, __current_mask, __current_segment):
         return gradio_imageslider.ImageSlider(render=True)
 
 
-def __generate_multi_mask_output(
-    image, auto_list, auto_mode, auto_bbox_mode, masks, bboxes
-):
+def __generate_auto_mask(image, auto_list, auto_mode, auto_bbox_mode, masks, bboxes):
     global segment_anything2assist
 
     # When value from gallery is called, it is a tuple
@@ -235,7 +235,7 @@ def __generate_multi_mask_output(
 
 
 @spaces.GPU(duration=60)
-def __generate_auto_mask(
+def generate_auto_mask(
     image,
     points_per_side,
     points_per_batch,
@@ -255,7 +255,7 @@ def __generate_auto_mask(
 ):
     global segment_anything2assist
     if VERBOSE:
-        print("SegmentAnything2AssistApp::__generate_auto_mask::Called.")
+        print("SegmentAnything2AssistApp::generate_auto_mask::Called.")
 
     __auto_masks, masks, bboxes = segment_anything2assist.generate_automatic_masks(
         image,
@@ -296,7 +296,7 @@ def __generate_auto_mask(
     else:
         choices = [str(i) for i in range(len(__auto_masks))]
 
-        returning_image = __generate_multi_mask_output(
+        returning_image = __generate_auto_mask(
             image, ["0"], output_mode, False, masks, bboxes
         )
         return (
@@ -318,6 +318,74 @@ def __generate_auto_mask(
         )
 
 
+def __generate_yolo_mask(
+    image,
+    yolo_mask,
+    output_mode,
+):
+    global segment_anything2assist
+    if VERBOSE:
+        print("SegmentAnything2AssistApp::generate_yolo_mask::Called.")
+
+    mask = yolo_mask[4]
+
+    if output_mode == "Mask":
+        return [image, mask]
+
+    mask, output_image = segment_anything2assist.apply_mask_to_image(image, mask)
+
+    if output_mode == "Segment":
+        return [image, output_image]
+
+
+@spaces.GPU(duration=60)
+def generate_yolo_mask(
+    image,
+    yolo_model_choice,
+    mask_threshold,
+    max_hole_area,
+    max_sprinkle_area,
+    output_mode,
+):
+    global segment_anything2assist
+    if VERBOSE:
+        print("SegmentAnything2AssistApp::generate_yolo_mask::Called.")
+
+    results = segment_anything2assist.generate_mask_from_image_with_yolo(
+        image,
+        YOLOv10ModelName=yolo_model_choice,
+        mask_threshold=mask_threshold,
+        max_hole_area=max_hole_area,
+        max_sprinkle_area=max_sprinkle_area,
+    )
+
+    if len(results) > 0:
+        if VERBOSE:
+            print("SegmentAnything2AssistApp::generate_yolo_mask::Masks generated.")
+
+        yolo_masks = []
+        for result in results:
+            yolo_mask = [
+                result["name"],
+                result["class"],
+                result["confidence"],
+                [result["box"]],
+                result["mask_chw"],
+                result["mask_iou"][0].item(),
+            ]
+            yolo_masks.append(yolo_mask)
+
+        return __generate_yolo_mask(image, yolo_masks[0], output_mode), gradio.Dataset(
+            label="YOLOv10 Assisted Masks", type="values", samples=yolo_masks
+        )
+
+    else:
+        if VERBOSE:
+            print("SegmentAnything2AssistApp::generate_yolo_mask::No masks generated.")
+
+        return gradio.ImageSlider(), gradio.Dataset()
+
+
 with gradio.Blocks() as base_app:
     gradio.Markdown(
         """
@@ -342,14 +410,14 @@ with gradio.Blocks() as base_app:
                 ["cpu", "cuda"], value="cuda", label="Device Choice"
             )
     base_model_choice.change(
-        __change_base_model, inputs=[base_model_choice, base_gpu_choice]
+        change_base_model, inputs=[base_model_choice, base_gpu_choice]
     )
     base_gpu_choice.change(
-        __change_base_model, inputs=[base_model_choice, base_gpu_choice]
+        change_base_model, inputs=[base_model_choice, base_gpu_choice]
     )
 
     # Image Segmentation
-    with gradio.Tab(label="Image Segmentation", id="image_tab") as image_tab:
+    with gradio.Tab(label="🌆 Image Segmentation", id="image_tab") as image_tab:
         gradio.Markdown("Image Segmentation", render=True)
         with gradio.Column():
             with gradio.Accordion("Image Annotation Documentation", open=False):
@@ -419,7 +487,7 @@ with gradio.Blocks() as base_app:
 
             # image_input.change(__post_process_annotator_inputs, inputs = [image_input])
             image_generate_mask_button.click(
-                __generate_mask,
+                generate_image_mask,
                 inputs=[
                     image_input,
                     image_generate_SAM_mask_threshold,
@@ -436,7 +504,7 @@ with gradio.Blocks() as base_app:
                 ],
             )
             image_output_mode.change(
-                __change_output_mode,
+                on_image_output_mode_change,
                 inputs=[
                     image_input,
                     image_output_mode,
@@ -447,7 +515,7 @@ with gradio.Blocks() as base_app:
             )
 
     # Auto Segmentation
-    with gradio.Tab(label="Auto Segmentation", id="auto_tab"):
+    with gradio.Tab(label="🤖 Auto Segmentation", id="auto_tab"):
         gradio.Markdown("Auto Segmentation", render=True)
         with gradio.Column():
             with gradio.Accordion("Auto Annotation Documentation", open=False):
@@ -558,7 +626,7 @@ with gradio.Blocks() as base_app:
                 )
 
             auto_generate_button.click(
-                __generate_auto_mask,
+                generate_auto_mask,
                 inputs=[
                     auto_input,
                     auto_generate_SAM_points_per_side,
@@ -586,7 +654,7 @@ with gradio.Blocks() as base_app:
                 ],
             )
             auto_output_list.change(
-                __generate_multi_mask_output,
+                __generate_auto_mask,
                 inputs=[
                     auto_input,
                     auto_output_list,
@@ -598,7 +666,7 @@ with gradio.Blocks() as base_app:
                 outputs=[auto_output],
             )
             auto_output_bbox.change(
-                __generate_multi_mask_output,
+                __generate_auto_mask,
                 inputs=[
                     auto_input,
                     auto_output_list,
@@ -610,7 +678,7 @@ with gradio.Blocks() as base_app:
                 outputs=[auto_output],
             )
             auto_output_mode.change(
-                __generate_multi_mask_output,
+                __generate_auto_mask,
                 inputs=[
                     auto_input,
                     auto_output_list,
@@ -622,6 +690,121 @@ with gradio.Blocks() as base_app:
                 outputs=[auto_output],
             )
 
+    # YOLOv10 assisted Segmentation.
+    with gradio.Tab("🤙 YOLOv10 assisted Segmentation"):
+        gradio.Markdown("YOLOv10 assisted Segmentation")
+        with gradio.Column():
+            with gradio.Accordion("YOLOv10 Documentation", open=False):
+                gradio.Markdown(
+                    """
+                    ### 🖼️ YOLOv10 Assisted Segmentation Documentation
+
+                    YOLOv10 assisted segmentation allows you to generate masks for an image using the YOLOv10 model. 
+                    In this app, you can configure various settings to control the mask generation process.
+
+                    **📝 How to Use YOLOv10 Assisted Segmentation:**
+                    - Upload or select an image.
+                    - Choose the desired YOLOv10 model from the dropdown.
+                    - Adjust the advanced settings to fine-tune the mask generation process.
+                    - Click the 'Generate YOLOv10 Mask' button to generate masks.
+
+                    **⚙️ Advanced Settings:**
+                    - **SAM Mask Threshold:** Threshold for the SAM mask generation.
+                    - **Max Hole Area:** Maximum area for holes in the mask.
+                    - **Max Sprinkle Area:** Maximum area for sprinkled regions in the mask.
+
+                    **🎨 Generating Masks:**
+                    - Once you have configured the settings, click the 'Generate YOLOv10 Mask' button.
+                    - The masks will be generated based on the selected parameters.
+                    - You can view the generated masks and adjust the settings if needed.
+                    """
+                )
+
+            yolo_input = gradio.Image("assets/cars.jpg")
+            yolo_model_choice = gradio.Dropdown(
+                choices=["nano", "small", "medium", "base", "large", "xlarge"],
+                value="nano",
+                label="YOLOv10 Model Choice",
+            )
+            with gradio.Accordion("Advanced Settings", open=False):
+                yolo_generate_SAM_mask_threshold = gradio.Slider(
+                    0.0, 1.0, 0.0, label="SAM Mask Threshold"
+                )
+                yolo_generate_SAM_max_hole_area = gradio.Slider(
+                    0, 1000, 0, label="SAM Max Hole Area"
+                )
+                yolo_generate_SAM_max_sprinkle_area = gradio.Slider(
+                    0, 1000, 0, label="SAM Max Sprinkle Area"
+                )
+
+            yolo_generate_mask_button = gradio.Button("Generate YOLOv10 Mask")
+            with gradio.Row():
+                with gradio.Column():
+                    yolo_output_mode = gradio.Radio(
+                        ["Segment", "Mask"], value="Segment", label="Output Mode"
+                    )
+                with gradio.Column(scale=3):
+                    yolo_output = gradio_imageslider.ImageSlider()
+
+            with gradio.Accordion("Debug 1", open=DEBUG, visible=DEBUG):
+                __yolo_name = gradio.Textbox(
+                    label="Name", interactive=DEBUG, visible=DEBUG
+                )
+                __yolo_class = gradio.Number(
+                    label="Class", interactive=DEBUG, visible=DEBUG
+                )
+                __yolo_confidence = gradio.Number(
+                    label="Confidence", interactive=DEBUG, visible=DEBUG
+                )
+                __yolo_box = gradio.DataFrame(
+                    value=[[1, 2, 3, 4]], label="Box", interactive=DEBUG, visible=DEBUG
+                )
+                __yolo_mask = gradio.Image(
+                    label="Mask", interactive=DEBUG, visible=DEBUG
+                )
+                __yolo_mask_iou = gradio.Number(
+                    label="Mask IOU", interactive=DEBUG, visible=DEBUG
+                )
+
+            with gradio.Row():
+                yolo_masks = gradio.Dataset(
+                    label="YOLOv10 Assisted Masks",
+                    type="values",
+                    components=[
+                        __yolo_name,
+                        __yolo_class,
+                        __yolo_confidence,
+                        __yolo_box,
+                        __yolo_mask,
+                        __yolo_mask_iou,
+                    ],
+                )
+
+            yolo_generate_mask_button.click(
+                generate_yolo_mask,
+                inputs=[
+                    yolo_input,
+                    yolo_model_choice,
+                    yolo_generate_SAM_mask_threshold,
+                    yolo_generate_SAM_max_hole_area,
+                    yolo_generate_SAM_max_sprinkle_area,
+                    yolo_output_mode,
+                ],
+                outputs=[yolo_output, yolo_masks],
+            )
+
+            yolo_masks.click(
+                __generate_yolo_mask,
+                inputs=[yolo_input, yolo_masks, yolo_output_mode],
+                outputs=[yolo_output],
+            )
+
+            yolo_output_mode.change(
+                __generate_yolo_mask,
+                inputs=[yolo_input, yolo_masks, yolo_output_mode],
+                outputs=[yolo_output],
+            )
+
 
 if __name__ == "__main__":
     base_app.launch()

requirements.txt

@@ -41,6 +41,7 @@ pandas==2.2.2
 pillow==10.4.0
 portalocker==2.10.1
 psutil==5.9.8
+py-cpuinfo==9.0.0
 pydantic==2.8.2
 pydantic_core==2.20.1
 pydub==0.25.1
@@ -53,7 +54,10 @@ PyYAML==6.0.2
 requests==2.32.3
 rich==13.7.1
 ruff==0.6.2
+safetensors==0.4.5
 SAM-2 @ git+https://github.com/facebookresearch/segment-anything-2.git@7e1596c0b6462eb1d1ba7e1492430fed95023598
+scipy==1.14.1
+seaborn==0.13.2
 semantic-version==2.10.0
 setuptools==73.0.1
 shellingham==1.5.4
@@ -62,11 +66,13 @@ sniffio==1.3.1
 spaces==0.29.3
 starlette==0.38.2
 sympy==1.13.2
+thop==0.1.1.post2209072238
 tomlkit==0.12.0
 tqdm==4.66.5
 typer==0.12.5
 typing_extensions==4.12.2
 tzdata==2024.1
+ultralytics @ git+https://github.com/THU-MIG/yolov10.git@cd2f79c70299c9041fb6d19617ef1296f47575b1
 urllib3==2.2.2
 uvicorn==0.30.6
 websockets==12.0

src/SegmentAnything2Assist.py

@@ -10,6 +10,8 @@ import pickle
 import sam2.build_sam
 import sam2.automatic_mask_generator
 
+from . import YOLOv10Plugin
+
 import cv2
 
 SAM2_MODELS = {
@@ -39,7 +41,7 @@ SAM2_MODELS = {
 class SegmentAnything2Assist:
     def __init__(
         self,
-        model_name: (
+        sam_model_name: (
             str
             | typing.Literal[
                 "sam2_hiera_tiny",
@@ -56,32 +58,35 @@ class SegmentAnything2Assist:
         download: bool = True,
         device: str | torch.device = torch.device("cpu"),
         verbose: bool = True,
+        YOLOv10Model: YOLOv10Plugin.YOLOv10Plugin | None = None,
     ) -> None:
         assert (
-            model_name in SAM2_MODELS.keys()
-        ), f"`model_name` should be either one of {list(SAM2_MODELS.keys())}"
+            sam_model_name in SAM2_MODELS.keys()
+        ), f"`sam_model_name` should be either one of {list(SAM2_MODELS.keys())}"
         assert configuration in ["Automatic Mask Generator", "Image"]
 
-        self.model_name = model_name
+        self.sam_model_name = sam_model_name
         self.configuration = configuration
-        self.config_file = SAM2_MODELS[model_name]["config_file"]
+        self.config_file = SAM2_MODELS[sam_model_name]["config_file"]
         self.device = device
 
         self.download_url = (
             download_url
             if download_url is not None
-            else SAM2_MODELS[model_name]["download_url"]
+            else SAM2_MODELS[sam_model_name]["download_url"]
         )
         self.model_path = (
             model_path
             if model_path is not None
-            else SAM2_MODELS[model_name]["model_path"]
+            else SAM2_MODELS[sam_model_name]["model_path"]
         )
         os.makedirs(os.path.dirname(self.model_path), exist_ok=True)
         self.verbose = verbose
 
         if self.verbose:
-            print(f"SegmentAnything2Assist::__init__::Model Name: {self.model_name}")
+            print(
+                f"SegmentAnything2Assist::__init__::Model Name: {self.sam_model_name}"
+            )
             print(
                 f"SegmentAnything2Assist::__init__::Configuration: {self.configuration}"
             )
@@ -109,6 +114,8 @@ class SegmentAnything2Assist:
             if self.verbose:
                 print("SegmentAnything2Assist::__init__::SAM2 is not loaded.")
 
+        self.YOLOv10Model = YOLOv10Model
+
     def is_model_available(self) -> bool:
         ret = os.path.exists(self.model_path)
         if self.verbose:
@@ -264,3 +271,43 @@ class SegmentAnything2Assist:
         all_masks = all_masks.astype(numpy.uint8)
         image_with_segments = cv2.bitwise_and(image, image, mask=all_masks)
         return image_with_bounding_boxes, all_masks, image_with_segments
+
+    def generate_mask_from_image_with_yolo(
+        self,
+        image,
+        YOLOv10Model: YOLOv10Plugin.YOLOv10Plugin | None = None,
+        YOLOv10ModelName: str | None = None,
+        mask_threshold=0.0,
+        max_hole_area=0.0,
+        max_sprinkle_area=0.0,
+    ):
+        if self.YOLOv10Model is None:
+            assert bool(YOLOv10Model) != bool(
+                YOLOv10ModelName
+            ), "Either YOLOv10Model or YOLOv10ModelName should be provided."
+
+        if YOLOv10Model is not None:
+            self.YOLOv10Model = self.YOLOv10Model
+
+        if YOLOv10ModelName is not None:
+            self.YOLOv10Model = YOLOv10Plugin.YOLOv10Plugin(
+                yolo_model_name=YOLOv10ModelName
+            )
+
+        results = self.YOLOv10Model.detect(image)
+
+        for _, result in enumerate(results):
+            mask_chw, mask_iou = self.generate_masks_from_image(
+                image,
+                point_coords=None,
+                point_labels=None,
+                box=result["box"],
+                mask_threshold=mask_threshold,
+                max_hole_area=max_hole_area,
+                max_sprinkle_area=max_sprinkle_area,
+            )
+
+            results[_]["mask_chw"] = numpy.squeeze(mask_chw, 0)
+            results[_]["mask_iou"] = mask_iou
+
+        return results

src/YOLOv10Plugin.py

@@ -0,0 +1,67 @@
+import typing
+import ultralytics
+
+YOLO_V10_MODELS = {
+    "nano": "jameslahm/yolov10n",
+    "small": "jameslahm/yolov10s",
+    "medium": "jameslahm/yolov10m",
+    "base": "jameslahm/yolov10b",
+    "large": "jameslahm/yolov10l",
+    "xlarge": "jameslahm/yolov10x",
+}
+
+
+class YOLOv10Plugin:
+    def __init__(
+        self,
+        yolo_model_name: (
+            str
+            | typing.Literal[
+                "nano",
+                "small",
+                "medium",
+                "base",
+                "large",
+                "xlarge",
+            ]
+        ) = "nano",
+        verbose: bool = True,
+    ):
+        assert (
+            yolo_model_name in YOLO_V10_MODELS.keys()
+        ), f"`yolo_model_name` should be either one of {list(YOLO_V10_MODELS.keys())}"
+        self.yolo_model_name = yolo_model_name
+        self.model = ultralytics.YOLOv10.from_pretrained(
+            YOLO_V10_MODELS[yolo_model_name]
+        )
+
+        self.verbose = verbose
+        if self.verbose:
+            print(f"YOLOv10Plugin::__init__::Model Name: {self.yolo_model_name}")
+
+    def detect(self, image):
+        results = self.model(image)
+        results = results[0].summary()
+
+        out = []
+        for result in results:
+            out.append(
+                {
+                    "name": result["name"],
+                    "class": result["class"],
+                    "confidence": result["confidence"],
+                    "box": [
+                        int(result["box"]["x1"]),
+                        int(result["box"]["y1"]),
+                        int(result["box"]["x2"]),
+                        int(result["box"]["y2"]),
+                    ],
+                }
+            )
+
+        return out
+
+
+if __name__ == "__main__":
+    yolo = YOLOv10Plugin()
+    yolo.detect("https://ultralytics.com/images/zidane.jpg")