face-detection ~ visual-mean

smartcrop.py

import os
import sys
import glob
import cv
import math
import argparse
from collections import OrderedDict

HAAR_CASCADE_FRONT = \
    "/usr/local/share/OpenCV/haarcascades/haarcascade_frontalface_alt2.xml"
HAAR_CASCADE_PROFILE = \
    "/usr/local/share/OpenCV/haarcascades/haarcascade_profileface.xml"

# 600 pixels is our current threshold for image height/width
SIZE_THRESHOLD = 600

MIN_WINDOW_SIZE_PROFILE = 70
MIN_WINDOW_SIZE_FRONT = 70
IMAGE_SCALE_FACTOR = 1
DEVIATIONS_TRESHOLD = 1


class SmartCrop(object):

    """
    Base class for SmartCrop (using face detection) for Image Cropping.

    """
    def __init__(self, image, **kwargs):
        self.image = image
        self.cvImage = cv.LoadImage(image)
        self.name = image.title().split('/')[-1].lower()

        self.faces = []
        self.pixel_crop = {}

        if 'out_dir' in kwargs:
            self.out_dir = kwargs['out_dir']

        self.get_point = self._detect_faces_get_point()

    @property
    def get_center(self):
        point = self.get_point
        # float the denominator in order to get float output
        return (point[0] / float(self.cvImage.width),
            point[1] / float(self.cvImage.height))

    @property
    def draw_face_points(self):
        """
        For human testing.
        See image's faces (green rects) and derived center (yellow dot).
        """

        origin = cv.LoadImage(self.image)

        if self.faces:
            for (x, y, w, h) in self.faces:
                cv.Rectangle(origin, (x, y), (x + w, y + h),
                    cv.RGB(0, 255, 0), 1)

        cv.Circle(origin, self.get_point, 8, cv.RGB(255, 255, 0), -1)
        cv.SaveImage('{}/{}'.format(self.out_dir, self.name), origin)

    @property
    def draw_bounding_box(self):
        """
        Assert pixel crop exists and draw bounding box for smart-crop.
        """
        for pos in ['left', 'bottom', 'top', 'right']:
            assert pos in self.pixel_crop

        origin = cv.LoadImage(self.image)
        cv.Circle(origin, self.get_point, 8, cv.RGB(255, 255, 0), -1)
        cv.Rectangle(origin,
            (int(self.pixel_crop['left']), int(self.pixel_crop['top'])),
            (int(self.pixel_crop['right']), int(self.pixel_crop['bottom'])),
            cv.RGB(255, 0, 0), 3)

        cv.SaveImage('{}/cropped.{}'.format(self.out_dir, self.name), origin)

    def get_positions(self, aspect_width=1, aspect_height=1):
        center_point = self.get_point

        aspect_ratio = int(aspect_width) / float(aspect_height)

        positions = OrderedDict()

        # set positions to center point
        positions['left'] = center_point[0]
        positions['top'] = center_point[1]
        positions['right'] = center_point[0]
        positions['bottom'] = center_point[1]

        positions.update(self._translate(positions, aspect_ratio))

        self.pixel_crop = positions.copy()

        # Return positions in terms of percentage
        for pos in positions:
            if pos == 'left' or pos == 'right':
                positions[pos] /= float(self.cvImage.width)
            else:
                positions[pos] /= float(self.cvImage.height)

        return positions

    def _translate(self, positions, aspect_ratio, do_transpose=True):
        if aspect_ratio >= 1:
            subjects = ['left', 'right', 'top', 'bottom']
        else:
            subjects = ['top', 'bottom', 'left', 'right']

        # New position math
        funcs = {
            'top': {
                'left': lambda sub: positions['left'] - (sub * aspect_ratio),
                'top': lambda sub: positions['top'] - sub,
                'right': lambda sub: positions['right'] + (sub * aspect_ratio),
                'bottom': lambda sub: positions['top'] + sub},
            'bottom': {
                'left': lambda sub: positions['left'] - (sub * aspect_ratio),
                'top': lambda sub: positions['top'] + sub,
                'right': lambda sub: positions['right'] + (sub * aspect_ratio),
                'bottom': lambda sub: self.cvImage.height},
            'left': {
                'left': lambda sub: positions['left'] - sub,
                'top': lambda sub: positions['top'] - (sub * aspect_ratio),
                'right': lambda sub: positions['right'] + sub,
                'bottom': lambda sub: positions['top'] + (sub * aspect_ratio)},
            'right': {
                'left': lambda sub: positions['left'] + sub,
                'top': lambda sub: positions['top'] - (sub * aspect_ratio),
                'right': lambda sub: self.cvImage.width,
                'bottom': lambda sub: positions['top'] + (sub * aspect_ratio)}}

        # Conditions that can non-fitted position values
        conditions = [
            lambda _val, _pos: False if _val < 0 else True,
            lambda _val, pos: False if (_pos == 'right' or _pos == 'left') \
                and _val > self.cvImage.width else True,
            lambda _val, pos: False if (_pos == 'bottom' or _pos == 'top') \
                and _val > self.cvImage.height else True
        ]

        # We assume that either the top or bottom edge will be 0
        for pos in subjects:
            sub = positions[pos]
            has_bounds = True
            new_pos = {}
            for _pos, _val in positions.items():
                if has_bounds:
                    resp = funcs[pos][_pos](sub)
                    new_pos[_pos] = resp
                    for cond in conditions:
                        if not cond(resp, _pos):
                            has_bounds = False
                            new_pos = {}
                            break
                else:
                    break
            if new_pos:
                return new_pos

        if not new_pos:
            if do_transpose:
                transpose_ratio = 1 / aspect_ratio
                new_pos = self._translate(positions, transpose_ratio,
                    do_transpose=False)
            else:
                # Fallback to entire image.
                new_pos['top'] = 0
                new_pos['left'] = 0
                new_pos['right'] = self.cvImage.width
                new_pos['bottom'] = self.cvImage.height

        return new_pos

    def _detect_faces_get_point(self):
        """
        Main method for detecting faces and getting the visual mean.
        """
        center_points = []
        detected = []

        image_scale = IMAGE_SCALE_FACTOR

        # These Window Sizes may be updated for "larger" than THRESHOLD images
        window_size_profile = MIN_WINDOW_SIZE_PROFILE
        window_size_front = MIN_WINDOW_SIZE_FRONT

        storage = cv.CreateMemStorage(0)
        cascade_profile = cv.Load(HAAR_CASCADE_PROFILE)
        cascade_front = cv.Load(HAAR_CASCADE_FRONT)

        # Convert to grayscale
        grayscale = cv.CreateImage((self.cvImage.width, self.cvImage.height),
            8, 1)
        cv.CvtColor(self.cvImage, grayscale, cv.CV_BGR2GRAY)

        # Downsample
        if grayscale.width >= SIZE_THRESHOLD and \
                grayscale.height >= SIZE_THRESHOLD:
            # Fixed Scale for "larger" than THRESHOLD images.
            image_scale = 2.5

        # Smooth grayscale image and Resize/Downsample
        cv.Smooth(grayscale, grayscale, cv.CV_GAUSSIAN, 3, 0)
        runner_image = cv.CreateImage((cv.Round(grayscale.width / image_scale),
            cv.Round(grayscale.height / image_scale)), 8, 1)
        cv.Resize(grayscale, runner_image, interpolation=cv.CV_INTER_LINEAR)

        # Perform algorithm and combine front-facing and profile lists
        detected_front = cv.HaarDetectObjects(runner_image, cascade_front,
            storage, 1.2, 2, cv.CV_HAAR_DO_CANNY_PRUNING, (window_size_front,
            window_size_front))
        detected_profile = cv.HaarDetectObjects(runner_image, cascade_profile,
            storage, 1.2, 2, cv.CV_HAAR_DO_CANNY_PRUNING, (window_size_profile,
            window_size_profile))

        detected = detected_front + detected_profile

        if detected:
            for (x, y, w, h), n in detected:
                _x = int(x * image_scale)
                _y = int(y * image_scale)
                _w = int(w * image_scale)
                _h = int(h * image_scale)
                self.faces.append((_x, _y, _w, _h))

            for (x, y, w, h) in self.faces:
                mid_center = (x + (w / 2), y + (h / 2))
                center_points.append(mid_center)

            return self._get_visual_center(center_points)

        return (self.cvImage.width / 2, self.cvImage.height / 2)

    def _get_visual_center(self, point_list):
        """
        Get the visual mean of all the mid-centers of our faces.
        """
        avg_point = tuple([sum(point) / len(point) for point in \
            zip(*point_list)])
        return self._remove_bad_faces(point_list, avg_point)

    def _getdeviations(self, point, avg, stddevs):
        """
        Gets number of standard deviations from the mean.
        """
        sd_point = [0, 0]
        for idx, sd in enumerate(stddevs):
            if sd <= 0:
                sd_point[idx] = 0
            else:
                sd_point[idx] = abs(point[idx] - avg[idx]) / sd

        return sd_point

    def _remove_bad_faces(self, points, avg):
        """
        Remove points that have both coords more than 1
        standard deviation from the mean.

        Also removes a "bad" point from self.faces.
        """
        new_points = []

        variances = map(lambda (x, y): ((x - avg[0]) ** 2, (y - avg[1]) ** 2),
            points)
        stddevs = [math.sqrt(sum(point) / len(point)) for point in \
            zip(*variances)]
        for idx, point in enumerate(points):
            check_point = self._getdeviations(point, avg, stddevs)
            if check_point[0] > DEVIATIONS_TRESHOLD and \
                    check_point[1] > DEVIATIONS_TRESHOLD:
                self.faces[idx] = None
            else:
                new_points.append(point)
        if new_points:
            self.faces = filter(None, self.faces)
            return tuple([sum(point) / len(point) for point in \
                zip(*new_points)])
        else:
            return avg


def runner(args):
    file_input = args.file_input
    file_output = args.file_output
    aspect_width = args.aspect_width
    aspect_height = args.aspect_height
    images = []

    if os.path.isdir(file_input):
        exts = ['jpg', 'jpeg', 'bmp', 'png', 'gif', 'Jpeg', 'Gif', 'Bmp',
            'Png']
        exts.extend([ext.upper() for ext in exts])
        for ext in exts:
            search_inside = os.path.join(file_input, '*.{}'.format(ext))
            images.extend(glob.glob(search_inside))
    else:
        images.append(file_input)

    if not os.path.exists(file_output):
        os.makedirs(file_output)

    for image in images:
        new = SmartCrop(image, out_dir=file_output)
        new.draw_face_points
        crop = new.get_positions(aspect_width, aspect_height)
        try:
            new.draw_bounding_box
        except AssertionError:
            print "Must call get_positions before drawing bounding_box"


def main():
    parser = argparse.ArgumentParser(description='Run SmartCrop program')
    parser.add_argument('-i', dest='file_input', default='.',
        help='A path to an image or directory for input')
    parser.add_argument('-o', dest='file_output', default='.',
        help='A path to an image or directory for output')
    parser.add_argument('-aw', dest='aspect_width', default=None,
        help='Aspect Width for Init Cropping Box'),
    parser.add_argument('-ah', dest='aspect_height', default=None,
        help='Aspect Height for Init Cropping Box')
    args = parser.parse_args()
    runner(args)

if __name__ == "__main__":
    main()