Add text mode

.gitignore

@@ -0,0 +1,7 @@
+*.pyc
+*.pyo
+*.pyd
+.DS_Store
+.idea
+weights
+gradio_cached_examples

README.md

@@ -5,7 +5,7 @@ colorFrom: pink
 colorTo: indigo
 sdk: gradio
 sdk_version: 3.35.2
-app_file: app.py
+app_file: app_gradio.py
 pinned: false
 license: apache-2.0
 ---

app.py → app_gradio.py

@@ -1,24 +1,31 @@
 from ultralytics import YOLO
 import gradio as gr
 import torch
-from tools import fast_process, format_results, box_prompt, point_prompt
+from utils.tools_gradio import fast_process
+from utils.tools import format_results, box_prompt, point_prompt, text_prompt
 from PIL import ImageDraw
 import numpy as np
 
 # Load the pre-trained model
-model = YOLO('checkpoints/FastSAM.pt')
+model = YOLO('./weights/FastSAM.pt')
 
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+device = torch.device(
+    "cuda"
+    if torch.cuda.is_available()
+    else "mps"
+    if torch.backends.mps.is_available()
+    else "cpu"
+)
 
 # Description
 title = "<center><strong><font size='8'>🏃 Fast Segment Anything 🤗</font></strong></center>"
 
 news = """ # 📖 News
+        🔥 2023/06/29: Support the text mode (Thanks for [gaoxinge](https://github.com/CASIA-IVA-Lab/FastSAM/pull/47)).
 
-        🔥 2023/06/24: Add the 'Advanced options" in Everything mode to get a more detailed adjustment.
-        
         🔥 2023/06/26: Support the points mode. (Better and faster interaction will come soon!)
-        
+
+        🔥 2023/06/24: Add the 'Advanced options" in Everything mode to get a more detailed adjustment.        
         """  
 
 description_e = """This is a demo on Github project 🏃 [Fast Segment Anything Model](https://github.com/CASIA-IVA-Lab/FastSAM). Welcome to give a star ⭐️ to it.
@@ -52,8 +59,8 @@ description_p = """ # 🎯 Instructions for points mode
                 
               """
 
-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"]]
+examples = [["examples/sa_8776.jpg"], ["examples/sa_414.jpg"], ["examples/sa_1309.jpg"], ["examples/sa_11025.jpg"],
+            ["examples/sa_561.jpg"], ["examples/sa_192.jpg"], ["examples/sa_10039.jpg"], ["examples/sa_862.jpg"]]
 
 default_example = examples[0]
 
@@ -68,10 +75,10 @@ def segment_everything(
     better_quality=False,
     withContours=True,
     use_retina=True,
+    text="",
     mask_random_color=True,
-    ):
+):
     input_size = int(input_size)  # 确保 imgsz 是整数
-
     # Thanks for the suggestion by hysts in HuggingFace.
     w, h = input.size
     scale = input_size / max(w, h)
@@ -85,18 +92,26 @@ def segment_everything(
                     iou=iou_threshold,
                     conf=conf_threshold,
                     imgsz=input_size,)
+
+    if len(text) > 0:
+        results = format_results(results[0], 0)
+        annotations, _ = text_prompt(results, text, input, device=device)
+        annotations = np.array([annotations])
+    else:
+        annotations = results[0].masks.data
     
-    fig = fast_process(annotations=results[0].masks.data,
-                        image=input,
-                        device=device,
-                        scale=(1024 // input_size),
-                        better_quality=better_quality,
-                        mask_random_color=mask_random_color,
-                        bbox=None,
-                        use_retina=use_retina,
-                        withContours=withContours,)
+    fig = fast_process(annotations=annotations,
+                       image=input,
+                       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(
     input,
     input_size=1024, 
@@ -104,9 +119,9 @@ def segment_with_points(
     conf_threshold=0.25,
     better_quality=False,
     withContours=True,
-    mask_random_color=True,
     use_retina=True,
-    ):    
+    mask_random_color=True,
+):
     global global_points
     global global_point_label
     
@@ -128,56 +143,51 @@ def segment_with_points(
                     imgsz=input_size,)
     
     results = format_results(results[0], 0)
-    
     annotations, _ = point_prompt(results, scaled_points, global_point_label, new_h, new_w)
     annotations = np.array([annotations])
-        
+
     fig = fast_process(annotations=annotations,
-                        image=input,
-                        device=device,
-                        scale=(1024 // input_size),
-                        better_quality=better_quality,
-                        mask_random_color=mask_random_color,
-                        bbox=None,
-                        use_retina=use_retina,
-                        withContours=withContours,)
+                       image=input,
+                       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
 
+
 def get_points_with_draw(image, label, evt: gr.SelectData):
-    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 global_points
     global global_point_label
-    print((x, y))
+
+    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
-    
 
-# input_size=1024
-# high_quality_visual=True
-# inp = 'assets/sa_192.jpg'
-# input = Image.open(inp)
-# device = 'cuda' if torch.cuda.is_available() else 'cpu'
-# input_size = int(input_size)  # 确保 imgsz 是整数
-# results = model(input, device=device, retina_masks=True, iou=0.7, conf=0.25, imgsz=input_size)
-# pil_image = fast_process(annotations=results[0].masks.data,
-#                             image=input, high_quality=high_quality_visual, device=device)
 
 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')
+cond_img_t = gr.Image(label="Input with text", value="examples/dogs.jpg", 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')
+segm_img_t = gr.Image(label="Segmented Image with text", interactive=False, type='pil')
 
 global_points = []
-global_point_label = [] # TODO:Clear points each image
+global_point_label = []
 
 input_size_slider = gr.components.Slider(minimum=512,
                                          maximum=1024,
@@ -188,14 +198,14 @@ input_size_slider = gr.components.Slider(minimum=512,
 
 with gr.Blocks(css=css, title='Fast Segment Anything') as demo:
     with gr.Row():
-            with gr.Column(scale=1):
-                # Title
-                gr.Markdown(title)
-        
-            with gr.Column(scale=1):
-                # News
-                gr.Markdown(news)
-                
+        with gr.Column(scale=1):
+            # Title
+            gr.Markdown(title)
+
+        with gr.Column(scale=1):
+            # News
+            gr.Markdown(news)
+
     with gr.Tab("Everything mode"):
         # Images
         with gr.Row(variant="panel"):
@@ -227,13 +237,13 @@ with gr.Blocks(css=css, title='Fast Segment Anything') as demo:
 
             with gr.Column():
                 with gr.Accordion("Advanced options", open=False):
+                    # text_box = gr.Textbox(label="text prompt")
                     iou_threshold = gr.Slider(0.1, 0.9, 0.7, step=0.1, label='iou', info='iou threshold for filtering the annotations')
                     conf_threshold = gr.Slider(0.1, 0.9, 0.25, step=0.05, label='conf', info='object confidence threshold')
                     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)
 
@@ -259,30 +269,102 @@ with gr.Blocks(css=css, title='Fast Segment Anything') as demo:
                 gr.Markdown("Try some of the examples below ⬇️")
                 gr.Examples(examples=examples,
                             inputs=[cond_img_p],
-                            outputs=segm_img_p,
-                            fn=segment_with_points,
+                            # outputs=segm_img_p,
+                            # fn=segment_with_points,
                             # cache_examples=True,
                             examples_per_page=4)
 
             with gr.Column():
                 # Description
                 gr.Markdown(description_p)
+                
+    with gr.Tab("Text mode"):
+        # Images
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1):
+                cond_img_t.render()
+
+            with gr.Column(scale=1):
+                segm_img_t.render()
+
+        # Submit & Clear
+        with gr.Row():
+            with gr.Column():
+                input_size_slider_t = 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.Row():
+                    with gr.Column():
+                        contour_check = gr.Checkbox(value=True, label='withContours', info='draw the edges of the masks')
+                        text_box = gr.Textbox(label="text prompt", value="a black dog")
+
+                    with gr.Column():
+                        segment_btn_t = gr.Button("Segment with text", variant='primary')
+                        clear_btn_t = gr.Button("Clear", variant="secondary")
+
+                gr.Markdown("Try some of the examples below ⬇️")
+                gr.Examples(examples=["examples/dogs.jpg"],
+                            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):
+                    iou_threshold = gr.Slider(0.1, 0.9, 0.7, step=0.1, label='iou', info='iou threshold for filtering the annotations')
+                    conf_threshold = gr.Slider(0.1, 0.9, 0.25, step=0.05, label='conf', info='object confidence threshold')
+                    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)
         
     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, iou_threshold, conf_threshold, mor_check, contour_check, retina_check],
-                    outputs=segm_img_e)
-    
+                        inputs=[
+                            cond_img_e,
+                            input_size_slider,
+                            iou_threshold,
+                            conf_threshold,
+                            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])
+                        inputs=[cond_img_p],
+                        outputs=[segm_img_p, cond_img_p])
     
+    segment_btn_t.click(segment_everything,
+                        inputs=[
+                            cond_img_t,
+                            input_size_slider_t,
+                            iou_threshold,
+                            conf_threshold,
+                            mor_check,
+                            contour_check,
+                            retina_check,
+                            text_box,
+                        ],
+                        outputs=segm_img_t)
+
     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])
+    clear_btn_t.click(clear_text, outputs=[cond_img_p, segm_img_p, text_box])
 
 demo.queue()
 demo.launch()

requirements.txt

@@ -2,6 +2,7 @@
 matplotlib==3.2.2
 numpy
 opencv-python
+clip>=0.2.0
 # Pillow>=7.1.2
 # PyYAML>=5.3.1
 # requests>=2.23.0

tools.py → utils/tools.py

@@ -3,7 +3,9 @@ from PIL import Image
 import matplotlib.pyplot as plt
 import cv2
 import torch
-# import clip
+import os
+import sys
+import clip
 
 
 def convert_box_xywh_to_xyxy(box):
@@ -49,7 +51,7 @@ def format_results(result, filter=0):
     return annotations
 
 
-def filter_masks(annotations):  # filte the overlap mask
+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)):
@@ -86,126 +88,171 @@ def get_bbox_from_mask(mask):
         w = x2 - x1
     return [x1, y1, x2, y2]
 
+
 def fast_process(
-    annotations,
-    image,
-    device,
-    scale,
-    better_quality=False,
-    mask_random_color=True,
-    bbox=None,
-    use_retina=True,
-    withContours=True,
-    ):
+    annotations, args, mask_random_color, bbox=None, points=None, edges=False
+):
     if isinstance(annotations[0], dict):
-        annotations = [annotation['segmentation'] for annotation in annotations]
-
-    original_h = image.height
-    original_w = image.width
-    if better_quality:
+        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 device == 'cpu':
+            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)
-        inner_mask = fast_show_mask(
+        fast_show_mask(
             annotations,
             plt.gca(),
             random_color=mask_random_color,
             bbox=bbox,
-            retinamask=use_retina,
+            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)
-        inner_mask = fast_show_mask_gpu(
+        fast_show_mask_gpu(
             annotations,
             plt.gca(),
-            random_color=mask_random_color,
+            random_color=args.randomcolor,
             bbox=bbox,
-            retinamask=use_retina,
+            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 withContours:
+    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']
+                mask = mask["segmentation"]
             annotation = mask.astype(np.uint8)
-            if use_retina == False:
+            if args.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)
+            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 // scale)
-        color = np.array([0 / 255, 0 / 255, 255 / 255, 0.9])
+        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)
-
-    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
+        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,
 ):
-    mask_sum = annotation.shape[0]
+    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)[::1]
+    sorted_indices = np.argsort(areas)
     annotation = annotation[sorted_indices]
 
     index = (annotation != 0).argmax(axis=0)
     if random_color == True:
-        color = np.random.random((mask_sum, 1, 1, 3))
+        color = np.random.random((msak_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
+        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
 
-    mask = np.zeros((height, weight, 4))
-
-    h_indices, w_indices = np.meshgrid(np.arange(height), np.arange(weight), indexing='ij')
+    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))
-
-    mask[h_indices, w_indices, :] = mask_image[indices]
+    # 使用向量化索引更新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))
+        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:
-        mask = cv2.resize(mask, (target_width, target_height), interpolation=cv2.INTER_NEAREST)
-
-    return mask
+        show = cv2.resize(
+            show, (target_width, target_height), interpolation=cv2.INTER_NEAREST
+        )
+    ax.imshow(show)
 
 
 def fast_show_mask_gpu(
@@ -213,12 +260,13 @@ def fast_show_mask_gpu(
     ax,
     random_color=False,
     bbox=None,
+    points=None,
+    point_label=None,
     retinamask=True,
     target_height=960,
     target_width=960,
 ):
-    device = annotation.device
-    mask_sum = annotation.shape[0]
+    msak_sum = annotation.shape[0]
     height = annotation.shape[1]
     weight = annotation.shape[2]
     areas = torch.sum(annotation, dim=(1, 2))
@@ -227,21 +275,23 @@ def fast_show_mask_gpu(
     # 找每个位置第一个非零值下标
     index = (annotation != 0).to(torch.long).argmax(dim=0)
     if random_color == True:
-        color = torch.rand((mask_sum, 1, 1, 3)).to(device)
+        color = torch.rand((msak_sum, 1, 1, 3)).to(annotation.device)
     else:
-        color = torch.ones((mask_sum, 1, 1, 3)).to(device) * torch.tensor(
+        color = torch.ones((msak_sum, 1, 1, 3)).to(annotation.device) * torch.tensor(
             [30 / 255, 144 / 255, 255 / 255]
-        ).to(device)
-    transparency = torch.ones((mask_sum, 1, 1, 1)).to(device) * 0.6
+        ).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的形式
-    mask = torch.zeros((height, weight, 4)).to(device)
-    h_indices, w_indices = torch.meshgrid(torch.arange(height), torch.arange(weight))
+    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的值
-    mask[h_indices, w_indices, :] = mask_image[indices]
-    mask_cpu = mask.cpu().numpy()
+    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(
@@ -249,31 +299,48 @@ def fast_show_mask_gpu(
                 (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:
-        mask_cpu = cv2.resize(
-            mask_cpu, (target_width, target_height), interpolation=cv2.INTER_NEAREST
+        show_cpu = cv2.resize(
+            show_cpu, (target_width, target_height), interpolation=cv2.INTER_NEAREST
         )
-    return mask_cpu
-
-
-# # clip
-# @torch.no_grad()
-# def retriev(
-#     model, preprocess, elements, search_text: str, device
-# ) -> int:
-#     preprocessed_images = [preprocess(image).to(device) for image in elements]
-#     tokenized_text = clip.tokenize([search_text]).to(device)
-#     stacked_images = torch.stack(preprocessed_images)
-#     image_features = model.encode_image(stacked_images)
-#     text_features = model.encode_text(tokenized_text)
-#     image_features /= image_features.norm(dim=-1, keepdim=True)
-#     text_features /= text_features.norm(dim=-1, keepdim=True)
-#     probs = 100.0 * image_features @ text_features.T
-#     return probs[:, 0].softmax(dim=0)
-
-
-def crop_image(annotations, image_path):
-    image = Image.open(image_path)
+    ax.imshow(show_cpu)
+
+
+# clip
+@torch.no_grad()
+def retriev(
+    model, preprocess, elements, search_text: str, device
+) -> int:
+    preprocessed_images = [preprocess(image).to(device) for image in elements]
+    tokenized_text = clip.tokenize([search_text]).to(device)
+    stacked_images = torch.stack(preprocessed_images)
+    image_features = model.encode_image(stacked_images)
+    text_features = model.encode_text(tokenized_text)
+    image_features /= image_features.norm(dim=-1, keepdim=True)
+    text_features /= text_features.norm(dim=-1, keepdim=True)
+    probs = 100.0 * image_features @ text_features.T
+    return probs[:, 0].softmax(dim=0)
+
+
+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:
@@ -324,7 +391,7 @@ def box_prompt(masks, bbox, target_height, target_width):
     return masks[max_iou_index].cpu().numpy(), max_iou_index
 
 
-def point_prompt(masks, points, pointlabel, target_height, target_width):  # numpy 处理
+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:
@@ -339,23 +406,23 @@ def point_prompt(masks, points, pointlabel, target_height, target_width):  # num
         else:
             mask = annotation
         for i, point in enumerate(points):
-            if mask[point[1], point[0]] == 1 and pointlabel[i] == 1:
+            if mask[point[1], point[0]] == 1 and point_label[i] == 1:
                 onemask += mask
-            if mask[point[1], point[0]] == 1 and pointlabel[i] == 0:
+            if mask[point[1], point[0]] == 1 and point_label[i] == 0:
                 onemask -= mask
     onemask = onemask >= 1
     return onemask, 0
 
 
-# def text_prompt(annotations, args):
-#     cropped_boxes, cropped_images, not_crop, filter_id, annotaions = crop_image(
-#         annotations, args.img_path
-#     )
-#     clip_model, preprocess = clip.load("ViT-B/32", device=args.device)
-#     scores = retriev(
-#         clip_model, preprocess, cropped_boxes, args.text_prompt, device=args.device
-#     )
-#     max_idx = scores.argsort()
-#     max_idx = max_idx[-1]
-#     max_idx += sum(np.array(filter_id) <= int(max_idx))
-#     return annotaions[max_idx]["segmentation"], max_idx
+def text_prompt(annotations, text, img_path, device):
+    cropped_boxes, cropped_images, not_crop, filter_id, annotations_ = crop_image(
+        annotations, img_path
+    )
+    clip_model, preprocess = clip.load("./weights/CLIP_ViT_B_32.pt", device=device)
+    scores = retriev(
+        clip_model, preprocess, cropped_boxes, text, device=device
+    )
+    max_idx = scores.argsort()
+    max_idx = max_idx[-1]
+    max_idx += sum(np.array(filter_id) <= int(max_idx))
+    return annotations_[max_idx]["segmentation"], max_idx

utils/tools_gradio.py

@@ -0,0 +1,175 @@
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+import cv2
+import torch
+
+
+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 = 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