OSSForks/faceswap
1
0
Fork 0
mirror of https://github.com/zebrajr/faceswap.git synced 2026-01-15 12:15:15 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Adding new tool effmpeg ("easy"-ffmpeg) with gui support. Extend gui functionality to support filetypes. Re-opening PR. (#373)

Browse Source 
* 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.
This commit is contained in:
Lev Velykoivanenko
2018-05-09 19:47:17 +02:00
committed by torzdf
parent 958493a64a
commit 80cde77a6d
15 changed files with 1247 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
faceswap.py
View File

@@ -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()

336
lib/cli.py
View File

@@ -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",

19
lib/utils.py
View File

@@ -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)

1
requirements-gpu-python35-cuda8.txt
View File

@@ -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

1
requirements-gpu-python36-cuda9.txt
View File

@@ -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

1
requirements-python35.txt
View File

@@ -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

1
requirements-python36.txt
View File

@@ -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

21
scripts/convert.py
View File

@@ -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 ###

12
scripts/extract.py
View File

@@ -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)

24
scripts/fsmedia.py
View File

@@ -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

156
scripts/gui.py
View File

@@ -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 """

25
scripts/train.py
View File

@@ -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()

33
tools.py
View File

@@ -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)

777
tools/effmpeg.py Normal file
View File

@@ -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)

13
tools/sort.py
View File

@@ -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"