import cv2 as cv
import torch
import torch.nn as nn
from torchvision import transforms
import numpy as np
from Model import DeePixBiS
from Loss import PixWiseBCELoss
from Metrics import predict, test_accuracy, test_loss
from rknnlite.api import RKNNLite
rknn_lite = RKNNLite()
ret = rknn_lite.load_rknn('/home/orangepi/Desktop/test-face-id/cpp/test_py/DeePixBiS.rknn')
ret = rknn_lite.init_runtime(core_mask=RKNNLite.NPU_CORE_0_1_2)
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))
])

faceClassifier = cv.CascadeClassifier('/home/orangepi/Desktop/test-face-id/cpp/test_py/Classifiers/haarface.xml')

camera = cv.VideoCapture(0)

while cv.waitKey(1) & 0xFF != ord('q'):
    _, img = camera.read()
    grey = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    faces = faceClassifier.detectMultiScale(
        grey, scaleFactor=1.1, minNeighbors=4)

    for x, y, w, h in faces:
        faceRegion = img[y:y + h, x:x + w]
        faceRegion = cv.cvtColor(faceRegion, cv.COLOR_BGR2RGB)
        faceRegion = tfms(faceRegion)
        faceRegion = faceRegion.unsqueeze(0).detach().cpu().numpy()

        mask, binary = rknn_lite.inference([faceRegion.astype(np.float16)])
        res = np.mean(mask.flatten())
        print(res)
        cv.rectangle(img, (x, y), (x + w, y + h), (0, 0, 255), 2)

        if res < 0.5:
            cv.putText(img, 'Fake', (x, y + h + 30),
                       cv.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255))
        else:
            cv.putText(img, 'Real', (x, y + h + 30),
                       cv.FONT_HERSHEY_COMPLEX, 1, (0, 0, 255))

    cv.imshow('Deep Pixel-wise Binary Supervision Anti-Spoofing', img)