From 80cde77a6de03aaef6bd5e6e4691341f04290227 Mon Sep 17 00:00:00 2001 From: Lev Velykoivanenko Date: Wed, 9 May 2018 19:47:17 +0200 Subject: [PATCH] Adding new tool effmpeg ("easy"-ffmpeg) with gui support. Extend gui functionality to support filetypes. Re-opening PR. (#373) * Pre push commit. Add filetypes support to gui through new classes in lib/cli.py Add various new functions to tools/effmpeg.py * Finish developing basic effmpeg functionality. Ready for public alpha test. * Add ffmpy to requirements. Fix gen-vid to allow specifying a new file in GUI. Fix extract throwing an error when supplied with a valid directory. Add two new gui user pop interactions: save (allows you to create new files/directories) and nothing (disables the prompt button when it's not needed). Improve logic and argument processing in effmpeg. * Fix post merge bugs. Reformat tools.py to match the new style of faceswap.py Fix some whitespace issues. * Fix matplotlib.use() being called after pyplot was imported. * Fix various effmpeg bugs and add ability do terminate nested subprocess to GUI. effmpeg changes: Fix get-fps not printing to terminal. Fix mux-audio not working. Add verbosity option. If verbose is not specified than ffmpeg output is reduced with the -hide_banner flag. scripts/gui.py changes: Add ability to terminate nested subprocesses, i.e. the following type of process tree should now be terminated safely: gui -> command -> command-subprocess -> command-subprocess -> command-sub-subprocess * Add functionality to tools/effmpeg.py, fix some docstring and print statement issues in some files. tools/effmpeg.py: Transpose choices now display detailed name in GUI, while in cli they can still be entered as a number or the full command name. Add quiet option to effmpeg that only shows critical ffmpeg errors. Improve user input handling. lib/cli.py; scripts/convert.py; scripts/extract.py; scripts/train.py: Fix some line length issues and typos in docstrings, help text and print statements. Fix some whitespace issues. lib/cli.py: Add filetypes to '--alignments' argument. Change argument action to DirFullPaths where appropriate. * Bug fixes and improvements to tools/effmpeg.py Fix bug where duration would not be used even when end time was not set. Add option to specify output filetype for extraction. Enchance gen-vid to be able to generate a video from images that were zero padded to any arbitrary number, and not just 5. Enchance gen-vid to be able to use any of the image formats that a video can be extracted into. Improve gen-vid output video quality. Minor code quality improvements and ffmpeg argument formatting improvements. * Remove dependency on psutil in scripts/gui.py and various small improvements. lib/utils.py: Add _image_extensions and _video_extensions as global variables to make them easily portable across all of faceswap. Fix lack of new lines between function and class declarions to conform to PEP8. Fix some typos and line length issues in doctsrings and comments. scripts/convert.py: Make tqdm print to stdout. scripts/extract.py: Make tqdm print to stdout. Apply workaround for occasional TqdmSynchronisationWarning being thrown. Fix some typos and line length issues in doctsrings and comments. scripts/fsmedia.py: Did TODO in scripts/fsmedia.py in Faces.load_extractor(): TODO Pass extractor_name as argument Fix lack of new lines between function and class declarions to conform to PEP8. Fix some typos and line length issues in doctsrings and comments. Change 2 print statements to use format() for string formatting instead of the old '%'. scripts/gui.py: Refactor subprocess generation and termination to remove dependency on psutil. Fix some typos and line length issues in comments. tools/effmpeg.py Refactor DataItem class to use new lib/utils.py global media file extensions. Improve ffmpeg subprocess termination handling. --- faceswap.py | 3 +- lib/cli.py | 336 +++++++++--- lib/utils.py | 19 +- requirements-gpu-python35-cuda8.txt | 1 + requirements-gpu-python36-cuda9.txt | 1 + requirements-python35.txt | 1 + requirements-python36.txt | 1 + scripts/convert.py | 21 +- scripts/extract.py | 12 +- scripts/fsmedia.py | 24 +- scripts/gui.py | 156 ++++-- scripts/train.py | 25 +- tools.py | 33 +- tools/effmpeg.py | 777 ++++++++++++++++++++++++++++ tools/sort.py | 13 +- 15 files changed, 1247 insertions(+), 176 deletions(-) create mode 100644 tools/effmpeg.py diff --git a/faceswap.py b/faceswap.py index 0519928..0921777 100755 --- a/faceswap.py +++ b/faceswap.py @@ -4,7 +4,7 @@ import sys import lib.cli as cli -if sys.version_info[0] < 3: +if sys.version_info[0] < 3: raise Exception("This program requires at least python3.2") if sys.version_info[0] == 3 and sys.version_info[1] < 2: raise Exception("This program requires at least python3.2") @@ -15,6 +15,7 @@ def bad_args(args): PARSER.print_help() exit(0) + if __name__ == "__main__": PARSER = cli.FullHelpArgumentParser() SUBPARSER = PARSER.add_subparsers() diff --git a/lib/cli.py b/lib/cli.py index ca054c4..b108f3c 100644 --- a/lib/cli.py +++ b/lib/cli.py @@ -1,16 +1,17 @@ #!/usr/bin python3 """ Command Line Arguments """ - import argparse import os import sys from plugins.PluginLoader import PluginLoader + class ScriptExecutor(object): """ Loads the relevant script modules and executes the script. - This class is initialised in each of the argparsers for the relevant command, - then execute script is called within their set_default function """ + This class is initialised in each of the argparsers for the relevant + command, then execute script is called within their set_default + function. """ def __init__(self, command, subparsers=None): self.command = command.lower() @@ -37,11 +38,134 @@ class ScriptExecutor(object): process = script(*args) process.process() + class FullPaths(argparse.Action): - """Expand user- and relative-paths""" + """ 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))) + os.path.expanduser(values))) + + +class DirFullPaths(FullPaths): + """ Class that gui uses to determine if you need to open a directory """ + pass + + +class FileFullPaths(FullPaths): + """ + Class that gui uses to determine if you need to open a file. + + Filetypes added as an argparse argument must be an iterable, i.e. a + list of lists, tuple of tuples, list of tuples etc... formatted like so: + [("File Type", ["*.ext", "*.extension"])] + A more realistic example: + [("Video File", ["*.mkv", "mp4", "webm"])] + + If the file extensions are not prepended with '*.', use the + prep_filetypes() method to format them in the arguments_list. + """ + def __init__(self, option_strings, dest, nargs=None, filetypes=None, + **kwargs): + super(FileFullPaths, self).__init__(option_strings, dest, **kwargs) + if nargs is not None: + raise ValueError("nargs not allowed") + self.filetypes = filetypes + + @staticmethod + def prep_filetypes(filetypes): + all_files = ("All Files", "*.*") + filetypes_l = list() + for i in range(len(filetypes)): + filetypes_l.append(FileFullPaths._process_filetypes(filetypes[i])) + filetypes_l.append(all_files) + return tuple(filetypes_l) + + @staticmethod + def _process_filetypes(filetypes): + """ """ + if filetypes is None: + return None + + filetypes_name = filetypes[0] + filetypes_l = filetypes[1] + if (type(filetypes_l) == list or type(filetypes_l) == tuple) \ + and all("*." in i for i in filetypes_l): + return filetypes # assume filetypes properly formatted + + if type(filetypes_l) != list and type(filetypes_l) != tuple: + raise ValueError("The filetypes extensions list was " + "neither a list nor a tuple: " + "{}".format(filetypes_l)) + + filetypes_list = list() + for i in range(len(filetypes_l)): + filetype = filetypes_l[i].strip("*.") + filetype = filetype.strip(';') + filetypes_list.append("*." + filetype) + return filetypes_name, filetypes_list + + def _get_kwargs(self): + names = [ + 'option_strings', + 'dest', + 'nargs', + 'const', + 'default', + 'type', + 'choices', + 'help', + 'metavar', + 'filetypes' + ] + return [(name, getattr(self, name)) for name in names] + + +class ComboFullPaths(FileFullPaths): + """ + Class that gui uses to determine if you need to open a file or a + directory based on which action you are choosing + """ + def __init__(self, option_strings, dest, nargs=None, filetypes=None, + actions_open_type=None, **kwargs): + if nargs is not None: + raise ValueError("nargs not allowed") + super(ComboFullPaths, self).__init__(option_strings, dest, + filetypes=None, **kwargs) + + self.actions_open_type = actions_open_type + self.filetypes = filetypes + + @staticmethod + def prep_filetypes(filetypes): + all_files = ("All Files", "*.*") + filetypes_d = dict() + for k, v in filetypes.items(): + filetypes_d[k] = () + if v is None: + filetypes_d[k] = None + continue + filetypes_l = list() + for i in range(len(v)): + filetypes_l.append(ComboFullPaths._process_filetypes(v[i])) + filetypes_d[k] = (tuple(filetypes_l), all_files) + return filetypes_d + + def _get_kwargs(self): + names = [ + 'option_strings', + 'dest', + 'nargs', + 'const', + 'default', + 'type', + 'choices', + 'help', + 'metavar', + 'filetypes', + 'actions_open_type' + ] + return [(name, getattr(self, name)) for name in names] + class FullHelpArgumentParser(argparse.ArgumentParser): """ Identical to the built-in argument parser, but on error it @@ -51,6 +175,7 @@ class FullHelpArgumentParser(argparse.ArgumentParser): args = {"prog": self.prog, "message": message} self.exit(2, "%(prog)s: error: %(message)s\n" % args) + class FaceSwapArgs(object): """ Faceswap argument parser functions that are universal to all commands. Should be the parent function of all @@ -67,16 +192,16 @@ class FaceSwapArgs(object): @staticmethod def get_argument_list(): - """ Put the arguments in a list so that they are accessible from both argparse and gui - overide for command specific arguments """ + """ Put the arguments in a list so that they are accessible from both + argparse and gui override for command specific arguments """ argument_list = [] return argument_list @staticmethod def get_optional_arguments(): - """ Put the arguments in a list so that they are accessible from both argparse and gui - This is used for when there are sub-children (e.g. convert and extract) - Override this for custom arguments """ + """ Put the arguments in a list so that they are accessible from both + argparse and gui. This is used for when there are sub-children + (e.g. convert and extract) Override this for custom arguments """ argument_list = [] return argument_list @@ -109,37 +234,48 @@ class ExtractConvertArgs(FaceSwapArgs): @staticmethod def get_argument_list(): - """ Put the arguments in a list so that they are accessible from both argparse and gui """ - argument_list = [] + """ Put the arguments in a list so that they are accessible from both + argparse and gui """ + alignments_filetypes = [["Serializers", ['json', 'yaml', 'p']], + ["JSON", ["json"]], + ["Pickle", ["p"]], + ["YAML", ["yaml"]]] + alignments_filetypes = FileFullPaths.prep_filetypes(alignments_filetypes) + argument_list = list() argument_list.append({"opts": ("-i", "--input-dir"), - "action": FullPaths, + "action": DirFullPaths, "dest": "input_dir", "default": "input", - "help": "Input directory. A directory containing the files " - "you wish to process. Defaults to 'input'"}) + "help": "Input directory. A directory " + "containing the files you wish to " + "process. Defaults to 'input'"}) argument_list.append({"opts": ("-o", "--output-dir"), - "action": FullPaths, + "action": DirFullPaths, "dest": "output_dir", "default": "output", - "help": "Output directory. This is where the converted files will " - "be stored. Defaults to 'output'"}) + "help": "Output directory. This is where the " + "converted files will be stored. " + "Defaults to 'output'"}) argument_list.append({"opts": ("--alignments", ), + "action": FileFullPaths, + "filetypes": alignments_filetypes, "type": str, "dest": "alignments_path", "help": "optional path to alignments file"}) argument_list.append({"opts": ("--serializer", ), "type": str.lower, "dest": "serializer", - "choices": ("yaml", "json", "pickle"), + "choices": ("json", "pickle", "yaml"), "help": "serializer for alignments file"}) argument_list.append({"opts": ("-D", "--detector"), "type": str, # case sensitive because this is used to load a plugin. "choices": ("hog", "cnn", "all"), "default": "hog", - "help": "Detector to use. 'cnn' detects many more angles but " - "will be much more resource intensive and may fail " - "on large files"}) + "help": "Detector to use. 'cnn' detects many " + "more angles but will be much more " + "resource intensive and may fail on " + "large files"}) argument_list.append({"opts": ("-l", "--ref_threshold"), "type": float, "dest": "ref_threshold", @@ -150,17 +286,19 @@ class ExtractConvertArgs(FaceSwapArgs): "dest": "nfilter", "nargs": "+", "default": None, - "help": "Reference image for the persons you do not want to " - "process. Should be a front portrait. Multiple images" - "can be added space separated"}) + "help": "Reference image for the persons you do " + "not want to process. Should be a front " + "portrait. Multiple images can be added " + "space separated"}) argument_list.append({"opts": ("-f", "--filter"), "type": str, "dest": "filter", "nargs": "+", "default": None, - "help": "Reference images for the person you want to process. " - "Should be a front portrait. Multiple images" - "can be added space separated"}) + "help": "Reference images for the person you " + "want to process. Should be a front " + "portrait. Multiple images can be added " + "space separated"}) argument_list.append({"opts": ("-v", "--verbose"), "action": "store_true", "dest": "verbose", @@ -168,6 +306,7 @@ class ExtractConvertArgs(FaceSwapArgs): "help": "Show verbose output"}) return argument_list + class ExtractArgs(ExtractConvertArgs): """ Class to parse the command line arguments for extraction. Inherits base options from ExtractConvertArgs where arguments @@ -175,30 +314,37 @@ class ExtractArgs(ExtractConvertArgs): @staticmethod def get_optional_arguments(): - """ Put the arguments in a list so that they are accessible from both argparse and gui """ + """ Put the arguments in a list so that they are accessible from both + argparse and gui """ argument_list = [] argument_list.append({"opts": ("-r", "--rotate-images"), "type": str, "dest": "rotate_images", "default": None, - "help": "If a face isn't found, rotate the images to try to " - "find a face. Can find more faces at the cost of extraction " - "speed. Pass in a single number to use increments of that " - "size up to 360, or pass in a list of numbers to enumerate " + "help": "If a face isn't found, rotate the " + "images to try to find a face. Can find " + "more faces at the cost of extraction " + "speed. Pass in a single number to use " + "increments of that size up to 360, or " + "pass in a list of numbers to enumerate " "exactly what angles to check"}) argument_list.append({"opts": ("-bt", "--blur-threshold"), "type": int, "dest": "blur_thresh", "default": None, - "help": "Automatically discard images blurrier than the specified " - "threshold. Discarded images are moved into a \"blurry\" " - "sub-folder. Lower values allow more blur"}) + "help": "Automatically discard images blurrier " + "than the specified threshold. " + "Discarded images are moved into a " + "\"blurry\" sub-folder. Lower values " + "allow more blur"}) argument_list.append({"opts": ("-j", "--processes"), "type": int, "default": 1, - "help": "Number of CPU processes to use. WARNING: ONLY USE THIS " - " IF YOU ARE NOT EXTRACTING ON A GPU. Anything above 1 " - " process on a GPU will run out of memory and will crash"}) + "help": "Number of CPU processes to use. " + "WARNING: ONLY USE THIS IF YOU ARE NOT " + "EXTRACTING ON A GPU. Anything above 1 " + "process on a GPU will run out of " + "memory and will crash"}) argument_list.append({"opts": ("-s", "--skip-existing"), "action": "store_true", "dest": "skip_existing", @@ -213,10 +359,12 @@ class ExtractArgs(ExtractConvertArgs): "action": "store_true", "dest": "align_eyes", "default": False, - "help": "Perform extra alignment to ensure left/right eyes " - "are at the same height"}) + "help": "Perform extra alignment to ensure " + "left/right eyes are at the same " + "height"}) return argument_list + class ConvertArgs(ExtractConvertArgs): """ Class to parse the command line arguments for conversion. Inherits base options from ExtractConvertArgs where arguments @@ -224,23 +372,27 @@ class ConvertArgs(ExtractConvertArgs): @staticmethod def get_optional_arguments(): - """ Put the arguments in a list so that they are accessible from both argparse and gui """ + """ Put the arguments in a list so that they are accessible from both + argparse and gui """ argument_list = [] argument_list.append({"opts": ("-m", "--model-dir"), - "action": FullPaths, + "action": DirFullPaths, "dest": "model_dir", "default": "models", - "help": "Model directory. A directory containing the trained model " - "you wish to process. Defaults to 'models'"}) + "help": "Model directory. A directory " + "containing the trained model you wish " + "to process. Defaults to 'models'"}) argument_list.append({"opts": ("-a", "--input-aligned-dir"), - "action": FullPaths, + "action": DirFullPaths, "dest": "input_aligned_dir", "default": None, - "help": "Input \"aligned directory\". A directory that should " - "contain the aligned faces extracted from the input files. " - "If you delete faces from this folder, they'll be skipped " - "during conversion. If no aligned dir is specified, all " - "faces will be converted"}) + "help": "Input \"aligned directory\". A " + "directory that should contain the " + "aligned faces extracted from the input " + "files. If you delete faces from this " + "folder, they'll be skipped during " + "conversion. If no aligned dir is " + "specified, all faces will be converted"}) argument_list.append({"opts": ("-t", "--trainer"), "type": str, # case sensitive because this is used to load a plug-in. @@ -261,12 +413,14 @@ class ConvertArgs(ExtractConvertArgs): "dest": "erosion_kernel_size", "type": int, "default": None, - "help": "Erosion kernel size. Positive values apply erosion " - "which reduces the edge of the swapped face. Negative " - "values apply dilation which allows the swapped face " - "to cover more space. (Masked converter only)"}) + "help": "Erosion kernel size. Positive values " + "apply erosion which reduces the edge " + "of the swapped face. Negative values " + "apply dilation which allows the " + "swapped face to cover more space. " + "(Masked converter only)"}) argument_list.append({"opts": ("-M", "--mask-type"), - #lowercase this, because its just a string later on. + # lowercase this, because it's just a string later on. "type": str.lower, "dest": "mask_type", "choices": ["rect", "facehull", "facehullandrect"], @@ -277,8 +431,9 @@ class ConvertArgs(ExtractConvertArgs): "dest": "sharpen_image", "choices": ["bsharpen", "gsharpen"], "default": None, - "help": "Use Sharpen Image.bsharpen for Box Blur, gsharpen for " - "Gaussian Blur (Masked converter only)"}) + "help": "Use Sharpen Image. bsharpen for Box " + "Blur, gsharpen for Gaussian Blur " + "(Masked converter only)"}) argument_list.append({"opts": ("-g", "--gpus"), "type": int, "default": 1, @@ -286,16 +441,18 @@ class ConvertArgs(ExtractConvertArgs): argument_list.append({"opts": ("-fr", "--frame-ranges"), "nargs": "+", "type": str, - "help": "frame ranges to apply transfer to e.g. For frames 10 to " - "50 and 90 to 100 use --frame-ranges 10-50 90-100. Files " - "must have the frame-number as the last number in the " - "name!"}) + "help": "frame ranges to apply transfer to e.g. " + "For frames 10 to 50 and 90 to 100 use " + "--frame-ranges 10-50 90-100. Files " + "must have the frame-number as the last " + "number in the name!"}) argument_list.append({"opts": ("-d", "--discard-frames"), "action": "store_true", "dest": "discard_frames", "default": False, - "help": "When used with --frame-ranges discards frames that are " - "not processed instead of writing them out unchanged"}) + "help": "When used with --frame-ranges discards " + "frames that are not processed instead " + "of writing them out unchanged"}) argument_list.append({"opts": ("-s", "--swap-model"), "action": "store_true", "dest": "swap_model", @@ -323,42 +480,49 @@ class ConvertArgs(ExtractConvertArgs): "help": "Average color adjust. (Adjust converter only)"}) return argument_list + class TrainArgs(FaceSwapArgs): """ Class to parse the command line arguments for training """ @staticmethod def get_argument_list(): - """ Put the arguments in a list so that they are accessible from both argparse and gui """ - argument_list = [] + """ Put the arguments in a list so that they are accessible from both + argparse and gui """ + argument_list = list() argument_list.append({"opts": ("-A", "--input-A"), - "action": FullPaths, + "action": DirFullPaths, "dest": "input_A", "default": "input_A", - "help": "Input directory. A directory containing training images " - "for face A. Defaults to 'input'"}) + "help": "Input directory. A directory " + "containing training images for face A. " + "Defaults to 'input'"}) argument_list.append({"opts": ("-B", "--input-B"), - "action": FullPaths, + "action": DirFullPaths, "dest": "input_B", "default": "input_B", - "help": "Input directory. A directory containing training images " - "for face B Defaults to 'input'"}) + "help": "Input directory. A directory " + "containing training images for face B. " + "Defaults to 'input'"}) argument_list.append({"opts": ("-m", "--model-dir"), - "action": FullPaths, + "action": DirFullPaths, "dest": "model_dir", "default": "models", - "help": "Model directory. This is where the training data will " - "be stored. Defaults to 'model'"}) + "help": "Model directory. This is where the " + "training data will be stored. " + "Defaults to 'model'"}) argument_list.append({"opts": ("-s", "--save-interval"), "type": int, "dest": "save_interval", "default": 100, - "help": "Sets the number of iterations before saving the model"}) + "help": "Sets the number of iterations before " + "saving the model"}) argument_list.append({"opts": ("-t", "--trainer"), "type": str, "choices": PluginLoader.get_available_models(), "default": PluginLoader.get_default_model(), - "help": "Select which trainer to use, Use LowMem for cards with " - " less than 2GB of VRAM"}) + "help": "Select which trainer to use, Use " + "LowMem for cards with less than 2GB of " + "VRAM"}) argument_list.append({"opts": ("-bs", "--batch-size"), "type": int, "default": 64, @@ -375,14 +539,14 @@ class TrainArgs(FaceSwapArgs): "action": "store_true", "dest": "preview", "default": False, - "help": "Show preview output. If not specified, write progress " - "to file"}) + "help": "Show preview output. If not specified, " + "write progress to file"}) argument_list.append({"opts": ("-w", "--write-image"), "action": "store_true", "dest": "write_image", "default": False, - "help": "Writes the training result to a file even on " - "preview mode"}) + "help": "Writes the training result to a file " + "even on preview mode"}) argument_list.append({"opts": ("-pl", "--use-perceptual-loss"), "action": "store_true", "dest": "perceptual_loss", @@ -392,8 +556,8 @@ class TrainArgs(FaceSwapArgs): "action": "store_true", "dest": "allow_growth", "default": False, - "help": "Sets allow_growth option of Tensorflow to spare memory " - "on some configs"}) + "help": "Sets allow_growth option of Tensorflow " + "to spare memory on some configs"}) argument_list.append({"opts": ("-v", "--verbose"), "action": "store_true", "dest": "verbose", @@ -408,12 +572,14 @@ class TrainArgs(FaceSwapArgs): "help": argparse.SUPPRESS}) return argument_list + class GuiArgs(FaceSwapArgs): """ Class to parse the command line arguments for training """ @staticmethod def get_argument_list(): - """ Put the arguments in a list so that they are accessible from both argparse and gui """ + """ Put the arguments in a list so that they are accessible from both + argparse and gui """ argument_list = [] argument_list.append({"opts": ("-d", "--debug"), "action": "store_true", diff --git a/lib/utils.py b/lib/utils.py index 94ed269..893c84e 100644 --- a/lib/utils.py +++ b/lib/utils.py @@ -9,15 +9,22 @@ import warnings from pathlib import Path + +# Global variables +_image_extensions = ['.bmp', '.jpeg', '.jpg', '.png', '.tif', '.tiff'] +_video_extensions = ['.avi', '.flv', '.mkv', '.mov', '.mp4', '.mpeg', '.webm'] + + def get_folder(path): """ Return a path to a folder, creating it if it doesn't exist """ output_dir = Path(path) output_dir.mkdir(parents=True, exist_ok=True) return output_dir + def get_image_paths(directory, exclude=list(), debug=False): """ Return a list of images that reside in a folder """ - image_extensions = [".jpg", ".jpeg", ".png", ".tif", ".tiff"] + image_extensions = _image_extensions exclude_names = [basename(Path(x).stem[:Path(x).stem.rfind('_')] + Path(x).suffix) for x in exclude] dir_contents = list() @@ -37,6 +44,7 @@ def get_image_paths(directory, exclude=list(), debug=False): return dir_contents + def backup_file(directory, filename): """ Backup a given file by appending .bk to the end """ origfile = join(directory, filename) @@ -46,6 +54,7 @@ def backup_file(directory, filename): if exists(origfile): os.rename(origfile, backupfile) + def set_system_verbosity(loglevel): """ Set the verbosity level of tensorflow and suppresses future and deprecation warnings from any modules @@ -63,10 +72,11 @@ def set_system_verbosity(loglevel): for warncat in (FutureWarning, DeprecationWarning): warnings.simplefilter(action='ignore', category=warncat) + class BackgroundGenerator(threading.Thread): """ Run a queue in the background. From: https://stackoverflow.com/questions/7323664/python-generator-pre-fetch """ - def __init__(self, generator, prefetch=1): #See below why prefetch count is flawed + def __init__(self, generator, prefetch=1): # See below why prefetch count is flawed threading.Thread.__init__(self) self.queue = Queue.Queue(prefetch) self.generator = generator @@ -75,8 +85,9 @@ class BackgroundGenerator(threading.Thread): def run(self): """ Put until queue size is reached. - Note: put blocks only if put is called while queue has already reached max size - => this makes 2 prefetched items! One in the queue, one waiting for insertion! """ + Note: put blocks only if put is called while queue has already + reached max size => this makes 2 prefetched items! One in the + queue, one waiting for insertion! """ for item in self.generator: self.queue.put(item) self.queue.put(None) diff --git a/requirements-gpu-python35-cuda8.txt b/requirements-gpu-python35-cuda8.txt index 8e3ad09..1728b34 100644 --- a/requirements-gpu-python35-cuda8.txt +++ b/requirements-gpu-python35-cuda8.txt @@ -4,6 +4,7 @@ h5py==2.7.1 Keras==2.1.2 opencv-python==3.3.0.10 tensorflow-gpu==1.4.0 +ffmpy==0.2.2 scikit-image dlib face_recognition diff --git a/requirements-gpu-python36-cuda9.txt b/requirements-gpu-python36-cuda9.txt index fe1066d..4093ff5 100644 --- a/requirements-gpu-python36-cuda9.txt +++ b/requirements-gpu-python36-cuda9.txt @@ -4,6 +4,7 @@ h5py==2.7.1 Keras==2.1.2 opencv-python==3.3.0.10 tensorflow-gpu==1.5.0 +ffmpy==0.2.2 scikit-image dlib face_recognition diff --git a/requirements-python35.txt b/requirements-python35.txt index e0de3ff..1ae0e61 100644 --- a/requirements-python35.txt +++ b/requirements-python35.txt @@ -4,6 +4,7 @@ h5py==2.7.1 Keras==2.1.2 opencv-python==3.3.0.10 tensorflow==1.4.1 +ffmpy==0.2.2 scikit-image dlib face_recognition diff --git a/requirements-python36.txt b/requirements-python36.txt index cd846f6..a8dfba4 100644 --- a/requirements-python36.txt +++ b/requirements-python36.txt @@ -4,6 +4,7 @@ h5py==2.7.1 Keras==2.1.2 opencv-python==3.3.0.10 tensorflow==1.5.0 +ffmpy==0.2.2 scikit-image dlib face_recognition diff --git a/scripts/convert.py b/scripts/convert.py index 397abac..d7c6c6b 100644 --- a/scripts/convert.py +++ b/scripts/convert.py @@ -3,6 +3,7 @@ import re import os +import sys from pathlib import Path from tqdm import tqdm @@ -13,6 +14,7 @@ from lib.utils import BackgroundGenerator, get_folder, get_image_paths from plugins.PluginLoader import PluginLoader + class Convert(object): """ The convert process. """ def __init__(self, arguments): @@ -27,7 +29,9 @@ class Convert(object): def process(self): """ Original & LowMem models go with Adjust or Masked converter - Note: GAN prediction outputs a mask + an image, while other predicts only an image """ + + Note: GAN prediction outputs a mask + an image, while other + predicts only an image. """ Utils.set_verbosity(self.args.verbose) if not self.alignments.have_alignments_file: @@ -89,7 +93,7 @@ class Convert(object): def prepare_images(self): """ Prepare the images for conversion """ filename = "" - for filename in tqdm(self.images.input_images): + for filename in tqdm(self.images.input_images, file=sys.stdout): if not self.check_alignments(filename): continue image = Utils.cv2_read_write('read', filename) @@ -157,18 +161,19 @@ class OptionalActions(object): input_aligned_dir = self.args.input_aligned_dir if input_aligned_dir is None: - print("Aligned directory not specified. All faces listed in the alignments file " - "will be converted") + print("Aligned directory not specified. All faces listed in the " + "alignments file will be converted") elif not os.path.isdir(input_aligned_dir): - print("Aligned directory not found. All faces listed in the alignments file " - "will be converted") + print("Aligned directory not found. All faces listed in the " + "alignments file will be converted") else: faces_to_swap = [Path(path) for path in get_image_paths(input_aligned_dir)] if not faces_to_swap: print("Aligned directory is empty, no faces will be converted!") elif len(faces_to_swap) <= len(self.input_images) / 3: - print("Aligned directory contains an amount of images much less than the input, \ - are you sure this is the right directory?") + print("Aligned directory contains an amount of images much " + "less than the input, are you sure this is the right " + "directory?") return faces_to_swap ### SKIP FRAME RANGES ### diff --git a/scripts/extract.py b/scripts/extract.py index e5cc1a5..530b9d4 100644 --- a/scripts/extract.py +++ b/scripts/extract.py @@ -2,13 +2,16 @@ """ The script to run the extract process of faceswap """ import os +import sys from pathlib import Path from tqdm import tqdm +tqdm.monitor_interval = 0 # workaround for TqdmSynchronisationWarning from lib.multithreading import pool_process from scripts.fsmedia import Alignments, Faces, Images, Utils + class Extract(object): """ The extract process. """ @@ -43,7 +46,7 @@ class Extract(object): def extract_single_process(self): """ Run extraction in a single process """ - for filename in tqdm(self.images.input_images): + for filename in tqdm(self.images.input_images, file=sys.stdout): filename, faces = self.process_single_image(filename) self.faces.faces_detected[os.path.basename(filename)] = faces @@ -52,7 +55,8 @@ class Extract(object): for filename, faces in tqdm(pool_process(self.process_single_image, self.images.input_images, processes=self.args.processes), - total=self.images.images_found): + total=self.images.images_found, + file=sys.stdout): self.faces.num_faces_detected += 1 self.faces.faces_detected[os.path.basename(filename)] = faces @@ -60,8 +64,8 @@ class Extract(object): """ Detect faces in an image. Rotate the image the specified amount until at least one face is found, or until image rotations are depleted. - Once at least one face has been detected, pass to process_single_face - to process the individual faces """ + Once at least one face has been detected, pass to + process_single_face to process the individual faces """ retval = filename, list() try: image = Utils.cv2_read_write('read', filename) diff --git a/scripts/fsmedia.py b/scripts/fsmedia.py index a27f9c5..fa4ec02 100644 --- a/scripts/fsmedia.py +++ b/scripts/fsmedia.py @@ -1,6 +1,6 @@ #!/usr/bin python3 -""" Holds the classes for the 3 main Faceswap 'media' objects for input (extract) - and output (convert) tasks. Those being: +""" Holds the classes for the 3 main Faceswap 'media' objects for + input (extract) and output (convert) tasks. Those being: Images Faces Alignments""" @@ -18,6 +18,7 @@ from lib.FaceFilter import FaceFilter from lib.utils import get_folder, get_image_paths, set_system_verbosity from plugins.PluginLoader import PluginLoader + class Utils(object): """ Holds utility functions that are required by more than one media object """ @@ -75,6 +76,7 @@ class Utils(object): print("Done!") return images_found, num_faces_detected + class Images(object): """ Holds the full frames/images """ def __init__(self, arguments): @@ -88,12 +90,12 @@ class Images(object): self.rotation_height = 0 def get_rotation_angles(self): - """ Set the rotation angles. Includes backwards compatibility for the 'on' - and 'off' options: + """ Set the rotation angles. Includes backwards compatibility for the + 'on' and 'off' options: - 'on' - increment 90 degrees - 'off' - disable - - 0 is prepended to the list, as whatever happens, we want to scan the image - in it's upright state """ + - 0 is prepended to the list, as whatever happens, we want to + scan the image in it's upright state """ rotation_angles = [0] if (not hasattr(self.args, 'rotate_images') @@ -152,6 +154,7 @@ class Images(object): rotated_height=self.rotation_height) return image + class Faces(object): """ Holds the faces """ def __init__(self, arguments): @@ -166,10 +169,8 @@ class Faces(object): self.verify_output = False @staticmethod - def load_extractor(): + def load_extractor(extractor_name="Align"): """ Load the requested extractor for extraction """ - # TODO Pass as argument - extractor_name = "Align" extractor = PluginLoader.get_extractor(extractor_name)() return extractor @@ -237,7 +238,7 @@ class Faces(object): self.num_faces_detected += 1 faces_count += 1 if faces_count > 1 and self.args.verbose: - print("Note: Found more than one face in an image! File: %s" % filename) + print("Note: Found more than one face in an image! File: {}".format(filename)) self.verify_output = True def draw_landmarks_on_face(self, face, image): @@ -267,6 +268,7 @@ class Faces(object): Path(filename).stem return blurry_file + class Alignments(object): """ Holds processes pertaining to the alignments file """ def __init__(self, arguments): @@ -328,5 +330,5 @@ class Alignments(object): if val: faces_detected[key] = val else: - print("Existing alignments file '%s' not found." % alignfile) + print("Existing alignments file '{}' not found.".format(alignfile)) return faces_detected diff --git a/scripts/gui.py b/scripts/gui.py index 6a93b73..b5144e5 100644 --- a/scripts/gui.py +++ b/scripts/gui.py @@ -14,6 +14,7 @@ from threading import Thread from time import time import matplotlib +matplotlib.use('TkAgg') import matplotlib.animation as animation from matplotlib import pyplot as plt from matplotlib import style @@ -21,11 +22,9 @@ from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg import numpy -from lib.cli import FullPaths +from lib.cli import FullPaths, ComboFullPaths, DirFullPaths, FileFullPaths from lib.Serializer import JSONSerializer -matplotlib.use('TkAgg') - PATHSCRIPT = os.path.realpath(os.path.dirname(sys.argv[0])) @@ -45,6 +44,7 @@ except ImportError: messagebox = None TclError = None + class Utils(object): """ Inter-class object holding items that are required across classes """ @@ -107,7 +107,7 @@ class Utils(object): self.change_action_button() def clear_display_panel(self): - ''' Clear the preview window and graph ''' + """ Clear the preview window and graph """ self.delete_preview() self.lossdict = dict() @@ -199,6 +199,9 @@ class Utils(object): if os.path.exists(self.previewloc): os.remove(self.previewloc) + def get_chosen_action(self, task_name): + return self.opts[task_name][0]['value'].get() + class Tooltip: """ @@ -355,6 +358,7 @@ class Tooltip: topwidget.destroy() self.topwidget = None + class FaceswapGui(object): """ The Graphical User Interface """ @@ -462,6 +466,7 @@ class FaceswapGui(object): self.gui.quit() exit() + class ConsoleOut(object): """ The Console out tab of the Display section """ @@ -508,7 +513,6 @@ class SysOutRouter(object): class CommandTab(object): - """ Tabs to hold the command options """ def __init__(self, utils, notebook, command): @@ -534,6 +538,7 @@ class CommandTab(object): sep = ttk.Frame(self.page, height=2, relief=tk.RIDGE) sep.pack(fill=tk.X, pady=(5, 0), side=tk.BOTTOM) + class OptionsFrame(object): """ Options Frame - Holds the Options for each command """ @@ -573,6 +578,7 @@ class OptionsFrame(object): canvas_width = event.width self.canvas.itemconfig(self.optscanvas, width=canvas_width) + class OptionControl(object): """ Build the correct control for the option parsed and place it on the frame """ @@ -625,7 +631,10 @@ class OptionControl(object): var.set(default) if sysbrowser is not None: + # if sysbrowser in "load file": self.add_browser_buttons(frame, sysbrowser, var) + # elif sysbrowser == "combo": + # self.add_browser_combo_button(frame, sysbrowser, var) ctlkwargs = {'variable': var} if control == ttk.Checkbutton else { 'textvariable': var} @@ -642,26 +651,80 @@ class OptionControl(object): def add_browser_buttons(self, frame, sysbrowser, filepath): """ Add correct file browser button for control """ - img = self.utils.icons[sysbrowser] + if sysbrowser == "combo": + img = self.utils.icons['load'] + else: + img = self.utils.icons[sysbrowser] action = getattr(self, 'ask_' + sysbrowser) + filetypes = self.option['filetypes'] fileopn = ttk.Button(frame, image=img, - command=lambda cmd=action: cmd(filepath)) + command=lambda cmd=action: cmd(filepath, + filetypes)) fileopn.pack(padx=(0, 5), side=tk.RIGHT) @staticmethod - def ask_folder(filepath): - """ Pop-up to get path to a folder """ + def ask_folder(filepath, filetypes=None): + """ + Pop-up to get path to a directory + :param filepath: tkinter StringVar object that will store the path to a + directory. + :param filetypes: Unused argument to allow filetypes to be given in + ask_load(). + """ dirname = filedialog.askdirectory() if dirname: filepath.set(dirname) @staticmethod - def ask_load(filepath): + def ask_load(filepath, filetypes=None): """ Pop-up to get path to a file """ - filename = filedialog.askopenfilename() + if filetypes is None: + filename = filedialog.askopenfilename() + else: + # In case filetypes were not configured properly in the + # arguments_list + try: + filename = filedialog.askopenfilename(filetypes=filetypes) + except TclError as te1: + filetypes = FileFullPaths.prep_filetypes(filetypes) + filename = filedialog.askopenfilename(filetypes=filetypes) + except TclError as te2: + filename = filedialog.askopenfilename() if filename: filepath.set(filename) + @staticmethod + def ask_save(filepath, filetypes=None): + """ Pop-up to get path to save a new file """ + if filetypes is None: + filename = filedialog.asksaveasfilename() + else: + # In case filetypes were not configured properly in the + # arguments_list + try: + filename = filedialog.asksaveasfilename(filetypes=filetypes) + except TclError as te1: + filetypes = FileFullPaths.prep_filetypes(filetypes) + filename = filedialog.asksaveasfilename(filetypes=filetypes) + except TclError as te2: + filename = filedialog.asksaveasfilename() + if filename: + filepath.set(filename) + + @staticmethod + def ask_nothing(filepath, filetypes=None): + """ Method that does nothing, used for disabling open/save pop up """ + return + + def ask_combo(self, filepath, filetypes): + actions_open_type = self.option['actions_open_type'] + task_name = actions_open_type['task_name'] + chosen_action = self.utils.get_chosen_action(task_name) + action = getattr(self, "ask_" + actions_open_type[chosen_action]) + filetypes = filetypes[chosen_action] + action(filepath, filetypes) + + class ActionFrame(object): """Action Frame - Displays information and action controls """ @@ -708,6 +771,7 @@ class ActionFrame(object): btnutl.pack(padx=2, side=tk.LEFT) Tooltip(btnutl, text=utl.capitalize() + ' ' + self.title + ' config', wraplength=200) + class DisplayTab(object): """ The display tabs """ @@ -792,6 +856,7 @@ class GraphDisplay(object): for child in self.graphpane.panes(): self.graphpane.remove(child) + class Graph(object): """ Each graph to be displayed. Until training is run it is not known how many graphs will be required, so they sit in their own class @@ -857,7 +922,7 @@ class Graph(object): self.trend_plot(xrng, loss) def recalculate_axes(self, loss): - ''' Recalculate the latest x and y axes limits from latest data ''' + """ Recalculate the latest x and y axes limits from latest data """ ymin = floor(min([min(lossvals) for lossvals in loss]) * 100) / 100 ymax = ceil(max([max(lossvals) for lossvals in loss]) * 100) / 100 @@ -872,17 +937,18 @@ class Graph(object): return xlim def raw_plot(self, x_range, loss): - ''' Raw value plotting ''' + """ Raw value plotting """ for idx, lossvals in enumerate(loss): self.losslines[idx].set_data(x_range, lossvals) def trend_plot(self, x_range, loss): - ''' Trend value plotting ''' + """ Trend value plotting """ for idx, lossvals in enumerate(loss): fit = numpy.polyfit(x_range, lossvals, 3) poly = numpy.poly1d(fit) self.trndlines[idx].set_data(x_range, poly(x_range)) + class PreviewDisplay(object): """ The Preview tab of the Display section """ @@ -927,6 +993,7 @@ class PreviewDisplay(object): class FaceswapControl(object): """ Control the underlying Faceswap tasks """ + __group_processes = ["effmpeg"] def __init__(self, utils, calling_file="faceswap.py"): self.pathexecscript = os.path.join(PATHSCRIPT, calling_file) @@ -968,6 +1035,8 @@ class FaceswapControl(object): 'stderr': PIPE, 'bufsize': 1, 'universal_newlines': True} + if self.command in self.__group_processes: + kwargs['preexec_fn'] = os.setsid if os.name == 'nt': kwargs['creationflags'] = subprocess.CREATE_NEW_PROCESS_GROUP self.process = Popen(self.args, **kwargs) @@ -1016,7 +1085,7 @@ class FaceswapControl(object): def capture_loss(self, string): """ Capture loss values from stdout """ - #TODO: Remove this hideous hacky fix. When the subprocess is terminated and + # TODO: Remove this hideous hacky fix. When the subprocess is terminated and # the loss dictionary is reset, 1 set of loss values ALWAYS slips through # and appends to the lossdict AFTER the subprocess has closed meaning that # checks on whether the dictionary is empty fail. @@ -1027,7 +1096,7 @@ class FaceswapControl(object): # sys.exit() on the stdout/err threads (no effect) # sys.stdout/stderr.flush (no effect) # thread.join (locks the whole process up, because the stdout thread - # stubbonly refuses to release it's last line) + # stubbornly refuses to release its last line) currentlenloss = len(self.utils.lossdict) if self.lenloss > currentlenloss: @@ -1066,15 +1135,27 @@ class FaceswapControl(object): return except ValueError as err: print(err) - print('Terminating Process...') - try: - self.process.terminate() - self.process.wait(timeout=10) - print('Terminated') - except TimeoutExpired: - print('Termination timed out. Killing Process...') - self.process.kill() - print('Killed') + elif self.command in self.__group_processes: + print('Terminating Process Group...') + pgid = os.getpgid(self.process.pid) + try: + os.killpg(pgid, signal.SIGINT) + self.process.wait(timeout=10) + print('Terminated') + except TimeoutExpired: + print('Termination timed out. Killing Process Group...') + os.killpg(pgid, signal.SIGKILL) + print('Killed') + else: + print('Terminating Process...') + try: + self.process.terminate() + self.process.wait(timeout=10) + print('Terminated') + except TimeoutExpired: + print('Termination timed out. Killing Process...') + self.process.kill() + print('Killed') def set_final_status(self, returncode): """ Set the status bar output based on subprocess return code """ @@ -1101,6 +1182,10 @@ class Gui(object): cmd = sys.argv + # If not running in gui mode return before starting to create a window + if 'gui' not in cmd: + return + self.args = arguments self.opts = self.extract_options(subparsers) self.utils = Utils(self.opts, calling_file=cmd[0]) @@ -1147,11 +1232,13 @@ class Gui(object): for opt in command: if opt.get('help', '') == SUPPRESS: command.remove(opt) - ctl, sysbrowser = self.set_control(opt) + ctl, sysbrowser, filetypes, actions_open_types = self.set_control(opt) opt['control_title'] = self.set_control_title( opt.get('opts', '')) opt['control'] = ctl opt['filesystem_browser'] = sysbrowser + opt['filetypes'] = filetypes + opt['actions_open_types'] = actions_open_types return opts @staticmethod @@ -1165,16 +1252,25 @@ class Gui(object): def set_control(option): """ Set the control and filesystem browser to use for each option """ sysbrowser = None + filetypes = None + actions_open_type = None ctl = ttk.Entry - if option.get('dest', '') == 'alignments_path': - sysbrowser = 'load' - elif option.get('action', '') == FullPaths: + if option.get('action', '') == FullPaths: sysbrowser = 'folder' + elif option.get('action', '') == DirFullPaths: + sysbrowser = 'folder' + elif option.get('action', '') == FileFullPaths: + sysbrowser = 'load' + filetypes = option.get('filetypes', None) + elif option.get('action', '') == ComboFullPaths: + sysbrowser = 'combo' + actions_open_type = option['actions_open_type'] + filetypes = option.get('filetypes', None) elif option.get('choices', '') != '': ctl = ttk.Combobox elif option.get('action', '') == 'store_true': ctl = ttk.Checkbutton - return ctl, sysbrowser + return ctl, sysbrowser, filetypes, actions_open_type def process(self): """ Builds the GUI """ diff --git a/scripts/train.py b/scripts/train.py index 96be29f..2f411de 100644 --- a/scripts/train.py +++ b/scripts/train.py @@ -18,7 +18,6 @@ class Train(object): def __init__(self, arguments): self.args = arguments self.images = self.get_images() - self.stop = False self.save_now = False self.preview_buffer = dict() @@ -31,10 +30,8 @@ class Train(object): def process(self): """ Call the training process object """ print("Training data directory: {}".format(self.args.model_dir)) - lvl = '0' if self.args.verbose else '2' set_system_verbosity(lvl) - thread = self.start_thread() if self.args.preview: @@ -45,7 +42,8 @@ class Train(object): self.end_thread(thread) def get_images(self): - """ Check the image dirs exist, contain images and return the image objects """ + """ Check the image dirs exist, contain images and return the image + objects """ images = [] for image_dir in [self.args.input_A, self.args.input_B]: if not os.path.isdir(image_dir): @@ -69,9 +67,10 @@ class Train(object): def end_thread(self, thread): """ On termination output message and join thread back to main """ - print("Exit requested! The trainer will complete its current cycle, save " - "the models and quit (it can take up a couple of seconds depending " - "on your training speed). If you want to kill it now, press Ctrl + c") + print("Exit requested! The trainer will complete its current cycle, " + "save the models and quit (it can take up a couple of seconds " + "depending on your training speed). If you want to kill it now, " + "press Ctrl + c") self.stop = True thread.join() sys.stdout.flush() @@ -106,7 +105,7 @@ class Train(object): return model def load_trainer(self, model): - """ Load the trainer requested for traning """ + """ Load the trainer requested for training """ images_a, images_b = self.images trainer = PluginLoader.get_trainer(self.trainer_name) @@ -137,7 +136,8 @@ class Train(object): """ Generate the preview window and wait for keyboard input """ print("Using live preview.\n" "Press 'ENTER' on the preview window to save and quit.\n" - "Press 'S' on the preview window to save model weights immediately") + "Press 'S' on the preview window to save model weights " + "immediately") while True: try: with self.lock: @@ -158,11 +158,12 @@ class Train(object): def monitor_console(): """ Monitor the console for any input followed by enter or ctrl+c """ # TODO: how to catch a specific key instead of Enter? - # there isnt a good multiplatform solution: + # there isn't a good multiplatform solution: # https://stackoverflow.com/questions/3523174 # TODO: Find a way to interrupt input() if the target iterations are reached. - # At the moment, setting a target iteration and using the -p flag is the only guaranteed - # way to exit the training loop on hitting target iterations. """ + # At the moment, setting a target iteration and using the -p flag is + # the only guaranteed way to exit the training loop on hitting target + # iterations. print("Starting. Press 'ENTER' to stop training and save model") try: input() diff --git a/tools.py b/tools.py index 54e5fde..3453d11 100755 --- a/tools.py +++ b/tools.py @@ -1,9 +1,9 @@ #!/usr/bin/env python3 import sys -from lib.cli import FullHelpArgumentParser # Importing the various tools from tools.sort import SortProcessor -from scripts.gui import TKGui +from tools.effmpeg import Effmpeg +import lib.cli as cli # Python version check if sys.version_info[0] < 3: @@ -13,7 +13,7 @@ if sys.version_info[0] == 3 and sys.version_info[1] < 2: def bad_args(args): - parser.print_help() + PARSER.print_help() exit(0) @@ -23,15 +23,18 @@ if __name__ == "__main__": _tools_warning += "understand how it works." print(_tools_warning) - parser = FullHelpArgumentParser() - subparser = parser.add_subparsers() - sort = SortProcessor( - subparser, "sort", "This command lets you sort images using various " - "methods.") - guiparsers = {'sort': sort} - print(__file__) - gui = TKGui( - subparser, guiparsers, "gui", "Launch the Faceswap Tools Graphical User Interface.") - parser.set_defaults(func=bad_args) - arguments = parser.parse_args() - arguments.func(arguments) + PARSER = cli.FullHelpArgumentParser() + SUBPARSER = PARSER.add_subparsers() + EFFMPEG = Effmpeg( + SUBPARSER, "effmpeg", + "This command allows you to easily execute common ffmpeg tasks.") + SORT = SortProcessor( + SUBPARSER, "sort", + "This command lets you sort images using various methods.") + GUIPARSERS = {'effmpeg': EFFMPEG, 'sort': SORT} + GUI = cli.GuiArgs( + SUBPARSER, "gui", + "Launch the Faceswap Tools Graphical User Interface.", GUIPARSERS) + PARSER.set_defaults(func=bad_args) + ARGUMENTS = PARSER.parse_args() + ARGUMENTS.func(ARGUMENTS) diff --git a/tools/effmpeg.py b/tools/effmpeg.py new file mode 100644 index 0000000..92f32e4 --- /dev/null +++ b/tools/effmpeg.py @@ -0,0 +1,777 @@ +#!/usr/bin/env python3 +# vim: set fileencoding=utf-8 : +""" +Created on 2018-03-16 15:14 + +@author: Lev Velykoivanenko (velykoivanenko.lev@gmail.com) +""" +import argparse +import os +import sys +import subprocess +import datetime + +from ffmpy import FFprobe, FFmpeg, FFRuntimeError + +# faceswap imports +from lib.cli import FileFullPaths, ComboFullPaths +from lib.utils import _image_extensions, _video_extensions + + +if sys.version_info[0] < 3: + raise Exception("This program requires at least python3.2") +if sys.version_info[0] == 3 and sys.version_info[1] < 2: + raise Exception("This program requires at least python3.2") + + +class DataItem(object): + """ + A simple class used for storing the media data items and directories that + Effmpeg uses for 'input', 'output' and 'ref_vid'. + """ + vid_ext = _video_extensions + # future option in effmpeg to use audio file for muxing + audio_ext = ['.aiff', '.flac', '.mp3', '.wav'] + img_ext = _image_extensions + + def __init__(self, path=None, name=None, item_type=None, ext=None, + fps=None): + self.path = path + self.name = name + self.type = item_type + self.ext = ext + self.fps = fps + self.dirname = None + self.set_type_ext(path) + self.set_dirname(self.path) + self.set_name(name) + if self.is_type("vid") and self.fps is None: + self.set_fps() + + def set_name(self, name=None): + if name is None and self.path is not None: + self.name = os.path.basename(self.path) + elif name is not None and self.path is None: + self.name = os.path.basename(name) + elif name is not None and self.path is not None: + self.name = os.path.basename(name) + else: + self.name = None + + def set_type_ext(self, path=None): + if path is not None: + self.path = path + if self.path is not None: + item_ext = os.path.splitext(self.path)[1] + if item_ext in DataItem.vid_ext: + item_type = 'vid' + elif item_ext in DataItem.audio_ext: + item_type = 'audio' + else: + item_type = 'dir' + self.type = item_type + self.ext = item_ext + else: + return + + def set_dirname(self, path=None): + if path is None and self.path is not None: + self.dirname = os.path.dirname(self.path) + elif path is not None and self.path is None: + self.dirname = os.path.dirname(path) + elif path is not None and self.path is not None: + self.dirname = os.path.dirname(path) + else: + self.dirname = None + + def is_type(self, item_type=None): + if item_type == "media": + return self.type in "vid audio" + elif item_type == "dir": + return self.type == "dir" + elif item_type == "vid": + return self.type == "vid" + elif item_type == "audio": + return self.type == "audio" + elif item_type.lower() == "none": + return self.type is None + else: + return False + + def set_fps(self): + try: + self.fps = Effmpeg.get_fps(self.path) + except FFRuntimeError: + self.fps = None + + +class Effmpeg(object): + """ + Class that allows for "easy" ffmpeg use. It provides a nice cli interface + for common video operations. + """ + + _actions_req_fps = ["extract", "gen_vid"] + _actions_req_ref_video = ["mux_audio"] + _actions_can_use_ref_video = ["gen_vid"] + _actions_have_dir_output = ["extract"] + _actions_have_vid_output = ["gen_vid", "mux_audio", "rescale", "rotate", + "slice"] + _actions_have_print_output = ["get_fps", "get_info"] + _actions_have_dir_input = ["gen_vid"] + _actions_have_vid_input = ["extract", "get_fps", "get_info", "rescale", + "rotate", "slice"] + + # Class variable that stores the common ffmpeg arguments based on verbosity + __common_ffmpeg_args_dict = {"normal": "-hide_banner ", + "quiet": "-loglevel panic -hide_banner ", + "verbose": ''} + + # _common_ffmpeg_args is the class variable that will get used by various + # actions and it will be set by the process_arguments() method based on + # passed verbosity + _common_ffmpeg_args = '' + + def __init__(self, subparser, command, description='default'): + self.argument_list = self.get_argument_list() + self.optional_arguments = list() + self.args = None + self.input = DataItem() + self.output = DataItem() + self.ref_vid = DataItem() + self.start = "" + self.end = "" + self.duration = "" + self.print_ = False + self.parse_arguments(description, subparser, command) + + @staticmethod + def get_argument_list(): + vid_files = FileFullPaths.prep_filetypes([["Video Files", + DataItem.vid_ext]]) + arguments_list = list() + arguments_list.append({"opts": ('-a', '--action'), + "dest": "action", + "choices": ("extract", "gen-vid", "get-fps", + "get-info", "mux-audio", "rescale", + "rotate", "slice"), + "default": "extract", + "help": """Choose which action you want ffmpeg + ffmpeg to do. + 'slice' cuts a portion of the video + into a separate video file. + 'get-fps' returns the chosen video's + fps."""}) + + arguments_list.append({"opts": ('-i', '--input'), + "action": ComboFullPaths, + "dest": "input", + "default": "input", + "help": "Input file.", + "required": True, + "actions_open_type": { + "task_name": "effmpeg", + "extract": "load", + "gen-vid": "folder", + "get-fps": "load", + "get-info": "load", + "mux-audio": "load", + "rescale": "load", + "rotate": "load", + "slice": "load", + }, + "filetypes": { + "extract": vid_files, + "gen-vid": None, + "get-fps": vid_files, + "get-info": vid_files, + "mux-audio": vid_files, + "rescale": vid_files, + "rotate": vid_files, + "slice": vid_files + }}) + + arguments_list.append({"opts": ('-o', '--output'), + "action": ComboFullPaths, + "dest": "output", + "default": "", + "help": """Output file. If no output is + specified then: if the output is + meant to be a video then a video + called 'out.mkv' will be created in + the input directory; if the output is + meant to be a directory then a + directory called 'out' will be + created inside the input + directory. + Note: the chosen output file + extension will determine the file + encoding.""", + "actions_open_type": { + "task_name": "effmpeg", + "extract": "save", + "gen-vid": "save", + "get-fps": "nothing", + "get-info": "nothing", + "mux-audio": "save", + "rescale": "save", + "rotate": "save", + "slice": "save" + }, + "filetypes": { + "extract": None, + "gen-vid": vid_files, + "get-fps": None, + "get-info": None, + "mux-audio": vid_files, + "rescale": vid_files, + "rotate": vid_files, + "slice": vid_files + }}) + + arguments_list.append({"opts": ('-r', '--reference-video'), + "action": ComboFullPaths, + "dest": "ref_vid", + "default": "None", + "help": """Path to reference video if 'input' + was not a video.""", + "actions_open_type": { + "task_name": "effmpeg", + "extract": "nothing", + "gen-vid": "load", + "get-fps": "nothing", + "get-info": "nothing", + "mux-audio": "load", + "rescale": "nothing", + "rotate": "nothing", + "slice": "nothing" + }, + "filetypes": { + "extract": None, + "gen-vid": vid_files, + "get-fps": None, + "get-info": None, + "mux-audio": vid_files, + "rescale": None, + "rotate": None, + "slice": None + }}) + + arguments_list.append({"opts": ('-fps', '--fps'), + "type": str, + "dest": "fps", + "default": "-1.0", + "help": """Provide video fps. Can be an integer, + float or fraction. Negative values + will make the program try to get the + fps from the input or reference + videos."""}) + + arguments_list.append({"opts": ("-ef", "--extract-filetype"), + "choices": DataItem.img_ext, + "dest": "extract_ext", + "default": ".png", + "help": """Image format that extracted images + should be saved as. '.bmp' will offer + the fastest extraction speed, but + will take the most storage space. + '.png' will be slower but will take + less storage."""}) + + arguments_list.append({"opts": ('-s', '--start'), + "type": str, + "dest": "start", + "default": "00:00:00", + "help": """Enter the start time from which an + action is to be applied. + Default: 00:00:00, in HH:MM:SS + format. You can also enter the time + with or without the colons, e.g. + 00:0000 or 026010."""}) + + arguments_list.append({"opts": ('-e', '--end'), + "type": str, + "dest": "end", + "default": "00:00:00", + "help": """Enter the end time to which an action + is to be applied. If both an end time + and duration are set, then the end + time will be used and the duration + will be ignored. + Default: 00:00:00, in HH:MM:SS."""}) + + arguments_list.append({"opts": ('-d', '--duration'), + "type": str, + "dest": "duration", + "default": "00:00:00", + "help": """Enter the duration of the chosen + action, for example if you enter + 00:00:10 for slice, then the first 10 + seconds after and including the start + time will be cut out into a new + video. + Default: 00:00:00, in HH:MM:SS + format. You can also enter the time + with or without the colons, e.g. + 00:0000 or 026010."""}) + + arguments_list.append({"opts": ('-m', '--mux-audio'), + "action": "store_true", + "dest": "mux_audio", + "default": False, + "help": """Mux the audio from the reference + video into the input video. This + option is only used for the 'gen-vid' + action. 'mux-audio' action has this + turned on implicitly."""}) + + arguments_list.append({"opts": ('-tr', '--transpose'), + "choices": ("(0, 90CounterClockwise&VerticalFlip)", + "(1, 90Clockwise)", + "(2, 90CounterClockwise)", + "(3, 90Clockwise&VerticalFlip)", + "None"), + "type": lambda v: Effmpeg.__parse_transpose(v), + "dest": "transpose", + "default": "None", + "help": """Transpose the video. If transpose is + set, then degrees will be ignored. For + cli you can enter either the number + or the long command name, + e.g. to use (1, 90Clockwise) + -tr 1 or -tr 90Clockwise"""}) + + arguments_list.append({"opts": ('-de', '--degrees'), + "type": str, + "dest": "degrees", + "default": "None", + "help": """Rotate the video clockwise by the + given number of degrees."""}) + + arguments_list.append({"opts": ('-sc', '--scale'), + "type": str, + "dest": "scale", + "default": "1920x1080", + "help": """Set the new resolution scale if the + chosen action is 'rescale'."""}) + + arguments_list.append({"opts": ('-q', '--quiet'), + "action": "store_true", + "dest": "quiet", + "default": False, + "help": """Reduces output verbosity so that only + serious errors are printed. If both + quiet and verbose are set, verbose + will override quiet."""}) + + arguments_list.append({"opts": ('-v', '--verbose'), + "action": "store_true", + "dest": "verbose", + "default": False, + "help": """Increases output verbosity. If both + quiet and verbose are set, verbose + will override quiet."""}) + + return arguments_list + + def parse_arguments(self, description, subparser, command): + parser = subparser.add_parser( + command, + help="This command lets you easily invoke" + "common ffmpeg commands.", + description=description, + epilog="Questions and feedback: \ + https://github.com/deepfakes/faceswap-playground" + ) + + for option in self.argument_list: + args = option['opts'] + kwargs = {key: option[key] for key in option.keys() if key != 'opts'} + parser.add_argument(*args, **kwargs) + + parser.set_defaults(func=self.process_arguments) + + def process_arguments(self, arguments): + self.args = arguments + + # Format action to match the method name + self.args.action = self.args.action.replace('-', '_') + + # Instantiate input DataItem object + self.input = DataItem(path=self.args.input) + + # Instantiate output DataItem object + if self.args.action in self._actions_have_dir_output: + self.output = DataItem(path=self.__get_default_output()) + elif self.args.action in self._actions_have_vid_output: + if self.__check_have_fps(self.args.fps) > 0: + self.output = DataItem(path=self.__get_default_output(), + fps=self.args.fps) + else: + self.output = DataItem(path=self.__get_default_output()) + + if self.args.ref_vid.lower() == "none" or self.args.ref_vid == '': + self.args.ref_vid = None + + # Instantiate ref_vid DataItem object + self.ref_vid = DataItem(path=self.args.ref_vid) + + # Check that correct input and output arguments were provided + if self.args.action in self._actions_have_dir_input and not self.input.is_type("dir"): + raise ValueError("The chosen action requires a directory as its " + "input, but you entered: " + "{}".format(self.input.path)) + if self.args.action in self._actions_have_vid_input and not self.input.is_type("vid"): + raise ValueError("The chosen action requires a video as its " + "input, but you entered: " + "{}".format(self.input.path)) + if self.args.action in self._actions_have_dir_output and not self.output.is_type("dir"): + raise ValueError("The chosen action requires a directory as its " + "output, but you entered: " + "{}".format(self.output.path)) + if self.args.action in self._actions_have_vid_output and not self.output.is_type("vid"): + raise ValueError("The chosen action requires a video as its " + "output, but you entered: " + "{}".format(self.output.path)) + + # Check that ref_vid is a video when it needs to be + if self.args.action in self._actions_req_ref_video: + if self.ref_vid.is_type("none"): + raise ValueError("The file chosen as the reference video is " + "not a video, either leave the field blank " + "or type 'None': " + "{}".format(self.ref_vid.path)) + elif self.args.action in self._actions_can_use_ref_video: + if self.ref_vid.is_type("none"): + print("Warning: no reference video was supplied, even though " + "one may be used with the chosen action. If this is " + "intentional then ignore this warning.", file=sys.stderr) + + # Process start and duration arguments + self.start = self.parse_time(self.args.start) + self.end = self.parse_time(self.args.end) + if not self.__check_equals_time(self.args.end, "00:00:00"): + self.duration = self.__get_duration(self.start, self.end) + else: + self.duration = self.parse_time(str(self.args.duration)) + + # If fps was left blank in gui, set it to default -1.0 value + if self.args.fps == '': + self.args.fps = str(-1.0) + + # Try to set fps automatically if needed and not supplied by user + if self.args.action in self._actions_req_fps \ + and self.__convert_fps(self.args.fps) <= 0: + if self.__check_have_fps(['r', 'i']): + _error_str = "No fps, input or reference video was supplied, " + _error_str += "hence it's not possible to " + _error_str += "'{}'.".format(self.args.action) + raise ValueError(_error_str) + elif self.output.fps is not None and self.__check_have_fps(['r', 'i']): + self.args.fps = self.output.fps + elif self.ref_vid.fps is not None and self.__check_have_fps(['i']): + self.args.fps = self.ref_vid.fps + elif self.input.fps is not None and self.__check_have_fps(['r']): + self.args.fps = self.input.fps + + # Processing transpose + if self.args.transpose.lower() == "none": + self.args.transpose = None + else: + self.args.transpose = self.args.transpose[1] + + # Processing degrees + if self.args.degrees.lower() == "none" or self.args.degrees == '': + self.args.degrees = None + elif self.args.transpose is None: + try: + int(self.args.degrees) + except ValueError as ve: + print("You have entered an invalid value for degrees: " + "{}".format(self.args.degrees), file=sys.stderr) + exit(1) + + # Set verbosity of output + self.__set_verbosity(self.args.quiet, self.args.verbose) + + # Set self.print_ to True if output needs to be printed to stdout + if self.args.action in self._actions_have_print_output: + self.print_ = True + + self.process() + + def process(self): + kwargs = {"input_": self.input, + "output": self.output, + "ref_vid": self.ref_vid, + "fps": self.args.fps, + "extract_ext": self.args.extract_ext, + "start": self.start, + "duration": self.duration, + "mux_audio": self.args.mux_audio, + "degrees": self.args.degrees, + "transpose": self.args.transpose, + "scale": self.args.scale, + "print_": self.print_} + action = getattr(self, self.args.action) + action(**kwargs) + + @staticmethod + def extract(input_=None, output=None, fps=None, extract_ext=None, + **kwargs): + _input_opts = Effmpeg._common_ffmpeg_args[:] + _input = {input_.path: _input_opts} + _output_opts = '-y -vf fps="' + str(fps) + '"' + _output_path = output.path + "/" + input_.name + "%05d" + extract_ext + _output = {_output_path: _output_opts} + ff = FFmpeg(inputs=_input, outputs=_output) + os.makedirs(output.path, exist_ok=True) + Effmpeg.__run_ffmpeg(ff) + + @staticmethod + def gen_vid(input_=None, output=None, fps=None, mux_audio=False, + ref_vid=None, **kwargs): + filename = Effmpeg.__get_extracted_filename(input_.path) + _input_opts = Effmpeg._common_ffmpeg_args[:] + _input_path = os.path.join(input_.path, filename) + _output_opts = '-y -c:v libx264 -vf fps="' + str(fps) + '" ' + if mux_audio: + _ref_vid_opts = '-c copy -map 0:0 -map 1:1' + _output_opts = _ref_vid_opts + ' ' + _output_opts + _inputs = {_input_path: _input_opts, ref_vid.path: None} + else: + _inputs = {_input_path: _input_opts} + _outputs = {output.path: _output_opts} + ff = FFmpeg(inputs=_inputs, outputs=_outputs) + Effmpeg.__run_ffmpeg(ff) + + @staticmethod + def get_fps(input_=None, print_=False, **kwargs): + _input_opts = '-v error -select_streams v -of ' + _input_opts += 'default=noprint_wrappers=1:nokey=1 ' + _input_opts += '-show_entries stream=r_frame_rate' + if type(input_) == str: + _inputs = {input_: _input_opts} + else: + _inputs = {input_.path: _input_opts} + ff = FFprobe(inputs=_inputs) + _fps = ff.run(stdout=subprocess.PIPE)[0].decode("utf-8") + _fps = _fps.strip() + if print_: + print("Video fps:", _fps) + else: + return _fps + + @staticmethod + def get_info(input_=None, print_=False, **kwargs): + _input_opts = Effmpeg._common_ffmpeg_args[:] + _inputs = {input_.path: _input_opts} + ff = FFprobe(inputs=_inputs) + out = ff.run(stdout=subprocess.PIPE, + stderr=subprocess.STDOUT)[0].decode('utf-8') + if print_: + print(out) + else: + return out + + @staticmethod + def rescale(input_=None, output=None, scale=None, **kwargs): + _input_opts = Effmpeg._common_ffmpeg_args[:] + _output_opts = '-y -vf scale="' + str(scale) + '"' + _inputs = {input_.path: _input_opts} + _outputs = {output.path: _output_opts} + ff = FFmpeg(inputs=_inputs, outputs=_outputs) + Effmpeg.__run_ffmpeg(ff) + + @staticmethod + def rotate(input_=None, output=None, degrees=None, transpose=None, + **kwargs): + if transpose is None and degrees is None: + raise ValueError("You have not supplied a valid transpose or " + "degrees value:\ntranspose: {}\ndegrees: " + "{}".format(transpose, degrees)) + + _input_opts = Effmpeg._common_ffmpeg_args[:] + _output_opts = '-y -c:a copy -vf ' + _bilinear = '' + if transpose is not None: + _output_opts += 'transpose="' + str(transpose) + '"' + elif int(degrees) != 0: + if int(degrees) % 90 == 0 and int(degrees) != 0: + _bilinear = ":bilinear=0" + _output_opts += 'rotate="' + str(degrees) + '*(PI/180)' + _output_opts += _bilinear + '" ' + + _inputs = {input_.path: _input_opts} + _outputs = {output.path: _output_opts} + ff = FFmpeg(inputs=_inputs, outputs=_outputs) + Effmpeg.__run_ffmpeg(ff) + + @staticmethod + def mux_audio(input_=None, output=None, ref_vid=None, **kwargs): + _input_opts = Effmpeg._common_ffmpeg_args[:] + _ref_vid_opts = None + _output_opts = '-y -c copy -map 0:0 -map 1:1 -shortest' + _inputs = {input_.path: _input_opts, ref_vid.path: _ref_vid_opts} + _outputs = {output.path: _output_opts} + ff = FFmpeg(inputs=_inputs, outputs=_outputs) + Effmpeg.__run_ffmpeg(ff) + + @staticmethod + def slice(input_=None, output=None, start=None, duration=None, **kwargs): + _input_opts = Effmpeg._common_ffmpeg_args[:] + _input_opts += "-ss " + start + _output_opts = "-y -t " + duration + " " + _output_opts += "-vcodec copy -acodec copy" + _inputs = {input_.path: _input_opts} + _output = {output.path: _output_opts} + ff = FFmpeg(inputs=_inputs, outputs=_output) + Effmpeg.__run_ffmpeg(ff) + + # Various helper methods + @classmethod + def __set_verbosity(cls, quiet, verbose): + if verbose: + cls._common_ffmpeg_args = cls.__common_ffmpeg_args_dict["verbose"] + elif quiet: + cls._common_ffmpeg_args = cls.__common_ffmpeg_args_dict["quiet"] + else: + cls._common_ffmpeg_args = cls.__common_ffmpeg_args_dict["normal"] + + def __get_default_output(self): + # Set output to the same directory as input + # if the user didn't specify it. + if self.args.output == "": + if self.args.action in self._actions_have_dir_output: + return os.path.join(self.input.dirname, 'out') + elif self.args.action in self._actions_have_vid_output: + if self.input.is_type("media"): + # Using the same extension as input leads to very poor + # output quality, hence the default is mkv for now + return os.path.join(self.input.dirname, + "out.mkv") # + self.input.ext) + else: # case if input was a directory + return os.path.join(self.input.dirname, 'out.mkv') + else: + return self.args.output + + def __check_have_fps(self, items): + items_to_check = list() + for i in items: + if i == 'r': + items_to_check.append('ref_vid') + elif i == 'i': + items_to_check.append('input') + elif i == 'o': + items_to_check.append('output') + + return all(getattr(self, i).fps is None for i in items_to_check) + + @staticmethod + def __run_ffmpeg(ff): + try: + ff.run(stderr=subprocess.STDOUT) + except FFRuntimeError as ffe: + # After receiving SIGINT ffmpeg has a 255 exit code + if ffe.exit_code == 255: + pass + else: + raise ValueError("An unexpected FFRuntimeError occurred: " + "{}".format(ffe)) + except KeyboardInterrupt: + pass # Do nothing if voluntary interruption + + @staticmethod + def __convert_fps(fps): + if '/' in fps: + _fps = fps.split('/') + return float(_fps[0]) / float(_fps[1]) + else: + return float(fps) + + @staticmethod + def __get_duration(start_time, end_time): + start = [int(i) for i in start_time.split(':')] + end = [int(i) for i in end_time.split(':')] + start = datetime.timedelta(hours=start[0], minutes=start[1], seconds=start[2]) + end = datetime.timedelta(hours=end[0], minutes=end[1], seconds=end[2]) + delta = end - start + s = delta.total_seconds() + return '{:02}:{:02}:{:02}'.format(int(s // 3600), int(s % 3600 // 60), int(s % 60)) + + @staticmethod + def __get_extracted_filename(path): + filename = '' + for file in os.listdir(path): + if any(i in file for i in DataItem.img_ext): + filename = file + break + filename = filename.split('.') + img_ext = filename[-1] + zero_pad = filename[-2] + name = '.'.join(filename[:-2]) + + vid_ext = '' + for ve in [ve.replace('.', '') for ve in DataItem.vid_ext]: + if ve in zero_pad: + vid_ext = ve + zero_pad = len(zero_pad.replace(ve, '')) + break + + zero_pad = str(zero_pad).zfill(2) + filename_list = [name, vid_ext + '%0' + zero_pad + 'd', img_ext] + return '.'.join(filename_list) + + @staticmethod + def __parse_transpose(value): + index = 0 + opts = ["(0, 90CounterClockwise&VerticalFlip)", + "(1, 90Clockwise)", + "(2, 90CounterClockwise)", + "(3, 90Clockwise&VerticalFlip)", + "None"] + if len(value) == 1: + index = int(value) + else: + for i in range(5): + if value in opts[i]: + index = i + break + return opts[index] + + @staticmethod + def __check_is_valid_time(value): + val = value.replace(':', '') + return val.isdigit() + + @staticmethod + def __check_equals_time(value, time): + v = value.replace(':', '') + t = time.replace(':', '') + return v.zfill(6) == t.zfill(6) + + @staticmethod + def parse_time(txt): + clean_txt = txt.replace(':', '') + hours = clean_txt[0:2] + minutes = clean_txt[2:4] + seconds = clean_txt[4:6] + return hours + ':' + minutes + ':' + seconds + + +def bad_args(args): + parser.print_help() + exit(0) + + +if __name__ == "__main__": + print('"Easy"-ffmpeg wrapper.\n') + + parser = argparse.ArgumentParser() + subparser = parser.add_subparsers() + sort = Effmpeg( + subparser, "effmpeg", "Wrapper for various common ffmpeg commands.") + + parser.set_defaults(func=bad_args) + arguments = parser.parse_args() + arguments.func(arguments) diff --git a/tools/sort.py b/tools/sort.py index db2b2a3..194a526 100644 --- a/tools/sort.py +++ b/tools/sort.py @@ -8,7 +8,7 @@ import cv2 from tqdm import tqdm from shutil import copyfile import json -from lib.cli import FullPaths +from lib.cli import DirFullPaths, FileFullPaths # DLIB is a GPU Memory hog, so the following modules should only be imported # when required @@ -43,13 +43,14 @@ class SortProcessor(object): def get_argument_list(): arguments_list = list() arguments_list.append({"opts": ('-i', '--input'), - "action": FullPaths, + "action": DirFullPaths, "dest": "input_dir", "default": "input_dir", "help": "Input directory of aligned faces.", "required": True}) arguments_list.append({"opts": ('-o', '--output'), + "action": DirFullPaths, "dest": "output_dir", "default": "output_dir", "help": "Output directory for sorted aligned " @@ -127,6 +128,8 @@ class SortProcessor(object): "will be created in the input directory."}) arguments_list.append({"opts": ('-lf', '--log-file'), + "action": FileFullPaths, + "filetypes": ("JSON", "*.json"), "dest": 'log_file_path', "default": 'sort_log.json', "help": "Specify a log file to use for saving the " @@ -170,10 +173,8 @@ class SortProcessor(object): def parse_arguments(self, description, subparser, command): parser = subparser.add_parser( command, - help="This command lets you sort images using various methods." - " Please backup your data and/or test this tool with a " - "smaller data set to make sure you understand how it" - "works.", + help="This command lets you sort images using various " + "methods.", description=description, epilog="Questions and feedback: \ https://github.com/deepfakes/faceswap-playground"