import gradio as gr
import numpy as np
import torch
from PIL import Image
from transformers import SamModel, SamProcessor
from gradio_image_prompter import ImagePrompter
import spaces

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
sam_model = SamModel.from_pretrained("facebook/sam-vit-huge").to("cuda")
sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
slimsam_model = SamModel.from_pretrained("nielsr/slimsam-50-uniform").to("cuda")
slimsam_processor = SamProcessor.from_pretrained("nielsr/slimsam-50-uniform")

def get_processor_and_model(slim: bool):
    if slim:
        return slimsam_processor, slimsam_model
    return sam_processor, sam_model

@spaces.GPU
def sam_box_inference(image, x_min, y_min, x_max, y_max, *, slim=False):

    processor, model = get_processor_and_model(slim)
    
    inputs = processor(
        Image.fromarray(image),
        input_boxes=[[[[x_min, y_min, x_max, y_max]]]],
        return_tensors="pt"
    ).to(device)

    with torch.no_grad():
        outputs = model(**inputs)

    mask = processor.image_processor.post_process_masks(
        outputs.pred_masks.cpu(),
        inputs["original_sizes"].cpu(),
        inputs["reshaped_input_sizes"].cpu()
    )[0][0][0].numpy()
    mask = mask[np.newaxis, ...]
    print(mask)
    print(mask.shape)
    return [(mask, "mask")]

@spaces.GPU
def sam_point_inference(image, x, y, *, slim=False):

    processor, model = get_processor_and_model(slim)
    
    inputs = processor(
        image,
        input_points=[[[x, y]]],
        return_tensors="pt").to(device)

    with torch.no_grad():
        outputs = model(**inputs)

    mask = processor.post_process_masks(
        outputs.pred_masks.cpu(),
        inputs["original_sizes"].cpu(),
        inputs["reshaped_input_sizes"].cpu()
    )[0][0][0].numpy()
    mask = mask[np.newaxis, ...]
    print(type(mask))
    print(mask.shape)
    return [(mask, "mask")]

def infer_point(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"]
    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]))
    print("Point inference returned.")
    return ((image, sam_point_inference(image, center_x, center_y, slim=True)),
    (image, sam_point_inference(image, center_x, center_y)))

def infer_box(prompts):
    # background (original image) layers[0] ( point prompt) composite (total image)
    image = prompts["image"]
    if image is None:
      gr.Error("Please upload an image and draw a box before submitting")
    points = prompts["points"][0]
    if points is None:
      gr.Error("Please draw a box before submitting.")
    print(points)

    # x_min = points[0] x_max = points[3] y_min = points[1] y_max = points[4]
    return ((image, sam_box_inference(image, points[0], points[1], points[3], points[4], slim=True)),
    (image, sam_box_inference(image, points[0], points[1], points[3], points[4])))
with gr.Blocks(title="SlimSAM") as demo:
  gr.Markdown("# SlimSAM")
  gr.Markdown("SlimSAM is the pruned-distilled version of SAM that is smaller.")
  gr.Markdown("In this demo, you can compare SlimSAM and SAM outputs in point and box prompts.")
  
  with gr.Tab("Box Prompt"):
    with gr.Row():
        with gr.Column(scale=1):
            # Title
            gr.Markdown("To try box prompting, simply upload and image and draw a box on it.")
    with gr.Row():
        with gr.Column():
            im = ImagePrompter()
            btn = gr.Button("Submit")
        with gr.Column():
          output_box_slimsam = gr.AnnotatedImage(label="SlimSAM Output")
          output_box_sam = gr.AnnotatedImage(label="SAM Output")


    btn.click(infer_box, inputs=im, outputs=[output_box_slimsam, output_box_sam])

  with gr.Tab("Point Prompt"):
    with gr.Row():
        with gr.Column(scale=1):
            # Title
            gr.Markdown("To try point prompting, simply upload and image and leave a dot on it.")
    with gr.Row():
        with gr.Column():
            im = gr.ImageEditor(
                type="pil",
            )
        with gr.Column():
          output_slimsam = gr.AnnotatedImage(label="SlimSAM Output")
          output_sam = gr.AnnotatedImage(label="SAM Output")

    im.change(infer_point, inputs=im, outputs=[output_slimsam, output_sam])
demo.launch(debug=True)