import gradio as gr
import cv2
import numpy as np
from PIL import Image
from ultralytics import SAM

model = SAM('sam_b.pt')


def process_image(input_image):
    if input_image is None:
        return None, "Can't find photo", 0

    input_image = np.array(input_image)
    bgr_image = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)

    r = model(bgr_image)[0]
    masks = r.masks.xy

    # point = (x, y)
    # segment_index = point_in_mask(masks, point)

    im_array = r.plot().astype(np.uint8)
    im_array = cv2.cvtColor(im_array, cv2.COLOR_BGR2RGB)
    im_pil = Image.fromarray(im_array)

    # cv2.circle(im_array, point, radius=5, color=(255, 255, 0), thickness=2)

    return im_pil


def point_in_mask(masks, point):
    x, y = point
    for i in range(len(masks) - 1, -1, -1):
        mask = np.array(masks[i])
        min_x = min(mask[:, 0])
        max_x = max(mask[:, 0])
        min_y = min(mask[:, 1])
        max_y = max(mask[:, 1])
        if min_x <= x <= max_x and min_y <= y <= max_y:
            return i
    return None


iface = gr.Interface(
    fn=process_image,
    inputs=[gr.Image()],
    outputs=[gr.Image(type="numpy")],
    title="Image Segmenter",
    description="Upload an image to segment."
)

iface.launch()