import cv2
from PIL import Image
from insightface.app import FaceAnalysis
import os


# 获取图片中人脸有哪些，返回人脸信息
# ctx_id:    如果有多块GPU，那么ctx_id=0表示使用第一块显卡,如果是负数则是使用CPU执行预测
# det_thresh: 配置的是人脸检测的阈值
# det_size:   检测模型图片大小
# det_thresh = 0.50
def get_faces(v_image_path, v_ctx_id, v_det_thresh, v_det_size, v_image_path_flag=None):
    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)
        # 在图片上标识出人脸区域
        if v_image_path_flag is not None and len(v_faces) > 0:
            rimg = app.draw_on(img, v_faces)
            cv2.imwrite(v_image_path_flag, rimg)
        return v_faces
    except Exception as err:
        return None


# 获取面积最大的人脸
def get_first_face(v_faces, output_path):
    # 计算每个面部的面积
    face_areas = [(face['bbox'][2] - face['bbox'][0]) * (face['bbox'][3] - face['bbox'][1]) for face in v_faces]
    # 根据面积降序排序面部
    sorted_faces = [face for _, face in sorted(zip(face_areas, v_faces), reverse=True)]
    # 输出排序后的面部数组
    bbox = sorted_faces[0]['bbox']  # 最大的人脸
    # 把一个图片的这个区域截取出来，另存在一个图片文件
    img = Image.open(image_path)
    # 将图像转换为RGB模式
    img_rgb = img.convert('RGB')
    # 截取矩形区域
    cropped_image = img_rgb.crop(bbox)
    # 保存截取的区域到新的图片文件
    cropped_image.save(output_path)
    # 关闭图像
    cropped_image.close()
    # 关闭图像
    img_rgb.close()


if __name__ == '__main__':
    # 声明 FaceAnalysis
    app = FaceAnalysis(allowed_modules=['detection'], providers=['CUDAExecutionProvider', 'CPUExecutionProvider'],
                       download=False)
    # 遍历目录下所有文件
    directory = r'C:\Users\Administrator\Desktop\Upload'
    # 对于多个人脸，打标记框
    flag_directory = r'C:\Users\Administrator\Desktop\Flag'
    # 第一个人脸
    first_directory = r'C:\Users\Administrator\Desktop\FirstFace'
    # 不存在就创建
    if not os.path.exists(flag_directory):
        os.makedirs(flag_directory)
    if not os.path.exists(first_directory):
        os.makedirs(first_directory)

    with os.scandir(directory) as entries:  # 使用os.scandir()获取目录条目
        for entry in entries:
            if entry.is_file():  # 检查是否是文件
                image_path = os.path.join(directory, entry.name)
                # 打标记框的图片
                v_image_path_flag = os.path.join(flag_directory, entry.name)
                # 第一个人脸图片
                v_first_face_path = os.path.join(first_directory, entry.name)
                # 获取人脸信息
                faces = get_faces(v_image_path=image_path, v_ctx_id=0, v_det_thresh=0.2, v_det_size=(640, 640),
                                  v_image_path_flag=v_image_path_flag)

                if faces is not None and len(faces) > 0:
                    # 获取第一个人脸图片,只截取第一个人脸的方法是不可以的
                    get_first_face(faces, output_path=v_first_face_path)
                else:
                    print("文件名称" + entry.name + ",人脸数量:" + str(len(faces)))