import os.path

import cv2
from insightface.app import FaceAnalysis
import shutil


# 准备好模型
def prepare_model():
    # 默认的模型目录是否存在
    directory = 'C:/Users/Administrator/.insightface/models/buffalo_l'
    if not os.path.exists(directory):  # 不存在就创建
        os.makedirs(directory)
    # 模型文件是不是存在
    files = ['1k3d68.onnx', '2d106det.onnx', 'det_10g.onnx', 'genderage.onnx', 'w600k_r50.onnx']
    # 云存储的前缀
    prefix = r'https://dsideal.obs.cn-north-1.myhuaweicloud.com/HuangHai/insightface/'

    # windows版本的wget
    # https://eternallybored.org/misc/wget/
    source_file = './Tools/wget.exe'
    target_file = "c:/Windows/System32/wget.exe"
    if not os.path.exists(target_file):
        if os.path.exists(source_file):
            shutil.copy(source_file, target_file)
        elif os.path.exists('.' + source_file):
            shutil.copy('.' + source_file, target_file)
    # 下载缺失的模型文件
    for file in files:
        if not os.path.exists(directory + '/' + file):
            os.system('wget ' + prefix + file + ' -P ' + directory)


# 获取图片中人脸有哪些，返回人脸个数
# ctx_id:    如果有多块GPU，那么ctx_id=0表示使用第一块显卡,如果是负数则是使用CPU执行预测
# det_size:   检测模型图片大小
# det_thresh = 0.50 配置的是人脸检测的阈值
def get_faces_count(v_image_path, v_ctx_id=0, v_det_thresh=0.5, v_det_size=(640, 640)):
    try:
        app.prepare(ctx_id=v_ctx_id, det_thresh=v_det_thresh, det_size=v_det_size)
        img = cv2.imread(v_image_path)  # 检测图片中人脸的个数
        v_faces = app.get(img)
        return len(v_faces)
    except Exception as err:
        print(err)
        return 0


if __name__ == '__main__':
    # 准备好模型
    prepare_model()
    # 声明 FaceAnalysis
    app = FaceAnalysis(allowed_modules=['detection'], providers=['CUDAExecutionProvider', 'CPUExecutionProvider'],
                       download=False)
    # 0
    image_path = r'C:\Users\Administrator\Desktop\Upload\Cat.png'
    # 1
    # image_path = r'C:\Users\Administrator\Desktop\Upload\d825cc1e-6518-afd5-1b02-ed88a7a72806.jpg'
    # 2
    # image_path=r'C:\Users\Administrator\Desktop\Upload\ff781209-7368-bb99-4b0c-02c1406519d5.jpg'
    # 获取人脸个数信息
    face_count = get_faces_count(v_image_path=image_path)
    print(face_count)