# -*- coding: utf-8 -*-
# @Organization  : insightface.ai
# @Author        : Jia Guo
# @Time          : 2021-05-04
# @Function      : 


from __future__ import division

import onnxruntime

__all__ = ['FaceAnalysis']

from utils.common import Face
from models.arcface_onnx import ArcFaceONNX
from models.attribute import Attribute
from models.landmark import Landmark
from models.retinaface import RetinaFace
from huggingface_hub import hf_hub_download

REPO_ID = "leonelhs/insightface"

model_detector_path = hf_hub_download(repo_id=REPO_ID, filename="det_10g.onnx")
model_landmark_3d_68_path = hf_hub_download(repo_id=REPO_ID, filename="1k3d68.onnx")
model_landmark_2d_106_path = hf_hub_download(repo_id=REPO_ID, filename="2d106det.onnx")
model_genderage_path = hf_hub_download(repo_id=REPO_ID, filename="genderage.onnx")
model_recognition_path = hf_hub_download(repo_id=REPO_ID, filename="w600k_r50.onnx")

class FaceAnalysis:
    def __init__(self):
        onnxruntime.set_default_logger_severity(3)

        self.detector = RetinaFace(model_file=model_detector_path, input_size=(640, 640), det_thresh=0.5)
        self.landmark_3d_68 = Landmark(model_file=model_landmark_3d_68_path)
        self.landmark_2d_106 = Landmark(model_file=model_landmark_2d_106_path)
        self.genderage = Attribute(model_file=model_genderage_path)
        self.recognition = ArcFaceONNX(model_file=model_recognition_path)

    def get(self, img, max_num=0):
        bboxes, kpss = self.detector.detect(img,
                                            max_num=max_num,
                                            metric='default')
        if bboxes.shape[0] == 0:
            return []
        ret = []
        for i in range(bboxes.shape[0]):
            bbox = bboxes[i, 0:4]
            det_score = bboxes[i, 4]
            kps = None
            if kpss is not None:
                kps = kpss[i]
            face = Face(bbox=bbox, kps=kps, det_score=det_score)
            self.landmark_3d_68.get(img, face)
            self.landmark_2d_106.get(img, face)
            self.genderage.get(img, face)
            self.recognition.get(img, face)
            ret.append(face)
        return ret