import gradio as gr
import cv2 as cv
import torch
from torchvision import transforms
from Model import DeePixBiS


labels = ['Live', 'Spoof']
thresh = 0.45
examples = [
    'examples/1_1_21_2_33_scene_fake.jpg', 'examples/frame150_real.jpg',
    'examples/1_2.avi_125_real.jpg', 'examples/1_3.avi_25_fake.jpg']
device = torch.device("cpu")
faceClassifier = cv.CascadeClassifier('Classifiers/haarface.xml')
tfms = transforms.Compose([
    transforms.ToPILImage(),
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
model = DeePixBiS(pretrained=False)
model.load_state_dict(torch.load('./DeePixBiS.pth'))
model.eval()


def find_largest_face(faces):
    largest_face = None
    largest_area = 0

    for (x, y, w, h) in faces:
        area = w * h
        if area > largest_area:
            largest_area = area
            largest_face = (x, y, w, h)
    return largest_face


def inference(img):
    grey = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    faces = faceClassifier.detectMultiScale(
        grey, scaleFactor=1.1, minNeighbors=4)
    face = find_largest_face(faces)
    confidences = {}
    if face is not None:
        x, y, w, h = face
        faceRegion = img[y:y + h, x:x + w]
        faceRegion = cv.cvtColor(faceRegion, cv.COLOR_BGR2RGB)
        faceRegion = tfms(faceRegion)
        faceRegion = faceRegion.unsqueeze(0)
        mask, binary = model.forward(faceRegion)
        res = torch.mean(mask).item()
        if res < thresh:
            cls = 'Spoof'
            color = (0, 0, 255)
            res = 1 - res
        else:
            cls = 'Real'
            color = (0, 255, 0)
        label = f'{cls} {res:.2f}'
        cv.rectangle(img, (x, y), (x + w, y + h), color, 2)
        cv.putText(img, label, (x, y + h + 30),
                   cv.FONT_HERSHEY_COMPLEX, 1, color)
        confidences = {label: res}
    return img, confidences


if __name__ == '__main__':
    demo = gr.Interface(
        fn=inference,
        inputs=[gr.Image(source='webcam', shape=None, type='numpy')],
        outputs=["image", gr.Label(num_top_classes=2)],
        examples=examples).queue(concurrency_count=2)
    demo.launch(share=False)