import os
import sys
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(BASE_DIR)
import torch
from src.dataset.face_align.yoloface import YoloFace

class AlignImage(object):
    def __init__(self, device='cuda', det_path='checkpoints/yoloface_v5m.pt'):
        self.facedet = YoloFace(pt_path=det_path, confThreshold=0.5, nmsThreshold=0.45, device=device)

    @torch.no_grad()
    def __call__(self, im, maxface=False):
        bboxes, kpss, scores = self.facedet.detect(im)
        face_num = bboxes.shape[0]

        five_pts_list = []
        scores_list = []
        bboxes_list = []
        for i in range(face_num):
            five_pts_list.append(kpss[i].reshape(5,2))
            scores_list.append(scores[i])
            bboxes_list.append(bboxes[i])

        if maxface and face_num>1:
            max_idx = 0
            max_area = (bboxes[0, 2])*(bboxes[0, 3])
            for i in range(1, face_num):
                area = (bboxes[i,2])*(bboxes[i,3])
                if area>max_area:
                    max_idx = i
            five_pts_list = [five_pts_list[max_idx]]
            scores_list = [scores_list[max_idx]]
            bboxes_list = [bboxes_list[max_idx]]

        return five_pts_list, scores_list, bboxes_list