from huggingface_hub import snapshot_download
import gradio as gr
import numpy as np
import torch
import sys
import os
from tinysam import sam_model_registry, SamPredictor

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

model_type = "vit_t"
sam = sam_model_registry[model_type](checkpoint="./tinysam/tinysam.pth")

predictor = SamPredictor(sam)

examples = [
    [os.path.join(os.path.dirname(__file__), "assets/1.jpg")],
    ["assets/2.jpg"],
    ["assets/3.jpg"],
    ["assets/4.jpeg"],
    ["assets/5.jpg"],
    ["assets/6.jpeg"]
]

default_example = examples[0]
# 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 = """# Interactive Instance Segmentation
                - Point-prompt instruction
                <ol>
                <li> Click on the left image (point input), visualizing the point on the right image </li>
                <li> Click the button of Segment with Point Prompt </li>
                </ol>
                - Box-prompt instruction
                <ol>
                <li> Click on the left image (one point input), visualizing the point on the right image </li>
                <li> Click on the left image (another point input), visualizing the point and the box on the right image</li>
                <li> Click the button of Segment with Box Prompt </li>
                </ol>
                - Github [link](https://github.com/xinghaochen/TinySAM)
              """
css = "h1 { text-align: center } .about { text-align: justify; padding-left: 10%; padding-right: 10%; }"

def infer(img):
    if img is None:
        gr.Error("Please upload an image and select a point.")
    if img["background"] is None:
        gr.Error("Please upload an image and select a point.")
    # background (original image) layers[0] ( point prompt) composite (total image)
    image = img["background"].convert("RGB")
    point_prompt = img["layers"][0]
    total_image = img["composite"]
    predictor.set_image(np.array(image))
    print("point_prompt : ", point_prompt)

    # get point prompt
    img_arr = np.array(point_prompt)
    if not np.any(img_arr):
        gr.Error("Please select a point on top of the image.")
    else:
        nonzero_indices = np.nonzero(img_arr)
        img_arr = np.array(point_prompt)
        nonzero_indices = np.nonzero(img_arr)
        center_x = int(np.mean(nonzero_indices[1]))
        center_y = int(np.mean(nonzero_indices[0]))
        input_point = np.array([[center_x, center_y]])
        input_label = np.array([1])
        masks, scores, logits = predictor.predict(
            point_coords=input_point,
            point_labels=input_label,
        )
        result_label = [(masks[scores.argmax(), :, :], "mask")]
        return image, result_label


with gr.Blocks(css=css, title="TinySAM") as demo:
    with gr.Row():
        with gr.Column(scale=1):
            # Title
            gr.Markdown(title)
    with gr.Row():
        with gr.Column():
            im = gr.ImageEditor(
                type="pil",
            )
        output = gr.AnnotatedImage()

    im.change(infer, inputs=im, outputs=output)

demo.launch(debug=True)