initial commit

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/cars.jpg filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1 @@
+.tmp/

README.md

@@ -0,0 +1,14 @@
+---
+title: Segment Anything 2 Assist
+emoji: 👁
+colorFrom: indigo
+colorTo: red
+sdk: gradio
+sdk_version: 4.42.0
+app_file: SegmentAnything2AssistApp.py
+pinned: true
+license: bsd-3-clause
+short_description: A tool to use SAM2 on images.
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

SegmentAnything2AssistApp.py

@@ -0,0 +1,312 @@
+import gradio
+import gradio_image_annotation 
+import gradio_imageslider
+import spaces
+import torch
+
+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}]
+}
+
+VERBOSE = True
+
+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
+    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)
+    except:
+        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
+    new_boxes = []
+    __image_point_coords = []
+    __image_point_labels = []
+    __image_box = []
+    
+    b_has_box = False
+    for box in value["boxes"]:
+        if box['label'] == '':
+            if not b_has_box:
+                new_box = box.copy()
+                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'] == '-':
+            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_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
+
+    if VERBOSE:
+        print("SegmentAnything2AssistApp::____post_process_annotator_inputs::Done.")
+
+
+
+@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
+    global segment_anything2assist
+    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,
+        mask_threshold,
+        max_hole_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])
+
+    if VERBOSE:
+        print("SegmentAnything2AssistApp::__generate_mask::Masks and Segments created.")
+
+    if image_output_mode == "Mask":
+        return [value["image"], __current_mask]
+    elif image_output_mode == "Segment":
+        return [value["image"], __current_segment]
+    else:
+        gradio.Warning("This is an issue, please report the problem!", duration=5)
+        return gradio_imageslider.ImageSlider(render = True)
+
+def __change_output_mode(image_input, radio):
+    global __current_mask, __current_segment
+    global __image_point_coords, __image_point_labels, __image_box
+    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)
+    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):
+    global segment_anything2assist
+    image_with_bbox, mask, segment = segment_anything2assist.apply_auto_mask_to_image(image, [int(i) - 1 for i in auto_list])
+    
+    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)
+def __generate_auto_mask(
+    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,
+    output_mode
+    ):
+    global segment_anything2assist
+    if VERBOSE:
+        print("SegmentAnything2AssistApp::__generate_auto_mask::Called.")    
+    
+    __auto_masks = 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:
+        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)
+    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)    
+    
+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"
+                )
+        with gradio.Column():
+            base_gpu_choice = gradio.Dropdown(
+                ['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)
+        with gradio.Column():
+            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 '-'. 
+                    To annotate an image, simply click and drag to draw a box around the desired region. 
+                    You can add multiple boxes with different labels. 
+                    Once you have annotated the image, click the 'Generate Mask' button to generate a mask based on the annotations. 
+                    The mask can be either a binary mask or a segmented mask, depending on the selected output mode. 
+                    You can switch between the output modes using the radio buttons. 
+                    If you make any changes to the annotations or the output mode, you need to regenerate the mask by clicking the button again. 
+                    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_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.Column():
+            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)
+            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 = gradio_imageslider.ImageSlider()
+            
+            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
+                ],
+                outputs = [
+                    auto_output,
+                    auto_output_list,
+                    auto_output_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])
+            
+            
+if __name__ == "__main__":
+    base_app.launch()
+

requirements.txt

@@ -0,0 +1,72 @@
+aiofiles==23.2.1
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+anyio==4.4.0
+certifi==2024.7.4
+charset-normalizer==3.3.2
+click==8.1.7
+contourpy==1.2.1
+cycler==0.12.1
+fastapi==0.112.2
+ffmpy==0.4.0
+filelock==3.15.4
+fonttools==4.53.1
+fsspec==2024.6.1
+gradio==4.42.0
+gradio_client==1.3.0
+gradio_image_annotation==0.2.3
+gradio_imageslider==0.0.20
+h11==0.14.0
+httpcore==1.0.5
+httpx==0.27.0
+huggingface-hub==0.24.6
+hydra-core==1.3.2
+idna==3.7
+importlib_resources==6.4.4
+iopath==0.1.10
+Jinja2==3.1.4
+kiwisolver==1.4.5
+markdown-it-py==3.0.0
+MarkupSafe==2.1.5
+matplotlib==3.9.2
+mdurl==0.1.2
+mpmath==1.3.0
+networkx==3.2.1
+numpy==2.1.0
+omegaconf==2.3.0
+opencv-python==4.10.0.84
+orjson==3.10.7
+packaging==24.1
+pandas==2.2.2
+pillow==10.4.0
+portalocker==2.10.1
+psutil==5.9.8
+pydantic==2.8.2
+pydantic_core==2.20.1
+pydub==0.25.1
+Pygments==2.18.0
+pyparsing==3.1.4
+python-dateutil==2.9.0.post0
+python-multipart==0.0.9
+pytz==2024.1
+PyYAML==6.0.2
+requests==2.32.3
+rich==13.7.1
+ruff==0.6.2
+SAM-2 @ git+https://github.com/facebookresearch/segment-anything-2.git@7e1596c0b6462eb1d1ba7e1492430fed95023598
+semantic-version==2.10.0
+setuptools==73.0.1
+shellingham==1.5.4
+six==1.16.0
+sniffio==1.3.1
+spaces==0.29.3
+starlette==0.38.2
+sympy==1.13.2
+tomlkit==0.12.0
+tqdm==4.66.5
+typer==0.12.5
+typing_extensions==4.12.2
+tzdata==2024.1
+urllib3==2.2.2
+uvicorn==0.30.6
+websockets==12.0

src/SegmentAnything2Assist.py

@@ -0,0 +1,217 @@
+import typing
+import os
+import sam2.sam2_image_predictor
+import tqdm
+import requests
+import torch
+import numpy
+import pickle
+
+import sam2.build_sam
+import sam2.automatic_mask_generator
+
+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"
+    },
+    "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"
+    },
+    "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"
+    },
+    "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"
+    },
+}
+      
+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("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