import argparse
import os

import numpy
import cv2

HAARCASCADE_FRONTALFACE_PATH = '/opt/local/share/OpenCV/haarcascades/haarcascade_frontalface_alt.xml'

if not os.access(HAARCASCADE_FRONTALFACE_PATH, os.R_OK):
    sys.exit("error: unable to open HAARCASCADE_FRONTALFACE_PATH; please check and set correctly")

def extract_faces(image, face_cascade, padding_fraction=0.4, resize=True, new_size=(224, 224)):
    """
    Return as a list the faces found in the image.
    Parameters
    ----------
    image : numpy uint8 (BGR-pixel order) format
    face_cascade : an initialized OpenCV2 CascadeClassifier
    padding_fraction: how much padding (as a fraction of found face size)
    resize : whether to resize the face once found
    new_size : the size of the resized face (with padding)
    Returns
    -------
    a list of OpenCV2 images (as numpy arrays), one per face
    """
    (height, width, _) = image.shape

    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    faces = face_cascade.detectMultiScale(gray, 1.1, 4)
    result = []
    for (i, (x,y,w,h)) in enumerate(faces):
        # margins, and padding (by replication) at border if necessary
        xpad = int(w * padding_fraction)
        ypad = int(h * padding_fraction)

        ymin = y - ypad
        if ymin < 0:
            yminpad = abs(ymin)
            ymin = 0
        else:
            yminpad = 0
        ymax = y + h + ypad
        if ymax >= height:
            ymax = height - 1
            ymaxpad = ymax - height + 1
        else:
            ymaxpad = 0
        xmin = x - xpad
        if xmin < 0:
            xminpad = abs(xmin)
            xmin = 0
        else:
            xminpad = 0
        xmax = x + w + xpad
        if xmax >= width:
            xmax = width - 1
            xmaxpad = xmax - width + 1
        else:
            xmaxpad = 0
        face = image[ymin:ymax, xmin:xmax]
        face_padded = cv2.copyMakeBorder(face,ymaxpad,yminpad,xminpad,xmaxpad,cv2.BORDER_REPLICATE)
        if resize:
            face_padded = cv2.resize(face_padded, new_size, interpolation=cv2.INTER_CUBIC)
        result.append(face_padded)
    return result

def make_classifier(cascade_path=HAARCASCADE_FRONTALFACE_PATH):
    return cv2.CascadeClassifier(cascade_path)

def main():
    parser = argparse.ArgumentParser(description='Extract face(s) from images; write results to FILE-face-0.png, FILE-face-1.png, etc.')
    parser.add_argument('files', metavar='FILE', nargs='+',
                        help='a file to perform image recognition on')
    args = parser.parse_args()

    if not os.access(HAARCASCADE_FRONTALFACE_PATH, os.R_OK):
        sys.exit("error: unable to open HAARCASCADE_FRONTALFACE_PATH; please check and set correctly")
    face_cascade = make_classifier()
    for f in args.files:
        image = cv2.imread(f)
        faces = extract_faces(image, face_cascade)
        for i, face in enumerate(faces):
            cv2.imwrite('{}-face-{}.png'.format(f, i), face)

if __name__ == '__main__':
    main()

import sys
if sys.version_info[0] != 2 or sys.version_info[1] != 7:
    sys.exit("This script written for Python 2.7; disable this check at your own risk.")
import argparse
import os

import caffe
import numpy as np

DATA_DIR = os.path.join('/Users/liberato/Research/2016-research-automated-image-detection', 'data')

# model from Rothe et al.; appears to be the VGG-16 architecture
MODEL_STRUCTURE = os.path.join(DATA_DIR, 'age.prototxt')

# two sets of model weights from Rothe, et al.
# first was trained on the IMDB-WIKI data set
# second was then fine-tuned on data from the ICCV competition
#MODEL_WEIGHTS = os.path.join(DATA_DIR, 'dex_imdb_wiki.caffemodel')
MODEL_WEIGHTS = os.path.join(DATA_DIR, 'dex_chalearn_iccv2015.caffemodel')

# CNNs often subtract the mean pixel value (over the set of training image pixels) to "center" pixel values; this is an established technique to improve accuracy
MEAN_IMAGE = os.path.join(DATA_DIR, 'imagenet_mean.binaryproto')

def load_net():
    net = caffe.Net(MODEL_STRUCTURE,
                    MODEL_WEIGHTS,
                    caffe.TEST) # test (as opposed to TRAIN) mode
    return net

def load_mean_pixels():
    # load mean image, determine mean pixel values
    blob = caffe.proto.caffe_pb2.BlobProto()
    data = open(MEAN_IMAGE, 'rb' ).read()
    blob.ParseFromString(data)
    mu = np.array(caffe.io.blobproto_to_array(blob)).reshape(3,256,256)
    mu = mu.mean(1).mean(1)  # average over pixels to obtain the mean (BGR) pixel values
    return mu

def construct_transformer(net, mu, rb_swap=True):
    # construct a transformer for input data
    transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
    transformer.set_transpose('data', (2,0,1))  # move image channels to outermost dimension
    transformer.set_mean('data', mu)            # subtract the dataset-mean value in each channel
    transformer.set_raw_scale('data', 255)      # rescale from [0, 1] to [0, 255]
    if rb_swap:
        transformer.set_channel_swap('data', (2,1,0))  # swap channels from RGB to BGR
    return transformer

def transform_image(image, transformer):
    """Transform and crop an input image in memory."""
    if image.shape != (224,224,3):
        sys.stderr.write('warning: resizing input to 224x224\n')
        sys.stderr.flush()
        image = caffe.io.resize_image(image, (224,224))
    transformed_image = transformer.preprocess('data', image)
    return transformed_image

def transform_file(path, transformer):
    """Transform and crop an input image stored in a file."""
    image = caffe.io.load_image(path)
    return transform_image(image, transformer)

def estimate_probs(transformed_images, net, progress_tracker=None):
    """
    Estimate the age of the face in each transformed_image. Return a length 100 list of probabilities. list[0] is the
    probability of age=0 (years), list[1] is the probability of age=1, etc.
    """

    # prepare data (input) buffer in the network
    net.blobs['data'].reshape(1,       # batch size
                              3,       # three-channel (BGR) images
                              224,224) # image size

    result = []

    for image in transformed_images:
        # copy image into the network
        net.blobs['data'].data[...] = image
        if progress_tracker: progress_tracker.update()

        # run the classifier
        output = net.forward()
        if progress_tracker: progress_tracker.update()

        output_prob = output['prob']

        result.append(np.ndarray.tolist(output_prob)[0])

    return result

def estimate_ages(transformed_images, net):
    """
    Estimate the age of the face in each transformed_image. Return a list of ordered ages lists; each age list is in order from most-to-least likely.
    """

    probabilities_list = estimate_probs(transformed_images, net)

    result = []

    for probs in probabilities_list:
        probs_ages = [(prob, age) for age, prob in enumerate(probs)]
        probs_ages.sort()
        probs_ages.reverse()
        sorted_ages = [age for _, age in probs_ages]
        result.append(sorted_ages)

    return result

def main():
    parser = argparse.ArgumentParser(description='Estimate ages of faces in FILEs; write results to STDOUT.')
    parser.add_argument('files', metavar='FILE', nargs='+',
                        help='a file to perform age estimation on')
    args = parser.parse_args()

    net = load_net()
    mean_px = load_mean_pixels()
    transformer = construct_transformer(net, mean_px)

    faces = [transform_file(f, transformer) for f in args.files]
    results = estimate_ages(faces, net)

    for (face, result) in zip(args.files, results):
        print face, result[0:5]

if __name__ == '__main__':
    main()