diff --git a/main.py b/main.py index 5f41fa2..7bd888b 100755 --- a/main.py +++ b/main.py @@ -1,10 +1,11 @@ import os.path import datetime -import subprocess +import pickle import tkinter as tk import cv2 from PIL import Image, ImageTk +import face_recognition import util @@ -52,12 +53,8 @@ def process_webcam(self): self._label.after(20, self.process_webcam) def login(self): - unknown_img_path = './.tmp.jpg' - cv2.imwrite(unknown_img_path, self.most_recent_capture_arr) - - output = str(subprocess.check_output(['face_recognition', self.db_dir, unknown_img_path])) - name = output.split(',')[1][:-3] + name = util.recognize(self.most_recent_capture_arr, self.db_dir) if name in ['unknown_person', 'no_persons_found']: util.msg_box('Ups...', 'Unknown user. Please register new user or try again.') @@ -67,8 +64,6 @@ def login(self): f.write('{},{}\n'.format(name, datetime.datetime.now())) f.close() - os.remove(unknown_img_path) - def register_new_user(self): self.register_new_user_window = tk.Toplevel(self.main_window) self.register_new_user_window.geometry("1200x520+370+120") @@ -106,7 +101,10 @@ def start(self): def accept_register_new_user(self): name = self.entry_text_register_new_user.get(1.0, "end-1c") - cv2.imwrite(os.path.join(self.db_dir, '{}.jpg'.format(name)), self.register_new_user_capture) + embeddings = face_recognition.face_encodings(self.register_new_user_capture)[0] + + file = open(os.path.join(self.db_dir, '{}.pickle'.format(name)), 'wb') + pickle.dump(embeddings, file) util.msg_box('Success!', 'User was registered successfully !') diff --git a/util.py b/util.py index d7086d1..f9a9248 100755 --- a/util.py +++ b/util.py @@ -1,5 +1,9 @@ +import os +import pickle + import tkinter as tk from tkinter import messagebox +import face_recognition def get_button(window, text, color, command, fg='white'): @@ -40,3 +44,32 @@ def get_entry_text(window): def msg_box(title, description): messagebox.showinfo(title, description) + + +def recognize(img, db_path): + # it is assumed there will be at most 1 match in the db + + embeddings_unknown = face_recognition.face_encodings(img) + if len(embeddings_unknown) == 0: + return 'no_persons_found' + else: + embeddings_unknown = embeddings_unknown[0] + + db_dir = sorted(os.listdir(db_path)) + + match = False + j = 0 + while not match and j < len(db_dir): + path_ = os.path.join(db_path, db_dir[j]) + + file = open(path_, 'rb') + embeddings = pickle.load(file) + + match = face_recognition.compare_faces([embeddings], embeddings_unknown)[0] + j += 1 + + if match: + return db_dir[j - 1][:-7] + else: + return 'unknown_person' +