#!/usr/bin/env python

import sys
import cv2
import numpy as np
import face_recognition

def record(path, dev):
    def draw_face_rec(name, frame):
        rgb_frame = cv2.cvtColor(frame, cv2.COLOR_GRAY2RGB)
        framelocs = face_recognition.face_locations(rgb_frame)
        frameencs = face_recognition.face_encodings(rgb_frame, framelocs)
        for (top, right, bottom, left), frameenc in zip(framelocs, frameencs):
            cv2.rectangle(rgb_frame, (left, top), (right, bottom), (0, 0, 255), 2)

        cv2.imshow(name, rgb_frame)
        if len(frameencs) == 0:
            return None
        else:
            return [ frameencs[0] ]

    cap = cv2.VideoCapture(dev)

    avg = 128
    while True:
        ret, frame = cap.read()
        mean = cv2.mean(frame)[0]
        avg = (avg + mean) / 2
        if mean < avg:
            continue

        cv2.imshow("frame", frame)

        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'):
            break
        elif key == ord('\r'):
            cv2.imshow("frame", frame)
            cv2.waitKey(1)
            encs = draw_face_rec("frame", frame)
            if encs == None:
                continue
            while True:
                key = cv2.waitKey(0) & 0xFF
                if key == ord('\r'):
                    np.save(path, encs[0])
                    return
                elif key == 27: # esc
                    break

if len(sys.argv) != 2:
    print(f"Usage: {sys.argv[0]} <path>")
    sys.exit(1)

record(sys.argv[1], 2)