import torch
import gradio as gr
from torchvision import transforms as T
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
from PIL import Image
import numpy as np

classes=["Other","Anger","Contempt","Happiness","Surprise"]


# load a resnet18 model pretrained on ImageNet
# and turn off autograd on model's parameters
def load_model(idx):
  model = torch.jit.load('model_2_60acc.pt',map_location=torch.device('cpu')).eval()
  for param in model.parameters():
    param.requires_grad = False
  return model
model=load_model(0)

# preprocess data
pretrained_std = torch.Tensor([0.229, 0.224, 0.225])
pretrained_mean = torch.Tensor([0.485, 0.456, 0.406])

optical_flow_t = T.Compose([
                       T.Resize((224,224)),
                       T.ToTensor(),
                       T.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225]),
                       ])  




    
ogpic=['ogtest0.png','ogtest1.png','ogtest2.png','ogtest3.png','ogtest4.png','ogtest5.png']
ofpic=['oftest0.jpg','oftest1.jpg','oftest2.jpg','oftest3.jpg','oftest4.jpg','oftest5.jpg']
diffpic=['difftest0.png','difftest1.png','difftest2.png','difftest3.png','difftest4.png','difftest5.png']
exp=[['ogtest0.png'],['ogtest1.png'],['ogtest2.png'],['ogtest3.png'],['ogtest4.png'],['ogtest5.png']]
vid=["vidtest0.mp4","vidtest1.mp4","vidtest2.mp4","vidtest3.mp4","vidtest4.mp4","vidtest5.mp4"]
actual=["Contempt","Other","Happiness","Anger","Other","Contempt"]
def main():

    with gr.Blocks() as demo:
        aa=gr.Variable(value=0)
        def set_example_image(img):
          aa.value=img
          return gr.Image.update(value=exp[img][0])

        def predss(img):
          #print(Image.open(ofpic[a]).shape())
          processed_img = optical_flow_t(Image.open(ofpic[aa.value]))
          tb = torch.unsqueeze(processed_img, dim=0)
          loaded_test = DataLoader(tb, batch_size=1,shuffle=False)
          # get predictions
          for i, inputs in enumerate(loaded_test):
            with torch.no_grad():
                output = model(inputs.to(torch.device('cpu'))) # Feed Network
                probs = torch.nn.functional.softmax(output[0], dim=0)
            top5_prob, top5_idx = torch.topk(probs, 5)
            preds = {classes[idx]: prob.item() for idx, prob in zip(top5_idx, top5_prob)}
            return ogpic[aa.value],ofpic[aa.value],diffpic[aa.value],vid[aa.value], preds,actual[aa.value]
        gr.Markdown('''## Micro-expression recognition
                    
                    ''')
                                                 
        with gr.Box():
            input_image = gr.Image(type="pil", label="Input Image")
            example_images = gr.Dataset(components=[input_image],
                                        samples=[['ogtest0.png'],['ogtest1.png'],['ogtest2.png'],['ogtest3.png'],['ogtest4.png'],['ogtest5.png']]
                                        ,type="index")

            with gr.Row():
                btn = gr.Button("Process")
        gr.Markdown('''### Original Image''')
        with gr.Box():
            with gr.Row():
                img_before = gr.Image(label="Original Image")
                img_after1 = gr.Image(label="Different frame")
            with gr.Row():
                img_after = gr.Image(label="Optical flow") 
                label_predict = gr.Label(label="Model Prediction")
        with gr.Box():
            with gr.Row():
                video = gr.Video(label="Original Video")
            with gr.Row():
                label_actual=gr.Label(label="Actual Emotion")

        
        # events
        btn.click(fn=predss, 
                  inputs=[input_image], 
                  outputs=[img_before,img_after,img_after1,video,label_predict,label_actual])
        
        example_images.click(fn=set_example_image,
                                 inputs=example_images,
                                 outputs=example_images.components)

    demo.launch()

if __name__ == '__main__':
    main()