faceaiapi.py

#!/usr/bin/python3
import mtcnn
import numpy as np
from keras import backend as K
from keras.utils.data_utils import get_file
from scipy import spatial
import os
import sys
import matplotlib.pyplot as plt
import cv2
import numpy as np
import PIL
import shutil
import tensorflow as tf



model=tf.keras.models.load_model("./faceai.h5")



def preprocess_input(x, data_format=None, version=1):
    x_temp = np.copy(x)
    if data_format is None:
        data_format = K.image_data_format()
    assert data_format in {'channels_last', 'channels_first'}

    if version == 1:
        if data_format == 'channels_first':
            x_temp = x_temp[:, ::-1, ...]
            x_temp[:, 0, :, :] -= 93.5940
            x_temp[:, 1, :, :] -= 104.7624
            x_temp[:, 2, :, :] -= 129.1863
        else:
            x_temp = x_temp[..., ::-1]
            x_temp[..., 0] -= 93.5940
            x_temp[..., 1] -= 104.7624
            x_temp[..., 2] -= 129.1863

    elif version == 2:
        if data_format == 'channels_first':
            x_temp = x_temp[:, ::-1, ...]
            x_temp[:, 0, :, :] -= 91.4953
            x_temp[:, 1, :, :] -= 103.8827
            x_temp[:, 2, :, :] -= 131.0912
        else:
            x_temp = x_temp[..., ::-1]
            x_temp[..., 0] -= 91.4953
            x_temp[..., 1] -= 103.8827
            x_temp[..., 2] -= 131.0912
    else:
        raise NotImplementedError

    return x_temp


def face_extractor(path):
    face_arr=[]
    detect=mtcnn.MTCNN()
    img=cv2.imread(path)
    n=len(detect.detect_faces(img))
#    plt.figure(figsize=(10,10))
    for i in range(n):
        x,y,w,h=detect.detect_faces(img)[i]['box']
        dis=img[y:y+h,x:x+w]
        dis=PIL.Image.fromarray(dis).resize((224,224))
        dis=np.asarray(dis,dtype='float32')
        face_arr.append(dis)
    return face_arr

def generate_embeddings(fpath='./encode/'):
    files=os.listdir(fpath)
    files=[fpath+f for f in files]
    faces=[]
    for i in files:
        faces.extend(face_extractor(i))
    faces=preprocess_input(faces,version=2)
    embeddings_known=model.predict(faces)
    return embeddings_known


def compare(embedding_known,embedding_unknown,limit=0.45):
    dist=spatial.distance.cosine(embedding_known.flatten(),embedding_unknown.flatten())
    if dist>limit:
        return 0
    else:
        return 1

def group(embeddings_known,fpath):
    files=os.listdir(fpath)
    files=[fpath+f for f in files]
    faces=[]
    for i in files:
        l=face_extractor(i)
        for j in l:
            j=preprocess_input(j,version=2)
            j=model.predict(np.expand_dims(j,axis=0))
            for k in embeddings_known:
                result=compare(j,k)
                if result==1:
                    if not os.path.exists('grouped'):
                        os.makedirs('grouped')
                    shutil.copy(i,'grouped')
                    break;                    
#def main():
#    embeddings_known=generate_embeddings('./encode/')
#    group(embeddings_known,'./source/')


#if __name__=="__main__":
#    main()