import gradio as gr
import numpy as np
import cv2
from mmseg.apis import init_model, inference_model
import torch


def process_single_img(img_bgr, model_name):
    print(type(img_bgr))
    palette = [
        ['background', [0, 0, 0]],
        ['red', [255, 0, 0]]
    ]

    palette_dict = {}
    for idx, each in enumerate(palette):
        palette_dict[idx] = each[1]

    if model_name == 'Mask2Former':
        config_file = 'CVRP_configs/CVRP_mask2former.py'
        checkpoint_file = 'checkpoint/Mask2Former.pth'
    elif model_name == 'KNet':
        config_file = 'CVRP_configs/CVRP_knet.py'
        checkpoint_file = 'checkpoint/KNet.pth'
    elif model_name == 'DeepLabV3+':
        config_file = 'CVRP_configs/CVRP_deeplabv3plus.py'
        checkpoint_file = 'checkpoint/DeepLabV3plus.pth'
    elif model_name == 'Segformer':
        config_file = 'CVRP_configs/CVRP_segformer.py'
        checkpoint_file = 'checkpoint/Segformer.pth'
    else:
        return None, None

    device = 'cuda:0'

    model = init_model(config_file, checkpoint_file, device=device)

    result = inference_model(model, img_bgr)
    pred_mask = result.pred_sem_seg.data[0].cpu().numpy()

    pred_mask_bgr = np.zeros((pred_mask.shape[0], pred_mask.shape[1], 3))
    for idx in palette_dict.keys():
        pred_mask_bgr[np.where(pred_mask == idx)] = palette_dict[idx]
    pred_mask_bgr = pred_mask_bgr.astype('uint8')

    pred_viz = cv2.addWeighted(img_bgr, 1, pred_mask_bgr, 1, 0)

    torch.cuda.empty_cache()

    return pred_viz, pred_mask_bgr
def run_segmentation(image_input, model_select):
    if model_select not in ["Mask2Former", "KNet", "DeepLabV3+", "Segformer"]:
        return None, None, [("No implementa", "Error"), ("", "")]
    else:
        color_img, binary_img = process_single_img(image_input, model_select)
        return color_img, binary_img, [("", ""), ("Segmentation Finished", "normal")]

title = """<p><h1 align="center">CVRP</h1></p>"""
# 设置SAM参数

with gr.Blocks() as iface:
    gr.Markdown(title)
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(interactive=True, visible=True, label="Input Image", height=360)
            with gr.Row():
                model_select = gr.Dropdown(choices=["Mask2Former", "KNet", "DeepLabV3+", "Segformer"], value="Mask2Former", label="Select model", visible=True)
                run_button = gr.Button(value="Run", interactive=True, visible=True)
            with gr.Row():
                gr.Examples(
                    examples=[['assets/T42_1220.jpg', 'Mask2Former'], ['assets/T99_799.jpg', 'Mask2Former'], ['assets/T92_10336.jpg', 'Mask2Former']],
                    inputs=[image_input, model_select])
                    
        with gr.Column():
            color_output = gr.Image(interactive=False, visible=True, label="Color Image", height=360)
            binary_output = gr.Image(interactive=False, visible=True, label="Binary Image", height=360)
            run_status = gr.HighlightedText(
                value=[("Text", "Error"), ("to be", "Label 2"), ("highlighted", "Label 3")], visible=True)        

    run_button.click(
        fn=run_segmentation,
        inputs=[image_input, model_select],
        outputs=[color_output, binary_output, run_status]
    )


iface.launch(debug=True, server_port=6006, server_name="127.0.0.1")