face recognition model

main.py

import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

from PIL import Image
import numpy as np
from matplotlib import pyplot as plt

import tensorflow as tf
from tensorflow.keras.models import Model

import cv2

from mtcnn import MTCNN
from keras_vggface.utils import preprocess_input
from keras_vggface.utils import decode_predictions

from model import get_model

model = get_model()
print(model.summary())


def crop_face(filename, required_size=(224, 224)):
    img = cv2.imread(filename)
    detector = MTCNN()
    results = detector.detect_faces(img)
    x, y, width, height = results[0]['box']
    face = img[y:y+height, x:x+width]
    image = Image.fromarray(face)
    image = image.resize(required_size)
    face_array = np.asarray(image)
    return face_array, face

def who_is_this(img, vgg_face_descriptor):
    face_array, face = crop_face(img)
    face_array = face_array.astype('float32')
    input_sample = np.expand_dims(face_array, axis=0)
    img_prediction = vgg_face_descriptor.predict(preprocess_input(input_sample))
    results = decode_predictions(img_prediction)
    prediction = results[0][0][0].replace("b'", "").replace("'","")

    return prediction

def get_prediction(image_path):
    model.load_weights(os.path.join(os.getcwd(), "weight", "vgg_face_weights.h5"))
    vgg_face_descriptor = Model(inputs=model.layers[0].input,\
                                outputs=model.layers[-2].output)

    return who_is_this(image_path, vgg_face_descriptor)

if __name__ == "__main__":
    print(get_prediction('test/Amir-Khan.jpg'))