采用多进程异步加速进行人脸图片提取

# -*- codeing: utf-8 -*-
# 异步处理是同步处理性能的4倍左右
import sys
import os
import cv2
import dlib
from multiprocessing import Pool
import os
import time

# input_dir = './images'
input_dir = './lfw'
output_dir = './tmp'

size = 180

index = 1
if not os.path.exists(output_dir):
    os.makedirs(output_dir)

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


def reversed_cmp(x, y):
    x1 = x.split('_')[1].split(".")[0]
    y1 = y.split('_')[1].split(".")[0]
    if x1 > y1:
        return -1
    if x1 < y1:
        return 1
    return 0


def process_LFW():
    global index
    for (path, dirnames, filenames) in os.walk(input_dir):
        for filename in filenames:
            if filename.endswith('.jpg'):
                print('Being processed picture %s' % index)
                img_path = path + '/' + filename
                # 从文件读取图片
                img = cv2.imread(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):
                    x1 = d.top() if d.top() > 0 else 0
                    y1 = d.bottom() if d.bottom() > 0 else 0
                    x2 = d.left() if d.left() > 0 else 0
                    y2 = d.right() if d.right() > 0 else 0
                    # img[y:y+h,x:x+w]
                    face = img[x1:y1, x2:y2]
                    # 调整图片的尺寸
                    face = cv2.resize(face, (size, size))
                    cv2.imshow('image', face)
                    # 保存图片
                    cv2.imwrite(output_dir + '/' + str(index) + '.jpg', face)
                    index += 1

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


def process_images(input_dir, output_dir):
    global index
    for (path, dirnames, filenames) in os.walk(input_dir):
        bb = sorted(filenames, key=lambda x: int(x.split('_')[1].split('.')[0]))
        # list.sort(cmp=None, key=None, reverse=False)
        for filename in bb:
            if filename.endswith('.jpg'):
                print('Being processed picture %s' % index)
                img_path = path + '/' + filename
                # 从文件读取图片
                img = cv2.imread(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):
                    x1 = d.top() if d.top() > 0 else 0
                    y1 = d.bottom() if d.bottom() > 0 else 0
                    x2 = d.left() if d.left() > 0 else 0
                    y2 = d.right() if d.right() > 0 else 0
                    # img[y:y+h,x:x+w]
                    face = img[x1:y1, x2:y2]
                    # 调整图片的尺寸
                    face = cv2.resize(face, (size, size))
                    cv2.imshow('image', face)

                    if valid_image(face):
                        # 保存图片
                        cv2.imwrite(output_dir + str(index) + '.jpg', face)
                        index += 1

                key = cv2.waitKey(30) & 0xff
                if key == 27:
                    break


def ProcessFaceImage(image, filename, index):
    print("in process", os.getpid())
    gray_img = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    # 使用detector进行人脸检测 dets为返回的结果
    dets = detector(gray_img, 1)
    # 使用enumerate 函数遍历序列中的元素以及它们的下标
    # 下标i即为人脸序号
    # left：人脸左边距离图片左边界的距离 ；right：人脸右边距离图片左边界的距离
    # top：人脸上边距离图片上边界的距离 ；bottom：人脸下边距离图片上边界的距离
    for i, d in enumerate(dets):
        x1 = d.top() if d.top() > 0 else 0
        y1 = d.bottom() if d.bottom() > 0 else 0
        x2 = d.left() if d.left() > 0 else 0
        y2 = d.right() if d.right() > 0 else 0
        # img[y:y+h,x:x+w]
        face = image[x1:y1, x2:y2]
        # 调整图片的尺寸
        face = cv2.resize(face, (size, size))
        # cv2.imshow('image', face)

        if valid_image(face):
            # 保存图片
            name = filename.split("_")[0]
            cv2.imwrite(output_dir + name + str(index) + '.jpg', face)
    return dets, index


def CallBackProcess(arg):
    # 在main进程中执行的，回调函数
    print("-->exec done in main:{} and get call back result {} ".format(os.getpid(), arg))


def process_ones_images_mult_process(input_dir, output_dir):
    global index
    for (path, dirnames, filenames) in os.walk(input_dir):
        pool = Pool(processes=7)  # 允许进程池同时放入7个进程
        print("主进程：", os.getpid())
        results = []
        for i in range(len(filenames)):
            filename = filenames[i]
            if filename.endswith('.jpg'):
                # print('Being processed picture %s' % i)
                img_path = path + '/' + filename
                # 从文件读取图片
                img = cv2.imread(img_path)
                result = pool.apply_async(func=ProcessFaceImage, args=(img, filename, i), callback=CallBackProcess)  # callback = 回调
                results.append(result)
                # print("result is : ", result)
        pool.close()  # 关闭进程池，表示不能在往进程池中添加进程
        pool.join()  # 等待进程池中的所有进程执行完毕，必须在close()之后调用
        print("Sub-process(es) done.")

        # 这里回调的函数是主进程去回调的（生产中若所有进程完毕后将结果写入数据库，只需要写个回调就行了，不必每个进程中写入数据库）
        # pool.apply(func=Foo,args=(1,))         # 串行
        # pool.apply_async(func=Foo,args=(1,))   # 并行
        # for i in range(len(filenames)):
        #     result = results[i]
        #     print("result get i :", result.get(i))


def process_ones_images_mult_process_syn(input_dir, output_dir):
    global index
    for (path, dirnames, filenames) in os.walk(input_dir):
        pool = Pool(processes=7)  # 允许进程池同时放入7个进程
        print("主进程：", os.getpid())
        results = []
        for i in range(len(filenames)):
            filename = filenames[i]
            if filename.endswith('.jpg'):
                # print('Being processed picture %s' % i)
                img_path = path + '/' + filename
                # 从文件读取图片
                img = cv2.imread(img_path)
                result = pool.apply(func=ProcessFaceImage, args=(img, filename, i))  # callback = 回调
                results.append(result)
                print("result is : ", result)
        pool.close()  # 关闭进程池，表示不能在往进程池中添加进程
        pool.join()  # 等待进程池中的所有进程执行完毕，必须在close()之后调用
        print("Sub-process(es) done.")


def get_img_var_score(image):
    """
     imageVar就是模糊度，大约100为分界100以内为模糊，100以上为清晰。
    :param image:
    :return: image_var
    """
    image_var = cv2.Laplacian(image, cv2.CV_64F).var()
    # print("image_var:", image_var)
    return image_var


def valid_image(image, var=100):
    image_var = get_img_var_score(image)
    return image_var > var


if __name__ == '__main__':
    # process_LFW()
    start = time.time()

    name = "dapai"
    output_dir = './videos/'  # .format(name)
    # process_images(input_dir="videos/{}".format(name), output_dir=output_dir)
    # 23.1s
    process_ones_images_mult_process(input_dir="videos/{}".format(name), output_dir=output_dir)

    # 92.3s
    # process_ones_images_mult_process_syn(input_dir="videos/{}".format(name), output_dir=output_dir)
    print("total cost time ", (time.time() - start))
    sys.exit(0)