diff --git a/examples/face_recognition_knn.py b/examples/face_recognition_knn.py index 307233b500b0394094d057fcbfd87f6344367b06..d4d01f8ff8e8cab3755a7e44c63edddcac1bb206 100644 --- a/examples/face_recognition_knn.py +++ b/examples/face_recognition_knn.py @@ -1,42 +1,49 @@ """ -This is an example of using the k-nearest-neighbors(knn) algorithm for face recognition. +This is an example of using the k-nearest-neighbors (KNN) algorithm for face recognition. When should I use this example? -This example is useful when you whish to recognize a large set of known people, -and make a prediction for an unkown person in a feasible computation time. +This example is useful when you wish to recognize a large set of known people, +and make a prediction for an unknown person in a feasible computation time. Algorithm Description: -The knn classifier is first trained on a set of labeled(known) faces, and can then predict the person -in an unkown image by finding the k most similar faces(images with closet face-features under eucledian distance) in its training set, -and performing a majority vote(possibly weighted) on their label. -For example, if k=3, and the three closest face images to the given image in the training set are one image of Biden and two images of Obama, -The result would be 'Obama'. -*This implemententation uses a weighted vote, such that the votes of closer-neighbors are weighted more heavily. +The knn classifier is first trained on a set of labeled (known) faces and can then predict the person +in an unknown image by finding the k most similar faces (images with closet face-features under eucledian distance) +in its training set, and performing a majority vote (possibly weighted) on their label. + +For example, if k=3, and the three closest face images to the given image in the training set are one image of Biden +and two images of Obama, The result would be 'Obama'. + +* This implementation uses a weighted vote, such that the votes of closer-neighbors are weighted more heavily. Usage: --First, prepare a set of images of the known people you want to recognize. - Organize the images in a single directory with a sub-directory for each known person. --Then, call the 'train' function with the appropriate parameters. - make sure to pass in the 'model_save_path' if you want to re-use the model without having to re-train it. --After training the model, you can call 'predict' to recognize the person in an unknown image. + +1. Prepare a set of images of the known people you want to recognize. Organize the images in a single directory + with a sub-directory for each known person. + +2. Then, call the 'train' function with the appropriate parameters. Make sure to pass in the 'model_save_path' if you + want to save the model to disk so you can re-use the model without having to re-train it. + +3. Call 'predict' and pass in your trained model to recognize the people in an unknown image. NOTE: This example requires scikit-learn to be installed! You can install it with pip: + $ pip3 install scikit-learn + """ -from math import sqrt +import math from sklearn import neighbors -from os import listdir -from os.path import isdir, join, isfile, splitext +import os +import os.path import pickle -from PIL import Image, ImageFont, ImageDraw, ImageEnhance +from PIL import Image, ImageDraw import face_recognition -from face_recognition import face_locations from face_recognition.cli import image_files_in_folder ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'} -def train(train_dir, model_save_path = "", n_neighbors = None, knn_algo = 'ball_tree', verbose=False): + +def train(train_dir, model_save_path=None, n_neighbors=None, knn_algo='ball_tree', verbose=False): """ Trains a k-nearest neighbors classifier for face recognition. @@ -54,100 +61,146 @@ def train(train_dir, model_save_path = "", n_neighbors = None, knn_algo = 'ball_ │ ├── .jpeg │ └── .jpeg └── ... - :param model_save_path: (optional) path to save model of disk - :param n_neighbors: (optional) number of neighbors to weigh in classification. Chosen automatically if not specified. + + :param model_save_path: (optional) path to save model on disk + :param n_neighbors: (optional) number of neighbors to weigh in classification. Chosen automatically if not specified :param knn_algo: (optional) underlying data structure to support knn.default is ball_tree :param verbose: verbosity of training :return: returns knn classifier that was trained on the given data. """ X = [] y = [] - for class_dir in listdir(train_dir): - if not isdir(join(train_dir, class_dir)): + + # Loop through each person in the training set + for class_dir in os.listdir(train_dir): + if not os.path.isdir(os.path.join(train_dir, class_dir)): continue - for img_path in image_files_in_folder(join(train_dir, class_dir)): + + # Loop through each training image for the current person + for img_path in image_files_in_folder(os.path.join(train_dir, class_dir)): image = face_recognition.load_image_file(img_path) - faces_bboxes = face_locations(image) - if len(faces_bboxes) != 1: - if verbose: - print("image {} not fit for training: {}".format(img_path, "didn't find a face" if len(faces_bboxes) < 1 else "found more than one face")) - continue - X.append(face_recognition.face_encodings(image, known_face_locations=faces_bboxes)[0]) - y.append(class_dir) + face_bounding_boxes = face_recognition.face_locations(image) + if len(face_bounding_boxes) != 1: + # If there are no people (or too many people) in a training image, skip the image. + if verbose: + print("Image {} not suitable for training: {}".format(img_path, "Didn't find a face" if len(face_bounding_boxes) < 1 else "Found more than one face")) + else: + # Add face encoding for current image to the training set + X.append(face_recognition.face_encodings(image, known_face_locations=face_bounding_boxes)[0]) + y.append(class_dir) + # Determine how many neighbors to use for weighting in the KNN classifier if n_neighbors is None: - n_neighbors = int(round(sqrt(len(X)))) + n_neighbors = int(round(math.sqrt(len(X)))) if verbose: - print("Chose n_neighbors automatically as:", n_neighbors) + print("Chose n_neighbors automatically:", n_neighbors) + # Create and train the KNN classifier knn_clf = neighbors.KNeighborsClassifier(n_neighbors=n_neighbors, algorithm=knn_algo, weights='distance') knn_clf.fit(X, y) - if model_save_path != "": + # Save the trained KNN classifier + if model_save_path is not None: with open(model_save_path, 'wb') as f: pickle.dump(knn_clf, f) + return knn_clf -def predict(X_img_path, knn_clf = None, model_save_path ="", DIST_THRESH = .5): + +def predict(X_img_path, knn_clf=None, model_path=None, distance_threshold=0.6): """ - recognizes faces in given image, based on a trained knn classifier + Recognizes faces in given image using a trained KNN classifier :param X_img_path: path to image to be recognized :param knn_clf: (optional) a knn classifier object. if not specified, model_save_path must be specified. - :param model_save_path: (optional) path to a pickled knn classifier. if not specified, model_save_path must be knn_clf. - :param DIST_THRESH: (optional) distance threshold in knn classification. the larger it is, the more chance of misclassifying an unknown person to a known one. + :param model_path: (optional) path to a pickled knn classifier. if not specified, model_save_path must be knn_clf. + :param distance_threshold: (optional) distance threshold for face classification. the larger it is, the more chance + of mis-classifying an unknown person as a known one. :return: a list of names and face locations for the recognized faces in the image: [(name, bounding box), ...]. - For faces of unrecognized persons, the name 'N/A' will be passed. + For faces of unrecognized persons, the name 'unknown' will be returned. """ + if not os.path.isfile(X_img_path) or os.path.splitext(X_img_path)[1][1:] not in ALLOWED_EXTENSIONS: + raise Exception("Invalid image path: {}".format(X_img_path)) - if not isfile(X_img_path) or splitext(X_img_path)[1][1:] not in ALLOWED_EXTENSIONS: - raise Exception("invalid image path: {}".format(X_img_path)) - - if knn_clf is None and model_save_path == "": - raise Exception("must supply knn classifier either thourgh knn_clf or model_save_path") + if knn_clf is None and model_path is None: + raise Exception("Must supply knn classifier either thourgh knn_clf or model_path") + # Load a trained KNN model (if one was passed in) if knn_clf is None: - with open(model_save_path, 'rb') as f: + with open(model_path, 'rb') as f: knn_clf = pickle.load(f) + # Load image file and find face locations X_img = face_recognition.load_image_file(X_img_path) - X_faces_loc = face_locations(X_img) - if len(X_faces_loc) == 0: - return [] + X_face_locations = face_recognition.face_locations(X_img) - faces_encodings = face_recognition.face_encodings(X_img, known_face_locations=X_faces_loc) + # If no faces are found in the image, return an empty result. + if len(X_face_locations) == 0: + return [] + # Find encodings for faces in the test iamge + faces_encodings = face_recognition.face_encodings(X_img, known_face_locations=X_face_locations) + # Use the KNN model to find the best matches for the test face closest_distances = knn_clf.kneighbors(faces_encodings, n_neighbors=1) + are_matches = [closest_distances[0][i][0] <= distance_threshold for i in range(len(X_face_locations))] - is_recognized = [closest_distances[0][i][0] <= DIST_THRESH for i in range(len(X_faces_loc))] + # Predict classes and remove classifications that aren't within the threshold + return [(pred, loc) if rec else ("unknown", loc) for pred, loc, rec in zip(knn_clf.predict(faces_encodings), X_face_locations, are_matches)] - # predict classes and cull classifications that are not with high confidence - return [(pred, loc) if rec else ("N/A", loc) for pred, loc, rec in zip(knn_clf.predict(faces_encodings), X_faces_loc, is_recognized)] -def draw_preds(img_path, preds): +def show_prediction_labels_on_image(img_path, predictions): """ - shows the face recognition results visually. + Shows the face recognition results visually. :param img_path: path to image to be recognized - :param preds: results of the predict function + :param predictions: results of the predict function :return: """ - source_img = Image.open(img_path).convert("RGBA") - draw = ImageDraw.Draw(source_img) - for pred in preds: - loc = pred[1] - name = pred[0] - # (top, right, bottom, left) => (left,top,right,bottom) - draw.rectangle(((loc[3], loc[0]), (loc[1],loc[2])), outline="red") - draw.text((loc[3], loc[0] - 30), name, font=ImageFont.truetype('Pillow/Tests/fonts/FreeMono.ttf', 30)) - source_img.show() + pil_image = Image.open(img_path).convert("RGB") + draw = ImageDraw.Draw(pil_image) + + for name, (top, right, bottom, left) in predictions: + # Draw a box around the face using the Pillow module + draw.rectangle(((left, top), (right, bottom)), outline=(0, 0, 255)) + + # There's a bug in Pillow where it blows up with non-UTF-8 text + # when using the default bitmap font + name = name.encode("UTF-8") + + # Draw a label with a name below the face + text_width, text_height = draw.textsize(name) + draw.rectangle(((left, bottom - text_height - 10), (right, bottom)), fill=(0, 0, 255), outline=(0, 0, 255)) + draw.text((left + 6, bottom - text_height - 5), name, fill=(255, 255, 255, 255)) + + # Remove the drawing library from memory as per the Pillow docs + del draw + + # Display the resulting image + pil_image.show() + if __name__ == "__main__": - knn_clf = train("knn_examples/train") - for img_path in listdir("knn_examples/test"): - preds = predict(join("knn_examples/test", img_path) ,knn_clf=knn_clf) - print(preds) - draw_preds(join("knn_examples/test", img_path), preds) + # STEP 1: Train the KNN classifier and save it to disk + # Once the model is trained and saved, you can skip this step next time. + print("Training KNN classifier...") + classifier = train("knn_examples/train", model_save_path="trained_knn_model.clf", n_neighbors=2) + print("Training complete!") + + # STEP 2: Using the trained classifier, make predictions for unknown images + for image_file in os.listdir("knn_examples/test"): + full_file_path = os.path.join("knn_examples/test", image_file) + + print("Looking for faces in {}".format(image_file)) + + # Find all people in the image using a trained classifier model + # Note: You can pass in either a classifier file name or a classifier model instance + predictions = predict(full_file_path, model_path="trained_knn_model.clf") + + # Print results on the console + for name, (top, right, bottom, left) in predictions: + print("- Found {} at ({}, {})".format(name, left, top)) + # Display results overlaid on an image + show_prediction_labels_on_image(os.path.join("knn_examples/test", image_file), predictions)