# import required packages
import cv2
import dlib
import argparse
import time
import numpy as np

# initialize hog + svm based face detector
hog_face_detector = dlib.get_frontal_face_detector()

def enhance_image(image):
    image_YCrCb = cv2.cvtColor(image, cv2.COLOR_BGR2YCR_CB)
    Y, Cr, Cb = cv2.split(image_YCrCb)
    Y = cv2.equalizeHist(Y)
    image_YCrCb = cv2.merge([Y, Cr, Cb])
    image = cv2.cvtColor(image_YCrCb, cv2.COLOR_YCR_CB2BGR)
    return image

def adjust_gamma(image, gamma=1.0):
    invGamma = 1.0 / gamma
    table = np.array([((i / 255.0) ** invGamma) * 255
                      for i in np.arange(0, 256)]).astype("uint8")
    return cv2.LUT(image, table)

def scale_faces(face_rects, down_scale=1.5):
    faces = []
    for face in face_rects:
        scaled_face = dlib.rectangle(int(face.left() * down_scale),
                                    int(face.top() * down_scale),
                                    int(face.right() * down_scale),
                                    int(face.bottom() * down_scale))
        faces.append(scaled_face)
    return faces

def detect_faces(srcimage, down_scale=1.5):
    image_scaled = cv2.resize(adjust_gamma(enhance_image(srcimage), gamma=1.0), None, fx=1.0/down_scale, fy=1.0/down_scale, interpolation=cv2.INTER_LINEAR)

    start = time.time()
    hog_faces = hog_face_detector(image_scaled, 1)
    end = time.time()
    hog_time = str(format(end - start, '.2f'))

    for face in scale_faces(hog_faces, down_scale):
        x,y,w,h = face.left(), face.top(), face.right(), face.bottom()
        cv2.rectangle(srcimage, (x,y), (w,h), (0,255,0), 2)

    # write at the top left corner of the image
    # for color identification
    img_height, img_width = srcimage.shape[:2]
    cv2.putText(srcimage, "HOG=" + str(len(hog_faces)) + "|" + hog_time, (img_width-125,20), cv2.FONT_HERSHEY_SIMPLEX, 0.5,(0,255,0), 2)
    return srcimage

# handle command line arguments
ap = argparse.ArgumentParser()
ap.add_argument('-i', '--image', required=True, help='path to image file')
ap.add_argument('-w', '--weights', default='./mmod_human_face_detector.dat',
                help='path to weights file')
args = ap.parse_args()
# load input image
image = cv2.imread(args.image)

if image is None:
    print("Could not read input image")
    exit()

image = detect_faces(image, down_scale=0.5)


# save output image
cv2.imwrite(args.image, image)
exit()