68个dlib的landmarks

目的：
获取指定文件夹里的所有图片的68个dlib的landmarks。
也可以间隔取图。
结果按照图片名字存放在字典中landmark_coords。
最后字典存放在josn或者txt里。

import dlib
import cv2
from tqdm import tqdm
import os

# Load the detector#加载探测器
detector = dlib.get_frontal_face_detector()

# Load the predictor#加载预测器
predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")

def generate_landmarks(img_dir, skip=1, valid=False):
    filenames = []
#获取路径，图片是jpg或png
    for root, dirs, files in os.walk(img_dir):
            for file in files:
                if file.endswith(".jpg") or file.endswith(".png"):
                    filenames.append(os.path.join(root, file))
    landmark_coords = {}#搞个存放的字典
    count = 0

    for i in tqdm(range(len(filenames))):
        if i%skip == 0:#按需间隔读取，只需要部分图片
            image_file = filenames[i] #'data/0000_00001.jpg'
            image = cv2.imread(image_file)
            image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            image_cropped = image_rgb

            # create a copy of the cropped image to be used later
            image_template = image_cropped.copy()
            # convert image to Grayscale
            image_gray = cv2.cvtColor(image_cropped, cv2.COLOR_BGR2GRAY)
            # Detect faces using dlib on the "grayscale image"

            try:
                faces = detector(image_gray)
                for face in faces:
                    landmarks = predictor(image=image_gray, box=face)
                    coord_list = []
                    for n in range(0, 68):
                        x = landmarks.part(n).x
                        y = landmarks.part(n).y
                        coord_list.append([x, y])#存放68个坐标
                    landmark_coords[image_file] = coord_list
            except Exception as e:
                print(e)

    import json
    if valid:
        save_name = 'landmarks_valid.json'
    else:
        save_name = 'landmarks2.json'
    with open(save_name, 'w') as fp:
        json.dump(landmark_coords, fp)


import time
start = time.process_time()
#generate_landmarks(img_dir='dataset/', skip=200)
generate_landmarks(img_dir='data/', skip=1)
#generate_landmarks(valid=True)

end =time.process_time()-start
print('time:%fs'%end)