# Code credit: [EdgeSAM Demo](https://huggingface.co/spaces/chongzhou/EdgeSAM).

import torch
import gradio as gr
from huggingface_hub import snapshot_download
import numpy as np
from tinysam import sam_model_registry, SamPredictor
from PIL import ImageDraw
from utils.tools_gradio import fast_process
import copy
import argparse

snapshot_download("merve/tinysam", local_dir="tinysam")

model_type = "vit_t"
sam = sam_model_registry[model_type](checkpoint="./tinysam/tinysam.pth")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
sam.to(device=device)
sam.eval()
predictor = SamPredictor(sam)

examples = [
    ["assets/1.jpg"],
    ["assets/2.jpg"],
    ["assets/3.jpg"],
    ["assets/4.jpeg"],
    ["assets/5.jpg"],
    ["assets/6.jpeg"]
]

# Description
title = "<center><strong><font size='8'>TinySAM<font></strong> <a href='https://github.com/xinghaochen/TinySAM'><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]
    print(f"x y: {x,y}")
    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']

    coord_np = np.array(session_state['coord_list'])
    label_np = np.array(session_state['label_list'])
    masks, scores, logits = predictor.predict(
        point_coords=coord_np,
        point_labels=label_np,
    )
    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(xy)
        masks, scores, _ = predictor.predict(
            point_coords=None,
            point_labels=None,
            box=box_np[None, :],
        )
        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
    return image, session_state


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

with gr.Blocks(css=css, title="TinySAM") 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=xinghaochen.tinysam' 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()