face_recognition.py 2.1 KB
import dlib
import cv2
import numpy as np
import os
import path
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import matplotlib.patheffects as path_effects


detector = dlib.get_frontal_face_detector()
sp = dlib.shape_predictor(
    'face_emotion_recognition/models/shape_predictor_68_face_landmarks.dat')
facerec = dlib.face_recognition_model_v1(
    'face_emotion_recognition/models/dlib_face_recognition_resnet_model_v1.dat')


def find_faces(img):
    dets = detector(img, 1)

    if len(dets) == 0:
        return np.empty(0), np.empty(0), np.empty(0)

    rects, shapes = [], []
    shapes_np = np.zeros((len(dets), 68, 2), dtype=np.int)
    for k, d in enumerate(dets):
        rect = ((d.left(), d.top()), (d.right(), d.bottom()))
        rects.append(rect)

        shape = sp(img, d)

        # convert dlib shape to numpy array
        for i in range(0, 68):
            shapes_np[k][i] = (shape.part(i).x, shape.part(i).y)

        shapes.append(shape)

        print("Part 0: {}, Part 1: {} ...".format(shape.part(0),
                                                  shape.part(1)))

    return rects, shapes, shapes_np


def encode_faces(img, shapes):
    face_descriptors = []
    for shape in shapes:
        face_descriptor = facerec.compute_face_descriptor(img, shape)
        face_descriptors.append(np.array(face_descriptor))

    return np.array(face_descriptors)


def face_to_npy():
    img_dict = 'static/img/'
    print(img_dict)
    img_paths = {}
    descs = {}
    for img_name in os.listdir(img_dict):
        if(img_name.rsplit('.')[1] == 'png' or img_name.rsplit('.')[1] == 'jpg'):
            img_paths[img_name.rsplit('.')[0]] = img_dict + img_name
    print(img_paths, descs)

    for name, img_path in img_paths.items():
        img_bgr = cv2.imread(img_path)
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

        _, img_shapes, _ = find_faces(img_rgb)
        print(_, img_shapes, _)
        if img_shapes:
            descs[name] = encode_faces(img_rgb, img_shapes)[0]
        else:
            os.remove(img_path)

    print(descs)
    np.save('static/img/descs.npy', descs)