REF: Uses internal variable for auto mask and image segmentation.

.gitignore

@@ -1 +1,2 @@
-.tmp/
+.tmp/
+.venv/

SegmentAnything2AssistApp.py

@@ -1,5 +1,5 @@
 import gradio
-import gradio_image_annotation 
+import gradio_image_annotation
 import gradio_imageslider
 import spaces
 import torch
@@ -8,29 +8,65 @@ import src.SegmentAnything2Assist as SegmentAnything2Assist
 
 example_image_annotation = {
     "image": "assets/cars.jpg",
-    "boxes": [{'label': '+', 'color': (0, 255, 0), 'xmin': 886, 'ymin': 551, 'xmax': 886, 'ymax': 551}, {'label': '-', 'color': (255, 0, 0), 'xmin': 1239, 'ymin': 576, 'xmax': 1239, 'ymax': 576}, {'label': '-', 'color': (255, 0, 0), 'xmin': 610, 'ymin': 574, 'xmax': 610, 'ymax': 574}, {'label': '', 'color': (0, 0, 255), 'xmin': 254, 'ymin': 466, 'xmax': 1347, 'ymax': 1047}]
+    "boxes": [
+        {
+            "label": "+",
+            "color": (0, 255, 0),
+            "xmin": 886,
+            "ymin": 551,
+            "xmax": 886,
+            "ymax": 551,
+        },
+        {
+            "label": "-",
+            "color": (255, 0, 0),
+            "xmin": 1239,
+            "ymin": 576,
+            "xmax": 1239,
+            "ymax": 576,
+        },
+        {
+            "label": "-",
+            "color": (255, 0, 0),
+            "xmin": 610,
+            "ymin": 574,
+            "xmax": 610,
+            "ymax": 574,
+        },
+        {
+            "label": "",
+            "color": (0, 0, 255),
+            "xmin": 254,
+            "ymin": 466,
+            "xmax": 1347,
+            "ymax": 1047,
+        },
+    ],
 }
 
+
 VERBOSE = True
+DEBUG = False
+
+
+segment_anything2assist = SegmentAnything2Assist.SegmentAnything2Assist(
+    model_name="sam2_hiera_tiny", device=torch.device("cpu")
+)
 
-segment_anything2assist = SegmentAnything2Assist.SegmentAnything2Assist(model_name = "sam2_hiera_tiny", device = torch.device("cuda"))
-__image_point_coords = None
-__image_point_labels = None
-__image_box = None
-__current_mask = None
-__current_segment = None
 
 def __change_base_model(model_name, device):
     global segment_anything2assist
+    gradio.Info(f"Changing model to {model_name} on {device}", duration=3)
     try:
-        segment_anything2assist = SegmentAnything2Assist.SegmentAnything2Assist(model_name = model_name, device = torch.device(device))
-        gradio.Info(f"Model changed to {model_name} on {device}", duration = 5)
+        segment_anything2assist = SegmentAnything2Assist.SegmentAnything2Assist(
+            model_name=model_name, device=torch.device(device)
+        )
+        gradio.Info(f"Model has been changed to {model_name} on {device}", duration=5)
     except:
-        gradio.Error(f"Model could not be changed", duration = 5)
+        gradio.Error(f"Model could not be changed", duration=5)
+
 
 def __post_process_annotator_inputs(value):
-    global __image_point_coords, __image_point_labels, __image_box
-    global __current_mask, __current_segment
     if VERBOSE:
         print("SegmentAnything2AssistApp::____post_process_annotator_inputs::Called.")
     __current_mask, __current_segment = None, None
@@ -38,111 +74,167 @@ def __post_process_annotator_inputs(value):
     __image_point_coords = []
     __image_point_labels = []
     __image_box = []
-    
+
     b_has_box = False
     for box in value["boxes"]:
-        if box['label'] == '':
+        if box["label"] == "":
             if not b_has_box:
                 new_box = box.copy()
-                new_box['color'] = (0, 0, 255)
+                new_box["color"] = (0, 0, 255)
                 new_boxes.append(new_box)
                 b_has_box = True
-            __image_box = [
-                box['xmin'],
-                box['ymin'],
-                box['xmax'],
-                box['ymax']
-            ]
-
-                
-        elif box['label'] == '+' or box['label'] == '-':
+            __image_box = [box["xmin"], box["ymin"], box["xmax"], box["ymax"]]
+
+        elif box["label"] == "+" or box["label"] == "-":
             new_box = box.copy()
-            new_box['color'] = (0, 255, 0) if box['label'] == '+' else (255, 0, 0)
-            new_box['xmin'] = int((box['xmin'] + box['xmax']) / 2)
-            new_box['ymin'] = int((box['ymin'] + box['ymax']) / 2)
-            new_box['xmax'] = new_box['xmin']
-            new_box['ymax'] = new_box['ymin']
+            new_box["color"] = (0, 255, 0) if box["label"] == "+" else (255, 0, 0)
+            new_box["xmin"] = int((box["xmin"] + box["xmax"]) / 2)
+            new_box["ymin"] = int((box["ymin"] + box["ymax"]) / 2)
+            new_box["xmax"] = new_box["xmin"]
+            new_box["ymax"] = new_box["ymin"]
             new_boxes.append(new_box)
-            
-            __image_point_coords.append([new_box['xmin'], new_box['ymin']])
-            __image_point_labels.append(1 if box['label'] == '+' else 0)
-        
-        if len(__image_box) == 0:
-            __image_box = None
-        
-        if len(__image_point_coords) == 0:
-            __image_point_coords = None
-        
-        if len(__image_point_labels) == 0:
-            __image_point_labels = None
+
+            __image_point_coords.append([new_box["xmin"], new_box["ymin"]])
+            __image_point_labels.append(1 if box["label"] == "+" else 0)
+
+    if len(__image_box) == 0:
+        __image_box = None
+
+    if len(__image_point_coords) == 0:
+        __image_point_coords = None
+
+    if len(__image_point_labels) == 0:
+        __image_point_labels = None
 
     if VERBOSE:
         print("SegmentAnything2AssistApp::____post_process_annotator_inputs::Done.")
 
+    return __image_point_coords, __image_point_labels, __image_box
 
 
-@spaces.GPU(duration = 60)
-def __generate_mask(value, mask_threshold, max_hole_area, max_sprinkle_area, image_output_mode):
-    global __current_mask, __current_segment
-    global __image_point_coords, __image_point_labels, __image_box
+@spaces.GPU(duration=60)
+def __generate_mask(
+    value,
+    mask_threshold,
+    max_hole_area,
+    max_sprinkle_area,
+    image_output_mode,
+):
     global segment_anything2assist
 
     # Force post processing of annotated image
-    __post_process_annotator_inputs(value)
+    image_point_coords, image_point_labels, image_box = __post_process_annotator_inputs(
+        value
+    )
 
     if VERBOSE:
         print("SegmentAnything2AssistApp::__generate_mask::Called.")
     mask_chw, mask_iou = segment_anything2assist.generate_masks_from_image(
         value["image"],
-        __image_point_coords,
-        __image_point_labels,
-        __image_box,
+        image_point_coords,
+        image_point_labels,
+        image_box,
         mask_threshold,
         max_hole_area,
-        max_sprinkle_area
+        max_sprinkle_area,
     )
 
     if VERBOSE:
         print("SegmentAnything2AssistApp::__generate_mask::Masks generated.")
 
-    __current_mask, __current_segment = segment_anything2assist.apply_mask_to_image(value["image"], mask_chw[0])
+    __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.")
 
+    __image_box = gradio.DataFrame(value=[[]])
+    __image_point_coords = gradio.DataFrame(value=[[]])
+    if DEBUG:
+        __image_box = gradio.DataFrame(
+            value=[image_box],
+            label="Box",
+            interactive=False,
+            headers=["XMin", "YMin", "XMax", "YMax"],
+        )
+        x = []
+        for i, _ in enumerate(image_point_coords):
+            x.append(
+                [
+                    image_point_labels[i],
+                    image_point_coords[i][0],
+                    image_point_coords[i][1],
+                ]
+            )
+        __image_point_coords = gradio.DataFrame(
+            value=x,
+            label="Point Coords",
+            interactive=False,
+            headers=["Label", "X", "Y"],
+        )
+
     if image_output_mode == "Mask":
-        return [value["image"], __current_mask]
+        return (
+            [value["image"], __current_mask],
+            __image_point_coords,
+            __image_box,
+            __current_mask,
+            __current_segment,
+        )
     elif image_output_mode == "Segment":
-        return [value["image"], __current_segment]
+        return (
+            [value["image"], __current_segment],
+            __image_point_coords,
+            __image_box,
+            __current_mask,
+            __current_segment,
+        )
     else:
         gradio.Warning("This is an issue, please report the problem!", duration=5)
-        return gradio_imageslider.ImageSlider(render = True)
+        return (
+            gradio_imageslider.ImageSlider(render=True),
+            __image_point_coords,
+            __image_box,
+            __current_mask,
+            __current_segment,
+        )
+
 
-def __change_output_mode(image_input, radio):
-    global __current_mask, __current_segment
-    global __image_point_coords, __image_point_labels, __image_box
+def __change_output_mode(image_input, radio, __current_mask, __current_segment):
     if VERBOSE:
         print("SegmentAnything2AssistApp::__generate_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)
+        return gradio_imageslider.ImageSlider(render=True)
     if radio == "Mask":
         return [image_input["image"], __current_mask]
     elif radio == "Segment":
         return [image_input["image"], __current_segment]
     else:
         gradio.Warning("This is an issue, please report the problem!", duration=5)
-        return gradio_imageslider.ImageSlider(render = True)
-    
-def __generate_multi_mask_output(image, auto_list, auto_mode, auto_bbox_mode):
+        return gradio_imageslider.ImageSlider(render=True)
+
+
+def __generate_multi_mask_output(
+    image, auto_list, auto_mode, auto_bbox_mode, masks, bboxes
+):
     global segment_anything2assist
-    image_with_bbox, mask, segment = segment_anything2assist.apply_auto_mask_to_image(image, [int(i) - 1 for i in auto_list])
-    
+
+    # When value from gallery is called, it is a tuple
+    if type(masks[0]) == tuple:
+        masks = [mask[0] for mask in masks]
+
+    image_with_bbox, mask, segment = segment_anything2assist.apply_auto_mask_to_image(
+        image, [int(i) - 1 for i in auto_list], masks, bboxes
+    )
+
     output_1 = image_with_bbox if auto_bbox_mode else image
     output_2 = mask if auto_mode == "Mask" else segment
     return [output_1, output_2]
-    
-@spaces.GPU(duration = 60)
+
+
+@spaces.GPU(duration=60)
 def __generate_auto_mask(
     image,
     points_per_side,
@@ -159,13 +251,13 @@ def __generate_auto_mask(
     min_mask_region_area,
     use_m2m,
     multimask_output,
-    output_mode
-    ):
+    output_mode,
+):
     global segment_anything2assist
     if VERBOSE:
-        print("SegmentAnything2AssistApp::__generate_auto_mask::Called.")    
-    
-    __auto_masks = segment_anything2assist.generate_automatic_masks(
+        print("SegmentAnything2AssistApp::__generate_auto_mask::Called.")
+
+    __auto_masks, masks, bboxes = segment_anything2assist.generate_automatic_masks(
         image,
         points_per_side,
         points_per_batch,
@@ -180,43 +272,84 @@ def __generate_auto_mask(
         crop_n_points_downscale_factor,
         min_mask_region_area,
         use_m2m,
-        multimask_output
+        multimask_output,
     )
-    
+
     if len(__auto_masks) == 0:
-        gradio.Warning("No masks generated, please tweak the advanced parameters.", duration = 5)
-        return gradio_imageslider.ImageSlider(), \
-            gradio.CheckboxGroup([], value = [], label = "Mask List", interactive = False), \
-            gradio.Checkbox(value = False, label = "Show Bounding Box", interactive = False)
+        gradio.Warning(
+            "No masks generated, please tweak the advanced parameters.", duration=5
+        )
+        return (
+            gradio_imageslider.ImageSlider(),
+            gradio.CheckboxGroup([], value=[], label="Mask List", interactive=False),
+            gradio.Checkbox(value=False, label="Show Bounding Box", interactive=False),
+            gradio.Gallery(
+                None, label="Output Gallery", interactive=False, type="numpy"
+            ),
+            gradio.DataFrame(
+                value=[[]],
+                label="Box",
+                interactive=False,
+                headers=["XMin", "YMin", "XMax", "YMax"],
+            ),
+        )
     else:
         choices = [str(i) for i in range(len(__auto_masks))]
-        returning_image = __generate_multi_mask_output(image, ["0"], output_mode, False)
-        return returning_image, \
-            gradio.CheckboxGroup(choices, value = ["0"], label = "Mask List", interactive = True), \
-            gradio.Checkbox(value = False, label = "Show Bounding Box", interactive = True)    
-    
+
+        returning_image = __generate_multi_mask_output(
+            image, ["0"], output_mode, False, masks, bboxes
+        )
+        return (
+            returning_image,
+            gradio.CheckboxGroup(
+                choices, value=["0"], label="Mask List", interactive=True
+            ),
+            gradio.Checkbox(value=False, label="Show Bounding Box", interactive=True),
+            gradio.Gallery(
+                masks, label="Output Gallery", interactive=True, type="numpy"
+            ),
+            gradio.DataFrame(
+                value=bboxes,
+                label="Box",
+                interactive=False,
+                headers=["XMin", "YMin", "XMax", "YMax"],
+                type="array",
+            ),
+        )
+
+
 with gradio.Blocks() as base_app:
     gradio.Markdown("# SegmentAnything2Assist")
     with gradio.Row():
         with gradio.Column():
             base_model_choice = gradio.Dropdown(
-                ['sam2_hiera_large', 'sam2_hiera_small', 'sam2_hiera_base_plus','sam2_hiera_tiny'], 
-                value = 'sam2_hiera_tiny',
-                label = "Model Choice"
-                )
+                [
+                    "sam2_hiera_large",
+                    "sam2_hiera_small",
+                    "sam2_hiera_base_plus",
+                    "sam2_hiera_tiny",
+                ],
+                value="sam2_hiera_tiny",
+                label="Model Choice",
+            )
         with gradio.Column():
             base_gpu_choice = gradio.Dropdown(
-                ['cpu', 'cuda'],
-                value = 'cuda',
-                label = "Device Choice"
+                ["cpu", "cuda"], value="cuda", label="Device Choice"
             )
-    base_model_choice.change(__change_base_model, inputs = [base_model_choice, base_gpu_choice])
-    base_gpu_choice.change(__change_base_model, inputs = [base_model_choice, base_gpu_choice])
-    with gradio.Tab(label = "Image Segmentation", id = "image_tab") as image_tab:
-        gradio.Markdown("Image Segmentation", render = True)
+    base_model_choice.change(
+        __change_base_model, inputs=[base_model_choice, base_gpu_choice]
+    )
+    base_gpu_choice.change(
+        __change_base_model, inputs=[base_model_choice, base_gpu_choice]
+    )
+
+    # Image Segmentation
+    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):
-                gradio.Markdown("""
+            with gradio.Accordion("Image Annotation Documentation", open=False):
+                gradio.Markdown(
+                    """
                     Image annotation allows you to mark specific regions of an image with labels. 
                     In this app, you can annotate an image by drawing boxes and assigning labels to them. 
                     The labels can be either '+' or '-'. 
@@ -229,88 +362,223 @@ with gradio.Blocks() as base_app:
                     Note that the advanced options allow you to adjust the SAM mask threshold, maximum hole area, and maximum sprinkle area. 
                     These options control the sensitivity and accuracy of the segmentation process. 
                     Experiment with different settings to achieve the desired results. 
-                """)
-            image_input = gradio_image_annotation.image_annotator(example_image_annotation)
-            with gradio.Accordion("Advanced Options", open = False):
-                image_generate_SAM_mask_threshold = gradio.Slider(0.0, 1.0, 0.0, label = "SAM Mask Threshold")
-                image_generate_SAM_max_hole_area = gradio.Slider(0, 1000, 0, label = "SAM Max Hole Area")
-                image_generate_SAM_max_sprinkle_area = gradio.Slider(0, 1000, 0, label = "SAM Max Sprinkle Area")
+                """
+                )
+            image_input = gradio_image_annotation.image_annotator(
+                example_image_annotation
+            )
+            with gradio.Accordion("Advanced Options", open=False):
+                image_generate_SAM_mask_threshold = gradio.Slider(
+                    0.0, 1.0, 0.0, label="SAM Mask Threshold"
+                )
+                image_generate_SAM_max_hole_area = gradio.Slider(
+                    0, 1000, 0, label="SAM Max Hole Area"
+                )
+                image_generate_SAM_max_sprinkle_area = gradio.Slider(
+                    0, 1000, 0, label="SAM Max Sprinkle Area"
+                )
             image_generate_mask_button = gradio.Button("Generate Mask")
-            image_output = gradio_imageslider.ImageSlider()
-            image_output_mode = gradio.Radio(["Segment", "Mask"], value = "Segment", label = "Output Mode")
-
-            image_input.change(__post_process_annotator_inputs, inputs = [image_input])
-            image_generate_mask_button.click(__generate_mask, inputs = [
-                image_input, 
-                image_generate_SAM_mask_threshold, 
-                image_generate_SAM_max_hole_area, 
-                image_generate_SAM_max_sprinkle_area, 
-                image_output_mode
-                ], 
-                outputs = [image_output])
-            image_output_mode.change(__change_output_mode, inputs = [image_input, image_output_mode], outputs = [image_output])
-    with gradio.Tab(label = "Auto Segmentation", id = "auto_tab"):
-        gradio.Markdown("Auto Segmentation", render = True)
+            with gradio.Row():
+                with gradio.Column():
+                    image_output_mode = gradio.Radio(
+                        ["Segment", "Mask"], value="Segment", label="Output Mode"
+                    )
+                with gradio.Column(scale=3):
+                    image_output = gradio_imageslider.ImageSlider()
+
+            with gradio.Accordion("Debug", open=DEBUG, visible=DEBUG):
+                __image_point_coords = gradio.DataFrame(
+                    value=[["+", 886, 551], ["-", 1239, 576]],
+                    label="Point Coords",
+                    interactive=False,
+                    headers=["Label", "X", "Y"],
+                )
+                __image_box = gradio.DataFrame(
+                    value=[[254, 466, 1347, 1047]],
+                    label="Box",
+                    interactive=False,
+                    headers=["XMin", "YMin", "XMax", "YMax"],
+                )
+                __current_mask = gradio.Image(label="Current Mask", interactive=False)
+                __current_segment = gradio.Image(
+                    label="Current Segment", interactive=False
+                )
+
+            # image_input.change(__post_process_annotator_inputs, inputs = [image_input])
+            image_generate_mask_button.click(
+                __generate_mask,
+                inputs=[
+                    image_input,
+                    image_generate_SAM_mask_threshold,
+                    image_generate_SAM_max_hole_area,
+                    image_generate_SAM_max_sprinkle_area,
+                    image_output_mode,
+                ],
+                outputs=[
+                    image_output,
+                    __image_point_coords,
+                    __image_box,
+                    __current_mask,
+                    __current_segment,
+                ],
+            )
+            image_output_mode.change(
+                __change_output_mode,
+                inputs=[
+                    image_input,
+                    image_output_mode,
+                    __current_mask,
+                    __current_segment,
+                ],
+                outputs=[image_output],
+            )
+
+    # Auto Segmentation
+    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):
-                gradio.Markdown("""
-                """)
+            with gradio.Accordion("Auto Annotation Documentation", open=False):
+                gradio.Markdown(
+                    """
+                """
+                )
             auto_input = gradio.Image("assets/cars.jpg")
-            with gradio.Accordion("Advanced Options", open = False):
-                auto_generate_SAM_points_per_side = gradio.Slider(1, 64, 32, 1, label = "Points Per Side", interactive = True)
-                auto_generate_SAM_points_per_batch = gradio.Slider(1, 64, 32, 1, label = "Points Per Batch", interactive = True)
-                auto_generate_SAM_pred_iou_thresh = gradio.Slider(0.0, 1.0, 0.8, 1, label = "Pred IOU Threshold", interactive = True)
-                auto_generate_SAM_stability_score_thresh = gradio.Slider(0.0, 1.0, 0.95, label = "Stability Score Threshold", interactive = True)
-                auto_generate_SAM_stability_score_offset = gradio.Slider(0.0, 1.0, 1.0, label = "Stability Score Offset", interactive = True)
-                auto_generate_SAM_mask_threshold = gradio.Slider(0.0, 1.0, 0.0, label = "Mask Threshold", interactive = True)
-                auto_generate_SAM_box_nms_thresh = gradio.Slider(0.0, 1.0, 0.7, label = "Box NMS Threshold", interactive = True)
-                auto_generate_SAM_crop_n_layers = gradio.Slider(0, 10, 0, 1, label = "Crop N Layers", interactive = True)
-                auto_generate_SAM_crop_nms_thresh = gradio.Slider(0.0, 1.0, 0.7, label = "Crop NMS Threshold", interactive = True)
-                auto_generate_SAM_crop_overlay_ratio = gradio.Slider(0.0, 1.0, 512 / 1500, label = "Crop Overlay Ratio", interactive = True)
-                auto_generate_SAM_crop_n_points_downscale_factor = gradio.Slider(1, 10, 1, label = "Crop N Points Downscale Factor", interactive = True)
-                auto_generate_SAM_min_mask_region_area = gradio.Slider(0, 1000, 0, label = "Min Mask Region Area", interactive = True)
-                auto_generate_SAM_use_m2m = gradio.Checkbox(label = "Use M2M", interactive = True)
-                auto_generate_SAM_multimask_output = gradio.Checkbox(value = True, label = "Multi Mask Output", interactive = True)
+            with gradio.Accordion("Advanced Options", open=False):
+                auto_generate_SAM_points_per_side = gradio.Slider(
+                    1, 64, 12, 1, label="Points Per Side", interactive=True
+                )
+                auto_generate_SAM_points_per_batch = gradio.Slider(
+                    1, 64, 32, 1, label="Points Per Batch", interactive=True
+                )
+                auto_generate_SAM_pred_iou_thresh = gradio.Slider(
+                    0.0, 1.0, 0.8, 1, label="Pred IOU Threshold", interactive=True
+                )
+                auto_generate_SAM_stability_score_thresh = gradio.Slider(
+                    0.0, 1.0, 0.95, label="Stability Score Threshold", interactive=True
+                )
+                auto_generate_SAM_stability_score_offset = gradio.Slider(
+                    0.0, 1.0, 1.0, label="Stability Score Offset", interactive=True
+                )
+                auto_generate_SAM_mask_threshold = gradio.Slider(
+                    0.0, 1.0, 0.0, label="Mask Threshold", interactive=True
+                )
+                auto_generate_SAM_box_nms_thresh = gradio.Slider(
+                    0.0, 1.0, 0.7, label="Box NMS Threshold", interactive=True
+                )
+                auto_generate_SAM_crop_n_layers = gradio.Slider(
+                    0, 10, 0, 1, label="Crop N Layers", interactive=True
+                )
+                auto_generate_SAM_crop_nms_thresh = gradio.Slider(
+                    0.0, 1.0, 0.7, label="Crop NMS Threshold", interactive=True
+                )
+                auto_generate_SAM_crop_overlay_ratio = gradio.Slider(
+                    0.0, 1.0, 512 / 1500, label="Crop Overlay Ratio", interactive=True
+                )
+                auto_generate_SAM_crop_n_points_downscale_factor = gradio.Slider(
+                    1, 10, 1, label="Crop N Points Downscale Factor", interactive=True
+                )
+                auto_generate_SAM_min_mask_region_area = gradio.Slider(
+                    0, 1000, 0, label="Min Mask Region Area", interactive=True
+                )
+                auto_generate_SAM_use_m2m = gradio.Checkbox(
+                    label="Use M2M", interactive=True
+                )
+                auto_generate_SAM_multimask_output = gradio.Checkbox(
+                    value=True, label="Multi Mask Output", interactive=True
+                )
             auto_generate_button = gradio.Button("Generate Auto Mask")
             with gradio.Row():
                 with gradio.Column():
-                    auto_output_mode = gradio.Radio(["Segment", "Mask"], value = "Segment", label = "Output Mode", interactive = True)
-                    auto_output_list = gradio.CheckboxGroup([], value = [], label = "Mask List", interactive = False)
-                    auto_output_bbox = gradio.Checkbox(value = False, label = "Show Bounding Box", interactive = False)
-                with gradio.Column(scale = 3):
+                    auto_output_mode = gradio.Radio(
+                        ["Segment", "Mask"],
+                        value="Segment",
+                        label="Output Mode",
+                        interactive=True,
+                    )
+                    auto_output_list = gradio.CheckboxGroup(
+                        [], value=[], label="Mask List", interactive=False
+                    )
+                    auto_output_bbox = gradio.Checkbox(
+                        value=False, label="Show Bounding Box", interactive=False
+                    )
+                with gradio.Column(scale=3):
                     auto_output = gradio_imageslider.ImageSlider()
-            
+            with gradio.Accordion("Debug", open=DEBUG, visible=DEBUG):
+                __auto_output_gallery = gradio.Gallery(
+                    None, label="Output Gallery", interactive=False, type="numpy"
+                )
+                __auto_bbox = gradio.DataFrame(
+                    value=[[]],
+                    label="Box",
+                    interactive=False,
+                    headers=["XMin", "YMin", "XMax", "YMax"],
+                )
+
             auto_generate_button.click(
-                __generate_auto_mask, 
-                inputs = [
-                    auto_input, 
-                    auto_generate_SAM_points_per_side, 
-                    auto_generate_SAM_points_per_batch, 
-                    auto_generate_SAM_pred_iou_thresh, 
-                    auto_generate_SAM_stability_score_thresh, 
-                    auto_generate_SAM_stability_score_offset, 
-                    auto_generate_SAM_mask_threshold, 
-                    auto_generate_SAM_box_nms_thresh, 
-                    auto_generate_SAM_crop_n_layers, 
-                    auto_generate_SAM_crop_nms_thresh, 
-                    auto_generate_SAM_crop_overlay_ratio, 
-                    auto_generate_SAM_crop_n_points_downscale_factor, 
-                    auto_generate_SAM_min_mask_region_area, 
-                    auto_generate_SAM_use_m2m, 
-                    auto_generate_SAM_multimask_output, 
-                    auto_output_mode
+                __generate_auto_mask,
+                inputs=[
+                    auto_input,
+                    auto_generate_SAM_points_per_side,
+                    auto_generate_SAM_points_per_batch,
+                    auto_generate_SAM_pred_iou_thresh,
+                    auto_generate_SAM_stability_score_thresh,
+                    auto_generate_SAM_stability_score_offset,
+                    auto_generate_SAM_mask_threshold,
+                    auto_generate_SAM_box_nms_thresh,
+                    auto_generate_SAM_crop_n_layers,
+                    auto_generate_SAM_crop_nms_thresh,
+                    auto_generate_SAM_crop_overlay_ratio,
+                    auto_generate_SAM_crop_n_points_downscale_factor,
+                    auto_generate_SAM_min_mask_region_area,
+                    auto_generate_SAM_use_m2m,
+                    auto_generate_SAM_multimask_output,
+                    auto_output_mode,
                 ],
-                outputs = [
+                outputs=[
                     auto_output,
                     auto_output_list,
-                    auto_output_bbox
-                ]
+                    auto_output_bbox,
+                    __auto_output_gallery,
+                    __auto_bbox,
+                ],
             )
-            auto_output_list.change(__generate_multi_mask_output, inputs = [auto_input, auto_output_list, auto_output_mode, auto_output_bbox], outputs = [auto_output])
-            auto_output_bbox.change(__generate_multi_mask_output, inputs = [auto_input, auto_output_list, auto_output_mode, auto_output_bbox], outputs = [auto_output])
-            auto_output_mode.change(__generate_multi_mask_output, inputs = [auto_input, auto_output_list, auto_output_mode, auto_output_bbox], outputs = [auto_output])
-            
-            
+            auto_output_list.change(
+                __generate_multi_mask_output,
+                inputs=[
+                    auto_input,
+                    auto_output_list,
+                    auto_output_mode,
+                    auto_output_bbox,
+                    __auto_output_gallery,
+                    __auto_bbox,
+                ],
+                outputs=[auto_output],
+            )
+            auto_output_bbox.change(
+                __generate_multi_mask_output,
+                inputs=[
+                    auto_input,
+                    auto_output_list,
+                    auto_output_mode,
+                    auto_output_bbox,
+                    __auto_output_gallery,
+                    __auto_bbox,
+                ],
+                outputs=[auto_output],
+            )
+            auto_output_mode.change(
+                __generate_multi_mask_output,
+                inputs=[
+                    auto_input,
+                    auto_output_list,
+                    auto_output_mode,
+                    auto_output_bbox,
+                    __auto_output_gallery,
+                    __auto_bbox,
+                ],
+                outputs=[auto_output],
+            )
+
+
 if __name__ == "__main__":
     base_app.launch()
-

src/SegmentAnything2Assist.py

@@ -14,204 +14,253 @@ import cv2
 
 SAM2_MODELS = {
     "sam2_hiera_tiny": {
-      "download_url": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt",
-      "model_path": ".tmp/checkpoints/sam2_hiera_tiny.pt",
-      "config_file": "sam2_hiera_t.yaml"
+        "download_url": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_tiny.pt",
+        "model_path": ".tmp/checkpoints/sam2_hiera_tiny.pt",
+        "config_file": "sam2_hiera_t.yaml",
     },
     "sam2_hiera_small": {
-      "download_url": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt",
-      "model_path": ".tmp/checkpoints/sam2_hiera_small.pt",
-      "config_file": "sam2_hiera_s.yaml"
+        "download_url": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt",
+        "model_path": ".tmp/checkpoints/sam2_hiera_small.pt",
+        "config_file": "sam2_hiera_s.yaml",
     },
     "sam2_hiera_base_plus": {
-      "download_url": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt",
-      "model_path": ".tmp/checkpoints/sam2_hiera_base_plus.pt",
-      "config_file": "sam2_hiera_b+.yaml"
+        "download_url": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_base_plus.pt",
+        "model_path": ".tmp/checkpoints/sam2_hiera_base_plus.pt",
+        "config_file": "sam2_hiera_b+.yaml",
     },
     "sam2_hiera_large": {
-      "download_url": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt",
-      "model_path": ".tmp/checkpoints/sam2_hiera_large.pt",
-      "config_file": "sam2_hiera_l.yaml"
+        "download_url": "https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt",
+        "model_path": ".tmp/checkpoints/sam2_hiera_large.pt",
+        "config_file": "sam2_hiera_l.yaml",
     },
 }
-      
+
+
 class SegmentAnything2Assist:
-  def __init__(
-    self,
-    model_name: str | typing.Literal["sam2_hiera_tiny", "sam2_hiera_small", "sam2_hiera_base_plus", "sam2_hiera_large"] = "sam2_hiera_small",
-    configuration: str |typing.Literal["Automatic Mask Generator", "Image"] = "Automatic Mask Generator",
-    download_url: str | None = None,
-    model_path: str | None = None,
-    download: bool = True,
-    device: str | torch.device = torch.device("cpu"),
-    verbose: bool = True
-  ) -> None:
-    assert model_name in SAM2_MODELS.keys(), f"`model_name` should be either one of {list(SAM2_MODELS.keys())}"
-    assert configuration in ["Automatic Mask Generator", "Image"]
-
-    self.model_name = model_name
-    self.configuration = configuration
-    self.config_file = SAM2_MODELS[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"]
-    self.model_path = model_path if model_path is not None else SAM2_MODELS[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__::Configuration: {self.configuration}")
-      print(f"SegmentAnything2Assist::__init__::Download URL: {self.download_url}")
-      print(f"SegmentAnything2Assist::__init__::Default Path: {self.model_path}")
-      print(f"SegmentAnything2Assist::__init__::Configuration File: {self.config_file}")
-
-    if download:
-      self.download_model()
-
-    if self.is_model_available():
-      self.sam2 = sam2.build_sam.build_sam2(config_file = self.config_file, ckpt_path = self.model_path, device = self.device)
-      if self.verbose:
-        print("SegmentAnything2Assist::__init__::SAM2 is loaded.")
-    else:
-      self.sam2 = None
-      if self.verbose:
-        print("SegmentAnything2Assist::__init__::SAM2 is not loaded.")
-
-
-  def is_model_available(self) -> bool:
-    ret = os.path.exists(self.model_path)
-    if self.verbose:
-      print(f"SegmentAnything2Assist::is_model_available::{ret}")
-    return ret
-
-  def load_model(self) -> None:
-    if self.is_model_available():
-      self.sam2 = sam2.build_sam(checkpoint = self.model_path)
-
-  def download_model(
-    self, 
-    force: bool = False
-  ) -> None:
-    if not force and self.is_model_available():
-        print(f"{self.model_path} already exists. Skipping download.")
-        return
-
-    response = requests.get(self.download_url, stream=True)
-    total_size = int(response.headers.get('content-length', 0))
-
-    with open(self.model_path, 'wb') as file, tqdm.tqdm(total = total_size, unit = 'B', unit_scale = True) as progress_bar:
-        for data in response.iter_content(chunk_size = 1024):
-            file.write(data)
-            progress_bar.update(len(data))
-
-  def generate_automatic_masks(
-    self,
-    image,
-    points_per_side = 32,
-    points_per_batch = 32,
-    pred_iou_thresh = 0.8,
-    stability_score_thresh = 0.95,
-    stability_score_offset = 1.0,
-    mask_threshold = 0.0,
-    box_nms_thresh = 0.7,
-    crop_n_layers = 0,
-    crop_nms_thresh = 0.7,
-    crop_overlay_ratio = 512 / 1500,
-    crop_n_points_downscale_factor = 1,
-    min_mask_region_area = 0,
-    use_m2m = False,
-    multimask_output = True
-  ):
-    if self.sam2 is None:
-      print("SegmentAnything2Assist::generate_automatic_masks::SAM2 is not loaded.")
-      return None
-    
-    generator = sam2.automatic_mask_generator.SAM2AutomaticMaskGenerator(
-      model = self.sam2,
-      points_per_side = points_per_side,
-      points_per_batch = points_per_batch,
-      pred_iou_thresh = pred_iou_thresh,
-      stability_score_thresh = stability_score_thresh,
-      stability_score_offset = stability_score_offset,
-      mask_threshold = mask_threshold,
-      box_nms_thresh = box_nms_thresh,
-      crop_n_layers = crop_n_layers,
-      crop_nms_thresh = crop_nms_thresh,
-      crop_overlay_ratio = crop_overlay_ratio,
-      crop_n_points_downscale_factor = crop_n_points_downscale_factor,
-      min_mask_region_area = min_mask_region_area,
-      use_m2m = use_m2m,
-      multimask_output = multimask_output
-      )
-    masks = generator.generate(image)
-    
-    pickle.dump(masks, open(".tmp/auto_masks.pkl", "wb"))
-    
-    return masks
-  
-  def generate_masks_from_image(
-    self,
-    image,
-    point_coords,
-    point_labels,
-    box,
-    mask_threshold = 0.0,
-    max_hole_area = 0.0,
-    max_sprinkle_area = 0.0
-  ):
-    generator = sam2.sam2_image_predictor.SAM2ImagePredictor(
-      self.sam2,
-      mask_threshold = mask_threshold,
-      max_hole_area = max_hole_area,
-      max_sprinkle_area = max_sprinkle_area
-    )
-    generator.set_image(image)
-    
-    masks_chw, mask_iou, mask_low_logits = generator.predict(
-      point_coords = numpy.array(point_coords) if point_coords is not None else None,
-      point_labels = numpy.array(point_labels) if point_labels is not None else None,
-      box = numpy.array(box) if box is not None else None,
-      multimask_output = False
-    )
-    
-    return masks_chw, mask_iou
-  
-  def apply_mask_to_image(
-    self,
-    image,
-    mask
-  ):
-    mask = numpy.array(mask)
-    mask = numpy.where(mask > 0, 255, 0).astype(numpy.uint8)
-    segment = cv2.bitwise_and(image, image, mask = mask)
-    return mask, segment
-
-  def apply_auto_mask_to_image(
-    self,
-    image,
-    auto_list
-  ):   
-    if not os.path.exists(".tmp/auto_masks.pkl"):
-      return
-
-    masks = pickle.load(open(".tmp/auto_masks.pkl", "rb"))
-      
-    image_with_bounding_boxes = image.copy()
-    all_masks = None
-    for _ in auto_list:
-      mask = numpy.array(masks[_]['segmentation'])
-      mask = numpy.where(mask == True, 255, 0).astype(numpy.uint8)
-      bbox = masks[_]["bbox"]
-      if all_masks is None:
-        all_masks = mask
-      else:
-        all_masks = cv2.bitwise_or(all_masks, mask)
-      
-      random_color = numpy.random.randint(0, 255, size = 3)
-      image_with_bounding_boxes = cv2.rectangle(image_with_bounding_boxes, (int(bbox[0]), int(bbox[1])), (int(bbox[0] + bbox[2]), int(bbox[1] + bbox[3])), random_color.tolist(), 2)
-      image_with_bounding_boxes = cv2.putText(image_with_bounding_boxes, f"{_ + 1}", (int(bbox[0]), int(bbox[1]) - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, random_color.tolist(), 2)
-    
-    all_masks = numpy.where(all_masks > 0, 255, 0).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 __init__(
+        self,
+        model_name: (
+            str
+            | typing.Literal[
+                "sam2_hiera_tiny",
+                "sam2_hiera_small",
+                "sam2_hiera_base_plus",
+                "sam2_hiera_large",
+            ]
+        ) = "sam2_hiera_small",
+        configuration: (
+            str | typing.Literal["Automatic Mask Generator", "Image"]
+        ) = "Automatic Mask Generator",
+        download_url: str | None = None,
+        model_path: str | None = None,
+        download: bool = True,
+        device: str | torch.device = torch.device("cpu"),
+        verbose: bool = True,
+    ) -> None:
+        assert (
+            model_name in SAM2_MODELS.keys()
+        ), f"`model_name` should be either one of {list(SAM2_MODELS.keys())}"
+        assert configuration in ["Automatic Mask Generator", "Image"]
+
+        self.model_name = model_name
+        self.configuration = configuration
+        self.config_file = SAM2_MODELS[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"]
+        )
+        self.model_path = (
+            model_path
+            if model_path is not None
+            else SAM2_MODELS[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__::Configuration: {self.configuration}"
+            )
+            print(
+                f"SegmentAnything2Assist::__init__::Download URL: {self.download_url}"
+            )
+            print(f"SegmentAnything2Assist::__init__::Default Path: {self.model_path}")
+            print(
+                f"SegmentAnything2Assist::__init__::Configuration File: {self.config_file}"
+            )
+
+        if download:
+            self.download_model()
+
+        if self.is_model_available():
+            self.sam2 = sam2.build_sam.build_sam2(
+                config_file=self.config_file,
+                ckpt_path=self.model_path,
+                device=self.device,
+            )
+            if self.verbose:
+                print("SegmentAnything2Assist::__init__::SAM2 is loaded.")
+        else:
+            self.sam2 = None
+            if self.verbose:
+                print("SegmentAnything2Assist::__init__::SAM2 is not loaded.")
+
+    def is_model_available(self) -> bool:
+        ret = os.path.exists(self.model_path)
+        if self.verbose:
+            print(f"SegmentAnything2Assist::is_model_available::{ret}")
+        return ret
+
+    def load_model(self) -> None:
+        if self.is_model_available():
+            self.sam2 = sam2.build_sam(checkpoint=self.model_path)
+
+    def download_model(self, force: bool = False) -> None:
+        if not force and self.is_model_available():
+            print(f"{self.model_path} already exists. Skipping download.")
+            return
+
+        response = requests.get(self.download_url, stream=True)
+        total_size = int(response.headers.get("content-length", 0))
+
+        with open(self.model_path, "wb") as file, tqdm.tqdm(
+            total=total_size, unit="B", unit_scale=True
+        ) as progress_bar:
+            for data in response.iter_content(chunk_size=1024):
+                file.write(data)
+                progress_bar.update(len(data))
+
+    def generate_automatic_masks(
+        self,
+        image,
+        points_per_side=32,
+        points_per_batch=32,
+        pred_iou_thresh=0.8,
+        stability_score_thresh=0.95,
+        stability_score_offset=1.0,
+        mask_threshold=0.0,
+        box_nms_thresh=0.7,
+        crop_n_layers=0,
+        crop_nms_thresh=0.7,
+        crop_overlay_ratio=512 / 1500,
+        crop_n_points_downscale_factor=1,
+        min_mask_region_area=0,
+        use_m2m=False,
+        multimask_output=True,
+    ):
+        if self.sam2 is None:
+            print(
+                "SegmentAnything2Assist::generate_automatic_masks::SAM2 is not loaded."
+            )
+            return None
+
+        generator = sam2.automatic_mask_generator.SAM2AutomaticMaskGenerator(
+            model=self.sam2,
+            points_per_side=points_per_side,
+            points_per_batch=points_per_batch,
+            pred_iou_thresh=pred_iou_thresh,
+            stability_score_thresh=stability_score_thresh,
+            stability_score_offset=stability_score_offset,
+            mask_threshold=mask_threshold,
+            box_nms_thresh=box_nms_thresh,
+            crop_n_layers=crop_n_layers,
+            crop_nms_thresh=crop_nms_thresh,
+            crop_overlay_ratio=crop_overlay_ratio,
+            crop_n_points_downscale_factor=crop_n_points_downscale_factor,
+            min_mask_region_area=min_mask_region_area,
+            use_m2m=use_m2m,
+            multimask_output=multimask_output,
+        )
+        masks = generator.generate(image)
+        segmentation_masks = [mask for mask in masks]
+        segmentation_masks = [
+            numpy.where(mask["segmentation"] == True, 255, 0).astype(numpy.uint8)
+            for mask in segmentation_masks
+        ]
+        segmentation_masks = [
+            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
+
+    def generate_masks_from_image(
+        self,
+        image,
+        point_coords,
+        point_labels,
+        box,
+        mask_threshold=0.0,
+        max_hole_area=0.0,
+        max_sprinkle_area=0.0,
+    ):
+        generator = sam2.sam2_image_predictor.SAM2ImagePredictor(
+            self.sam2,
+            mask_threshold=mask_threshold,
+            max_hole_area=max_hole_area,
+            max_sprinkle_area=max_sprinkle_area,
+        )
+        generator.set_image(image)
+
+        masks_chw, mask_iou, mask_low_logits = generator.predict(
+            point_coords=(
+                numpy.array(point_coords) if point_coords is not None else None
+            ),
+            point_labels=(
+                numpy.array(point_labels) if point_labels is not None else None
+            ),
+            box=numpy.array(box) if box is not None else None,
+            multimask_output=False,
+        )
+
+        return masks_chw, mask_iou
+
+    def apply_mask_to_image(self, image, mask):
+        mask = numpy.array(mask)
+        mask = numpy.where(mask > 0, 255, 0).astype(numpy.uint8)
+        segment = cv2.bitwise_and(image, image, mask=mask)
+        return mask, segment
+
+    def apply_auto_mask_to_image(self, image, auto_list, masks, bboxes):
+        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)
+
+            bbox = bboxes[_]
+            if all_masks is None:
+                all_masks = mask
+            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,
+                (int(bbox[0]), int(bbox[1])),
+                (int(bbox[0] + bbox[2]), int(bbox[1] + bbox[3])),
+                random_color.tolist(),
+                2,
+            )
+            image_with_bounding_boxes = cv2.putText(
+                image_with_bounding_boxes,
+                f"{_ + 1}",
+                (int(bbox[0]), int(bbox[1]) - 10),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.5,
+                random_color.tolist(),
+                2,
+            )
+
+        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