import sys
import os
import cv2
import dlib
import PIL.Image as Image


def cut_all_face(input_dir, output_dir):
    '''剪切所有的图片'''

    # 使用dlib自带的frontal_face_detector作为我们的特征提取器
    detector = dlib.get_frontal_face_detector()

    output_dir = output_dir + r'\all'
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    index = 1
    for path, dirnames, filenames in os.walk(input_dir):
        for filename in filenames:
            if filename.endswith('.jpg') or filename.endswith('.png'):
                print('正在处理图片 %s' % index)
            img_path = path + '/' + filename
            # 从文件读取图片
            img = cv2.imread(img_path)
            img1 = Image.open(img_path)
            # 转为灰度图片
            gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            # 使用detector进行人脸检测 dets为返回的结果
            dets = detector(gray_img, 1)

            # 使用enumerate 函数遍历序列中的元素以及它们的下标
            # 下标i即为人脸序号
            # left：人脸左边距离图片左边界的距离 ；right：人脸右边距离图片左边界的距离
            # top：人脸上边距离图片上边界的距离 ；bottom：人脸下边距离图片上边界的距离
            for i, d in enumerate(dets):
                # print(i, d)
                x1 = int(d.left()) if d.left() > 0 else 0
                y1 = int(d.top()) if d.top() > 0 else 0
                x2 = int(d.right()) if d.right() > 0 else 0
                y2 = int(d.bottom()) if d.bottom() > 0 else 0
                box = (x1 - 20, y1 - 20, x2 + 20, y2 + 20)
                # print(box)
                face = img1.crop(box)
                # 调整图片的尺寸
                face = face.resize((120, 120), Image.ANTIALIAS)
                # 保存图片
                face.save(output_dir + '\\' + str(index) + '.jpg', 'JPEG', quality=95)
                index += 1

            key = cv2.waitKey(30) & 0xff
            if key == 27:
                sys.exit(0)
    return output_dir


input_dir = r'C:\Users\Desktop\face'
output_dir = r'C:\Users\Desktop\test'
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

cut_all_face(input_dir=input_dir, output_dir=output_dir)