funlocker.py

import os
import glob
import cv2
from PIL import Image as Img
from PIL import ImageTk
import time
import numpy as np
import faiss

from tkinter import *
import tkinter.font as tkFont
import tkinter.ttk as ttk

from utils.doorlock import DoorlockController
from utils.faceid import (FaceDetection, FaceRecognition)

CAMERA_WIDTH, CAMERA_HEIGHT = 640, 480
REFRESH_TIME_MS = 3
EMBEDDINGS_DIM = 512
FACE_THRESHOLD = 0.6
FACEDB_PATH = 'faces/embeddings.npy'

# Design Resources for Window Form
tk_resources = dict()
faiss_index = None
embeddings_db = None

# TPU Global Instances
face_detector = FaceDetection()
face_recognizer = FaceRecognition()


class ModalDialog(Toplevel):
    def __init__(self, parent, body_class, *args, title=None, on_destroy=None, **kwargs):
        super().__init__(parent)
        self.transient(parent)
        self.title(title or "Modal")
        self.on_destroy = on_destroy

        self.body = body_class(self, *args, **kwargs)
        self.body.pack(padx=5, pady=5)

        self.button = Button(self, text="Close", command=self.on_ok)
        self.button.pack(side="bottom", pady=5)

        self.grab_set()
        self.protocol("WM_DELETE_WINDOW", self.on_ok)

        self.update_idletasks()
        x = parent.winfo_x() + parent.winfo_width()//2 - self.winfo_width()//2
        y = parent.winfo_y() + parent.winfo_height()//2 - self.winfo_height()//2
        self.geometry(f"+{x}+{y}")

        self.wait_window(self)

    def on_ok(self):
        if self.on_destroy is not None:
            self.on_destroy()
        self.destroy()


class LabeledToggle(Frame):
    def __init__(self, form, text, button_value=True, on_click=None):
        super(LabeledToggle, self).__init__()
        self.form = form
        self.on_click = on_click
        self.button_value = button_value
        rowframe = Frame(self.form)

        # Label
        label = Label(rowframe, text=text, anchor="center", padx=3, font=tk_resources['FONT3'])
        label.pack(side="left")

        # Button
        button_image = self.get_button_image(self.button_value)
        self.button = Button(rowframe, image=button_image, padx=3, bd=0, activebackground="#d9d9d9", command=self.button_on_click)
        self.button.pack(side="right")

        rowframe.pack(fill='x', anchor='n', pady=5, expand=False)
    
    def button_on_click(self):
        self.button_value = not self.button_value
        self.button.config(image=self.get_button_image(self.button_value))
        if self.on_click is not None:
            self.on_click()
    
    def get_button_image(self, button_value):
        return tk_resources['BUTTON_ON'] if button_value else tk_resources['BUTTON_OFF']
    
    def redraw_button(self, button_value):
        self.button_value = button_value
        self.button.config(image=self.get_button_image(self.button_value))


class AddFace(Frame):
    def __init__(self, parent, camera):
        super(AddFace, self).__init__(parent)
        self.cap = camera
        self.face_image = None
        self.captured = False

        # Label
        self.label_subtitle_0 = Label(parent, text="Add Face", anchor="center", font=tk_resources['FONT2'])
        self.label_subtitle_0.pack(fill="x", pady=10, expand=False)

        # CameraView
        self.cam_view = Label(parent, width=224, height=224, bd=0)
        self.cam_view.pack(fill="none", expand=False)
        self.cam_view.after(REFRESH_TIME_MS, self.refresh_frames)
        
        # Button
        self.cpature_button = Button(self, text="Register", command=self.capture)
        self.cpature_button.pack(side="top", padx=10)

    def capture(self):
        if self.face_image is None:
            return
        self.captured = True

        face_transform = cv2.cvtColor(self.face_image.astype(np.float32) / 255, cv2.COLOR_BGR2RGB)
        embed = face_recognizer.get_emb(np.asarray([face_transform]))[0].flatten()
        idx, confidence = recognize_faces(embed, thresh=FACE_THRESHOLD)
        if idx != -1 and confidence > 0.1:
            append_embedding(embed, self.face_image)
            self.label_subtitle_0.configure(text="Successfully added Face", fg="blue")
            print("[Success] Registration of facial data")
        else:
            self.label_subtitle_0.configure(text="Aleady added the face", fg="red")
            self.face_image = None
            self.captured = False
    
    def refresh_frames(self):
        if self.captured:
            return
        
        raw_image = cv2.flip(self.cap.read()[1], 1)

        detections = face_detector.detect_faces(raw_image)
        if not detections:
            self.face_image = None
            imgtk = tk_resources['NO_FACE']
        else:
            _, points, _ = detections[0]
            face_image = face_detector.get_face(raw_image, points)
            self.face_image = face_image

            face_image_upscale = cv2.resize(face_image, (224, 224))
            rgb_image = cv2.cvtColor(face_image_upscale, cv2.COLOR_BGR2RGB)
            imgtk = ImageTk.PhotoImage(image=Img.fromarray(rgb_image))
        
        self.cam_view.photo_image = imgtk
        self.cam_view.configure(image=imgtk)
        self.cam_view.after(REFRESH_TIME_MS, self.refresh_frames)


class MainWindow:
    def __init__(self, form):
        self.form = form
        self.running_cam = False
        self.modal_count = 0

        self.doorlock = DoorlockController(on_change=self.on_change_door)

        # Initialize Camera
        self.cap = cv2.VideoCapture(0)
        if self.cap.isOpened():
            self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, CAMERA_WIDTH)
            self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, CAMERA_HEIGHT)
        else:
            print("[ERROR] Camera could not be opened.")

        ### CameraView
        self.cam_view = Label(self.form, width=CAMERA_WIDTH, height=CAMERA_HEIGHT, bd=0)
        self.cam_view.pack(side="left", fill="both", expand=False)
        self.cam_view.after(REFRESH_TIME_MS, self.refresh_frames)

        ### Right Panel
        self.r_panel = PanedWindow(self.form)
        self.r_panel.pack(side="left", fill="both", expand=True)

        ###### Title on Right Panel
        label_title = Label(self.r_panel, text="FUnLocker v1.0", anchor="center", font=tk_resources['FONT1'])
        label_title.pack(fill="x", pady=10, expand=False)

        ###### Separator on Right Panel
        separator_0 = ttk.Separator(self.r_panel, orient='horizontal')
        separator_0.pack(fill='x')

        ###### SubTitle9 on Right Panel
        label_subtitle_0 = Label(self.r_panel, text="Control Panel", anchor="center", font=tk_resources['FONT2'])
        label_subtitle_0.pack(fill="x", pady=10, expand=False)

        ###### Camera on Right Panel
        use_camera = True if self.cap.isOpened() else False
        self.running_cam = use_camera
        self.camera_btn = LabeledToggle(self.r_panel, text="Camera", button_value=use_camera, on_click=self.camera_btn_clicked)

        ###### Lock Button on Right Panel
        self.lock_btn = LabeledToggle(self.r_panel, text="Door\nLocked", on_click=self.lock_btn_clicked)

        ###### Door Status on Right Panel
        door_status, door_color = ("Door Closed", "blue") if self.doorlock.door_closed else ("Door Opened", "red")
        self.label_door = Label(self.r_panel, text=door_status, font=tk_resources['FONT2'], fg=door_color)
        self.label_door.pack(fill="x", anchor="n", pady=2, expand=False)

        ### Separator
        separator_1 = ttk.Separator(self.r_panel, orient="horizontal")
        separator_1.pack(fill='x', pady=5)

        ###### SubTitle1 on Right Panel
        label_subtitle_1 = Label(self.r_panel, text="Face KEYs", anchor="center", font=tk_resources['FONT2'])
        label_subtitle_1.pack(fill="x", pady=10, expand=False)

        ###### Add Face on Right Panel
        add_face = Button(self.r_panel, text="Add Face", width=13, overrelief="solid", command=self.add_face)
        add_face.pack(fill="none", pady=5, expand=False)

        return
    
    def refresh_frames(self):
        if self.running_cam:
            raw_image = cv2.flip(self.cap.read()[1], 1)
            rgb_image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB)

            # Face Detection
            detections = face_detector.detect_faces(raw_image)
            embed = None
            if detections:
                _, points, _ = detections[0]
                face_image = face_detector.get_face(raw_image, points)
                face_transform = cv2.cvtColor(face_image.astype(np.float32) / 255, cv2.COLOR_BGR2RGB)
                embed = face_recognizer.get_emb(np.asarray([face_transform]))[0].flatten()

            # Face Recognition
            if embed is not None:
                idx, confidence = recognize_faces(embed, thresh=FACE_THRESHOLD)
                if idx != -1 and confidence > 0.1 and self.doorlock.door_locked and self.doorlock.door_closed:
                    print(f"Detected!! [{idx:02d}] Confidence:{confidence:04f}")
                    self.lock_btn_clicked()

            imgtk = ImageTk.PhotoImage(image=Img.fromarray(rgb_image))
            refresh_time = REFRESH_TIME_MS
        else:
            imgtk = tk_resources['NO_CAMERA']
            refresh_time = 100
        
        self.cam_view.photo_image = imgtk
        self.cam_view.configure(image=imgtk)
        self.cam_view.after(refresh_time, self.refresh_frames)
    
    def camera_btn_clicked(self):
        if self.cap.isOpened() == False:
            self.cap = cv2.VideoCapture(0)
            if self.cap.isOpened():
                self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, CAMERA_WIDTH)
                self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, CAMERA_HEIGHT)
            else:
                self.camera_btn.redraw_button(False)
        self.running_cam = self.camera_btn.button_value

    def lock_btn_clicked(self):
        self.doorlock.toggle_door()
    
    def add_face(self):
        if self.modal_count == 0 and self.running_cam:
            self.modal_count += 1
            self.running_cam = False

            dialog = ModalDialog(self.form, AddFace,
                                 camera=self.cap,
                                 title="Add", on_destroy=self.close_modal)
    
    def close_modal(self):
        self.modal_count -= 1
        self.running_cam = True
    
    def on_change_door(self):
        door_status, door_color = ("Door Closed", "blue") if self.doorlock.door_closed else ("Door Opened", "red")
        self.label_door.config(text=door_status, fg=door_color)

        self.lock_btn.redraw_button(self.doorlock.door_locked)


def init_faiss():
    global faiss_index
    try:
        # Cosine
        faiss_index = faiss.IndexFlatIP(EMBEDDINGS_DIM)
        faiss.normalize_L2(embeddings_db)

        # L2
        # faiss_index = faiss.IndexFlatL2(EMBEDDINGS_DIM)

        faiss_index.add(embeddings_db)
    except Exception as e:
        print(f"Faiss initialize failed. [{repr(e)}]")


def load_embeddings(npy_file=None):
    global embeddings_db

    if os.path.exists(npy_file or FACEDB_PATH) == False:
        folder_path = './faces'
        file_pattern = f"{folder_path}/face_*.png"
        file_paths = glob.glob(file_pattern)
        file_names = sorted([os.path.basename(path) for path in file_paths])

        embeddings_db = []
        for idx, f in enumerate(file_names):
            print(f"{folder_path}/{f}")
            raw_image = cv2.imread(f"{folder_path}/{f}")
            detections = face_detector.detect_faces(raw_image)
            
            if not detections:
                file_names.pop(idx)
                print(f"Face not detected! [{f}]")
            else:
                _, points, _ = detections[0]
                face_image = face_detector.get_face(raw_image, points)
                face_transform = cv2.cvtColor(face_image.astype(np.float32) / 255, cv2.COLOR_BGR2RGB)
                embed = face_recognizer.get_emb(np.asarray([face_transform]))[0].flatten()
                embeddings_db.append(embed)
        if len(embeddings_db) > 0:
            print(f"{len(embeddings_db)} images loaded!")
            np.save(FACEDB_PATH, embeddings_db)

    try:
        embeddings_db = np.load(FACEDB_PATH)
        embeddings_db = embeddings_db.astype("float32")
        print(f"Loaded faces: {embeddings_db.shape[0]:02d}")
    except Exception as e:
        print(f"Faiss initialize failed. [{repr(e)}]")
    init_faiss()


def append_embedding(embed, face_image=None):
    global embeddings_db
    if embeddings_db is not None:
        embeddings_db = np.vstack((embeddings_db, embed[np.newaxis, :]))
    else:
        embeddings_db = embed[np.newaxis, :]
    embeddings_db = embeddings_db.astype("float32")
    np.save(FACEDB_PATH, embeddings_db)
    init_faiss()

    if face_image is not None:
        cv2.imwrite(f'faces/face_{embeddings_db.shape[0]:02d}.png', face_image)


def recognize_faces(embed, thresh=0.6):
    if embeddings_db is None:
        return -1, -1
    
    embed = embed.reshape(1, -1)
    faiss.normalize_L2(embed)  # Cosine

    distance, idx = faiss_index.search(embed, 2)
    distance, idx = distance[0], idx[0]

    if thresh and distance[0] > thresh:
        idx = idx[0]
        conf = (distance[0] - thresh) / (1.4 - thresh)
        dist = distance[0]
    elif len(distance) >= 2:
        idx = idx[0]
        conf = (distance[1] - distance[0]) / 1.4
        dist = distance[0]
    else:
        idx = -1
        conf = -1
    print("[Debug] ---> ", distance, dist, idx, conf)
    return idx, conf


def main():
    global tk_resources

    winform = Tk()

    winform.attributes('-fullscreen', True)
    winform.geometry("800x480")

    def close_win(e):
        winform.destroy()

    # Bind shortcut keys
    winform.bind('<Escape>', lambda e: close_win(e))
    winform.bind('q', lambda e: close_win(e))

    # Load resources
    tk_resources['FONT1'] = tkFont.Font(family="DejaVu Sans", size=12, weight="bold")
    tk_resources['FONT2'] = tkFont.Font(family="DejaVu Sans", size=11, weight="bold")
    tk_resources['FONT3'] = tkFont.Font(family="DejaVu Sans", size=9)
    tk_resources['BUTTON_ON'] = PhotoImage(file="resources/toggle_btn_on.png")
    tk_resources['BUTTON_OFF'] = PhotoImage(file="resources/toggle_btn_off.png")
    tk_resources['NO_CAMERA'] = PhotoImage(file="resources/no_camera.png")
    tk_resources['NO_FACE'] = PhotoImage(file="resources/not_detected.png")

    load_embeddings()

    mainWindow = MainWindow(winform)
    winform.mainloop()


if __name__ == "__main__":
    main()