app.py

"""
Copyright (c) 2024-present Naver Cloud Corp.
This source code is based on code from the Segment Anything Model (SAM)
(https://github.com/facebookresearch/segment-anything).

This source code is licensed under the license found in the
LICENSE file in the root directory of this source tree.
"""
import os, sys
sys.path.append(os.getcwd())

# Gradio demo, comparison SAM vs ZIM
import os
import torch
import gradio as gr
from gradio_image_prompter import ImagePrompter
import numpy as np
import cv2
from zim import zim_model_registry, ZimPredictor, ZimAutomaticMaskGenerator
from segment_anything import sam_model_registry, SamPredictor, SamAutomaticMaskGenerator
from zim.utils import show_mat_anns

from huggingface_hub import hf_hub_download

def get_shortest_axis(image):
    h, w, _ = image.shape
    return h if h < w else w

def reset_image(image, prompts):
    if image is None:
        image = np.zeros((1024, 1024, 3), dtype=np.uint8)
    else:
        image = image['image']
    zim_predictor.set_image(image)
    sam_predictor.set_image(image)
    prompts = dict()
    black = np.zeros(image.shape[:2], dtype=np.uint8)

    return (image, image, image, image, black, black, black, black, prompts)

def reset_example_image(image, prompts):
    if image is None:
        image = np.zeros((1024, 1024, 3), dtype=np.uint8)

    zim_predictor.set_image(image)
    sam_predictor.set_image(image)
    prompts = dict()
    black = np.zeros(image.shape[:2], dtype=np.uint8)

    image_dict = {}
    image_dict['image'] = image
    image_dict['prompts'] = prompts

    return (image, image_dict, image, image, image, black, black, black, black, prompts)

def run_amg(image):
    gr.Info('Checkout ZIM Auto Mask tab.', duration=3)
    zim_masks = zim_mask_generator.generate(image)
    zim_masks_vis = show_mat_anns(image, zim_masks)
    
    sam_masks = sam_mask_generator.generate(image)
    sam_masks_vis = show_mat_anns(image, sam_masks)
    
    return zim_masks_vis, sam_masks_vis
    
    
def run_model(image, prompts):
    if not prompts:
        raise gr.Error(f'Please input any point or BBox')
    gr.Info('Checkout ZIM Mask tab.', duration=3)
    point_coords = None
    point_labels = None
    boxes = None

    if "point" in prompts:
        point_coords, point_labels = [], []

        for type, pts in prompts["point"]:
            point_coords.append(pts)
            point_labels.append(type)
        point_coords = np.array(point_coords)
        point_labels = np.array(point_labels)

    if "bbox" in prompts:
        boxes = prompts['bbox']
        boxes = np.array(boxes)

    if "scribble" in prompts:
        point_coords, point_labels = [], []

        for pts in prompts["scribble"]:
            point_coords.append(np.flip(pts))
            point_labels.append(1)
        if len(point_coords) == 0:
            raise gr.Error("Please input any scribbles.")
        point_coords = np.array(point_coords)
        point_labels = np.array(point_labels)

    # run ZIM
    zim_mask, _, _ = zim_predictor.predict(
        point_coords=point_coords,
        point_labels=point_labels,
        box=boxes,
        multimask_output=False,
    )
    zim_mask = np.squeeze(zim_mask, axis=0)
    zim_mask = np.uint8(zim_mask * 255)

    # run SAM
    sam_mask, _, _ = sam_predictor.predict(
        point_coords=point_coords,
        point_labels=point_labels,
        box=boxes,
        multimask_output=False,
        )
    sam_mask = np.squeeze(sam_mask, axis=0)
    sam_mask = np.uint8(sam_mask * 255)
    
    return zim_mask, sam_mask

def reset_scribble(image, scribble, prompts):
    # scribble = dict()
    for k in prompts.keys():
        prompts[k] = []

    for k, v in scribble.items():
        scribble[k] = None

    black = np.zeros(image.shape[:3], dtype=np.uint8)

    return scribble, black, black

def update_scribble(image, scribble, prompts):
    if "point" in prompts:
        del prompts["point"]

    if "bbox" in prompts:
        del prompts["bbox"]
    
    prompts = dict() # reset prompt
    scribble_mask = scribble["layers"][0][..., -1] > 0

    scribble_coords = np.argwhere(scribble_mask)
    n_points = min(len(scribble_coords), 24)
    indices = np.linspace(0, len(scribble_coords)-1, n_points, dtype=int)
    scribble_sampled = scribble_coords[indices]

    prompts["scribble"] = scribble_sampled
    
    zim_mask, sam_mask = run_model(image, prompts)

    return zim_mask, sam_mask, prompts


def draw_point(img, pt, size, color):
    # draw circle with white boundary region
    cv2.circle(img, (int(pt[0]), int(pt[1])), int(size * 1.3), (255, 255, 255), -1)
    cv2.circle(img, (int(pt[0]), int(pt[1])), int(size * 0.9), color, -1)


def draw_images(image, mask, prompts):
    if len(prompts) == 0 or mask.shape[1] == 1:
        return image, image, image

    minor = get_shortest_axis(image)
    size = int(minor / 80)

    image = np.float32(image)

    def blending(image, mask):
        mask = np.float32(mask) / 255
        blended_image = np.zeros_like(image, dtype=np.float32)
        blended_image[:, :, :] = [108, 0, 192]
        blended_image = (image * 0.5) + (blended_image * 0.5)
    
        img_with_mask = mask[:, :, None] * blended_image + (1 - mask[:, :, None]) * image
        img_with_mask = np.uint8(img_with_mask)

        return img_with_mask

    img_with_mask = blending(image, mask)
    img_with_point = img_with_mask.copy()

    if "point" in prompts:
        for type, pts in prompts["point"]:
            if type == "Positive":
                color = (0, 0, 255)
                draw_point(img_with_point, pts, size, color)
            elif type == "Negative":
                color = (255, 0, 0)
                draw_point(img_with_point, pts, size, color)

    size = int(minor / 200)

    return (
        img,
        img_with_mask,
    )

def get_point_or_box_prompts(img, prompts):
    image, img_prompts = img['image'], img['points']
    point_prompts = []
    box_prompts = []
    for prompt in img_prompts:
        for p in range(len(prompt)):
            prompt[p] = int(prompt[p])
        if prompt[2] == 2 and prompt[5] == 3:  # box prompt
            if len(box_prompts) != 0:
                raise gr.Error("Please input only one BBox.", duration=3)
            box_prompts.append([prompt[0], prompt[1], prompt[3], prompt[4]])
        elif prompt[2] == 1 and prompt[5] == 4:  # Positive point prompt
            point_prompts.append((1, (prompt[0], prompt[1])))
        elif prompt[2] == 0 and prompt[5] == 4:  # Negative point prompt
            point_prompts.append((0, (prompt[0], prompt[1])))

    if "scribble" in prompts:
        del prompts["scribble"]

    if len(point_prompts) > 0:
        prompts['point'] = point_prompts
    elif 'point' in prompts:
        del prompts['point']

    if len(box_prompts) > 0:
        prompts['bbox'] = box_prompts
    elif 'bbox' in prompts:
        del prompts['bbox']

    zim_mask, sam_mask = run_model(image, prompts)

    return image, zim_mask, sam_mask, prompts

def get_examples():
    assets_dir = os.path.join(os.path.dirname(__file__), 'examples')
    images = os.listdir(assets_dir)
    return [os.path.join(assets_dir, img) for img in images]

def download_onnx_weights(repo_id="naver-iv/zim-anything-vitb", file_dir="zim_vit_b_2043"):
    hf_hub_download(repo_id=repo_id, filename=f"{file_dir}/encoder.onnx")
    filepath = hf_hub_download(repo_id=repo_id, filename=f"{file_dir}/decoder.onnx")

    return os.path.dirname(filepath)

    
if __name__ == "__main__":
    backbone = "vit_b"
    
    # load ZIM
    zim = zim_model_registry[backbone](checkpoint=download_onnx_weights())
    if torch.cuda.is_available():
        zim.cuda()
    zim_predictor = ZimPredictor(zim)
    zim_mask_generator = ZimAutomaticMaskGenerator(
        zim, 
        pred_iou_thresh=0.7, 
        points_per_batch=8,
        stability_score_thresh=0.9, 
    )

    # load SAM
    ckpt_sam = "ckpts/sam_vit_b_01ec64.pth"
    sam = sam_model_registry[backbone](checkpoint=ckpt_sam)
    if torch.cuda.is_available():
        sam.cuda()
    sam_predictor = SamPredictor(sam)
    sam_mask_generator = SamAutomaticMaskGenerator(
        sam,
        points_per_batch=8,
    )
    
    with gr.Blocks() as demo:
        gr.Markdown("# <center> [Demo] ZIM: Zero-Shot Image Matting for Anything")

        prompts = gr.State(dict())
        img = gr.Image(visible=False)
        example_image = gr.Image(visible=False)
        
        with gr.Row():
            with gr.Column():
                # Point and Bbox prompt
                with gr.Tab(label="Point or Box"):
                    img_with_point_or_box = ImagePrompter(
                        label="query image", 
                        sources="upload"
                    )
                    interactions = "Left Click (Pos) | Middle/Right Click (Neg) | Press Move (Box)"
                    gr.Markdown("<h3 style='text-align: center'> {} </h3>".format(interactions))
                    run_bttn = gr.Button("Run")
                    amg_bttn = gr.Button("Automatic Mask Generation")

                # Scribble prompt
                with gr.Tab(label="Scribble"):
                    img_with_scribble = gr.ImageEditor(
                        label="Scribble", 
                        brush=gr.Brush(colors=["#00FF00"], default_size=15),
                        sources="upload", 
                        transforms=None, 
                        layers=False
                    )
                    interactions = "Press Move (Scribble)"
                    gr.Markdown("<h3 style='text-align: center'> Step 1. Select Draw button </h3>")
                    gr.Markdown("<h3 style='text-align: center'> Step 2. {} </h3>".format(interactions))
                    scribble_bttn = gr.Button("Run")
                    scribble_reset_bttn = gr.Button("Reset Scribbles")
                    amg_scribble_bttn = gr.Button("Automatic Mask Generation")
                
                # Example image
                gr.Examples(get_examples(), inputs=[example_image])

            # with gr.Row():
            with gr.Column():
                with gr.Tab(label="ZIM Image"):
                    img_with_zim_mask = gr.Image(
                        label="ZIM Image", 
                        interactive=False
                    )

                with gr.Tab(label="ZIM Mask"):
                    zim_mask = gr.Image(
                        label="ZIM Mask", 
                        image_mode="L", 
                        interactive=False
                    )
                with gr.Tab(label="ZIM Auto Mask"):
                    zim_amg = gr.Image(
                        label="ZIM Auto Mask", 
                        interactive=False
                    )
                    
            with gr.Column():
                with gr.Tab(label="SAM Image"):
                    img_with_sam_mask = gr.Image(
                        label="SAM image", 
                        interactive=False
                    )

                with gr.Tab(label="SAM Mask"):
                    sam_mask = gr.Image(
                        label="SAM Mask", 
                        image_mode="L", 
                        interactive=False
                    )
                    
                with gr.Tab(label="SAM Auto Mask"):
                    sam_amg = gr.Image(
                        label="SAM Auto Mask", 
                        interactive=False
                    )
                    
        example_image.change(
            reset_example_image,
            [example_image, prompts],
            [
                img,
                img_with_point_or_box,
                img_with_scribble,
                img_with_zim_mask,
                img_with_sam_mask,
                zim_amg,
                sam_amg,
                zim_mask,
                sam_mask,
                prompts,
            ]
        )

        img_with_point_or_box.upload(
            reset_image,
            [img_with_point_or_box, prompts],
            [
                img,
                img_with_scribble,
                img_with_zim_mask,
                img_with_sam_mask,
                zim_amg,
                sam_amg,
                zim_mask,
                sam_mask,
                prompts,
            ],
        )

        amg_bttn.click(
            run_amg,
            [img],
            [zim_amg, sam_amg]
        )
        amg_scribble_bttn.click(
            run_amg,
            [img],
            [zim_amg, sam_amg]
        )
        
        run_bttn.click(
            get_point_or_box_prompts,
            [img_with_point_or_box, prompts],
            [img, zim_mask, sam_mask, prompts]
        )

        zim_mask.change(
            draw_images,
            [img, zim_mask, prompts],
            [
                img, img_with_zim_mask, 
            ],
        )
        sam_mask.change(
            draw_images,
            [img, sam_mask, prompts],
            [
                img, img_with_sam_mask, 
            ],
        )
        
        scribble_reset_bttn.click(
            reset_scribble,
            [img, img_with_scribble, prompts],
            [img_with_scribble, zim_mask, sam_mask],
        )
        scribble_bttn.click(
            update_scribble,
            [img, img_with_scribble, prompts],
            [zim_mask, sam_mask, prompts],
        )

    demo.queue()
    demo.launch()