# !git clone https://github.com/polimi-ispl/icpr2020dfdc
# !pip install efficientnet-pytorch
# !pip install -U git+https://github.com/albu/albumentations > /dev/null

# %cd icpr2020dfdc/notebook

import torch
from torch.utils.model_zoo import load_url
from PIL import Image
import matplotlib.pyplot as plt

import sys
sys.path.append('./icpr2020dfdc/')

from blazeface import FaceExtractor, BlazeFace
from architectures import fornet,weights
from isplutils import utils

import gradio as gr



"""
Choose an architecture between
- EfficientNetB4
- EfficientNetB4ST
- EfficientNetAutoAttB4
- EfficientNetAutoAttB4ST
- Xception
"""
net_model = 'EfficientNetAutoAttB4'

"""
Choose a training dataset between
- DFDC
- FFPP
"""
train_db = 'DFDC'

device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')
face_policy = 'scale'
face_size = 224

model_url = weights.weight_url['{:s}_{:s}'.format(net_model,train_db)]
net = getattr(fornet,net_model)().eval().to(device)
net.load_state_dict(load_url(model_url,map_location=device,check_hash=True))

transf = utils.get_transformer(face_policy, face_size, net.get_normalizer(), train=False)

facedet = BlazeFace().to(device)
facedet.load_weights("./icpr2020dfdc/blazeface/blazeface.pth")
facedet.load_anchors("./icpr2020dfdc/blazeface/anchors.npy")
face_extractor = FaceExtractor(facedet=facedet)

title = "Face Manipulation Detection Through Ensemble of CNNs"

def inference(img):
  # im_original = Image.open(img)
  im_faces = face_extractor.process_image(img=img)
  im_face = im_faces['faces'][0]
  faces_t = torch.stack( [ transf(image=im)['image'] for im in [im_face] ] )
  
  with torch.no_grad():
      faces_pred = torch.sigmoid(net(faces_t.to(device))).cpu().numpy().flatten()
  # print(faces_pred[0])
  if faces_pred[0] >= 0.5:
    return "./Labels/Fake.jpg", f"{faces_pred[0]*100:.2f}%"
  else:
    return "./Labels/Real.jpg", f"{faces_pred[0]*100:.2f}%"

demo = gr.Interface(
    fn=inference, 
    inputs=[gr.inputs.Image(type="pil")],
    outputs=[gr.outputs.Image(type="pil"),gr.outputs.Label(type="text", label="Score")],
    title=title
    )
demo.launch()