78 lines
2.2 KiB
Python
78 lines
2.2 KiB
Python
import os
|
|
import dlib
|
|
import cv2
|
|
import sys
|
|
import numpy as np
|
|
|
|
try:
|
|
import cPickle # Python2.
|
|
except ImportError:
|
|
import _pickle as cPickle # Python3.
|
|
|
|
def enroll_face_dataset():
|
|
pwd = sys.path[0]
|
|
PREDICTOR_PATH = pwd + '/Facial_models/shape_predictor_68_face_landmarks.dat'
|
|
FACE_RECOGNITION_MODEL_PATH = pwd + '/Facial_models/dlib_face_recognition_resnet_model_v1.dat'
|
|
|
|
faceDetector = dlib.get_frontal_face_detector()
|
|
shapePredictor = dlib.shape_predictor(PREDICTOR_PATH)
|
|
faceRecognizer = dlib.face_recognition_model_v1(FACE_RECOGNITION_MODEL_PATH)
|
|
|
|
|
|
faceDatasetFolder = pwd + '/Facial_images/face_rec/train/'
|
|
|
|
subfolders = []
|
|
for x in os.listdir(faceDatasetFolder):
|
|
xpath = os.path.join(faceDatasetFolder, x)
|
|
if os.path.isdir(xpath):
|
|
subfolders.append(xpath)
|
|
|
|
|
|
nameLabelMap = {}
|
|
labels = []
|
|
imagePaths = []
|
|
for i, subfolder in enumerate(subfolders):
|
|
for x in os.listdir(subfolder):
|
|
xpath = os.path.join(subfolder, x)
|
|
if x.endswith('jpg'):
|
|
imagePaths.append(xpath)
|
|
labels.append(i)
|
|
nameLabelMap[xpath] = subfolder.split('/')[-1]
|
|
|
|
index = {}
|
|
i = 0
|
|
faceDescriptors = None
|
|
for imagePath in imagePaths:
|
|
print("processing: {}".format(imagePath))
|
|
img = cv2.imread(imagePath)
|
|
|
|
faces = faceDetector(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
|
|
|
|
print("{} Face(s) found".format(len(faces)))
|
|
|
|
for k, face in enumerate(faces):
|
|
|
|
shape = shapePredictor(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), face)
|
|
|
|
landmarks = [(p.x, p.y) for p in shape.parts()]
|
|
|
|
faceDescriptor = faceRecognizer.compute_face_descriptor(img, shape)
|
|
|
|
|
|
faceDescriptorList = [x for x in faceDescriptor]
|
|
faceDescriptorNdarray = np.asarray(faceDescriptorList, dtype=np.float64)
|
|
faceDescriptorNdarray = faceDescriptorNdarray[np.newaxis, :]
|
|
|
|
|
|
if faceDescriptors is None:
|
|
faceDescriptors = faceDescriptorNdarray
|
|
else:
|
|
faceDescriptors = np.concatenate((faceDescriptors, faceDescriptorNdarray), axis=0)
|
|
|
|
index[i] = nameLabelMap[imagePath]
|
|
i += 1
|
|
|
|
# Write descriors and index to disk
|
|
np.save(pwd+'/Facial_models/descriptors.npy', faceDescriptors)
|
|
with open(pwd+'/Facial_models/index.pkl', 'wb') as f:
|
|
cPickle.dump(index, f)
|