import cv2, dlib, os
from scipy.spatial import distance
import pickle

def save(name, list):
    with open(str(name) + '.pkl', 'wb') as filehandle:
        pickle.dump(list, filehandle)


def load(name):
    with open(str(name) + '.pkl', 'rb') as filehandle:
        list = pickle.load(filehandle)
    return list

predictor_path = 'datasets/shape_predictor_5_face_landmarks.dat'
face_rec_model_path = 'datasets/dlib_face_recognition_resnet_model_v1.dat'
weights = 'datasets/mmod_human_face_detector.dat'
path = 'photo_pairs/'
files = os.listdir(path)

detector = dlib.get_frontal_face_detector()
sp = dlib.shape_predictor(predictor_path)
facerec = dlib.face_recognition_model_v1(face_rec_model_path)

pics = {}
fail_pics = []

for img in files:
    image = cv2.imread(path + img)
    cnn_face_detector = dlib.cnn_face_detection_model_v1(weights)
    faces_cnn = cnn_face_detector(image, 1)

    if len(faces_cnn) > 0:
        for ind, face in enumerate(faces_cnn):
            destRGB = dlib.load_rgb_image(path + img)
            faceBox = dlib.rectangle(left=face.rect.left(), top=face.rect.top(), right=face.rect.right(),
                                     bottom=face.rect.bottom())
            shape = sp(destRGB, faceBox)
            face_descriptor = facerec.compute_face_descriptor(destRGB, shape)
            newname = '123/rez/' + str(ind) + '_' + img.split('.')[0]
            if len(face_descriptor) > 0:
                print('got vectors', img)
                save(newname, face_descriptor)
                if img in pics:
                    pics[img].append(face_descriptor)
                else:
                    pics[img] = []
                    pics[img].append(face_descriptor)
            else:
                print('no vectors', img)
                fail_pics.append(img)
    else:
        print('no face', img)
        fail_pics.append(img)

save('photo_pairs', pics)
save('fail_photo_pairs', fail_pics)