Clean up KNN example and fix PIL font bug in it

examples/face_recognition_knn.py

 """
-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_
         │   ├── <somename1>.jpeg
         │   └── <somename2>.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)