face_detection.py

from typing import List
import face_recognition
from PIL import Image
from os import path as p
import os
import shutil
import numpy as np

def getCroppedImages(image: Image.Image, cap = -1):
    """Takes a PIL image, and returns a list of PIL images with the cropped faces"""
    image = image.convert("RGB")
    face_locations = face_recognition.face_locations(np.array(image))
    
    outputs = []
    num = 0
    for f in face_locations:
        # allow for capping the number of faces detected, helps prevent multiple faces being detected
        # when there should only be one
        if(cap != -1 and num >= cap):
            break
        
        outputs.append(image.crop([f[3], f[0], f[1], f[2]]))
        num += 1
        
    return face_locations, outputs

    
def getFaceLocationsAndFiles(input_path: str, output_dir: str, output_name: str,):
    """
    input_path: local or absolute path to the input image\n
    output_dir: folder (local or absolute path) to save the output image to\n
    output_name: name for the new image (do not include extension)\n
    returns: list of face locations and list of paths to the cropped face images 
    """
    locs, imgs = getCroppedImages(Image.open(input_path))
    output_paths = []
    num = 0
    for i in imgs:
        path = p.join(output_dir, f"{output_name}_{num}.png")
        i.save(path)
        output_paths.append(p.abspath(path))
        num += 1
        
    return locs, output_paths


def createCroppedSetFromImage(input_path: str, output_dir: str, output_name: str, cap = -1):
    """
    input_path: local or absolute path to the input image\n
    output_dir: folder (local or absolute path) to save the output image to\n
    output_name: name for the new image (do not include extension)\n
    cap: set to -1 to process all faces detected, otherwise will limit faces to value\n
    returns: list of absolute paths to the images
    """
    imgs = getCroppedImages(Image.open(input_path))[1]
    
    num = 0
    for i in imgs:
        path = p.join(output_dir, f"{output_name}_{num}.png")
        i.save(path)
        num += 1
    

def createCroppedSets(input_path: str, output_path: str):
    """
    Iterates through subdirectories in the input directory, making a duplicate of it in the output directory.
    Then iterates through all images in the subdirectory, cropping each face and saving it to the output directory. 
    
    For example for params "input" and "output", input/Dwayne Johnson/img.png will be cropped and
    saved to output/Dwayne Johnson/0_0.png (repeated for all folder and the images they contain) 
    
    WARNING: IF A DIRECTORY WITH THE SAME NAME AS THE OUTPUT DIRECTORY, IT WILL BE DELETED
    """
    # delete previous folder so there's no conflicts
    if(p.exists(output_path)):
        shutil.rmtree(output_path)
    os.mkdir(output_path)
    
    # iterate through all subdirectories in input directory
    for dir in [name for name in os.listdir(input_path) if p.isdir(p.join(input_path, name))]:
        sub_out = p.join(output_path, dir)
        sub_in = p.join(input_path, dir)
        os.mkdir(sub_out)
        
        n = 0
        # iterate through all files in subdirectory
        for img in [name for name in os.listdir(sub_in) if p.isfile(p.join(sub_in, name))]:
            createCroppedSetFromImage(p.join(sub_in, img), sub_out, n, 1)
            n += 1