import numpy as np
import torch
import kornia as K
from kornia.core import Tensor
from kornia.contrib import FaceDetector, FaceDetectorResult, FaceKeypoint

print('Loading Face Detector...')
face_detection = FaceDetector()
print('DONE')

def detect_face(input):

    # Preprocessing
    img = K.image_to_tensor(np.array(input), keepdim=False)
    img = K.color.bgr_to_rgb(img.float())
    
    with torch.no_grad():
        dets = face_detection(img)
        
    return [FaceDetectorResult(o) for o in dets[0]]

def process_face(dets):
    vis_threshold = 0.8
    faces = []
    hairs = []
    
    for b in dets:
        if b.score  < vis_threshold:
            continue
    
        reye_kpt=b.get_keypoint(FaceKeypoint.EYE_RIGHT).int().tolist()
        leye_kpt=b.get_keypoint(FaceKeypoint.EYE_LEFT).int().tolist()
        rmou_kpt=b.get_keypoint(FaceKeypoint.MOUTH_RIGHT).int().tolist()
        lmou_kpt=b.get_keypoint(FaceKeypoint.MOUTH_LEFT).int().tolist()
        nose_kpt=b.get_keypoint(FaceKeypoint.NOSE).int().tolist()
    
        faces.append([nose_kpt,
                     rmou_kpt,
                     lmou_kpt,
                     reye_kpt,
                     leye_kpt
                    ])
    
        # point above
        top=((b.top_right + b.top_left)/2).int().tolist()
        bot=((b.bottom_right + b.bottom_left)/2).int().tolist()
        face_h = np.abs(top[1]-bot[1])
        top_margin=[top[0], top[1]-face_h*0.1]
    
        hairs.append([
                          top_margin
                    ])

    return faces, hairs