add files

.gitattributes

@@ -0,0 +1,2 @@
+efficient_sam_s_gpu.jit filter=lfs diff=lfs merge=lfs -text
+efficient_sam_s_cpu.jit filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,12 @@
+---
+title: SAM Arena
+emoji: 🐢
+colorFrom: red
+colorTo: green
+sdk: gradio
+sdk_version: 4.9.0
+app_file: app.py
+pinned: false
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,251 @@
+# Thanks to the following repos:
+# https://huggingface.co/spaces/An-619/FastSAM/blob/main/app_gradio.py
+# https://huggingface.co/spaces/SkalskiP/EfficientSAM
+from typing import Tuple
+
+from ultralytics import YOLO
+from PIL import ImageDraw
+from PIL import Image
+import gradio as gr
+import numpy as np
+import torch
+
+from transformers import SamModel, SamProcessor
+
+import supervision as sv
+from utils.tools_gradio import fast_process
+from utils.tools import format_results, point_prompt
+from utils.draw import draw_circle, calculate_dynamic_circle_radius
+from utils.efficient_sam import load, inference_with_box, inference_with_point
+
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Load the pre-trained models
+FASTSAM_MODEL = YOLO('FastSAM-s.pt')
+SAM_MODEL = SamModel.from_pretrained("facebook/sam-vit-huge").to(DEVICE)
+SAM_PROCESSOR = SamProcessor.from_pretrained("facebook/sam-vit-huge")
+EFFICIENT_SAM_MODEL = load(device=DEVICE)
+
+MASK_COLOR = sv.Color.from_hex("#FF0000")
+PROMPT_COLOR = sv.Color.from_hex("#D3D3D3")
+MASK_ANNOTATOR = sv.MaskAnnotator(
+    color=MASK_COLOR,
+    color_lookup=sv.ColorLookup.INDEX)
+
+title = "<center><strong><font size='8'>🤗 Segment Anything Model Arena ⚔️</font></strong></center>"
+
+description = "<center><font size='4'>This is a demo of the <strong>Segment Anything Model Arena</strong>, a collection of models for segmenting anything. "
+
+css = "h1 { text-align: center } .about { text-align: justify; padding-left: 10%; padding-right: 10%; }"
+
+#examples = [["examples/retail01.png"], ["examples/vend01.png"], ["examples/vend02.png"]]
+
+POINT_EXAMPLES = [
+    ['https://media.roboflow.com/efficient-sam/corgi.jpg', 1291, 751],
+    ['https://media.roboflow.com/efficient-sam/horses.jpg', 1168, 939],
+    ['https://media.roboflow.com/efficient-sam/bears.jpg', 913, 1051]
+]
+
+#default_example = examples[0]
+
+def annotate_image_with_point_prompt_result(
+    image: np.ndarray,
+    detections: sv.Detections,
+    x: int,
+    y: int
+) -> np.ndarray:
+    h, w, _ = image.shape
+    bgr_image = image[:, :, ::-1]
+    annotated_bgr_image = MASK_ANNOTATOR.annotate(
+        scene=bgr_image, detections=detections)
+    annotated_bgr_image = draw_circle(
+        scene=annotated_bgr_image,
+        center=sv.Point(x=x, y=y),
+        radius=calculate_dynamic_circle_radius(resolution_wh=(w, h)),
+        color=PROMPT_COLOR)
+    return annotated_bgr_image[:, :, ::-1]
+
+def SAM_points_inference(image: np.ndarray) -> np.ndarray:
+    global global_points
+    input_points = [[[float(num) for num in sublist]] for sublist in global_points]
+    print(input_points)
+    #input_points = [[[773.0, 167.0]]]
+    x = int(input_points[0][0][0])
+    y = int(input_points[0][0][1])
+    
+    inputs = SAM_PROCESSOR(
+        Image.fromarray(image),
+        input_points=[input_points],
+        return_tensors="pt"
+    ).to(DEVICE)
+
+    with torch.no_grad():
+        outputs = SAM_MODEL(**inputs)
+
+    mask = SAM_PROCESSOR.image_processor.post_process_masks(
+        outputs.pred_masks.cpu(),
+        inputs["original_sizes"].cpu(),
+        inputs["reshaped_input_sizes"].cpu()
+    )[0][0][0].numpy()
+    mask = mask[np.newaxis, ...]
+    detections = sv.Detections(xyxy=sv.mask_to_xyxy(masks=mask), mask=mask)
+
+    return annotate_image_with_point_prompt_result(
+        image=image, detections=detections, x=x, y=y)
+ 
+def FastSAM_points_inference(
+    input,
+    input_size=1024, 
+    iou_threshold=0.7,
+    conf_threshold=0.25,
+    better_quality=False,
+    withContours=True,
+    use_retina=True,
+    mask_random_color=True,
+):
+    global global_points
+    global global_point_label
+    input = Image.fromarray(input)
+    input_size = int(input_size)  # 确保 imgsz 是整数
+    # Thanks for the suggestion by hysts in HuggingFace.
+    w, h = input.size
+    scale = input_size / max(w, h)
+    new_w = int(w * scale)
+    new_h = int(h * scale)
+    input = input.resize((new_w, new_h))
+    
+    scaled_points = [[int(x * scale) for x in point] for point in global_points]
+
+    results = FASTSAM_MODEL(input,
+                    device=DEVICE,
+                    retina_masks=True,
+                    iou=iou_threshold,
+                    conf=conf_threshold,
+                    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,)
+    
+    global_points = []
+    global_point_label = []
+    
+    return fig
+
+def EfficientSAM_points_inference(image: np.ndarray):
+    x, y = int(global_points[0][0]), int(global_points[0][1])
+    point = np.array([[int(x), int(y)]])
+    mask = inference_with_point(image, point, EFFICIENT_SAM_MODEL, DEVICE)
+    mask = mask[np.newaxis, ...]
+    detections = sv.Detections(xyxy=sv.mask_to_xyxy(masks=mask), mask=mask)
+
+    return annotate_image_with_point_prompt_result(image=image, detections=detections, x=x, y=y)
+    
+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, 0, 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')
+    image = Image.fromarray(image)
+    draw = ImageDraw.Draw(image)
+    draw.ellipse([(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)], fill=point_color)
+    return image
+
+def clear(_: np.ndarray) -> Tuple[None, None, None, None]:
+    return None, None, None, None
+
+gr_input_image = gr.Image(label="Input", value='examples/fruits.jpg')
+
+fast_sam_segmented_image = gr.Image(label="Fast SAM", interactive=False, type='pil')
+
+edge_sam_segmented_imaged = gr.Image(label="Edge SAM", interactive=False, type='pil')
+
+
+global_points = []
+global_point_label = []
+
+with gr.Blocks() as demo:
+    with gr.Tab("Points prompt"):
+        # Input Image
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1, min_width="320", variant="compact"):
+                gr_input_image.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():
+                        inference_point_button = gr.Button("Segment", variant='primary')
+                        clear_button = gr.Button("Clear points", variant='secondary')
+        
+        # Segment Results Grid  
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1):
+                sam_segmented_image = gr.Image(label="SAM")
+            with gr.Column(scale=1):
+                efficient_sam_segmented_image = gr.Image(label="Efficient SAM")
+        
+        with gr.Row(variant="panel"):
+            with gr.Column(scale=1):
+                fast_sam_segmented_image.render()
+            with gr.Column(scale=1):
+                edge_sam_segmented_imaged.render()
+    
+    gr.Markdown("AI Generated Examples")
+    # gr.Examples(examples=examples,
+    #             inputs=[gr_input_image],
+    #             # outputs=sam_segmented_image,
+    #             # fn=segment_with_points,
+    #             # cache_examples=True,
+    #             examples_per_page=3)
+    
+    gr_input_image.select(get_points_with_draw, [gr_input_image, add_or_remove], gr_input_image)
+
+    inference_point_button.click(
+        SAM_points_inference,
+        inputs=[gr_input_image],
+        outputs=[sam_segmented_image]
+    )
+    
+    inference_point_button.click(
+        EfficientSAM_points_inference,
+        inputs=[gr_input_image],
+        outputs=[efficient_sam_segmented_image])
+    
+    inference_point_button.click(
+        FastSAM_points_inference,
+        inputs=[gr_input_image],
+        outputs=[fast_sam_segmented_image])
+    
+    # inference_point_button.click(
+    #     EdgeSAM_points_inference,
+    #     inputs=[gr_input_image],
+    #     outputs=[fast_sam_segmented_image, gr_input_image])
+    
+    gr_input_image.change(
+        clear,
+        inputs=gr_input_image,
+        outputs=[efficient_sam_segmented_image, sam_segmented_image, fast_sam_segmented_image]
+    )
+
+    clear_button.click(clear, outputs=[gr_input_image, efficient_sam_segmented_image, sam_segmented_image, fast_sam_segmented_image])
+
+
+demo.queue()
+demo.launch(debug=True, show_error=True)

efficient_sam_s_cpu.jit

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8b63ab268e9020b0fb7fc9f46e742644d4c9ea6e5d9caf56045f0afb6475db09
+size 106006979

efficient_sam_s_gpu.jit

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e47c589ead2c6a80d38050ce63083a551e288db27113d534e0278270fc7cba26
+size 106006979

requirements.txt

@@ -0,0 +1,10 @@
+torch
+torchvision
+
+pillow
+gradio==3.44.0
+transformers
+supervision
+ultralytics
+clip
+opencv-python

utils/draw.py

@@ -0,0 +1,33 @@
+#https://huggingface.co/spaces/SkalskiP/EfficientSAM
+from typing import Tuple
+
+import cv2
+import numpy as np
+import supervision as sv
+
+
+def draw_circle(
+    scene: np.ndarray, center: sv.Point, color: sv.Color, radius: int = 2
+) -> np.ndarray:
+    cv2.circle(
+        scene,
+        center=center.as_xy_int_tuple(),
+        radius=radius,
+        color=color.as_bgr(),
+        thickness=-1,
+    )
+    return scene
+
+
+def calculate_dynamic_circle_radius(resolution_wh: Tuple[int, int]) -> int:
+    min_dimension = min(resolution_wh)
+    if min_dimension < 480:
+        return 4
+    if min_dimension < 720:
+        return 8
+    if min_dimension < 1080:
+        return 8
+    if min_dimension < 2160:
+        return 16
+    else:
+        return 16

utils/efficient_sam.py

@@ -0,0 +1,80 @@
+import torch
+import numpy as np
+from torchvision.transforms import ToTensor
+
+GPU_EFFICIENT_SAM_CHECKPOINT = "efficient_sam_s_gpu.jit"
+CPU_EFFICIENT_SAM_CHECKPOINT = "efficient_sam_s_cpu.jit"
+
+
+def load(device: torch.device) -> torch.jit.ScriptModule:
+    if device.type == "cuda":
+        model = torch.jit.load(GPU_EFFICIENT_SAM_CHECKPOINT)
+    else:
+        model = torch.jit.load(CPU_EFFICIENT_SAM_CHECKPOINT)
+    model.eval()
+    return model
+
+
+def inference_with_box(
+    image: np.ndarray,
+    box: np.ndarray,
+    model: torch.jit.ScriptModule,
+    device: torch.device
+) -> np.ndarray:
+    bbox = torch.reshape(torch.tensor(box), [1, 1, 2, 2])
+    bbox_labels = torch.reshape(torch.tensor([2, 3]), [1, 1, 2])
+    img_tensor = ToTensor()(image)
+
+    predicted_logits, predicted_iou = model(
+        img_tensor[None, ...].to(device),
+        bbox.to(device),
+        bbox_labels.to(device),
+    )
+    predicted_logits = predicted_logits.cpu()
+    all_masks = torch.ge(torch.sigmoid(predicted_logits[0, 0, :, :, :]), 0.5).numpy()
+    predicted_iou = predicted_iou[0, 0, ...].cpu().detach().numpy()
+
+    max_predicted_iou = -1
+    selected_mask_using_predicted_iou = None
+    for m in range(all_masks.shape[0]):
+        curr_predicted_iou = predicted_iou[m]
+        if (
+                curr_predicted_iou > max_predicted_iou
+                or selected_mask_using_predicted_iou is None
+        ):
+            max_predicted_iou = curr_predicted_iou
+            selected_mask_using_predicted_iou = all_masks[m]
+    return selected_mask_using_predicted_iou
+
+
+def inference_with_point(
+    image: np.ndarray,
+    point: np.ndarray,
+    model: torch.jit.ScriptModule,
+    device: torch.device
+) -> np.ndarray:
+    pts_sampled = torch.reshape(torch.tensor(point), [1, 1, -1, 2])
+    max_num_pts = pts_sampled.shape[2]
+    pts_labels = torch.ones(1, 1, max_num_pts)
+    img_tensor = ToTensor()(image)
+
+    predicted_logits, predicted_iou = model(
+        img_tensor[None, ...].to(device),
+        pts_sampled.to(device),
+        pts_labels.to(device),
+    )
+    predicted_logits = predicted_logits.cpu()
+    all_masks = torch.ge(torch.sigmoid(predicted_logits[0, 0, :, :, :]), 0.5).numpy()
+    predicted_iou = predicted_iou[0, 0, ...].cpu().detach().numpy()
+
+    max_predicted_iou = -1
+    selected_mask_using_predicted_iou = None
+    for m in range(all_masks.shape[0]):
+        curr_predicted_iou = predicted_iou[m]
+        if (
+                curr_predicted_iou > max_predicted_iou
+                or selected_mask_using_predicted_iou is None
+        ):
+            max_predicted_iou = curr_predicted_iou
+            selected_mask_using_predicted_iou = all_masks[m]
+    return selected_mask_using_predicted_iou

utils/tools.py

@@ -0,0 +1,443 @@
+# https://huggingface.co/spaces/An-619/FastSAM/edit/main/utils/tools.py
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
+import cv2
+import torch
+import os
+import sys
+import clip
+
+
+def convert_box_xywh_to_xyxy(box):
+    if len(box) == 4:
+        return [box[0], box[1], box[0] + box[2], box[1] + box[3]]
+    else:
+        result = []
+        for b in box:
+            b = convert_box_xywh_to_xyxy(b)
+            result.append(b)               
+    return result
+
+
+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(result, filter=0):
+    annotations = []
+    n = len(result.masks.data)
+    for i in range(n):
+        annotation = {}
+        mask = result.masks.data[i] == 1.0
+
+        if torch.sum(mask) < filter:
+            continue
+        annotation["id"] = i
+        annotation["segmentation"] = mask.cpu().numpy()
+        annotation["bbox"] = result.boxes.data[i]
+        annotation["score"] = result.boxes.conf[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)
+
+
+# clip
+@torch.no_grad()
+def retriev(
+    model, preprocess, elements: [Image.Image], search_text: str, device
+):
+    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:
+        image = image.resize((mask_w, mask_h))
+    cropped_boxes = []
+    cropped_images = []
+    not_crop = []
+    origin_id = []
+    for _, mask in enumerate(annotations):
+        if np.sum(mask["segmentation"]) <= 100:
+            continue
+        origin_id.append(_)
+        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, origin_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))
+    masks = sorted(masks, key=lambda x: x['area'], reverse=True)
+    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] = 1
+            if mask[point[1], point[0]] == 1 and point_label[i] == 0:
+                onemask[mask] = 0
+    onemask = onemask >= 1
+    return onemask, 0
+
+
+def text_prompt(annotations, text, img_path, device, wider=False, threshold=0.9):
+    cropped_boxes, cropped_images, not_crop, origin_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 = origin_id[int(max_idx)]
+
+    # find the biggest mask which contains the mask with max score
+    if wider:
+        mask0 = annotations_[max_idx]["segmentation"]
+        area0 = np.sum(mask0)
+        areas = [(i, np.sum(mask["segmentation"])) for i, mask in enumerate(annotations_) if i in origin_id]
+        areas = sorted(areas, key=lambda area: area[1], reverse=True)
+        indices = [area[0] for area in areas]
+        for index in indices:
+            if index == max_idx or np.sum(annotations_[index]["segmentation"] & mask0) / area0 > threshold:
+                max_idx = index
+                break
+
+    return annotations_[max_idx]["segmentation"], max_idx

utils/tools_gradio.py

@@ -0,0 +1,176 @@
+# https://huggingface.co/spaces/An-619/FastSAM/edit/main/utils/tools_gradio.py
+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:
+        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 not retinamask:
+        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:
+        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 not retinamask:
+        mask_cpu = cv2.resize(
+            mask_cpu, (target_width, target_height), interpolation=cv2.INTER_NEAREST
+        )
+    return mask_cpu