# Small disclaimer
print("[INFO] Starting program... (might take a while to load)")

# Import stuff
import numpy as np
import imutils
import threading
import time
import cv2
import os
import argparse

# Define important variables
faceDetector = "Face Detector"
minConfidence = 0.7

ap = argparse.ArgumentParser()
ap.add_argument("-i", "--image", type=str,
    help="full path to image")
args = vars(ap.parse_args())

# load our serialized face detector model from disk
print("[INFO] loading face detector model...")
prototxtPath = os.path.sep.join([args["face"], "deploy.prototxt"])
weightsPath = os.path.sep.join([args["face"],
    "res10_300x300_ssd_iter_140000.caffemodel"])
faceNet = cv2.dnn.readNet(prototxtPath, weightsPath)

''' Load models '''

# Load the face detector model (https://github.com/gopinath-balu/computer_vision/tree/master/CAFFE_DNN)
print("[INFO] Loading face detector model...")
prototxtPath = faceDetector + "\\deploy.prototxt"
weightsPath = faceDetector + "\\facedetector.caffemodel"
faceNet = cv2.dnn.readNet(prototxtPath, weightsPath)

''' Define main functions '''

def FaceDetector(frame, faceNet):
    global detections

    # Grab webcam size and then construct a blob (group of connected pixels in an image that share some common property https://learnopencv.com/blob-detection-using-opencv-python-c/)
    (h, w) = frame.shape[:2]
    blob = cv2.dnn.blobFromImage(frame, 1.0, (300, 300),
        (104.0, 177.0, 123.0))

    # Pass the blob through face detector
    faceNet.setInput(blob)
    detections = faceNet.forward()

# Main mask detection function
def MaskDetector(frame, maskNet):
    (h, w) = frame.shape[:2]
    FaceDetector(frame, faceNet)

    # Init list of faces, their corresponding locations, and the list of predictions from face detector
    faces = []
    locs = []
    preds = []

    # Loop over each face detection
    for i in range(0, detections.shape[2]):

        # Extract the confidence associated with the detection
        confidence = detections[0, 0, i, 2]

        # Filter out weak detections
        if confidence > minConfidence:

            # Get the (x, y) coordinates of the bounding box for the object
            box = detections[0, 0, i, 3:7] * np.array([w, h, w, h])
            (startX, startY, endX, endY) = box.astype("int")

            # Ensure the bounding boxes fit inside the frame
            (startX, startY) = (max(0, startX), max(0, startY))
            (endX, endY) = (min(w - 1, endX), min(h - 1, endY))

            # Extract the face return value, convert it from BGR to RGB ordering, resize to 244x244 and process it
            face = frame[startY:endY, startX:endX]
            if face.any():
                face = cv2.cvtColor(face, cv2.COLOR_BGR2RGB)
                face = cv2.resize(face, (224, 224))
                face = img_to_array(face)
                face = preprocess_input(face)

                # Add the face and bounding boxes to their respective dicts
                faces.append(face)
                locs.append((startX, startY, endX, endY))

    # Only make a mask prediction if at least one face was detected
    if len(faces) > 0:
        print("1")
    else:
        print("0")

# Loop over each frame from video
while True:
    # Resize to 600x400
    frame = cv2.imread('[args["image"]')
    frame = imutils.resize(frame, width=600, height=400)

    # Run the faces through face detector
    MaskDetector(frame, maskNet)

    # Close window if `Q` is pressed
    if key == ord("q"):
        break

# Close window and stop webcam
vs.stop()