import gradio as gr
import numpy as np
import os
import torch
from model import SegmentationModel
import cv2

def segment(img, bbox_x, bbox_y, bbox_w, bbox_h):
    img = np.array(img)
    model_path = os.path.join(os.path.dirname(__file__), 'model_segmentation.pth')

    model = SegmentationModel(model_path)
    bounding_boxes = torch.tensor([
            [int(bbox_x), int(bbox_y), int(bbox_w), int(bbox_h)],
        ]).int()

    data_for_segmentation = dict()
    data_for_segmentation['bbox'] = bounding_boxes
    data_for_segmentation['conf'] = torch.ones(len(bounding_boxes))
    data_for_segmentation['labels'] = ['cat'] * len(bounding_boxes)

    segmentation = model.single_inference(img, data_for_segmentation)
    mask = segmentation[0]
    bounding_box = data_for_segmentation['bbox'][0]
    print(bounding_box)

    img = cv2.rectangle(img, (int(bounding_box[0]), int(bounding_box[1])), (int(bounding_box[2]), int(bounding_box[3])), (0, 255, 0), 2)
    alpha = 0.5
    color = [0.000, 0.447, 0.741]
    for c in range(3):
        img[:, :, c] = np.where(mask == 1,
                                    img[:, :, c] *
                                    (1 - alpha) + alpha * color[c] * 255,
                                    img[:, :, c])
    print(mask.max()) 
    return img


# Interface with input [image, bounding box] and output [image, segmentation]
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            img = gr.Image(label="Image", value="demo_img.jpg")
            with gr.Row():
                bbox_x = gr.Text(label="bbox_x1", value="265")
                bbox_y = gr.Text(label="bbox_y1", value="210")
                bbox_w = gr.Text(label="bbox_x2", value="350")
                bbox_h = gr.Text(label="bbox_y2", value="300")

        segm_img = gr.Image(type="numpy",label="Segmentation")
    button = gr.Button(label="Segment")
    button.click(segment, [img, bbox_x, bbox_y, bbox_w, bbox_h], segm_img)

demo.launch()