diff --git a/.gitignore b/.gitignore index 7bbc71c0..4b559a35 100644 --- a/.gitignore +++ b/.gitignore @@ -99,3 +99,9 @@ ENV/ # mypy .mypy_cache/ + +#pycharm config files +/.idea +.idea/ +#train images +images/ diff --git a/faceai/FaceRecognition/FaceRecognition.py b/faceai/FaceRecognition/FaceRecognition.py new file mode 100644 index 00000000..270df153 --- /dev/null +++ b/faceai/FaceRecognition/FaceRecognition.py @@ -0,0 +1,94 @@ +import dlib +import sys +sys.path.append("..") +from tool import compare, imtool, save_load + + +class FaceRecognition: + def __init__(self, shape_predict_path, recognition_model_path): + self.shape_detector = dlib.get_frontal_face_detector() + self.shape_predictor = dlib.shape_predictor(shape_predict_path) + self.face_recognition = dlib.face_recognition_model_v1(recognition_model_path) + self.train_data = [] + self.train_label = [] + + def train(self, train_data, is_debug=False): + if not isinstance(train_data, list): + raise TypeError("training_data required to be a list which contains labels and images") + + self.train_data = [] + self.train_label = [] + for data in train_data: + if is_debug: + print("Processing {}".format(data.get("label"))) + for img in data.get("data"): + encodings, bounds = self._get_face_encoding(img) + self.train_data = self.train_data + encodings + self.train_label = self.train_label + [data.get("label") for i in encodings] + + return self + + def _get_face_encoding(self, image): + bounds = [] + encodings = [] + detect_faces = self.shape_detector(image, 1) + for face in detect_faces: + shape = self.shape_predictor(image, face) + encodings.append(self.face_recognition.compute_face_descriptor(image, shape)) + bounds.append((face.left(), face.top(), face.right(), face.bottom())) + + return encodings, bounds + + def predict(self, image, threshold=0.6): + encodings, bounds = self._get_face_encoding(image) + predict_info = [] + for i, vector in enumerate(encodings): + sort_indexes, distances = compare.compare(self.train_data, vector) + distance = distances[sort_indexes[0]] + if distance < threshold: + predict_info.append((self.train_label[sort_indexes[0]], bounds[i])) + else: + predict_info.append((None, bounds[i])) + + return predict_info + + def save(self, file_name): + data = { + "train_data": self.train_data, + "train_label": self.train_label + } + save_load.save(data, file_name) + return self + + def load(self, file_name): + data = save_load.load(file_name) + self.train_data = data.get("train_data") + self.train_label = data.get("train_label") + return self + + +if __name__ == '__main__': + from tool import file + import argparse + + parse = argparse.ArgumentParser() + parse.add_argument("dlib_face_path", help="Please provide the dlib_face_recognition_model path, you can download it from http://dlib.net/files/dlib_face_recognition_resnet_model_v1.dat.bz2 and extract it.") + parse.add_argument("dlib_shape_path", help="Please provide the dlib_shape_detector_model path, you can download it from http://dlib.net/files/shape_predictor_5_face_landmarks.dat.bz2 and extract it.") + parse.add_argument("train_path", help="Please provide the training images folder path.") + parse.add_argument("test_path", help="Please provide the testing images folder path.") + args = parse.parse_args() + + recognition_path = args.dlib_face_path #"../dlib_model/dlib_face_recognition_resnet_model_v1.dat" + shape_path = args.dlib_shape_path #"../dlib_model/shape_predictor_5_face_landmarks.dat" + train_data = file.get_images_data(args.train_path) #("../../images/train") + test_data = file.get_images_data(args.test_path) #("../../images/test") + + face = FaceRecognition(shape_path, recognition_path) + face.train(train_data, is_debug=True) + + for data in test_data: + for image in data.get("data"): + for label, bounds in face.predict(image): + imtool.put_text(image, label if label is not None else "Unknown", bounds) + imtool.draw_rect(image, bounds) + imtool.show_image(image, width=400) diff --git a/faceai/FaceRecognition/__init__.py b/faceai/FaceRecognition/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/faceai/tool/__init__.py b/faceai/tool/__init__.py new file mode 100644 index 00000000..e19651ab --- /dev/null +++ b/faceai/tool/__init__.py @@ -0,0 +1,4 @@ +# +# coding=utf-8 +# description: tool package init file +# \ No newline at end of file diff --git a/faceai/tool/compare.py b/faceai/tool/compare.py new file mode 100644 index 00000000..59e04b74 --- /dev/null +++ b/faceai/tool/compare.py @@ -0,0 +1,26 @@ +# +# coding=utf-8 +# description: calculate and compare the distance of multiple 2-d vectors. +# + +import numpy as np + + +def distance(vector1, vector2): + if not isinstance(vector1, np.ndarray): + vector1 = np.array(vector1) + if not isinstance(vector2, np.ndarray): + vector2 = np.array(vector2) + + return np.linalg.norm(vector1 - vector2) + + +def compare(vectors, vector): + distances = [] + for v in vectors: + distances.append(distance(v, vector)) + return np.argsort(distances), distances + + +if __name__ == '__main__': + pass diff --git a/faceai/tool/file.py b/faceai/tool/file.py new file mode 100644 index 00000000..71356e2c --- /dev/null +++ b/faceai/tool/file.py @@ -0,0 +1,45 @@ +# +# coding=utf-8 +# description: file operation functions +# + +import os +import cv2 + + +IMAGE_SUFFIX = ["jpg", "jpeg", "png"] + + +def is_image(image_name): + return any([image_name.lower().endswith(suffix) for suffix in IMAGE_SUFFIX]) + + +def get_images_data(train_path, image_flag=cv2.IMREAD_COLOR): + """ + loading training images or test images. + :param train_path: + :param image_flag: + :return: + """ + if not os.path.isdir(train_path): + raise NotADirectoryError(train_path + " is not a dir path.") + + data = [] + + for path in os.listdir(train_path): + name = path + path = os.path.join(train_path, path) + images = [] + data.append({ + "label": name, + "data": images + }) + for image_name in os.listdir(path): + if not is_image(image_name): + continue + images.append(cv2.imread(os.path.join(path, image_name), image_flag)) + return data + + +if __name__ == '__main__': + pass diff --git a/faceai/tool/imtool.py b/faceai/tool/imtool.py new file mode 100644 index 00000000..17ce4067 --- /dev/null +++ b/faceai/tool/imtool.py @@ -0,0 +1,39 @@ +# +# coding=utf-8 +# description: +# +import cv2 + + +def show_image(img, width=None, height=None, win_name="Image", wait_time=0): + cv2.imshow(win_name, resize(img, width, height)) + cv2.waitKey(wait_time) & 0XFF + + +def resize(img, width, height, interpolation=cv2.INTER_AREA): + size = None + + if width is None and height is None: + size = (img.shape[1], img.shape[0]) + elif width is not None: + ratio = width / float(img.shape[1]) + height = int(img.shape[0] * ratio) + size = (width, height) + elif height is not None: + ratio = height / float(img.shape[0]) + width = int(img.shape[1] * ratio) + size = (width, height) + + return cv2.resize(img, size, interpolation=interpolation) + + +def draw_rect(img, rect, color=(0, 255, 0), line_width=1): + cv2.rectangle(img, rect[:2], rect[2:], color, line_width) + + +def put_text(img, text, rect, color=(0, 255, 0), thickness=0.8): + cv2.putText(img, text, rect[:2], cv2.FONT_HERSHEY_COMPLEX_SMALL, thickness, color) + + +if __name__ == '__main__': + pass diff --git a/faceai/tool/save_load.py b/faceai/tool/save_load.py new file mode 100644 index 00000000..7aed7104 --- /dev/null +++ b/faceai/tool/save_load.py @@ -0,0 +1,20 @@ +# +# coding=utf-8 +# description: +# +import pickle + + +def load(file_name): + with open(file_name, "rb") as fs: + data = pickle.load(fs) + return data + + +def save(data, file_name): + with open(file_name, "wb") as fs: + pickle.dump(data, fs) + + +if __name__ == '__main__': + pass