Initial Commit

README.md

@@ -1,13 +1,172 @@
----
-title: MobileSAM
-emoji: 🦀
-colorFrom: blue
-colorTo: gray
-sdk: gradio
-sdk_version: 3.35.2
-app_file: app.py
-pinned: false
-license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+<p float="center">
+  <img src="assets/logo.png?raw=true" width="99.1%" />
+</p>
+
+# Faster Segment Anything (MobileSAM)
+
+:pushpin: MobileSAM paper is available at [paper link](https://arxiv.org/pdf/2306.14289.pdf).
+
+![MobileSAM](assets/model_diagram.jpg?raw=true)
+
+<p float="left">
+  <img src="assets/mask_comparision.jpg?raw=true" width="99.1%" />
+</p>
+
+**MobileSAM** performs on par with the original SAM (at least visually) and keeps exactly the same pipeline as the original SAM except for a change on the image encoder. Specifically, we replace the original heavyweight ViT-H encoder (632M) with a much smaller Tiny-ViT (5M). On a single GPU, MobileSAM runs around 12ms per image: 8ms on the image encoder and 4ms on the mask decoder. 
+
+The comparison of ViT-based image encoder is summarzed as follows: 
+
+Image Encoder                                      | Original SAM | MobileSAM 
+:-----------------------------------------:|:---------|:-----:
+Paramters      |  611M   | 5M
+Speed      |  452ms  | 8ms
+
+Original SAM and MobileSAM have exactly the same prompt-guided mask decoder: 
+
+Mask Decoder                                      | Original SAM | MobileSAM 
+:-----------------------------------------:|:---------|:-----:
+Paramters      |  3.876M   | 3.876M
+Speed      |  4ms  | 4ms
+
+The comparison of the whole pipeline is summarzed as follows: 
+Whole Pipeline (Enc+Dec)                                      | Original SAM | MobileSAM 
+:-----------------------------------------:|:---------|:-----:
+Paramters      |  615M   | 9.66M
+Speed      |  456ms  | 12ms
+
+**Original SAM and MobileSAM with a (single) point as the prompt.**  
+
+<p float="left">
+  <img src="assets/mask_point.jpg?raw=true" width="99.1%" />
+</p>
+
+**Original SAM and MobileSAM with a box as the prompt.** 
+<p float="left">
+  <img src="assets/mask_box.jpg?raw=true" width="99.1%" />
+</p>
+
+**Is MobileSAM faster and smaller than FastSAM? Yes, to our knowledge!** 
+MobileSAM is around 7 times smaller and around 5 times faster than the concurrent FastSAM. 
+The comparison of the whole pipeline is summarzed as follows: 
+Whole Pipeline (Enc+Dec)                                      | FastSAM | MobileSAM 
+:-----------------------------------------:|:---------|:-----:
+Paramters      |  68M   | 9.66M
+Speed      |  64ms  |12ms
+
+**Is MobileSAM better than FastSAM for performance? Yes, to our knowledge!** 
+FastSAM cannot work with a single prompt as the original SAM or our MobileSAM. Therefore, we compare the mIoU with two prompt points (with different pixel distances) and show the resutls as follows. Our MobileSAM is much better than FastSAM under this setup. 
+mIoU                                     | FastSAM | MobileSAM 
+:-----------------------------------------:|:---------|:-----:
+100      |  0.27   | 0.73
+200      |  0.33  |0.71
+300      |  0.37  |0.74
+400      |  0.41  |0.73
+500      |  0.41  |0.73
+
+
+
+
+**How to Adapt from SAM to MobileSAM?** Since MobileSAM keeps exactly the same pipeline as the original SAM, we inherit pre-processing, post-processing, and all other interfaces from the original SAM. The users who use the original SAM can adapt to MobileSAM with zero effort, by assuming everything is exactly the same except for a smaller image encoder in the SAM.
+
+**How is MobileSAM trained?** MobileSAM is trained on a single GPU with 100k datasets (1% of the original images) for less than a day. The training code will be available soon.
+
+
+
+## Installation
+
+The code requires `python>=3.8`, as well as `pytorch>=1.7` and `torchvision>=0.8`. Please follow the instructions [here](https://pytorch.org/get-started/locally/) to install both PyTorch and TorchVision dependencies. Installing both PyTorch and TorchVision with CUDA support is strongly recommended.
+
+Install Mobile Segment Anything:
+
+```
+pip install git+https://github.com/ChaoningZhang/MobileSAM.git
+```
+
+or clone the repository locally and install with
+
+```
+git clone git@github.com:ChaoningZhang/MobileSAM.git
+cd MobileSAM; pip install -e .
+```
+
+
+## <a name="GettingStarted"></a>Getting Started
+The MobileSAM can be loaded in the following ways:
+
+```
+from mobile_encoder.setup_mobile_sam import setup_model
+checkpoint = torch.load('../weights/mobile_sam.pt')
+mobile_sam = setup_model()
+mobile_sam.load_state_dict(checkpoint,strict=True)
+```
+
+Then the model can be easily used in just a few lines to get masks from a given prompt:
+
+```
+from segment_anything import SamPredictor
+device = "cuda"
+mobile_sam.to(device=device)
+mobile_sam.eval()
+predictor = SamPredictor(mobile_sam)
+predictor.set_image(<your_image>)
+masks, _, _ = predictor.predict(<input_prompts>)
+```
+
+or generate masks for an entire image:
+
+```
+from segment_anything import SamAutomaticMaskGenerator
+
+mask_generator = SamAutomaticMaskGenerator(mobile_sam)
+masks = mask_generator.generate(<your_image>)
+```
+
+
+## BibTex of our MobileSAM
+If you use MobileSAM in your research, please use the following BibTeX entry. :mega: Thank you!
+
+```bibtex
+@article{mobile_sam,
+  title={Faster Segment Anything: Towards Lightweight SAM for Mobile Applications},
+  author={Zhang, Chaoning and Han, Dongshen and Qiao, Yu and Kim, Jung Uk and Bae, Sung Ho and Lee, Seungkyu and Hong, Choong Seon},
+  journal={arXiv preprint arXiv:2306.14289},
+  year={2023}
+}
+```
+
+## Acknowledgement
+
+<details>
+<summary>
+<a href="https://github.com/facebookresearch/segment-anything">SAM</a> (Segment Anything) [<b>bib</b>]
+</summary>
+
+```bibtex
+@article{kirillov2023segany,
+  title={Segment Anything}, 
+  author={Kirillov, Alexander and Mintun, Eric and Ravi, Nikhila and Mao, Hanzi and Rolland, Chloe and Gustafson, Laura and Xiao, Tete and Whitehead, Spencer and Berg, Alexander C. and Lo, Wan-Yen and Doll{\'a}r, Piotr and Girshick, Ross},
+  journal={arXiv:2304.02643},
+  year={2023}
+}
+```
+</details>
+
+
+
+<details>
+<summary>
+<a href="https://github.com/microsoft/Cream/tree/main/TinyViT">TinyViT</a> (TinyViT: Fast Pretraining Distillation for Small Vision Transformers) [<b>bib</b>]
+</summary>
+
+```bibtex
+@InProceedings{tiny_vit,
+  title={TinyViT: Fast Pretraining Distillation for Small Vision Transformers},
+  author={Wu, Kan and Zhang, Jinnian and Peng, Houwen and Liu, Mengchen and Xiao, Bin and Fu, Jianlong and Yuan, Lu},
+  booktitle={European conference on computer vision (ECCV)},
+  year={2022}
+```
+</details>
+
+
+
+

app.py

@@ -0,0 +1,348 @@
+import gradio as gr
+import numpy as np
+import torch
+import os
+from mobile_sam import SamAutomaticMaskGenerator, SamPredictor, sam_model_registry
+from PIL import ImageDraw
+from utils.tools import box_prompt, format_results, point_prompt
+from utils.tools_gradio import fast_process
+
+# Most of our demo code is from [FastSAM Demo](https://huggingface.co/spaces/An-619/FastSAM). Huge thanks for AN-619.
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Load the pre-trained model
+sam_checkpoint = "./mobile_sam.pt"
+model_type = "vit_t"
+
+mobile_sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+mobile_sam = mobile_sam.to(device=device)
+mobile_sam.eval()
+
+mask_generator = SamAutomaticMaskGenerator(mobile_sam)
+predictor = SamPredictor(mobile_sam)
+
+# Description
+title = "<center><strong><font size='8'>Faster Segment Anything(MobileSAM)<font></strong></center>"
+
+description_e = """This is a demo on Github project [Faster Segment Anything(MobileSAM) Model](https://github.com/ChaoningZhang/MobileSAM). Welcome to give a star ⭐️ to it.
+                
+                🎯 Upload an Image, segment it with Faster Segment Anything (Everything mode). The other modes will come soon.
+                
+                ⌛️ It takes about 5~ seconds to generate segment results. The concurrency_count of queue is 1, please wait for a moment when it is crowded.
+                
+                🚀 To get faster results, you can use a smaller input size and leave high_visual_quality unchecked.
+                
+                📣 You can also obtain the segmentation results of any Image through this Colab: [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://drive.google.com/file/d/1k6azd5wdOOYkFwi61uXoIHfP-qBzuoOu/view?usp=sharing)
+                
+                🏠 Check out our [Model Card 🏃](https://huggingface.co/dhkim2810/MobileSAM)
+
+                😚 Most of our demo code is from [FastSAM Demo](https://huggingface.co/spaces/An-619/FastSAM). Huge thanks for AN-619.
+                
+              """
+
+description_p = """ # 🎯 Instructions for points mode
+                This is a demo on Github project [Faster Segment Anything(MobileSAM) Model](https://github.com/ChaoningZhang/MobileSAM). Welcome to give a star ⭐️ to it.
+
+                🎯 Upload an Image, segment it with Faster Segment Anything (Everything mode). The other modes will come soon.
+                
+                ⌛️ It takes about 5~ seconds to generate segment results. The concurrency_count of queue is 1, please wait for a moment when it is crowded.
+
+                🚀 To get faster results, you can use a smaller input size and leave high_visual_quality unchecked.
+
+                📣 You can also obtain the segmentation results of any Image through this Colab: [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://drive.google.com/file/d/1jibN6HTQcC4C2okoaKLRzHIo_pS0Eeom/view?usp=sharing)
+
+                🏠 Check out our [Model Card 🏃](https://huggingface.co/dhkim2810/MobileSAM)
+
+
+                
+                1. Upload an image or choose an example.
+                
+                2. Choose the point label ('Add mask' means a positive point. 'Remove' Area means a negative point that is not segmented).
+                
+                3. Add points one by one on the image.
+                
+                4. Click the 'Segment with points prompt' button to get the segmentation results.
+                
+                **5. If you get Error, click the 'Clear points' button and try again may help.**
+
+              """
+
+examples = [
+    ["assets/sa_8776.jpg"],
+    ["assets/sa_414.jpg"],
+    ["assets/sa_1309.jpg"],
+    ["assets/sa_11025.jpg"],
+    ["assets/sa_561.jpg"],
+    ["assets/sa_192.jpg"],
+    ["assets/sa_10039.jpg"],
+    ["assets/sa_862.jpg"],
+]
+
+default_example = examples[0]
+
+css = "h1 { text-align: center } .about { text-align: justify; padding-left: 10%; padding-right: 10%; }"
+
+
+@torch.no_grad()
+def segment_everything(
+    image,
+    input_size=1024,
+    better_quality=False,
+    withContours=True,
+    use_retina=True,
+    mask_random_color=True,
+):
+    global mask_generator
+
+    input_size = int(input_size)
+    w, h = image.size
+    scale = input_size / max(w, h)
+    new_w = int(w * scale)
+    new_h = int(h * scale)
+    image = image.resize((new_w, new_h))
+
+    nd_image = np.array(image)
+    annotations = mask_generator.generate(nd_image)
+
+    fig = fast_process(
+        annotations=annotations,
+        image=image,
+        device=device,
+        scale=(1024 // input_size),
+        better_quality=better_quality,
+        mask_random_color=mask_random_color,
+        bbox=None,
+        use_retina=use_retina,
+        withContours=withContours,
+    )
+    return fig
+
+
+def segment_with_points(
+    image,
+    input_size=1024,
+    better_quality=False,
+    withContours=True,
+    use_retina=True,
+    mask_random_color=True,
+):
+    global global_points
+    global global_point_label
+
+    input_size = int(input_size)
+    w, h = image.size
+    scale = input_size / max(w, h)
+    new_w = int(w * scale)
+    new_h = int(h * scale)
+    image = image.resize((new_w, new_h))
+
+    scaled_points = np.array([[int(x * scale) for x in point] for point in global_points])
+    global_point_label = np.array(global_point_label)
+
+    nd_image = np.array(image)
+    predictor.set_image(nd_image)
+    masks, scores, logits = predictor.predict(
+        point_coords=scaled_points,
+        point_labels=global_point_label,
+        multimask_output=True,
+    )
+
+    results = format_results(masks, scores, logits, 0)
+
+    annotations, _ = point_prompt(
+        results, scaled_points, global_point_label, new_h, new_w
+    )
+    annotations = np.array([annotations])
+
+    fig = fast_process(
+        annotations=annotations,
+        image=image,
+        device=device,
+        scale=(1024 // input_size),
+        better_quality=better_quality,
+        mask_random_color=mask_random_color,
+        bbox=None,
+        use_retina=use_retina,
+        withContours=withContours,
+    )
+
+    global_points = []
+    global_point_label = []
+    # return fig, None
+    return fig, image
+
+
+def get_points_with_draw(image, label, evt: gr.SelectData):
+    global global_points
+    global global_point_label
+
+    x, y = evt.index[0], evt.index[1]
+    point_radius, point_color = 15, (255, 255, 0) if label == "Add Mask" else (
+        255,
+        0,
+        255,
+    )
+    global_points.append([x, y])
+    global_point_label.append(1 if label == "Add Mask" else 0)
+
+    print(x, y, label == "Add Mask")
+
+    # 创建一个可以在图像上绘图的对象
+    draw = ImageDraw.Draw(image)
+    draw.ellipse(
+        [(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)],
+        fill=point_color,
+    )
+    return image
+
+
+cond_img_e = gr.Image(label="Input", value=default_example[0], type="pil")
+cond_img_p = gr.Image(label="Input with points", value=default_example[0], type="pil")
+
+segm_img_e = gr.Image(label="Segmented Image", interactive=False, type="pil")
+segm_img_p = gr.Image(
+    label="Segmented Image with points", interactive=False, type="pil"
+)
+
+global_points = []
+global_point_label = []
+
+input_size_slider = gr.components.Slider(
+    minimum=512,
+    maximum=1024,
+    value=1024,
+    step=64,
+    label="Input_size",
+    info="Our model was trained on a size of 1024",
+)
+
+with gr.Blocks(css=css, title="Faster Segment Anything(MobileSAM)") as demo:
+    with gr.Row():
+        with gr.Column(scale=1):
+            # Title
+            gr.Markdown(title)
+
+    # with gr.Tab("Everything mode"):
+    #     # Images
+    #     with gr.Row(variant="panel"):
+    #         with gr.Column(scale=1):
+    #             cond_img_e.render()
+    #
+    #         with gr.Column(scale=1):
+    #             segm_img_e.render()
+    #
+    #     # Submit & Clear
+    #     with gr.Row():
+    #         with gr.Column():
+    #             input_size_slider.render()
+    #
+    #             with gr.Row():
+    #                 contour_check = gr.Checkbox(
+    #                     value=True,
+    #                     label="withContours",
+    #                     info="draw the edges of the masks",
+    #                 )
+    #
+    #                 with gr.Column():
+    #                     segment_btn_e = gr.Button(
+    #                         "Segment Everything", variant="primary"
+    #                     )
+    #                     clear_btn_e = gr.Button("Clear", variant="secondary")
+    #
+    #             gr.Markdown("Try some of the examples below ⬇️")
+    #             gr.Examples(
+    #                 examples=examples,
+    #                 inputs=[cond_img_e],
+    #                 outputs=segm_img_e,
+    #                 fn=segment_everything,
+    #                 cache_examples=True,
+    #                 examples_per_page=4,
+    #             )
+    #
+    #         with gr.Column():
+    #             with gr.Accordion("Advanced options", open=False):
+    #                 # text_box = gr.Textbox(label="text prompt")
+    #                 with gr.Row():
+    #                     mor_check = gr.Checkbox(
+    #                         value=False,
+    #                         label="better_visual_quality",
+    #                         info="better quality using morphologyEx",
+    #                     )
+    #                     with gr.Column():
+    #                         retina_check = gr.Checkbox(
+    #                             value=True,
+    #                             label="use_retina",
+    #                             info="draw high-resolution segmentation masks",
+    #                         )
+    #             # Description
+    #             gr.Markdown(description_e)
+    #
+    with gr.Tab("Points mode"):
+        # Images
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1):
+                cond_img_p.render()
+
+            with gr.Column(scale=1):
+                segm_img_p.render()
+
+        # Submit & Clear
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    add_or_remove = gr.Radio(
+                        ["Add Mask", "Remove Area"],
+                        value="Add Mask",
+                        label="Point_label (foreground/background)",
+                    )
+
+                    with gr.Column():
+                        segment_btn_p = gr.Button(
+                            "Segment with points prompt", variant="primary"
+                        )
+                        clear_btn_p = gr.Button("Clear points", variant="secondary")
+
+                gr.Markdown("Try some of the examples below ⬇️")
+                gr.Examples(
+                    examples=examples,
+                    inputs=[cond_img_p],
+                    # outputs=segm_img_p,
+                    # fn=segment_with_points,
+                    # cache_examples=True,
+                    examples_per_page=4,
+                )
+
+            with gr.Column():
+                # Description
+                gr.Markdown(description_p)
+
+    cond_img_p.select(get_points_with_draw, [cond_img_p, add_or_remove], cond_img_p)
+
+    # segment_btn_e.click(
+    #     segment_everything,
+    #     inputs=[
+    #         cond_img_e,
+    #         input_size_slider,
+    #         mor_check,
+    #         contour_check,
+    #         retina_check,
+    #     ],
+    #     outputs=segm_img_e,
+    # )
+
+    segment_btn_p.click(
+        segment_with_points, inputs=[cond_img_p], outputs=[segm_img_p, cond_img_p]
+    )
+
+    def clear():
+        return None, None
+
+    def clear_text():
+        return None, None, None
+
+    # clear_btn_e.click(clear, outputs=[cond_img_e, segm_img_e])
+    clear_btn_p.click(clear, outputs=[cond_img_p, segm_img_p])
+
+demo.queue()
+demo.launch()

requirements.txt

@@ -0,0 +1,4 @@
+torch
+torchvision
+timm
+git+https://github.com/dhkim2810/MobileSAM.git

utils/tools.py

@@ -0,0 +1,406 @@
+import os
+import sys
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from PIL import Image
+
+
+def convert_box_xywh_to_xyxy(box):
+    x1 = box[0]
+    y1 = box[1]
+    x2 = box[0] + box[2]
+    y2 = box[1] + box[3]
+    return [x1, y1, x2, y2]
+
+
+def segment_image(image, bbox):
+    image_array = np.array(image)
+    segmented_image_array = np.zeros_like(image_array)
+    x1, y1, x2, y2 = bbox
+    segmented_image_array[y1:y2, x1:x2] = image_array[y1:y2, x1:x2]
+    segmented_image = Image.fromarray(segmented_image_array)
+    black_image = Image.new("RGB", image.size, (255, 255, 255))
+    # transparency_mask = np.zeros_like((), dtype=np.uint8)
+    transparency_mask = np.zeros(
+        (image_array.shape[0], image_array.shape[1]), dtype=np.uint8
+    )
+    transparency_mask[y1:y2, x1:x2] = 255
+    transparency_mask_image = Image.fromarray(transparency_mask, mode="L")
+    black_image.paste(segmented_image, mask=transparency_mask_image)
+    return black_image
+
+
+def format_results(masks, scores, logits, filter=0):
+    annotations = []
+    n = len(scores)
+    for i in range(n):
+        annotation = {}
+
+        mask = masks[i]
+        tmp = np.where(mask != 0)
+        if np.sum(mask) < filter:
+            continue
+        annotation["id"] = i
+        annotation["segmentation"] = mask
+        annotation["bbox"] = [
+            np.min(tmp[0]),
+            np.min(tmp[1]),
+            np.max(tmp[1]),
+            np.max(tmp[0]),
+        ]
+        annotation["score"] = scores[i]
+        annotation["area"] = annotation["segmentation"].sum()
+        annotations.append(annotation)
+    return annotations
+
+
+def filter_masks(annotations):  # filter the overlap mask
+    annotations.sort(key=lambda x: x["area"], reverse=True)
+    to_remove = set()
+    for i in range(0, len(annotations)):
+        a = annotations[i]
+        for j in range(i + 1, len(annotations)):
+            b = annotations[j]
+            if i != j and j not in to_remove:
+                # check if
+                if b["area"] < a["area"]:
+                    if (a["segmentation"] & b["segmentation"]).sum() / b[
+                        "segmentation"
+                    ].sum() > 0.8:
+                        to_remove.add(j)
+
+    return [a for i, a in enumerate(annotations) if i not in to_remove], to_remove
+
+
+def get_bbox_from_mask(mask):
+    mask = mask.astype(np.uint8)
+    contours, hierarchy = cv2.findContours(
+        mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
+    )
+    x1, y1, w, h = cv2.boundingRect(contours[0])
+    x2, y2 = x1 + w, y1 + h
+    if len(contours) > 1:
+        for b in contours:
+            x_t, y_t, w_t, h_t = cv2.boundingRect(b)
+            # 将多个bbox合并成一个
+            x1 = min(x1, x_t)
+            y1 = min(y1, y_t)
+            x2 = max(x2, x_t + w_t)
+            y2 = max(y2, y_t + h_t)
+        h = y2 - y1
+        w = x2 - x1
+    return [x1, y1, x2, y2]
+
+
+def fast_process(
+    annotations, args, mask_random_color, bbox=None, points=None, edges=False
+):
+    if isinstance(annotations[0], dict):
+        annotations = [annotation["segmentation"] for annotation in annotations]
+    result_name = os.path.basename(args.img_path)
+    image = cv2.imread(args.img_path)
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    original_h = image.shape[0]
+    original_w = image.shape[1]
+    if sys.platform == "darwin":
+        plt.switch_backend("TkAgg")
+    plt.figure(figsize=(original_w / 100, original_h / 100))
+    # Add subplot with no margin.
+    plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
+    plt.margins(0, 0)
+    plt.gca().xaxis.set_major_locator(plt.NullLocator())
+    plt.gca().yaxis.set_major_locator(plt.NullLocator())
+    plt.imshow(image)
+    if args.better_quality == True:
+        if isinstance(annotations[0], torch.Tensor):
+            annotations = np.array(annotations.cpu())
+        for i, mask in enumerate(annotations):
+            mask = cv2.morphologyEx(
+                mask.astype(np.uint8), cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8)
+            )
+            annotations[i] = cv2.morphologyEx(
+                mask.astype(np.uint8), cv2.MORPH_OPEN, np.ones((8, 8), np.uint8)
+            )
+    if args.device == "cpu":
+        annotations = np.array(annotations)
+        fast_show_mask(
+            annotations,
+            plt.gca(),
+            random_color=mask_random_color,
+            bbox=bbox,
+            points=points,
+            point_label=args.point_label,
+            retinamask=args.retina,
+            target_height=original_h,
+            target_width=original_w,
+        )
+    else:
+        if isinstance(annotations[0], np.ndarray):
+            annotations = torch.from_numpy(annotations)
+        fast_show_mask_gpu(
+            annotations,
+            plt.gca(),
+            random_color=args.randomcolor,
+            bbox=bbox,
+            points=points,
+            point_label=args.point_label,
+            retinamask=args.retina,
+            target_height=original_h,
+            target_width=original_w,
+        )
+    if isinstance(annotations, torch.Tensor):
+        annotations = annotations.cpu().numpy()
+    if args.withContours == True:
+        contour_all = []
+        temp = np.zeros((original_h, original_w, 1))
+        for i, mask in enumerate(annotations):
+            if type(mask) == dict:
+                mask = mask["segmentation"]
+            annotation = mask.astype(np.uint8)
+            if args.retina == False:
+                annotation = cv2.resize(
+                    annotation,
+                    (original_w, original_h),
+                    interpolation=cv2.INTER_NEAREST,
+                )
+            contours, hierarchy = cv2.findContours(
+                annotation, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
+            )
+            for contour in contours:
+                contour_all.append(contour)
+        cv2.drawContours(temp, contour_all, -1, (255, 255, 255), 2)
+        color = np.array([0 / 255, 0 / 255, 255 / 255, 0.8])
+        contour_mask = temp / 255 * color.reshape(1, 1, -1)
+        plt.imshow(contour_mask)
+
+    save_path = args.output
+    if not os.path.exists(save_path):
+        os.makedirs(save_path)
+    plt.axis("off")
+    fig = plt.gcf()
+    plt.draw()
+
+    try:
+        buf = fig.canvas.tostring_rgb()
+    except AttributeError:
+        fig.canvas.draw()
+        buf = fig.canvas.tostring_rgb()
+
+    cols, rows = fig.canvas.get_width_height()
+    img_array = np.fromstring(buf, dtype=np.uint8).reshape(rows, cols, 3)
+    cv2.imwrite(
+        os.path.join(save_path, result_name), cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
+    )
+
+
+# CPU post process
+def fast_show_mask(
+    annotation,
+    ax,
+    random_color=False,
+    bbox=None,
+    points=None,
+    point_label=None,
+    retinamask=True,
+    target_height=960,
+    target_width=960,
+):
+    msak_sum = annotation.shape[0]
+    height = annotation.shape[1]
+    weight = annotation.shape[2]
+    # 将annotation 按照面积 排序
+    areas = np.sum(annotation, axis=(1, 2))
+    sorted_indices = np.argsort(areas)
+    annotation = annotation[sorted_indices]
+
+    index = (annotation != 0).argmax(axis=0)
+    if random_color == True:
+        color = np.random.random((msak_sum, 1, 1, 3))
+    else:
+        color = np.ones((msak_sum, 1, 1, 3)) * np.array(
+            [30 / 255, 144 / 255, 255 / 255]
+        )
+    transparency = np.ones((msak_sum, 1, 1, 1)) * 0.6
+    visual = np.concatenate([color, transparency], axis=-1)
+    mask_image = np.expand_dims(annotation, -1) * visual
+
+    show = np.zeros((height, weight, 4))
+    h_indices, w_indices = np.meshgrid(
+        np.arange(height), np.arange(weight), indexing="ij"
+    )
+    indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
+    # 使用向量化索引更新show的值
+    show[h_indices, w_indices, :] = mask_image[indices]
+    if bbox is not None:
+        x1, y1, x2, y2 = bbox
+        ax.add_patch(
+            plt.Rectangle(
+                (x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="b", linewidth=1
+            )
+        )
+    # draw point
+    if points is not None:
+        plt.scatter(
+            [point[0] for i, point in enumerate(points) if point_label[i] == 1],
+            [point[1] for i, point in enumerate(points) if point_label[i] == 1],
+            s=20,
+            c="y",
+        )
+        plt.scatter(
+            [point[0] for i, point in enumerate(points) if point_label[i] == 0],
+            [point[1] for i, point in enumerate(points) if point_label[i] == 0],
+            s=20,
+            c="m",
+        )
+
+    if retinamask == False:
+        show = cv2.resize(
+            show, (target_width, target_height), interpolation=cv2.INTER_NEAREST
+        )
+    ax.imshow(show)
+
+
+def fast_show_mask_gpu(
+    annotation,
+    ax,
+    random_color=False,
+    bbox=None,
+    points=None,
+    point_label=None,
+    retinamask=True,
+    target_height=960,
+    target_width=960,
+):
+    msak_sum = annotation.shape[0]
+    height = annotation.shape[1]
+    weight = annotation.shape[2]
+    areas = torch.sum(annotation, dim=(1, 2))
+    sorted_indices = torch.argsort(areas, descending=False)
+    annotation = annotation[sorted_indices]
+    # 找每个位置第一个非零值下标
+    index = (annotation != 0).to(torch.long).argmax(dim=0)
+    if random_color == True:
+        color = torch.rand((msak_sum, 1, 1, 3)).to(annotation.device)
+    else:
+        color = torch.ones((msak_sum, 1, 1, 3)).to(annotation.device) * torch.tensor(
+            [30 / 255, 144 / 255, 255 / 255]
+        ).to(annotation.device)
+    transparency = torch.ones((msak_sum, 1, 1, 1)).to(annotation.device) * 0.6
+    visual = torch.cat([color, transparency], dim=-1)
+    mask_image = torch.unsqueeze(annotation, -1) * visual
+    # 按index取数，index指每个位置选哪个batch的数，把mask_image转成一个batch的形式
+    show = torch.zeros((height, weight, 4)).to(annotation.device)
+    h_indices, w_indices = torch.meshgrid(
+        torch.arange(height), torch.arange(weight), indexing="ij"
+    )
+    indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
+    # 使用向量化索引更新show的值
+    show[h_indices, w_indices, :] = mask_image[indices]
+    show_cpu = show.cpu().numpy()
+    if bbox is not None:
+        x1, y1, x2, y2 = bbox
+        ax.add_patch(
+            plt.Rectangle(
+                (x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="b", linewidth=1
+            )
+        )
+    # draw point
+    if points is not None:
+        plt.scatter(
+            [point[0] for i, point in enumerate(points) if point_label[i] == 1],
+            [point[1] for i, point in enumerate(points) if point_label[i] == 1],
+            s=20,
+            c="y",
+        )
+        plt.scatter(
+            [point[0] for i, point in enumerate(points) if point_label[i] == 0],
+            [point[1] for i, point in enumerate(points) if point_label[i] == 0],
+            s=20,
+            c="m",
+        )
+    if retinamask == False:
+        show_cpu = cv2.resize(
+            show_cpu, (target_width, target_height), interpolation=cv2.INTER_NEAREST
+        )
+    ax.imshow(show_cpu)
+
+
+def crop_image(annotations, image_like):
+    if isinstance(image_like, str):
+        image = Image.open(image_like)
+    else:
+        image = image_like
+    ori_w, ori_h = image.size
+    mask_h, mask_w = annotations[0]["segmentation"].shape
+    if ori_w != mask_w or ori_h != mask_h:
+        image = image.resize((mask_w, mask_h))
+    cropped_boxes = []
+    cropped_images = []
+    not_crop = []
+    filter_id = []
+    # annotations, _ = filter_masks(annotations)
+    # filter_id = list(_)
+    for _, mask in enumerate(annotations):
+        if np.sum(mask["segmentation"]) <= 100:
+            filter_id.append(_)
+            continue
+        bbox = get_bbox_from_mask(mask["segmentation"])  # mask 的 bbox
+        cropped_boxes.append(segment_image(image, bbox))  # 保存裁剪的图片
+        # cropped_boxes.append(segment_image(image,mask["segmentation"]))
+        cropped_images.append(bbox)  # 保存裁剪的图片的bbox
+
+    return cropped_boxes, cropped_images, not_crop, filter_id, annotations
+
+
+def box_prompt(masks, bbox, target_height, target_width):
+    h = masks.shape[1]
+    w = masks.shape[2]
+    if h != target_height or w != target_width:
+        bbox = [
+            int(bbox[0] * w / target_width),
+            int(bbox[1] * h / target_height),
+            int(bbox[2] * w / target_width),
+            int(bbox[3] * h / target_height),
+        ]
+    bbox[0] = round(bbox[0]) if round(bbox[0]) > 0 else 0
+    bbox[1] = round(bbox[1]) if round(bbox[1]) > 0 else 0
+    bbox[2] = round(bbox[2]) if round(bbox[2]) < w else w
+    bbox[3] = round(bbox[3]) if round(bbox[3]) < h else h
+
+    # IoUs = torch.zeros(len(masks), dtype=torch.float32)
+    bbox_area = (bbox[3] - bbox[1]) * (bbox[2] - bbox[0])
+
+    masks_area = torch.sum(masks[:, bbox[1] : bbox[3], bbox[0] : bbox[2]], dim=(1, 2))
+    orig_masks_area = torch.sum(masks, dim=(1, 2))
+
+    union = bbox_area + orig_masks_area - masks_area
+    IoUs = masks_area / union
+    max_iou_index = torch.argmax(IoUs)
+
+    return masks[max_iou_index].cpu().numpy(), max_iou_index
+
+
+def point_prompt(masks, points, point_label, target_height, target_width):  # numpy 处理
+    h = masks[0]["segmentation"].shape[0]
+    w = masks[0]["segmentation"].shape[1]
+    if h != target_height or w != target_width:
+        points = [
+            [int(point[0] * w / target_width), int(point[1] * h / target_height)]
+            for point in points
+        ]
+    onemask = np.zeros((h, w))
+    for i, annotation in enumerate(masks):
+        if type(annotation) == dict:
+            mask = annotation["segmentation"]
+        else:
+            mask = annotation
+        for i, point in enumerate(points):
+            if mask[point[1], point[0]] == 1 and point_label[i] == 1:
+                onemask += mask
+            if mask[point[1], point[0]] == 1 and point_label[i] == 0:
+                onemask -= mask
+    onemask = onemask >= 1
+    return onemask, 0

utils/tools_gradio.py

@@ -0,0 +1,192 @@
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from PIL import Image
+
+
+def fast_process(
+    annotations,
+    image,
+    device,
+    scale,
+    better_quality=False,
+    mask_random_color=True,
+    bbox=None,
+    use_retina=True,
+    withContours=True,
+):
+    if isinstance(annotations[0], dict):
+        annotations = [annotation["segmentation"] for annotation in annotations]
+
+    original_h = image.height
+    original_w = image.width
+    if better_quality:
+        if isinstance(annotations[0], torch.Tensor):
+            annotations = np.array(annotations.cpu())
+        for i, mask in enumerate(annotations):
+            mask = cv2.morphologyEx(
+                mask.astype(np.uint8), cv2.MORPH_CLOSE, np.ones((3, 3), np.uint8)
+            )
+            annotations[i] = cv2.morphologyEx(
+                mask.astype(np.uint8), cv2.MORPH_OPEN, np.ones((8, 8), np.uint8)
+            )
+    if device == "cpu":
+        annotations = np.array(annotations)
+        inner_mask = fast_show_mask(
+            annotations,
+            plt.gca(),
+            random_color=mask_random_color,
+            bbox=bbox,
+            retinamask=use_retina,
+            target_height=original_h,
+            target_width=original_w,
+        )
+    else:
+        if isinstance(annotations[0], np.ndarray):
+            annotations = np.array(annotations)
+            annotations = torch.from_numpy(annotations)
+        inner_mask = fast_show_mask_gpu(
+            annotations,
+            plt.gca(),
+            random_color=mask_random_color,
+            bbox=bbox,
+            retinamask=use_retina,
+            target_height=original_h,
+            target_width=original_w,
+        )
+    if isinstance(annotations, torch.Tensor):
+        annotations = annotations.cpu().numpy()
+
+    if withContours:
+        contour_all = []
+        temp = np.zeros((original_h, original_w, 1))
+        for i, mask in enumerate(annotations):
+            if type(mask) == dict:
+                mask = mask["segmentation"]
+            annotation = mask.astype(np.uint8)
+            if use_retina == False:
+                annotation = cv2.resize(
+                    annotation,
+                    (original_w, original_h),
+                    interpolation=cv2.INTER_NEAREST,
+                )
+            contours, _ = cv2.findContours(
+                annotation, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
+            )
+            for contour in contours:
+                contour_all.append(contour)
+        cv2.drawContours(temp, contour_all, -1, (255, 255, 255), 2 // scale)
+        color = np.array([0 / 255, 0 / 255, 255 / 255, 0.9])
+        contour_mask = temp / 255 * color.reshape(1, 1, -1)
+
+    image = image.convert("RGBA")
+    overlay_inner = Image.fromarray((inner_mask * 255).astype(np.uint8), "RGBA")
+    image.paste(overlay_inner, (0, 0), overlay_inner)
+
+    if withContours:
+        overlay_contour = Image.fromarray((contour_mask * 255).astype(np.uint8), "RGBA")
+        image.paste(overlay_contour, (0, 0), overlay_contour)
+
+    return image
+
+
+# CPU post process
+def fast_show_mask(
+    annotation,
+    ax,
+    random_color=False,
+    bbox=None,
+    retinamask=True,
+    target_height=960,
+    target_width=960,
+):
+    mask_sum = annotation.shape[0]
+    height = annotation.shape[1]
+    weight = annotation.shape[2]
+    # 将annotation 按照面积 排序
+    areas = np.sum(annotation, axis=(1, 2))
+    sorted_indices = np.argsort(areas)[::1]
+    annotation = annotation[sorted_indices]
+
+    index = (annotation != 0).argmax(axis=0)
+    if random_color == True:
+        color = np.random.random((mask_sum, 1, 1, 3))
+    else:
+        color = np.ones((mask_sum, 1, 1, 3)) * np.array(
+            [30 / 255, 144 / 255, 255 / 255]
+        )
+    transparency = np.ones((mask_sum, 1, 1, 1)) * 0.6
+    visual = np.concatenate([color, transparency], axis=-1)
+    mask_image = np.expand_dims(annotation, -1) * visual
+
+    mask = np.zeros((height, weight, 4))
+
+    h_indices, w_indices = np.meshgrid(
+        np.arange(height), np.arange(weight), indexing="ij"
+    )
+    indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
+
+    mask[h_indices, w_indices, :] = mask_image[indices]
+    if bbox is not None:
+        x1, y1, x2, y2 = bbox
+        ax.add_patch(
+            plt.Rectangle(
+                (x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="b", linewidth=1
+            )
+        )
+
+    if retinamask == False:
+        mask = cv2.resize(
+            mask, (target_width, target_height), interpolation=cv2.INTER_NEAREST
+        )
+
+    return mask
+
+
+def fast_show_mask_gpu(
+    annotation,
+    ax,
+    random_color=False,
+    bbox=None,
+    retinamask=True,
+    target_height=960,
+    target_width=960,
+):
+    device = annotation.device
+    mask_sum = annotation.shape[0]
+    height = annotation.shape[1]
+    weight = annotation.shape[2]
+    areas = torch.sum(annotation, dim=(1, 2))
+    sorted_indices = torch.argsort(areas, descending=False)
+    annotation = annotation[sorted_indices]
+    # 找每个位置第一个非零值下标
+    index = (annotation != 0).to(torch.long).argmax(dim=0)
+    if random_color == True:
+        color = torch.rand((mask_sum, 1, 1, 3)).to(device)
+    else:
+        color = torch.ones((mask_sum, 1, 1, 3)).to(device) * torch.tensor(
+            [30 / 255, 144 / 255, 255 / 255]
+        ).to(device)
+    transparency = torch.ones((mask_sum, 1, 1, 1)).to(device) * 0.6
+    visual = torch.cat([color, transparency], dim=-1)
+    mask_image = torch.unsqueeze(annotation, -1) * visual
+    # 按index取数，index指每个位置选哪个batch的数，把mask_image转成一个batch的形式
+    mask = torch.zeros((height, weight, 4)).to(device)
+    h_indices, w_indices = torch.meshgrid(torch.arange(height), torch.arange(weight))
+    indices = (index[h_indices, w_indices], h_indices, w_indices, slice(None))
+    # 使用向量化索引更新show的值
+    mask[h_indices, w_indices, :] = mask_image[indices]
+    mask_cpu = mask.cpu().numpy()
+    if bbox is not None:
+        x1, y1, x2, y2 = bbox
+        ax.add_patch(
+            plt.Rectangle(
+                (x1, y1), x2 - x1, y2 - y1, fill=False, edgecolor="b", linewidth=1
+            )
+        )
+    if retinamask == False:
+        mask_cpu = cv2.resize(
+            mask_cpu, (target_width, target_height), interpolation=cv2.INTER_NEAREST
+        )
+    return mask_cpu