face_train.py

import os
from tkinter import messagebox

import cv2
import numpy as np


class FaceTrain:
    def __init__(self, class_name):
        self.class_name = class_name
        self.xml_file = os.path.join(os.getcwd(), "haarcascade_frontalface_default.xml")

    def face_detect(self, image):
        img = image.copy()
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        clf = cv2.CascadeClassifier(self.xml_file)
        faces = clf.detectMultiScale(gray, scaleFactor=1.2, minNeighbors=5)
        if len(faces) == 0:
            return None, None
        else:
            (x, y, w, h) = faces[0]
            return gray[y: y + w, x: x + h], faces[0]

    def prepare_training_data(self):
        training_data_dir = os.path.join(os.getcwd(), "images", self.class_name)
        dirs = os.listdir(training_data_dir)
        faces = []
        labels = []
        for dir_name in dirs:
            if not dir_name.startswith("s"):
                continue
            try:
                label = int(dir_name.replace("s", ""))
            except:
                messagebox.showerror(
                    "Error",
                    "Unable to prepare training data for "
                    + dir_name
                    + "\nfolder in images directory does not follow the naming scheme",
                )
                continue
            subject_dir_path = os.path.join(training_data_dir, dir_name)
            images = os.listdir(subject_dir_path)
            for image_name in images:
                if image_name.startswith("."):
                    continue
                img = cv2.imread(os.path.join(subject_dir_path, image_name))
                face, rect = self.face_detect(img)
                if face is not None:
                    faces.append(face)
                    labels.append(label)

        return faces, labels

    def main(self):
        # Preparing data...
        faces, labels = self.prepare_training_data()

        if 0 in (len(faces), len(labels)):
            messagebox.showerror("Training Unsuccessful", "No student image found!")
            return

        face_recognizer = cv2.face.LBPHFaceRecognizer_create()
        face_recognizer.train(faces, np.array(labels))

        face_recognizer_file = os.path.join(
            os.getcwd(), "extras", self.class_name, "face_recognizer_file.xml"
        )
        face_recognizer.save(face_recognizer_file)

        messagebox.showinfo(
            "Training Successful", "Trained for {} images".format(len(faces))
        )


if __name__ == "__main__":
    pass