import torch, torchvision
from monai.networks.nets import UNet
from monai.networks.layers import Norm
from monai.inferers import sliding_window_inference
import PIL
from torchvision.utils import save_image
import numpy as np

model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
    norm=Norm.BATCH,
)

model.load_state_dict(torch.load("weights/model.pt", map_location=torch.device('cpu')))

import gradio as gr

def load_image0():
    return load_image(0)

def load_image1():
    return load_image(1)

def load_image2():
    return load_image(2)

def load_image3():
    return load_image(3)

def load_image4():
    return load_image(4)

def load_image5():
    return load_image(5)

def load_image6():
    return load_image(6)

def load_image7():
    return load_image(7)

def load_image8():
    return load_image(8)

def load_image(index):
    return [index, f"thumbnails/val_image{index}.png", f"thumbnails_label/val_label{index}.png"]

def predict(index):
    val_data = torch.load(f"samples/val_data{index}.pt")
    
    model.eval()
    with torch.no_grad():
        roi_size = (160, 160, 160)
        sw_batch_size = 4
        val_outputs = sliding_window_inference(val_data, roi_size, sw_batch_size, model)

    meta_tsr = torch.argmax(val_outputs, dim=1)[0, :, :, 80]
    pil_image = torchvision.transforms.functional.to_pil_image(meta_tsr.to(torch.float32))

    return pil_image


with gr.Blocks(title="Spleen 3D segmentation with MONAI - ClassCat",
                    css=".gradio-container {background:azure;}"
               ) as demo:
    sample_index = gr.State([])

    gr.HTML("""<div style="font-family:'Times New Roman', 'Serif'; font-size:16pt; font-weight:bold; text-align:center; color:royalblue;">Spleen 3D segmentation with MONAI</div>""")

    gr.HTML("""<h4 style="color:navy;">1. Select an example, which includes input images and label images, by clicking "Example x" button.</h4>""")

    with gr.Row():
        input_image = gr.Image(label="a piece of input image data", type="filepath")
        label_image = gr.Image(label="label image", type="filepath")
        output_image = gr.Image(label="predicted image", type="pil")


    with gr.Row():
        with gr.Column():
            ex_btn0 = gr.Button("Example 1")
            ex_btn0.style(full_width=False, css="width:20px;")
            ex_image0 = gr.Image(value='thumbnails/val_image0.png', interactive=False, label='ex 1')
            ex_image0.style(width=128, height=128)

        with gr.Column():
            ex_btn1 = gr.Button("Example 2")
            ex_btn1.style(full_width=False, css="width:20px;")
            ex_image1 = gr.Image(value='thumbnails/val_image1.png', interactive=False, label='ex 2')
            ex_image1.style(width=128, height=128)

        with gr.Column():
            ex_btn2 = gr.Button("Example 3")
            ex_btn2.style(full_width=False, css="width:20px;")
            ex_image2 = gr.Image(value='thumbnails/val_image2.png', interactive=False, label='ex 3')
            ex_image2.style(width=128, height=128)

        with gr.Column():
            ex_btn3 = gr.Button("Example 4")
            ex_btn3.style(full_width=False, css="width:20px;")
            ex_image3 = gr.Image(value='thumbnails/val_image3.png', interactive=False, label='ex 4')
            ex_image3.style(width=128, height=128)

        with gr.Column():
            ex_btn4 = gr.Button("Example 5")
            ex_btn4.style(full_width=False, css="width:20px;")
            ex_image4 = gr.Image(value='thumbnails/val_image4.png', interactive=False, label='ex 5')
            ex_image4.style(width=128, height=128)

        with gr.Column():
            ex_btn5 = gr.Button("Example 6")
            ex_btn5.style(full_width=False, css="width:20px;")
            ex_image5 = gr.Image(value='thumbnails/val_image5.png', interactive=False, label='ex 6')
            ex_image5.style(width=128, height=128)

    ex_btn0.click(fn=load_image0, outputs=[sample_index, input_image, label_image])
    ex_btn1.click(fn=load_image1, outputs=[sample_index, input_image, label_image])
    ex_btn2.click(fn=load_image2, outputs=[sample_index, input_image, label_image])
    ex_btn3.click(fn=load_image3, outputs=[sample_index, input_image, label_image])
    ex_btn4.click(fn=load_image4, outputs=[sample_index, input_image, label_image])
    ex_btn5.click(fn=load_image5, outputs=[sample_index, input_image, label_image])


    gr.HTML("""<br/>""")
    gr.HTML("""<h4 style="color:navy;">2. Then, click "Infer" button to predict a segmentation image. It will take about 15 seconds (on cpu)</h4>""")

    send_btn = gr.Button("Infer")
    send_btn.click(fn=predict, inputs=[sample_index], outputs=[output_image])


#demo.queue()
demo.launch(debug=True)


### EOF ###