import os
import numpy as np
import cv2 as cv

# Ustawia ścieżkę do kaskadowych klasyfikatorów Haara do wykrywania twarzy.
cascade_path = "C:/Python39/Lib/site-packages/cv2/data/"
face_detector = cv.CascadeClassifier(cascade_path +
                                     'haarcascade_frontalface_default.xml')

# Ustawia ścieżkę do obrazów szkoleniowych oraz przygotowuje słownik
# nazwisk i etykiet.
# Gotowe zdjęcia kapitana Demminga znajdują się w katalogu demming_trainer.
# Zdjęcia własnej twarzy umieść w katalogu trainer.

# Użyj katalogu trainer w przypadku korzystania z obrazów własnych (odkomentuj).
# train_path = './trainer'
# Użyj katalogu demming_trainer w przypadku korzystania z gotowych obrazów.
train_path = './demming_trainer'

image_paths = [os.path.join(train_path, f) for f in os.listdir(train_path)]
images, labels = [], []

# Znajduje prostokąty z twarzami i przypisuje im etykiety w postaci liczb.
for image in image_paths:
    train_image = cv.imread(image, cv.IMREAD_GRAYSCALE)
    label = int(os.path.split(image)[-1].split('.')[1])
    name = os.path.split(image)[-1].split('.')[0]
    frame_num = os.path.split(image)[-1].split('.')[2]
    faces = face_detector.detectMultiScale(train_image)
    for (x, y, w, h) in faces:
        images.append(train_image[y:y + h, x:x + w])
        labels.append(label)
        print(f"Przetwarzanie danych szkoleniowych: plik {name}.{label}"
              f".{frame_num}.jpg")
        cv.imshow("Obraz szkoleniowy", train_image[y:y + h, x:x + w])
        cv.waitKey(50)

cv.destroyAllWindows()

# Przeprowadza szkolenie algorytmu.
recognizer = cv.face.LBPHFaceRecognizer_create()
recognizer.train(images, np.array(labels))
recognizer.write('lbph_trainer.yml')
print("Szkolenie zakonczone. Wyjscie z programu...")
