Commit 99587795 authored by Kyle Anderson's avatar Kyle Anderson
Browse files

Working reworked register user

The new register user, using custom methods, appears to work nicely from basic testing. Will continue with other things now.
parent d55e0090
import os
import cv2
DATABASE_LOC = "./dataset/faces.pickle"
DATA_DIR = "./data"
DATABASE_LOC = os.path.join(DATA_DIR, "faces.pickle")
CASCADE_DIR = "./cascades"
# Output location for pickle database files
OUTPUT_DIR = "./output"
# Directory for the face detection model.
FACE_DETECTION_MODEL_DIR = os.path.join(RES_DIRECTORY, "face_detection_model")
EMBEDDINGS_PROCESSOR_LOC = os.path.join(RES_DIRECTORY, "openface_nn4.small2.v1.t7")
def display_frame(frame):
......@@ -13,3 +24,41 @@ def start_video_stream(camera: int):
Also waits for the video stream to open before returning it."""
video_stream = cv2.VideoCapture(0)
return video_stream
def load_cascade(cascade_loc: str) -> cv2.CascadeClassifier:
Opens the cascade classifier at the given path.
:param cascade_loc: The file location of the cascade.
:return:The CascadeClassifier class.
return cv2.CascadeClassifier(cascade_loc)
def load_detector(proto_path: str, model_path: str):
Loads the caffe detector with the given proto text file and the model file.
:param proto_path: The path location of the prototext file.
:param model_path: The path to the caffe model.
:return: The detector.
return cv2.dnn.readNetFromCaffe(proto_path, model_path)
def load_embedding_model(model_path: str):
Loads the torch embedding model at the given location.
:param model_path: The path to the model.
:return: The embedding model
return cv2.dnn.readNetFromTorch(model_path)
CAFFE_MODEL_NAME = "res10_300x300_ssd_iter_140000.caffemodel"
PROTOTXT_NAME = "deploy.prototxt"
def load_detector_from_dir(detector_dir: str):
prototxt: str = os.path.join(detector_dir, PROTOTXT_NAME)
caffe_model: str = os.path.join(detector_dir, CAFFE_MODEL_NAME)
return load_detector(prototxt, caffe_model)
import os
import cv2
import imutils
import numpy
from imutils import paths as impaths
import common
import data_handler
def detect_faces(face_detector, image):
Detects faces in the provided image using the provided face_detector.
:param face_detector: The face_detector.
:param image: The image to be processed.
:return: The detected faces in the image.
image_blob = cv2.dnn.blobFromImage(cv2.resize(image, (300, 300)), scalefactor=1.0, size=(300, 300),
mean=(104.0, 177.0, 123.0), swapRB=False, crop=False)
return face_detector.forward()
def find_best_match(faces, image, min_confidence: float = 0.5):
Finds the single face in the given list of faces with the best match
:param image: The image from which the face was detected.
:param faces: The list of faces to go through.
:param min_confidence: The minimum percentage confidence in order for a face to be considered recognized.
:return: The best matched face, or None if none were matched or none were matched with a large enough confidence.
best_match = None
if len(faces) > 0:
# Assume that each image has only one face, so take the bounding box with the largest probability of being a face.
i = numpy.argmax(faces[0, 0, :, 2])
confidence = faces[0, 0, i, 2]
# Only continue if the confidence is enough
if confidence > min_confidence:
img_height, img_width = image.shape[0:2]
# Determine the bounding box for the face
box = faces[0, 0, i, 3:7] * numpy.array([img_width, img_height, img_width, img_height])
# Get start and end positions for the box.
startx, starty, endx, endy = box.astype("int")
# Extract face ROI and get dimensions for it
face = image[starty: endy, startx:endx]
face_height, face_width = face.shape[0:2]
# Don't match the face if it's too small.
if face_width >= 20 and face_height >= 20:
best_match = face
return best_match
def extract_face_embeddings(face, embedding_cnn):
Extracts the facial embeddings for the given face
:param face: The face for which embeddings should be created
:param embedding_cnn: The embedding cnn to be used.
:return: The embeddings for the face
# Construct a blob and pass it to the embedder to obtain a 128-d quantification for the face.
face_blob = cv2.dnn.blobFromImage(face, scalefactor=1.0 / 255, size=(96, 96), mean=(0, 0, 0), swapRB=True,
vec = embedding_cnn.forward()
return vec
def process_dataset(directory_location: str, detector_dir: str = common.FACE_DETECTION_MODEL_DIR,
embedding_model_path: str = common.EMBEDDINGS_PROCESSOR_LOC,
show_output: bool = False, file_output: str = None) -> dict:
Processes the images in the given directory for facial identification.
:param directory_location: The path to a directory full of a dataset of images for the same person.
Note that each subdirectory within this directory should be named the same as the user_id for the user.
--- some_user_id
-- image1.png
-- image2.png
--- some_other_user_id
-- image1.png
-- image2.png
:param detector_dir: String location of the detection file directory.
:param embedding_model_path: The path to the embedding model.
:param show_output: True to print progress, False otherwise.
:param file_output: The pickle file to which the embeddings should be outputted. None means it won't be saved.
:return: The processed dataset dictionary, with format { "user_id" : [encoding1, encoding2, ...] , ... }
# Dictionary with results.
result_database = {}
image_paths = list(impaths.list_images(directory_location))
face_detector = common.load_detector_from_dir(detector_dir)
embedding_cnn = common.load_embedding_model(embedding_model_path)
for (i, image_path) in enumerate(image_paths):
current_user_id: str = image_path.split(os.path.sep)[-2]
if show_output:
print(f"Processing image {i + 1} for user {current_user_id}.")
image = cv2.imread(image_path)
image = imutils.resize(image, width=IMAGE_RESIZE_WIDTH)
faces = detect_faces(face_detector, image)
face = find_best_match(faces, image)
if face is not None:
facial_embeddings = extract_face_embeddings(face, embedding_cnn)
if facial_embeddings is not None and len(facial_embeddings) > 0:
if current_user_id not in result_database:
result_database[current_user_id] = []
if file_output is not None:
data_handler.write_database(file_output, result_database)
return result_database
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(
description="Registers users' facial encodings from a dataset of images containing their face.")
parser.add_argument("dataset", type=str, help="Location of the dataset which should be processed.")
parser.add_argument("output", type=str,
help="Location of the output pickle database file to which the encodings should be written.")
args = parser.parse_args()
process_dataset(args.dataset, show_output=True, file_output=args.output)
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
This diff is collapsed.
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment