import face_recognition
import cv2
import os
from PIL import Image
import numpy as np


# Tune this this to memory.
BATCH_FRAMES = 25

def process_video(vid):
    print(f"Processing: {vid}")
    basename, _ = os.path.splitext(os.path.basename(vid))
    video_capture = cv2.VideoCapture(vid)
    faces = 0
    frames = []
    while video_capture.isOpened():
        # Grab a single frame of video
        ret, frame_cv = video_capture.read()

        # Bail out when the video file ends
        if not ret:
            break

        frame = cv2.cvtColor(frame_cv, cv2.COLOR_BGR2RGB)
        frames.append(frame)
        print(".", end="", flush=""),
		# Tune this to GPU Memory Size.
        if len(frames) == BATCH_FRAMES:
            print("#", end="", flush="")
            batch_of_face_locations = face_recognition.batch_face_locations(
                frames, number_of_times_to_upsample=0
            )

            for idx,face_locations in enumerate(batch_of_face_locations):
                if len(face_locations)==0:
                    continue

                face_encodings = face_recognition.face_encodings(
                    face_image=frames[idx],
                    known_face_locations=face_locations
                )
                for face_encoding, face_location in zip(face_encodings, face_locations):
                    print("$", end="", flush=True)
                    top, right, bottom, left = face_location

                    buffer = 0.5
                    max_bottom, max_right, _ = frame.shape
                    width = right - left
                    height = bottom - top
                    top2 = int(max([0, top - height * buffer]))
                    left2 = int(max([0, left - width * buffer]))

                    bottom2 = int(min([max_bottom, bottom + height * buffer]))
                    right2 = int(min([max_right, right + width * buffer]))

                    # You can access the actual face itself like this:
                    face_image = frame[top2:bottom2, left2:right2]
                    pil_image = Image.fromarray(face_image)

                    pil_image.save(f"faces/{basename}-{faces:08d}.jpg")

                    with open(f"faces/{basename}-{faces:08d}.np", "wb") as b:
                        np.save(b, face_encoding, allow_pickle=False, fix_imports=False)

                    faces += 1
            print("", end="\n", flush=True)
            frames = []

    print("#", end="", flush="")
    batch_of_face_locations = face_recognition.batch_face_locations(
        frames, number_of_times_to_upsample=0
    )

    for idx,face_locations in enumerate(batch_of_face_locations):
        if len(face_locations)==0:
            continue

        face_encodings = face_recognition.face_encodings(
            face_image=frames[idx],
            known_face_locations=face_locations
        )
        for face_encoding, face_location in zip(face_encodings, face_locations):
            print("$", end="", flush=True)
            top, right, bottom, left = face_location

            buffer = 0.1
            max_bottom, max_right, _ = frame.shape
            width = right - left
            height = bottom - top
            top2 = int(max([0, top - height * buffer]))
            left2 = int(max([0, left - width * buffer]))

            bottom2 = int(min([max_bottom, bottom + height * buffer]))
            right2 = int(min([max_right, right + width * buffer]))

            # You can access the actual face itself like this:
            face_image = frame[top2:bottom2, left2:right2]
            pil_image = Image.fromarray(face_image)

            pil_image.save(f"faces/{basename}-{faces:08d}.jpg")

            with open(f"faces/{basename}-{faces:08d}.np", "wb") as b:
                np.save(b, face_encoding, allow_pickle=False, fix_imports=False)
            faces += 1

# Open video file
import glob

vids = glob.glob("*.mp4")
for vid in vids:
    process_video(vid)