import gradio as gr
import numpy as np
import torch
from torchvision.transforms import ToTensor
from PIL import Image

# loading EfficientSAM model
model_path = "efficientsam_s_cpu.jit"
with open(model_path, "rb") as f:
    model = torch.jit.load(f)

# getting mask using points
def get_sam_mask_using_points(img_tensor, pts_sampled, model):
    pts_sampled = torch.reshape(torch.tensor(pts_sampled), [1, 1, -1, 2])
    max_num_pts = pts_sampled.shape[2]
    pts_labels = torch.ones(1, 1, max_num_pts)

    predicted_logits, predicted_iou = model(
        img_tensor[None, ...],
        pts_sampled,
        pts_labels,
    )
    predicted_logits = predicted_logits.cpu()
    all_masks = torch.ge(torch.sigmoid(predicted_logits[0, 0, :, :, :]), 0.5).numpy()
    predicted_iou = predicted_iou[0, 0, ...].cpu().detach().numpy()

    max_predicted_iou = -1
    selected_mask_using_predicted_iou = None
    for m in range(all_masks.shape[0]):
        curr_predicted_iou = predicted_iou[m]
        if (
            curr_predicted_iou > max_predicted_iou
            or selected_mask_using_predicted_iou is None
        ):
            max_predicted_iou = curr_predicted_iou
            selected_mask_using_predicted_iou = all_masks[m]
    return selected_mask_using_predicted_iou

# examples
examples = [["examples/image1.jpg"], ["examples/image2.jpg"], ["examples/image3.jpg"], ["examples/image4.jpg"],
            ["examples/image5.jpg"], ["examples/image6.jpg"], ["examples/image7.jpg"], ["examples/image8.jpg"],
            ["examples/image9.jpg"], ["examples/image10.jpg"], ["examples/image11.jpg"], ["examples/image12.jpg"]
            ["examples/image13.jpg"], ["examples/image14.jpg"]]


with gr.Blocks() as demo:
    with gr.Row():
        input_img = gr.Image(label="Input",height=512)
        output_img = gr.Image(label="Selected Segment",height=512)

    with gr.Row():
        gr.Markdown("Try some of the examples below ⬇️")
        gr.Examples(examples=examples,
                    inputs=[input_img])

    def get_select_coords(img, evt: gr.SelectData):
        img_tensor = ToTensor()(img)
        _, H, W = img_tensor.shape

        visited_pixels = set()
        pixels_in_queue = set()
        pixels_in_segment = set()

        mask = get_sam_mask_using_points(img_tensor, [[evt.index[0], evt.index[1]]], model)

        out = img.copy()

        out = out.astype(np.uint8)
        out *= mask[:,:,None]
        for pixel in pixels_in_segment:
            out[pixel[0], pixel[1]] = img[pixel[0], pixel[1]]
        print(out)
        return out

    input_img.select(get_select_coords, [input_img], output_img)

if __name__ == "__main__":
    demo.launch()