WIP

dragon.py

import pickle
import os
import mimetypes
import face_recognition
from heapq import nsmallest
import shutil

dump = 'data.dat'

allowedMimetypes = [
    'image/jpeg'
]

known_face_encodings = []
known_face_encoding_files = []
known_faces = []

if not os.path.isfile(dump):
    for root, dirs, files in os.walk('input'):
        for file in files:
            # MIME type check
            type, encoding = mimetypes.guess_type(file)

            if type not in allowedMimetypes:
                print("MIME type {} not allowed. Skipping...".format(type))
                continue

            filepath = os.path.join(root, file)

            # Face count check
            image = face_recognition.load_image_file(filepath)
            face_locations = face_recognition.face_locations(image, 0, 'cnn')

            if len(face_locations) == 0:
                print("No faces found in image {}. Skipping...".format(file))
                continue

            print("{} faces found in image {}".format(len(face_locations), file))

            face_encodings = face_recognition.face_encodings(image, face_locations)
            known_face_encodings.extend(face_encodings)

            for face_encoding in face_encodings:
                known_face_encoding_files.append(filepath) # Track face_encoding origin

        print("Dumping known face encodings")
        with open(dump, 'wb') as f:
            pickle.dump(known_face_encodings, f)
            pickle.dump(known_face_encoding_files, f)

# Load known face encodings dump
with open(dump, 'rb') as f:
    known_face_encodings = pickle.load(f)
    known_face_encoding_files = pickle.load(f)

for i, face_encoding in enumerate(known_face_encodings):
    face_distances = face_recognition.face_distance(known_face_encodings, face_encoding)
    smallest_face_distance = nsmallest(2, face_distances)[-1]

    origin_file = known_face_encoding_files[i]

    if smallest_face_distance < 0.6:
        matched_i = list(face_distances).index(smallest_face_distance)
        matched_file = known_face_encoding_files[matched_i]

        print("A face in image {} has already been found in image {}".format(origin_file, matched_file))
        print("Distance: {}".format(smallest_face_distance))

        try:
            j = [j for j in known_faces if matched_file in j][0]
            known_faces[known_faces.index(j)].append(origin_file)
        except IndexError:
            try:
                j = [j for j in known_faces if origin_file in j][0]
                known_faces[known_faces.index(j)].append(matched_file)
            except IndexError:
                known_faces.append([matched_file, origin_file])
                pass
            pass

print("Copying images for {} faces".format(len(known_faces)))
for i, known_face in enumerate(known_faces):
    dir = os.path.join('output', i)
    os.mkdir(dir)

    for file in known_face:
        shutil.copy(file, os.path.join(dir, os.path.basename(file)))