cli.py

import argparse
import os
import time
from tqdm import tqdm

from pathlib import Path
from lib.FaceFilter import FaceFilter
from lib.faces_detect import detect_faces
from lib.utils import get_image_paths, get_folder

class FullPaths(argparse.Action):
    """Expand user- and relative-paths"""

    def __call__(self, parser, namespace, values, option_string=None):
        setattr(namespace, self.dest, os.path.abspath(
            os.path.expanduser(values)))

class DirectoryProcessor(object):
    '''
    Abstract class that processes a directory of images
    and writes output to the specified folder
    '''
    arguments = None
    parser = None

    input_dir = None
    output_dir = None

    verify_output = False
    images_found = 0
    images_processed = 0
    faces_detected = 0

    def __init__(self, subparser, command, description='default'):
        self.create_parser(subparser, command, description)
        self.parse_arguments(description, subparser, command)

    def process_arguments(self, arguments):
        self.arguments = arguments
        print("Input Directory: {}".format(self.arguments.input_dir))
        print("Output Directory: {}".format(self.arguments.output_dir))
        print('Starting, this may take a while...')

        self.output_dir = get_folder(self.arguments.output_dir)
        try:
            self.input_dir = get_image_paths(self.arguments.input_dir)
        except:
            print('Input directory not found. Please ensure it exists.')
            exit(1)

        self.images_found = len(self.input_dir)
        self.filter = self.load_filter()
        self.process()
        self.finalize()
        
    def read_directory(self):
        for filename in tqdm(self.input_dir):
            if self.arguments.verbose:
                print('Processing: {}'.format(os.path.basename(filename)))

            yield filename
            self.images_processed = self.images_processed + 1
    
    def get_faces(self, image):
        faces_count = 0
        for face in detect_faces(image):
            if self.filter is not None and not self.filter.check(face):
                print('Skipping not recognized face!')
                continue

            yield faces_count, face

            self.faces_detected = self.faces_detected + 1
            faces_count +=1
        
        if faces_count > 0 and self.arguments.verbose:
            print('Note: Found more than one face in an image!')
            self.verify_output = True
    
    def load_filter(self):
        filter_file = "filter.jpg" # TODO Pass as argument
        if Path(filter_file).exists():
            print('Loading reference image for filtering')
            return FaceFilter(filter_file)

    # for now, we limit this class responsability to the read of files. images and faces are processed outside this class
    def process(self):
        # implement your image processing!
        raise NotImplementedError()

    def parse_arguments(self, description, subparser, command):
        self.parser.add_argument('-i', '--input-dir',
                            action=FullPaths,
                            dest="input_dir",
                            default="input",
                            help="Input directory. A directory containing the files \
                            you wish to process. Defaults to 'input'")
        self.parser.add_argument('-o', '--output-dir',
                            action=FullPaths,
                            dest="output_dir",
                            default="output",
                            help="Output directory. This is where the converted files will \
                                be stored. Defaults to 'output'")
        self.parser.add_argument('-v', '--verbose',
                            action="store_true",
                            dest="verbose",
                            default=False,
                            help="Show verbose output")
        self.parser = self.add_optional_arguments(self.parser)
        self.parser.set_defaults(func=self.process_arguments)

    def create_parser(self, subparser, command, description):
        parser = subparser.add_parser(
            command,
            description=description,
            epilog="Questions and feedback: \
            https://github.com/deepfakes/faceswap-playground"
        )
        return parser

    def add_optional_arguments(self, parser):
        # Override this for custom arguments
        return parser

    def finalize(self):
        print('-------------------------')
        print('Images found:        {}'.format(self.images_found))
        print('Images processed:    {}'.format(self.images_processed))
        print('Faces detected:      {}'.format(self.faces_detected))
        print('-------------------------')

        if self.verify_output:
            print('Note:')
            print('Multiple faces were detected in one or more pictures.')
            print('Double check your results.')
            print('-------------------------')
        print('Done!')