# Code credit: [FastSAM Demo](https://huggingface.co/spaces/An-619/FastSAM).

import torch
import gradio as gr
import numpy as np
from segment_anything import sam_model_registry, SamPredictor
from segment_anything.onnx import SamPredictorONNX
from PIL import ImageDraw
from utils.tools_gradio import fast_process
import copy
import argparse

# Use ONNX to speed up the inference.
ENABLE_ONNX = False

parser = argparse.ArgumentParser(
    description="Host EdgeSAM as a local web service."
)
parser.add_argument(
    "--checkpoint",
    default="weights/edge_sam_3x.pth",
    type=str,
    help="The path to the PyTorch checkpoint of EdgeSAM."
)
parser.add_argument(
    "--encoder-onnx-path",
    default="weights/edge_sam_3x_encoder.onnx",
    type=str,
    help="The path to the ONNX model of EdgeSAM's encoder."
)
parser.add_argument(
    "--decoder-onnx-path",
    default="weights/edge_sam_3x_decoder.onnx",
    type=str,
    help="The path to the ONNX model of EdgeSAM's decoder."
)
parser.add_argument(
    "--server-name",
    default="0.0.0.0",
    type=str,
    help="The server address that this demo will be hosted on."
)
parser.add_argument(
    "--port",
    default=8080,
    type=int,
    help="The port that this demo will be hosted on."
)
args = parser.parse_args()

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if ENABLE_ONNX:
    predictor = SamPredictorONNX(args.encoder_onnx_path, args.decoder_onnx_path)
else:
    sam = sam_model_registry["edge_sam"](checkpoint=args.checkpoint, upsample_mode="bicubic")
    sam = sam.to(device=device)
    sam.eval()
    predictor = SamPredictor(sam)

examples = [
    ["assets/1.jpeg"],
    ["assets/2.jpeg"],
    ["assets/3.jpeg"],
    ["assets/4.jpeg"],
    ["assets/5.jpeg"],
    ["assets/6.jpeg"],
    ["assets/7.jpeg"],
    ["assets/8.jpeg"],
    ["assets/9.jpeg"],
    ["assets/10.jpeg"],
    ["assets/11.jpeg"],
    ["assets/12.jpeg"],
    ["assets/13.jpeg"],
    ["assets/14.jpeg"],
    ["assets/15.jpeg"],
    ["assets/16.jpeg"]
]

# Description
title = "<center><strong><font size='8'>EdgeSAM<font></strong> <a href='https://github.com/chongzhou96/EdgeSAM'><font size='6'>[GitHub]</font></a> </center>"

description_p = """ # Instructions for point mode

                1. Upload an image or click one of the provided examples.
                2. Select the point type.
                3. Click once or multiple times on the image to indicate the object of interest.
                4. The Clear button clears all the points.
                5. The Reset button resets both points and the image.

              """

description_b = """ # Instructions for box mode

                1. Upload an image or click one of the provided examples.
                2. Click twice on the image (diagonal points of the box).
                3. The Clear button clears the box.
                4. The Reset button resets both the box and the image.

              """

css = "h1 { text-align: center } .about { text-align: justify; padding-left: 10%; padding-right: 10%; }"


def reset(session_state):
    session_state['coord_list'] = []
    session_state['label_list'] = []
    session_state['box_list'] = []
    session_state['ori_image'] = None
    session_state['image_with_prompt'] = None
    session_state['feature'] = None
    return None, session_state


def reset_all(session_state):
    session_state['coord_list'] = []
    session_state['label_list'] = []
    session_state['box_list'] = []
    session_state['ori_image'] = None
    session_state['image_with_prompt'] = None
    session_state['feature'] = None
    return None, None, session_state


def clear(session_state):
    session_state['coord_list'] = []
    session_state['label_list'] = []
    session_state['box_list'] = []
    session_state['image_with_prompt'] = copy.deepcopy(session_state['ori_image'])
    return session_state['ori_image'], session_state


def on_image_upload(
    image,
    session_state,
    input_size=1024
):
    session_state['coord_list'] = []
    session_state['label_list'] = []
    session_state['box_list'] = []

    input_size = int(input_size)
    w, h = image.size
    scale = input_size / max(w, h)
    new_w = int(w * scale)
    new_h = int(h * scale)
    image = image.resize((new_w, new_h))
    session_state['ori_image'] = copy.deepcopy(image)
    session_state['image_with_prompt'] = copy.deepcopy(image)
    print("Image changed")
    nd_image = np.array(image)
    session_state['feature'] = predictor.set_image(nd_image)

    return image, session_state


def convert_box(xyxy):
    min_x = min(xyxy[0][0], xyxy[1][0])
    max_x = max(xyxy[0][0], xyxy[1][0])
    min_y = min(xyxy[0][1], xyxy[1][1])
    max_y = max(xyxy[0][1], xyxy[1][1])
    xyxy[0][0] = min_x
    xyxy[1][0] = max_x
    xyxy[0][1] = min_y
    xyxy[1][1] = max_y
    return xyxy


def segment_with_points(
    label,
    session_state,
    evt: gr.SelectData,
    input_size=1024,
    better_quality=False,
    withContours=True,
    use_retina=True,
    mask_random_color=False,
):
    x, y = evt.index[0], evt.index[1]
    point_radius, point_color = 5, (97, 217, 54) if label == "Positive" else (237, 34, 13)
    session_state['coord_list'].append([x, y])
    session_state['label_list'].append(1 if label == "Positive" else 0)

    print(f"coord_list: {session_state['coord_list']}")
    print(f"label_list: {session_state['label_list']}")

    draw = ImageDraw.Draw(session_state['image_with_prompt'])
    draw.ellipse(
        [(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)],
        fill=point_color,
    )
    image = session_state['image_with_prompt']

    if ENABLE_ONNX:
        coord_np = np.array(session_state['coord_list'])[None]
        label_np = np.array(session_state['label_list'])[None]
        masks, scores, _ = predictor.predict(
            features=session_state['feature'],
            point_coords=coord_np,
            point_labels=label_np,
        )
        masks = masks.squeeze(0)
        scores = scores.squeeze(0)
    else:
        coord_np = np.array(session_state['coord_list'])
        label_np = np.array(session_state['label_list'])
        masks, scores, logits = predictor.predict(
            features=session_state['feature'],
            point_coords=coord_np,
            point_labels=label_np,
            num_multimask_outputs=4,
            use_stability_score=True
        )

    print(f'scores: {scores}')
    area = masks.sum(axis=(1, 2))
    print(f'area: {area}')

    annotations = np.expand_dims(masks[scores.argmax()], axis=0)

    seg = fast_process(
        annotations=annotations,
        image=image,
        device=device,
        scale=(1024 // input_size),
        better_quality=better_quality,
        mask_random_color=mask_random_color,
        bbox=None,
        use_retina=use_retina,
        withContours=withContours,
    )

    return seg, session_state


def segment_with_box(
        session_state,
        evt: gr.SelectData,
        input_size=1024,
        better_quality=False,
        withContours=True,
        use_retina=True,
        mask_random_color=False,
):
    x, y = evt.index[0], evt.index[1]
    point_radius, point_color, box_outline = 5, (97, 217, 54), 5
    box_color = (0, 255, 0)

    if len(session_state['box_list']) == 0:
        session_state['box_list'].append([x, y])
    elif len(session_state['box_list']) == 1:
        session_state['box_list'].append([x, y])
    elif len(session_state['box_list']) == 2:
        session_state['image_with_prompt'] = copy.deepcopy(session_state['ori_image'])
        session_state['box_list'] = [[x, y]]

    print(f"box_list: {session_state['box_list']}")

    draw = ImageDraw.Draw(session_state['image_with_prompt'])
    draw.ellipse(
        [(x - point_radius, y - point_radius), (x + point_radius, y + point_radius)],
        fill=point_color,
    )
    image = session_state['image_with_prompt']

    if len(session_state['box_list']) == 2:
        box = convert_box(session_state['box_list'])
        xy = (box[0][0], box[0][1], box[1][0], box[1][1])
        draw.rectangle(
            xy,
            outline=box_color,
            width=box_outline
        )

        box_np = np.array(box)
        if ENABLE_ONNX:
            point_coords = box_np.reshape(2, 2)[None]
            point_labels = np.array([2, 3])[None]
            masks, _, _ = predictor.predict(
                features=session_state['feature'],
                point_coords=point_coords,
                point_labels=point_labels,
            )
            annotations = masks[:, 0, :, :]
        else:
            masks, scores, _ = predictor.predict(
                features=session_state['feature'],
                box=box_np,
                num_multimask_outputs=1,
            )
            annotations = masks

        seg = fast_process(
            annotations=annotations,
            image=image,
            device=device,
            scale=(1024 // input_size),
            better_quality=better_quality,
            mask_random_color=mask_random_color,
            bbox=None,
            use_retina=use_retina,
            withContours=withContours,
        )
        return seg, session_state
    return image, session_state


img_p = gr.Image(label="Input with points", type="pil")
img_b = gr.Image(label="Input with box", type="pil")

with gr.Blocks(css=css, title="EdgeSAM") as demo:
    session_state = gr.State({
        'coord_list': [],
        'label_list': [],
        'box_list': [],
        'ori_image': None,
        'image_with_prompt': None,
        'feature': None
    })

    with gr.Row():
        with gr.Column(scale=1):
            # Title
            gr.Markdown(title)

    with gr.Tab("Point mode") as tab_p:
        # Images
        with gr.Row(variant="panel"):
            with gr.Column(scale=1):
                img_p.render()
            with gr.Column(scale=1):
                with gr.Row():
                    add_or_remove = gr.Radio(
                        ["Positive", "Negative"],
                        value="Positive",
                        label="Point Type"
                    )

                    with gr.Column():
                        clear_btn_p = gr.Button("Clear", variant="secondary")
                        reset_btn_p = gr.Button("Reset", variant="secondary")
                with gr.Row():
                    gr.Markdown(description_p)

        with gr.Row():
            with gr.Column():
                gr.Markdown("Try some of the examples below ⬇️")
                gr.Examples(
                    examples=examples,
                    inputs=[img_p, session_state],
                    outputs=[img_p, session_state],
                    examples_per_page=8,
                    fn=on_image_upload,
                    run_on_click=True
                )

    with gr.Tab("Box mode") as tab_b:
        # Images
        with gr.Row(variant="panel"):
            with gr.Column(scale=1):
                img_b.render()
            with gr.Row():
                with gr.Column():
                    clear_btn_b = gr.Button("Clear", variant="secondary")
                    reset_btn_b = gr.Button("Reset", variant="secondary")
                    gr.Markdown(description_b)

        with gr.Row():
            with gr.Column():
                gr.Markdown("Try some of the examples below ⬇️")
                gr.Examples(
                    examples=examples,
                    inputs=[img_b, session_state],
                    outputs=[img_b, session_state],
                    examples_per_page=8,
                    fn=on_image_upload,
                    run_on_click=True
                )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown(
                "<center><img src='https://visitor-badge.laobi.icu/badge?page_id=chongzhou/edgesam' alt='visitors'></center>")

    img_p.upload(on_image_upload, [img_p, session_state], [img_p, session_state])
    img_p.select(segment_with_points, [add_or_remove, session_state], [img_p, session_state])

    clear_btn_p.click(clear, [session_state], [img_p, session_state])
    reset_btn_p.click(reset, [session_state], [img_p, session_state])
    tab_p.select(fn=reset_all, inputs=[session_state], outputs=[img_p, img_b, session_state])

    img_b.upload(on_image_upload, [img_b, session_state], [img_b, session_state])
    img_b.select(segment_with_box, [session_state], [img_b, session_state])

    clear_btn_b.click(clear, [session_state], [img_b, session_state])
    reset_btn_b.click(reset, [session_state], [img_b, session_state])
    tab_b.select(fn=reset_all, inputs=[session_state], outputs=[img_p, img_b, session_state])

demo.queue()
# demo.launch(server_name=args.server_name, server_port=args.port)
demo.launch()