# Face Img Filter v0.1
# 创建人：曾逸夫
# 创建时间：2023-10-02

"""
分类模型字典
性别分类：
{0: 'man', 1: 'woman'}
年龄分类：
{0: 'Children', 1: 'Elderly', 2: 'Infants', 3: 'Middle', 4: 'Teenagers', 5: 'Youth'}
"""

import argparse
import os
import shutil
import time

from ultralytics import YOLO

from util.time_format import time_format


def parse_args(known=False):
    parser = argparse.ArgumentParser(description="Face Img Filter v0.1")
    parser.add_argument("--genderCls", "-gc", action="store_true", help="gender classify")
    parser.add_argument("--ageCls", "-ac", action="store_true", help="age classify")
    parser.add_argument("--genderName", "-gn", default="man", type=str, help="gender classify name")
    parser.add_argument("--ageName", "-an", default="Youth", type=str, help="age classify name")
    parser.add_argument("--imgRoot", "-ir", default="./sample", type=str, help="img root path")
    parser.add_argument("--imgTarget", "-it", default="./result", type=str, help="img target path")
    parser.add_argument("--conf", "-cf", default=0.5, type=float, help="Conf")
    parser.add_argument("--iou", "-iu", default=0.45, type=float, help="IoU")
    parser.add_argument("--inferSize", "-is", default=640, type=int, help="inference size")
    parser.add_argument("--faceModel", "-fm", default="./yolov8_widerface-s.pt", type=str, help="YOLOv8 face detect Model")
    # ---------- 分类模型 ----------
    parser.add_argument("--genderModel", "-gm", default="./yolov8s-cls_gender.pt", type=str, help="YOLOv8 gender classify Model")
    parser.add_argument("--ageModel", "-am", default="./yolov8s-cls_age.pt", type=str, help="YOLOv8 age classify Model")

    args = parser.parse_known_args()[0] if known else parser.parse_args()
    return args


# 模型加载
def model_loading(img_path, conf, iou, infer_size, face_model):
    model = YOLO(face_model)

    results = model(source=img_path, imgsz=infer_size, conf=conf, iou=iou)
    results = list(results)[0]
    return results


# 分类模型加载
def model_cls_loading(cls_model, cls_name, img_root_path, img_target_path, img_name):
    yolo_cls_model = YOLO(cls_model)
    yolo_cls_r = yolo_cls_model(f"{img_root_path}/{img_name}")
    for result in yolo_cls_r:
        cls_r_name = result.names[result.probs.top1]
    if (cls_r_name == cls_name):
        shutil.copy(f"{img_root_path}/{img_name}", img_target_path)


# 人脸过滤
def face_img_filter(gender_cls, age_cls, gender_name, age_name, img_root_path, img_target_path, conf, iou, infer_size, face_model, gender_model, age_model):
    is_dir(img_target_path)

    face_total_num = 0  # 人脸图片总数
    img_names = os.listdir(img_root_path)
    print(f"一共有 {len(img_names)} 张图片")

    for img_name in img_names:
        try:
            result = model_loading(f"{img_root_path}/{img_name}", conf, iou, infer_size, face_model)
        except Exception as e:
            print(e)
            continue

        face_num = len(result.boxes.xyxy.cpu().numpy().tolist())

        if (face_num == 0):
            # 未检测到人脸
            print(f"{img_name}: 没有检测到人脸")
        else:
            # 检测到人脸
            print(f"{img_name}: 检测到人脸")
            face_total_num += 1  # 人脸图片计数

            if (face_num == 1):  # 对于单一人脸图片启动分类模式
                if (gender_cls):  # 性别分类
                    model_cls_loading(gender_model, gender_name, img_root_path, img_target_path, img_name)
                elif (age_cls):  # 年龄分类
                    model_cls_loading(age_model, age_name, img_root_path, img_target_path, img_name)
                else:
                    shutil.copy(f"{img_root_path}/{img_name}", img_target_path)

            elif (face_num >= 2 and not gender_cls and not age_cls):
                # 复制到指定目录
                shutil.copy(f"{img_root_path}/{img_name}", img_target_path)

    print(f"一共检测到 {face_total_num} 张人脸图片（ {len(img_names)} ）")


# 目录检查
def is_dir(dir_name):
    if (not os.path.exists(dir_name)):
        os.makedirs(dir_name)
        print(f"{dir_name}目录不存在，创建完成")
    else:
        print(f"{dir_name}目录已存在")


def main(args):
    gender_cls = args.genderCls  # 性别分类
    age_cls = args.ageCls  # 年龄分类
    gender_name = args.genderName  # 性别分类名称
    age_name = args.ageName  # 年龄分类名称
    img_root_path = args.imgRoot
    img_target_path = args.imgTarget
    conf = args.conf
    iou = args.iou
    infer_size = args.inferSize
    face_model = args.faceModel
    gender_model = args.genderModel
    age_model = args.ageModel


    s_time = time.time()  # 起始时间
    print("-------------------- 程序启动 --------------------")

    face_img_filter(gender_cls, age_cls, gender_name, age_name, img_root_path, img_target_path, conf, iou, infer_size, face_model, gender_model, age_model)

    print("-------------------- 程序结束 --------------------")
    e_time = time.time()  # 终止时间
    total_time = e_time - s_time  # 程序用时

    # 格式化时间格式，便于观察
    outTimeMsg = f"用时：{time_format(total_time)}"
    print(outTimeMsg)  # 打印用时


if __name__ == "__main__":
    args = parse_args()
    main(args)