face_cropper.py

import dlib
from PIL import Image
from skimage import io
from skimage.transform import resize


def detect_faces(image):
    """Get face border positions for all the faces in the image.
    
    Params:
        image: np.ndarray - single image in rgb format.
    
    Return:
        face_frames: list(int) - list of left, right top, bottom 
                                 edges of faces present in image.
    """

    # Create a face detector.
    face_detector = dlib.get_frontal_face_detector()

    # Run detector and get bounding boxes of the faces on image.
    detected_faces = face_detector(image, 1)
    face_frames = [(x.left(), x.top(), x.right(), x.bottom()) 
                   for x in detected_faces]
    return face_frames
  
def clip_crop_rect(image, left, top, right, bottom):
    """Method to prevent the edges from being outside bounds of image.
    
    If an edge does go over the boundaries, the box is shifted in
    the opposite direction so that the edges can retain their dimensions
    and also not intersect with the boundaries of the image.
    
    Params:
        image: np.ndarray - the rgb image.
        left: int         - the left edge of the crop.
        top: int          - the top edge of the crop.
        right: int        - the right edge of the crop.
        bottom: int       - the bottom edge of the crop.
    """
    
    # The most extreme values the edges can take on without
    # adjusting the position of the other edges.
    left_bound = 0
    right_bound = image.width
    top_bound = 0
    bottom_bound = image.height

    if left < left_bound:
        difference = abs(left - left_bound)
        left = left_bound
        right += difference

    if right >= right_bound:
        difference = abs(right - right_bound)
        right = right_bound
        left -= difference

    if top < top_bound:
        difference = abs(top - top_bound)
        top = top_bound
        bottom += difference

    if bottom >= bottom_bound:
        difference = abs(bottom - bottom_bound)
        bottom = bottom_bound
        top -= difference
        
    # A set of edges that do not intersect with the borders of the image.
    return left, top, right, bottom
  

# Dimensions of the new dataset.
target_shape  = (256, 256)

# No face crops less than 256, 256 to be pushed to the dataset.
threshold = 1.0

# List the images in the folder.
images = os.listdir(images_folder)

# Get each image number and filename in the folder.
for image_index, image_name in enumerate(images):

    # The relative path to the image.
    image_path = os.path.join(images_folder, image_name)

    # Load the image.
    image_array = io.imread(image_path)

    # Get a list of lists containing the borders for the present 
    # faces.
    detected_faces = detect_faces(image_array)

    # Create PIL image.
    image = Image.fromarray(image_array)

    # Iterate over the faces in the current image.
    for i, (left, top, right, bottom) in enumerate(detected_faces):

        # Adjust the edges if they overlap with the border of the
        # image.
        if prevent_boundary_overlap:
            left, top, right, bottom = \
                clip_crop_rect(image, left, top, right, bottom)

        # Create a tuple.
        face_rect = left, top, right, bottom

        # Cropped face.
        face = image.crop(face_rect)

        # Don't add pixelated images into the dataset.
        width_big_enough = face.width / target_shape[0] >= threshold
        height_big_enough = face.height / target_shape[1] >= threshold
        if width_big_enough and height_big_enough:

            # Create array.
            face = np.array(face)

            # Resize to standard size.
            face = resize(face, target_shape)

            # Save crop to new folder.
            new_image_path = os.path.join(output_folder, image_name)
            io.imsave(new_image_path, face)
            counter += 1

    # Print progress.
    print('{}/{} done    '.format(image_index, len(images)), end='\r')