from huggingface_hub import snapshot_download
import gradio as gr
import numpy as np
import torch
import sys
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)

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))

  # get point prompt
  img_arr = np.array(point_prompt)
  nonzero_indices = np.nonzero(img_arr)
  if nonzero_indices is None:
      gr.Error("Please select a point.")
  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[0, :, :], "mask")]
  return image, result_label


with gr.Blocks() as demo:
    gr.Markdown("## TinySAM")
    gr.Markdown("**[TinySAM](https://arxiv.org/abs/2312.13789) is a framework to distill Segment Anything Model.**")
    gr.Markdown("**To try it out, simply upload an image and leave a point on what you would like to segment.**")
    with gr.Row():
        with gr.Column():
            im = gr.ImageEditor(
                type="pil"
            )
            submit_btn = gr.Button()
        output = gr.AnnotatedImage()
    submit_btn.click(infer, inputs=im, outputs=output)

demo.launch(debug=True)