|
|
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 |
|
|
import imageio |
|
|
|
|
|
|
|
|
classes=["Other","Anger","Contempt","Happiness","Surprise"] |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def load_model(idx): |
|
|
model = torch.jit.load('/content/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) |
|
|
|
|
|
|
|
|
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=['/content/ogtest0.png','/content/ogtest1.png','/content/ogtest2.png','/content/ogtest3.png','/content/ogtest4.png','/content/ogtest5.png'] |
|
|
ofpic=['/content/oftest0.jpg','/content/oftest1.jpg','/content/oftest2.jpg','/content/oftest3.jpg','/content/oftest4.jpg','/content/oftest5.jpg'] |
|
|
diffpic=['/content/difftest0.png','/content/difftest1.png','/content/difftest2.png','/content/difftest3.png','/content/difftest4.png','/content/difftest5.png'] |
|
|
exp=[['/content/ogtest0.png'],['/content/ogtest1.png'],['/content/ogtest2.png'],['/content/ogtest3.png'],['/content/ogtest4.png'],['/content/ogtest5.png']] |
|
|
vid=["/content/vidtest0.mp4","/content/vidtest1.mp4","/content/vidtest2.mp4","/content/vidtest3.mp4","/content/vidtest4.mp4","/content/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): |
|
|
|
|
|
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) |
|
|
|
|
|
for i, inputs in enumerate(loaded_test): |
|
|
with torch.no_grad(): |
|
|
output = model(inputs.to(torch.device('cpu'))) |
|
|
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=[['/content/ogtest0.png'],['/content/ogtest1.png'],['/content/ogtest2.png'],['/content/ogtest3.png'],['/content/ogtest4.png'],['/content/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="Prediction") |
|
|
with gr.Box(): |
|
|
with gr.Row(): |
|
|
video = gr.Video(label="Original Video") |
|
|
with gr.Row(): |
|
|
label_actual=gr.Label(label="Actual Emotion") |
|
|
|
|
|
|
|
|
|
|
|
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() |
