Face recognizer with opencv

# coding:utf-8
from constants import *
from os import path, getcwd, listdir
import cv2
import numpy as np
import time
import sys


def load_cascade():
    """
        Carregando XML com dados para detecção de rosto(s)
    """
    return cv2.CascadeClassifier(path.join(getcwd(), "haarcascade_frontalface_default.xml"))

def load_image(image=None):
    """
        Carrega a imagem do diretório atual (se houver), caso contrário, irá buscar a imagem dentro da pasta
    """
    if path.exists(image):
        return cv2.imread(image)
    elif path.exists(path.join(IMAGES_PATH, image)):
        return cv2.imread(path.join(IMAGES_PATH, image))
    else:
        raise FileNotFoundError("Imagem não existe")

def treatment(image=None):
    """
        Aqui fazemos o tratamento da imagem para uma melhor visão computacional
    """
    blur = cv2.GaussianBlur(image, (5,5), 0)
    image = cv2.cvtColor(blur, cv2.COLOR_BGR2GRAY)
    return image

def detect_face(image):
    """
        Função responsável por fazer a detecção do rost da pessoa na foto
    """
    face_cascade = load_cascade()
    faces = face_cascade.detectMultiScale(image, 1.1, 5)
    if len(faces) == 0:
        return None
    else:
        (x, y, w, h) = faces[0]
        (width, height) = ((x+w)-w), int(y*2.3)
        return image[y:y+w, x:x+h], faces[0]

def prepare_training_data(data_folder_path):
    """
        Treinamos os rostos conhecidos com os nomes que temos e fazemos comparação com os novos
    """
    dirs = listdir(data_folder_path)
    faces = []
    labels = []
    for dir_name in dirs:
        if not dir_name.startswith('s'):
            continue
        label = int(dir_name.replace('s', ''))
        subject_dir_path = path.join(data_folder_path, dir_name)
        subject_images_names = listdir(subject_dir_path)
        for image_name in subject_images_names:
            if image_name.startswith('.'):
                continue
            else:
                print("Treinando imagem %s " % (image_name))
                image_path = path.join(subject_dir_path, image_name)
                image = cv2.imread(image_path)
                face, rect = detect_face(image)
                if face is not None:
                    face = cv2.cvtColor(face, cv2.COLOR_BGR2GRAY)
                    faces.append(face)
                    labels.append(label)
    return faces, labels

def draw_rectangle(img, rect):
    (x,y,w,h) = rect
    cv2.rectangle(img, (x,y), (x+w, y+h), (0, 255, 0), 2)

def draw_text(img, text, x, y):
    cv2.putText(img, text, (x,y), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (255, 0, 0), 2)

def predict(fr, test_img):
    img = test_img.copy()
    face, rect = detect_face(img)
    label = fr.predict(face)
    label_text = SUBJECTS[label[0]]
    draw_rectangle(img, rect)
    draw_text(img, label_text, rect[0], rect[1]+40)
    return img, label_text, label[1]

def wait():
    sys.stdout.write('.')
    sys.stdout.flush()

def main():
    """
        Função principal de nosso programa que fará chamadas para outras
    """
    print("Treinando dados...")
    start = time.time()
    if not path.exists("dados_treino.xml"):
        print("Criando novo modelo de dados")
        faces, labels = prepare_training_data("known_faces")
        print("Total de rostos conhecidos {}\nTotal de etiquetas conhecidas: {}\n".format(len(faces), len(labels)))
        face_recognizer = cv2.face.LBPHFaceRecognizer_create()
        face_recognizer.train(faces, np.array(labels))
    else:
        face_recognizer = cv2.face.LBPHFaceRecognizer_create()
        face_recognizer.read('dados_treino.xml')
    end = time.time()
    print("Levou %f segundos" % (end - start))
    print("Salvando resultado do treino...")
    face_recognizer.save("dados_treino.xml")
    print("Comparando rostos...")
    s = raw_input("Quer continuar? (S/n): ")
    while s.lower() == 's':
        try:
            image = raw_input("Coloque a imagem aqui e compararemos com nosso banco de dados...\n")
            image = treatment(load_image(image))
            face, rect = detect_face(image)
            face_similar, name, probability = predict(face_recognizer, face)
            print("Tem {}% de chance de ser o {}".format(probability, name))
            cv2.imshow("Resultado", face_similar)
            cv2.waitKey(0)
            cv2.destroyAllWindows()
        except KeyboardInterrupt:
            print("Saindo...")
            exit(0)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("Saindo...")
        exit(0)