Upload files

.gitignore

@@ -0,0 +1,12 @@
+*.pyc
+*.pyo
+*.pyd
+__py
+**/__pycache__/
+data
+onnx
+results
+**.egg-info
+*.log
+*.onnx
+.hypothesis

app.py

@@ -0,0 +1,214 @@
+import numpy as np
+import os
+
+from nanosam import Predictor
+
+import gradio as gr
+import time
+from PIL import ImageDraw
+from utils import download_file_from_url, fast_process, format_results, point_prompt
+
+# Most of our demo code is from [FastSAM Demo](https://huggingface.co/spaces/An-619/FastSAM). Huge thanks for AN-619.
+
+if not os.path.exists("onnx/sam_hgv2_b4_ln_nonorm_image_encoder.onnx"):
+    download_file_from_url(
+        "https://huggingface.co/dragonSwing/nanosam/resolve/main/sam_hgv2_b4_ln_nonorm_image_encoder.onnx",
+        "onnx/sam_hgv2_b4_ln_nonorm_image_encoder.onnx",
+    )
+
+if not os.path.exists("onnx/efficientvit_l0_mask_decoder.onnx"):
+    download_file_from_url(
+        "https://huggingface.co/dragonSwing/nanosam/resolve/main/efficientvit_l0_mask_decoder.onnx",
+        "onnx/efficientvit_l0_mask_decoder.onnx",
+    )
+
+# Load the pre-trained model
+image_encoder_cfg = {
+    "path": "onnx/sam_hgv2_b4_ln_nonorm_image_encoder.onnx",
+    "provider": "cpu",
+    "normalize_input": False,
+}
+mask_decoder_cfg = {
+    "path": "onnx/efficientvit_l0_mask_decoder.onnx",
+    "provider": "cpu",
+}
+predictor = Predictor(image_encoder_cfg, mask_decoder_cfg)
+
+# Description
+title = "<center><strong><font size='8'>Faster Segment Anything(NanoSAM)<font></strong></center>"
+
+description_p = """ ## This is a demo of [Faster Segment Anything(NanoSAM) Model](https://github.com/binh234/nanosam).
+                # Instructions for point mode
+                0. Restart by click the Restart button
+                1. Select a point with Add Mask for the foreground (Must)
+                2. Select a point with Remove Area for the background (Optional)
+                3. Click the Start Segmenting.
+                - Github [link](https://github.com/binh234/nanosam)
+                - Model Card [link](https://huggingface.co/dragoswing/nanosam)
+                We will provide box mode soon. 
+                Enjoy!
+              """
+
+examples = [
+    ["assets/picture3.jpg"],
+    ["assets/picture4.jpg"],
+    ["assets/picture5.jpg"],
+    ["assets/picture6.jpg"],
+    ["assets/picture1.jpg"],
+    ["assets/picture2.jpg"],
+    ["assets/dogs.jpg"],
+]
+
+css = "h1 { text-align: center } .about { text-align: justify; padding-left: 10%; padding-right: 10%; }"
+
+
+def get_empty_state():
+    return {"points": [], "point_labels": [], "features": None}
+
+
+def clear():
+    return None, None, get_empty_state()
+
+
+def set_image(image):
+    state = get_empty_state()
+    start = time.perf_counter()
+    predictor.set_image(image)
+    end = time.perf_counter()
+    print(f"Encoder time: {end - start: .3f}s")
+    state["features"] = predictor.features
+    return state
+
+
+def segment_with_points(
+    image,
+    state,
+    better_quality=False,
+    withContours=True,
+    use_retina=True,
+    mask_random_color=True,
+):
+    global predictor
+
+    points = np.asarray(state["points"])
+    point_labels = np.asarray(state["point_labels"])
+    if len(points) == 0 and len(point_labels) == 0:
+        raise gr.Error("No points selected")
+    if len(points) != len(point_labels):
+        raise gr.Error("Mismatch length between points and point labels")
+    if state["features"] is None:
+        raise gr.Error(
+            "Image was not set correctly, please wait for a moment after uploading image before drawing points!"
+        )
+
+    predictor.features = state["features"]
+    img_w, img_h = image.size
+    predictor.original_size = (img_h, img_w)
+    start = time.perf_counter()
+    masks, scores, logits = predictor.predict(
+        points=points,
+        point_labels=point_labels,
+    )
+    end = time.perf_counter()
+    print(f"Decoder time: {end - start: .3f}s")
+
+    # results = format_results(masks[0], scores[0], logits[0], 0)
+
+    # annotations, _ = point_prompt(results, points, point_labels, img_h, img_w)
+    # annotations = np.array([annotations])
+
+    fig = fast_process(
+        annotations=[masks[0, scores.argmax()] > 0],
+        image=image,
+        scale=1,
+        better_quality=better_quality,
+        mask_random_color=mask_random_color,
+        bbox=None,
+        use_retina=use_retina,
+        withContours=withContours,
+    )
+
+    # return fig, None
+    return fig
+
+
+def get_points_with_draw(image, label, evt: gr.SelectData, state):
+    x, y = evt.index[0], evt.index[1]
+    point_radius, point_color = 15, (
+        (255, 255, 0)
+        if label == "Add Mask"
+        else (
+            255,
+            0,
+            255,
+        )
+    )
+    state["points"].append([x, y])
+    state["point_labels"].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, state
+
+
+cond_img_p = gr.Image(label="Input with points", type="pil", interactive=True)
+
+segm_img_p = gr.Image(label="Segmented Image with points", interactive=False, type="pil")
+
+global_points = []
+global_point_labels = []
+
+with gr.Blocks(css=css, title="Faster Segment Anything(NanoSAM)") as demo:
+    state = gr.State(value=get_empty_state())
+    with gr.Row():
+        with gr.Column(scale=1):
+            # Title
+            gr.Markdown(title)
+
+    with gr.Tab("Point 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",
+                    )
+
+                    with gr.Column():
+                        segment_btn_p = gr.Button("Start segmenting!", variant="primary")
+                        restart_btn_p = gr.Button("Restart", variant="secondary")
+
+                gr.Markdown("Try some of the examples below ⬇️")
+                gr.Examples(
+                    examples=examples,
+                    inputs=[cond_img_p],
+                    outputs=[state],
+                    fn=set_image,
+                    run_on_click=True,
+                    examples_per_page=4,
+                )
+
+            with gr.Column():
+                # Description
+                gr.Markdown(description_p)
+
+    cond_img_p.upload(set_image, inputs=[cond_img_p], outputs=[state])
+    cond_img_p.select(get_points_with_draw, [cond_img_p, add_or_remove, state], [cond_img_p, state])
+    segment_btn_p.click(segment_with_points, [cond_img_p, state], [segm_img_p])
+    restart_btn_p.click(clear, outputs=[cond_img_p, segm_img_p, state])
+
+demo.queue().launch()

requirements.txt

@@ -0,0 +1,5 @@
+matplotlib
+onnx>=1.14.0
+onnxruntime>=1.14.0
+opencv-python-headless
+git+https://github.com/binh234/nanosam.git

utils.py

@@ -0,0 +1,178 @@
+import os
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+import requests
+from PIL import Image
+
+
+def point_prompt(masks, points, point_label, target_height, target_width):
+    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:
+                if point_label[i] == 0:
+                    onemask -= mask
+                else:
+                    onemask += mask
+                break
+    onemask = onemask > 0
+    return onemask, 0
+
+
+def format_results(masks, scores, logits, filter=0):
+    annotations = []
+    n = len(scores)
+    for i in range(n):
+        annotation = {}
+
+        mask = masks[i] > 0
+        tmp = np.where(mask)
+        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"] = mask.sum()
+        annotations.append(annotation)
+    return annotations
+
+
+def fast_process(
+    annotations,
+    image,
+    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:
+        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)
+            )
+    annotations = np.asarray(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,
+    )
+
+    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]
+    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 download_file_from_url(url, output_file, chunk_size=8192):
+    output_dir = os.path.dirname(output_file)
+    os.makedirs(output_dir, exist_ok=True)
+    try:
+        with requests.get(url, stream=True) as response:
+            if response.status_code == 200:
+                with open(output_file, 'wb') as f:
+                    for chunk in response.iter_content(chunk_size=chunk_size):
+                        f.write(chunk)
+            else:
+                print(f"Failed to download file. Status code: {response.status_code}")
+    except Exception as e:
+        print(f"An error occurred: {e}")