Gui v3.0b (#436)

* GUI version 3 (#411)

GUI version 3.0a

* Required for Shaonlu mode (#416)

Added two modes - Original and Shaonlu.
The later requires this file to function.

* model update (#417)

New, functional Original 128 model

* OriginalHighRes 128 model update (#418)

Required for OriginalHighRes Model to function

* Add OriginalHighRes 128 update to gui branch (#421)

* Required for Shaonlu mode (#416)

Added two modes - Original and Shaonlu.
The later requires this file to function.

* model update (#417)

New, functional Original 128 model

* OriginalHighRes 128 model update (#418)

Required for OriginalHighRes Model to function

* Dev gui (#420)

* reduce singletons

* Fix tooltips and screen boundaries on popup

* Remove dpi fix. Fix context filebrowsers

* fix tools.py execution and context filebrowser bugs

* Bugfixes (#422)

* Bump matplotlib requirement. Fix polyfit. Fix TQDM on sort

* Fixed memory usage at 6GB cards. (#423)

- Switched default encoder to ORIGINAL
- Fixed memory consumption. Tested with geforce gtx 9800 ti with 6Gb; batch_size 8 no OOM or memory warnings now.

* Staging (#426)

* altered trainer (#425)

altered trainer to accommodate with model change

* Update Model.py (#424)

- Added saving state (currently only saved epoch number, to be extended in future)
- Changed saving to ThreadPoolExecutor

* Add DPI Scaling (#428)

* Add dpi scaling

* Hotfix for effmpeg. (#429)

effmpeg fixed so it works both in cli and gui.
Initial work done to add previewing feature to effmpeg (currently does nothing).
Some small spacing changes in other files to improve PEP8 conformity.

* PEP8 Linting (#430)

* pep8 linting

* Requirements version bump (#432)

* altered trainer (#425)

altered trainer to accommodate with model change

* Update Model.py (#424)

- Added saving state (currently only saved epoch number, to be extended in future)
- Changed saving to ThreadPoolExecutor

* Requirements version bump (#431)

This bumps the versions of:

    scandir
    h5py
    Keras
    opencv-python

to their latest vesions.

Virtual Environment will need to be setup again to make use of these.

* High DPI Fixes (#433)

* dpi scaling

* DPI Fixes

* Fix and improve context manager. (#434)

effmpeg tool:
Context manager for GUI fixed.

Context manager in general:
Functionality extended to allow configuring the context with both:
  command -> action
  command -> variable (cli argument) -> action

* Change epoch option to iterations

* Change epochs to iterations
This commit is contained in:
torzdf 2018-06-20 19:25:31 +02:00 committed by GitHub
parent 5275718365
commit be8e235eab
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
36 changed files with 4252 additions and 2073 deletions

7
.gitignore vendored
View File

@ -1,11 +1,16 @@
*
!*.keep
!*.py
!*.md
!*.txt
!*.png
!Dockerfile*
!requirements*
!icons
!lib
!lib/gui
!lib/gui/.cache
!lib/gui/.cache/preview
!lib/gui/.cache/icons
!scripts
!plugins
!tools

View File

@ -19,15 +19,18 @@ def bad_args(args):
if __name__ == "__main__":
PARSER = cli.FullHelpArgumentParser()
SUBPARSER = PARSER.add_subparsers()
EXTRACT = cli.ExtractArgs(
SUBPARSER, "extract", "Extract the faces from pictures")
TRAIN = cli.TrainArgs(
SUBPARSER, "train", "This command trains the model for the two faces A and B")
CONVERT = cli.ConvertArgs(
SUBPARSER, "convert", "Convert a source image to a new one with the face swapped")
GUIPARSERS = {'extract': EXTRACT, 'train': TRAIN, 'convert': CONVERT}
GUI = cli.GuiArgs(
SUBPARSER, "gui", "Launch the Faceswap Graphical User Interface", GUIPARSERS)
EXTRACT = cli.ExtractArgs(SUBPARSER,
"extract",
"Extract the faces from pictures")
TRAIN = cli.TrainArgs(SUBPARSER,
"train",
"This command trains the model for the two faces A and B")
CONVERT = cli.ConvertArgs(SUBPARSER,
"convert",
"Convert a source image to a new one with the face swapped")
GUI = cli.GuiArgs(SUBPARSER,
"gui",
"Launch the Faceswap Graphical User Interface")
PARSER.set_defaults(func=bad_args)
ARGUMENTS = PARSER.parse_args()
ARGUMENTS.func(ARGUMENTS)

View File

@ -22,7 +22,7 @@ def onExit():
for detector in dlib_detectors:
del detector
class TorchBatchNorm2D(keras.engine.topology.Layer):
class TorchBatchNorm2D(keras.engine.base_layer.Layer):
def __init__(self, axis=-1, momentum=0.99, epsilon=1e-3, **kwargs):
super(TorchBatchNorm2D, self).__init__(**kwargs)
self.supports_masking = True

View File

@ -1,7 +1,9 @@
#!/usr/bin/env python3
""" Command Line Arguments """
import argparse
from importlib import import_module
import os
import platform
import sys
from plugins.PluginLoader import PluginLoader
@ -19,23 +21,63 @@ class ScriptExecutor(object):
def import_script(self):
""" Only import a script's modules when running that script."""
if self.command == 'extract':
from scripts.extract import Extract as script
elif self.command == 'train':
from scripts.train import Train as script
elif self.command == 'convert':
from scripts.convert import Convert as script
elif self.command == 'gui':
from scripts.gui import Gui as script
else:
script = None
self.test_for_gui()
cmd = os.path.basename(sys.argv[0])
src = "tools" if cmd == "tools.py" else "scripts"
mod = ".".join((src, self.command.lower()))
module = import_module(mod)
script = getattr(module, self.command.title())
return script
def test_for_gui(self):
""" If running the gui, check the prerequisites """
if self.command != "gui":
return
self.test_tkinter()
self.check_display()
@staticmethod
def test_tkinter():
""" If the user is running the GUI, test whether the
tkinter app is available on their machine. If not
exit gracefully.
This avoids having to import every tk function
within the GUI in a wrapper and potentially spamming
traceback errors to console """
try:
import tkinter
except ImportError:
print(
"It looks like TkInter isn't installed for your OS, so "
"the GUI has been disabled. To enable the GUI please "
"install the TkInter application.\n\n"
"You can try:\n"
" Windows/macOS: Install ActiveTcl Community "
"Edition from "
"www.activestate.com\n"
" Ubuntu/Mint/Debian: sudo apt install python3-tk\n"
" Arch: sudo pacman -S tk\n"
" CentOS/Redhat: sudo yum install tkinter\n"
" Fedora: sudo dnf install python3-tkinter\n")
exit(1)
@staticmethod
def check_display():
""" Check whether there is a display to output the GUI. If running on
Windows then assume not running in headless mode """
if not os.environ.get("DISPLAY", None) and os.name != "nt":
print("No display detected. GUI mode has been disabled.")
if platform.system() == "Darwin":
print("macOS users need to install XQuartz. "
"See https://support.apple.com/en-gb/HT201341")
exit(1)
def execute_script(self, arguments):
""" Run the script for called command """
script = self.import_script()
args = (arguments, ) if self.command != 'gui' else (arguments, self.subparsers)
process = script(*args)
process = script(arguments)
process.process()
@ -43,7 +85,7 @@ class FullPaths(argparse.Action):
""" Expand user- and relative-paths """
def __call__(self, parser, namespace, values, option_string=None):
setattr(namespace, self.dest, os.path.abspath(
os.path.expanduser(values)))
os.path.expanduser(values)))
class DirFullPaths(FullPaths):
@ -55,14 +97,7 @@ 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.
see lib/gui/utils.py FileHandler for current GUI filetypes
"""
def __init__(self, option_strings, dest, nargs=None, filetypes=None,
**kwargs):
@ -71,98 +106,63 @@ class FileFullPaths(FullPaths):
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'
"option_strings",
"dest",
"nargs",
"const",
"default",
"type",
"choices",
"help",
"metavar",
"filetypes"
]
return [(name, getattr(self, name)) for name in names]
class ComboFullPaths(FileFullPaths):
class SaveFileFullPaths(FileFullPaths):
"""
Class that gui uses to determine if you need to save a file.
see lib/gui/utils.py FileHandler for current GUI filetypes
"""
pass
class ContextFullPaths(FileFullPaths):
"""
Class that gui uses to determine if you need to open a file or a
directory based on which action you are choosing
To use ContextFullPaths the action_option item should indicate which
cli option dictates the context of the filesystem dialogue
Bespoke actions are then set in lib/gui/utils.py FileHandler
"""
def __init__(self, option_strings, dest, nargs=None, filetypes=None,
actions_open_type=None, **kwargs):
def __init__(self, option_strings, dest, nargs=None, filetypes=None,
action_option=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
super(ContextFullPaths, self).__init__(option_strings, dest,
filetypes=None, **kwargs)
self.action_option = action_option
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'
"option_strings",
"dest",
"nargs",
"const",
"default",
"type",
"choices",
"help",
"metavar",
"filetypes",
"action_option"
]
return [(name, getattr(self, name)) for name in names]
@ -181,8 +181,12 @@ class FaceSwapArgs(object):
to all commands. Should be the parent function of all
subsequent argparsers """
def __init__(self, subparser, command, description="default", subparsers=None):
self.argument_list = self.get_argument_list()
self.optional_arguments = self.get_optional_arguments()
if not subparser:
return
self.parser = self.create_parser(subparser, command, description)
self.add_arguments()
@ -236,11 +240,6 @@ class ExtractConvertArgs(FaceSwapArgs):
def get_argument_list():
""" Put the arguments in a list so that they are accessible from both
argparse and gui """
alignments_filetypes = [["Serializers", ['json', 'p', 'yaml']],
["JSON", ["json"]],
["Pickle", ["p"]],
["YAML", ["yaml"]]]
alignments_filetypes = FileFullPaths.prep_filetypes(alignments_filetypes)
argument_list = list()
argument_list.append({"opts": ("-i", "--input-dir"),
"action": DirFullPaths,
@ -258,7 +257,7 @@ class ExtractConvertArgs(FaceSwapArgs):
"Defaults to 'output'"})
argument_list.append({"opts": ("--alignments", ),
"action": FileFullPaths,
"filetypes": alignments_filetypes,
"filetypes": 'alignments',
"type": str,
"dest": "alignments_path",
"help": "Optional path to an alignments file."})
@ -532,10 +531,10 @@ class TrainArgs(FaceSwapArgs):
"type": int,
"default": 64,
"help": "Batch size, as a power of 2 (64, 128, 256, etc)"})
argument_list.append({"opts": ("-ep", "--epochs"),
argument_list.append({"opts": ("-it", "--iterations"),
"type": int,
"default": 1000000,
"help": "Length of training in epochs"})
"help": "Length of training in iterations"})
argument_list.append({"opts": ("-g", "--gpus"),
"type": int,
"default": 1,

View File

Before

Width:  |  Height:  |  Size: 691 B

After

Width:  |  Height:  |  Size: 691 B

BIN
lib/gui/.cache/icons/graph.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 574 B

BIN
lib/gui/.cache/icons/move.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 550 B

View File

Before

Width:  |  Height:  |  Size: 406 B

After

Width:  |  Height:  |  Size: 406 B

View File

Before

Width:  |  Height:  |  Size: 263 B

After

Width:  |  Height:  |  Size: 263 B

View File

Before

Width:  |  Height:  |  Size: 773 B

After

Width:  |  Height:  |  Size: 773 B

View File

Before

Width:  |  Height:  |  Size: 530 B

After

Width:  |  Height:  |  Size: 530 B

BIN
lib/gui/.cache/icons/zoom.png Executable file

Binary file not shown.

After

Width:  |  Height:  |  Size: 642 B

View File

7
lib/gui/__init__.py Normal file
View File

@ -0,0 +1,7 @@
from lib.gui.command import CommandNotebook
from lib.gui.display import DisplayNotebook
from lib.gui.options import CliOptions, Config
from lib.gui.stats import CurrentSession
from lib.gui.statusbar import StatusBar
from lib.gui.utils import ConsoleOut, Images
from lib.gui.wrapper import ProcessWrapper

361
lib/gui/command.py Normal file
View File

@ -0,0 +1,361 @@
#!/usr/bin python3
""" The command frame for Faceswap GUI """
import tkinter as tk
from tkinter import ttk
from .options import Config
from .tooltip import Tooltip
from .utils import Images, FileHandler
class CommandNotebook(ttk.Notebook):
""" Frame to hold each individual tab of the command notebook """
def __init__(self, parent, cli_options, tk_vars, scaling_factor):
width = int(420 * scaling_factor)
height = int(500 * scaling_factor)
ttk.Notebook.__init__(self, parent, width=width, height=height)
parent.add(self)
self.cli_opts = cli_options
self.tk_vars = tk_vars
self.actionbtns = dict()
self.set_running_task_trace()
self.build_tabs()
def set_running_task_trace(self):
""" Set trigger action for the running task
to change the action buttons text and command """
self.tk_vars["runningtask"].trace("w", self.change_action_button)
def build_tabs(self):
""" Build the tabs for the relevant command """
for category in self.cli_opts.categories:
cmdlist = self.cli_opts.commands[category]
for command in cmdlist:
title = command.title()
commandtab = CommandTab(self, category, command)
self.add(commandtab, text=title)
def change_action_button(self, *args):
""" Change the action button to relevant control """
for cmd in self.actionbtns.keys():
btnact = self.actionbtns[cmd]
if self.tk_vars["runningtask"].get():
ttl = "Terminate"
hlp = "Exit the running process"
else:
ttl = cmd.title()
hlp = "Run the {} script".format(cmd.title())
btnact.config(text=ttl)
Tooltip(btnact, text=hlp, wraplength=200)
class CommandTab(ttk.Frame):
""" Frame to hold each individual tab of the command notebook """
def __init__(self, parent, category, command):
ttk.Frame.__init__(self, parent)
self.category = category
self.cli_opts = parent.cli_opts
self.actionbtns = parent.actionbtns
self.tk_vars = parent.tk_vars
self.command = command
self.build_tab()
def build_tab(self):
""" Build the tab """
OptionsFrame(self)
self.add_frame_separator()
ActionFrame(self)
def add_frame_separator(self):
""" Add a separator between top and bottom frames """
sep = ttk.Frame(self, height=2, relief=tk.RIDGE)
sep.pack(fill=tk.X, pady=(5, 0), side=tk.TOP)
class OptionsFrame(ttk.Frame):
""" Options Frame - Holds the Options for each command """
def __init__(self, parent):
ttk.Frame.__init__(self, parent)
self.pack(side=tk.TOP, fill=tk.BOTH, expand=True)
self.opts = parent.cli_opts
self.command = parent.command
self.canvas = tk.Canvas(self, bd=0, highlightthickness=0)
self.canvas.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
self.optsframe = ttk.Frame(self.canvas)
self.optscanvas = self.canvas.create_window((0, 0),
window=self.optsframe,
anchor=tk.NW)
self.chkbtns = self.checkbuttons_frame()
self.build_frame()
self.opts.set_context_option(self.command)
def checkbuttons_frame(self):
""" Build and format frame for holding the check buttons """
container = ttk.Frame(self.optsframe)
lbl = ttk.Label(container, text="Options", width=16, anchor=tk.W)
lbl.pack(padx=5, pady=5, side=tk.LEFT, anchor=tk.N)
chkframe = ttk.Frame(container)
chkframe.pack(side=tk.BOTTOM, expand=True)
chkleft = ttk.Frame(chkframe, name="leftFrame")
chkleft.pack(side=tk.LEFT, anchor=tk.N, expand=True)
chkright = ttk.Frame(chkframe, name="rightFrame")
chkright.pack(side=tk.RIGHT, anchor=tk.N, expand=True)
return container, chkframe
def build_frame(self):
""" Build the options frame for this command """
self.add_scrollbar()
self.canvas.bind("<Configure>", self.resize_frame)
for option in self.opts.gen_command_options(self.command):
optioncontrol = OptionControl(self.command,
option,
self.optsframe,
self.chkbtns[1])
optioncontrol.build_full_control()
if self.chkbtns[1].winfo_children():
self.chkbtns[0].pack(side=tk.BOTTOM, fill=tk.X, expand=True)
def add_scrollbar(self):
""" Add a scrollbar to the options frame """
scrollbar = ttk.Scrollbar(self, command=self.canvas.yview)
scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
self.canvas.config(yscrollcommand=scrollbar.set)
self.optsframe.bind("<Configure>", self.update_scrollbar)
def update_scrollbar(self, event):
""" Update the options frame scrollbar """
self.canvas.configure(scrollregion=self.canvas.bbox("all"))
def resize_frame(self, event):
""" Resize the options frame to fit the canvas """
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 """
def __init__(self, command, option, option_frame, checkbuttons_frame):
self.command = command
self.option = option
self.option_frame = option_frame
self.chkbtns = checkbuttons_frame
def build_full_control(self):
""" Build the correct control type for the option passed through """
ctl = self.option["control"]
ctltitle = self.option["control_title"]
sysbrowser = self.option["filesystem_browser"]
ctlhelp = " ".join(self.option.get("help", "").split())
ctlhelp = ". ".join(i.capitalize() for i in ctlhelp.split(". "))
ctlhelp = ctltitle + " - " + ctlhelp
dflt = self.option.get("default", "")
if ctl == ttk.Checkbutton:
dflt = self.option.get("default", False)
choices = self.option["choices"] if ctl == ttk.Combobox else None
ctlframe = self.build_one_control_frame()
if ctl != ttk.Checkbutton:
self.build_one_control_label(ctlframe, ctltitle)
ctlvars = (ctl, ctltitle, dflt, ctlhelp)
self.option["value"] = self.build_one_control(ctlframe,
ctlvars,
choices,
sysbrowser)
def build_one_control_frame(self):
""" Build the frame to hold the control """
frame = ttk.Frame(self.option_frame)
frame.pack(fill=tk.X, expand=True)
return frame
@staticmethod
def build_one_control_label(frame, control_title):
""" Build and place the control label """
lbl = ttk.Label(frame, text=control_title, width=16, anchor=tk.W)
lbl.pack(padx=5, pady=5, side=tk.LEFT, anchor=tk.N)
def build_one_control(self, frame, controlvars, choices, sysbrowser):
""" Build and place the option controls """
control, control_title, default, helptext = controlvars
default = default if default is not None else ""
var = tk.BooleanVar(
frame) if control == ttk.Checkbutton else tk.StringVar(frame)
var.set(default)
if sysbrowser is not None:
self.add_browser_buttons(frame, sysbrowser, var)
if control == ttk.Checkbutton:
self.checkbutton_to_checkframe(control,
control_title,
var,
helptext)
else:
self.control_to_optionsframe(control,
frame,
var,
choices,
helptext)
return var
def checkbutton_to_checkframe(self, control, control_title, var, helptext):
""" Add checkbuttons to the checkbutton frame """
leftframe = self.chkbtns.children["leftFrame"]
rightframe = self.chkbtns.children["rightFrame"]
chkbtn_count = len({**leftframe.children, **rightframe.children})
frame = leftframe if chkbtn_count % 2 == 0 else rightframe
ctl = control(frame, variable=var, text=control_title)
ctl.pack(side=tk.TOP, padx=5, pady=5, anchor=tk.W)
Tooltip(ctl, text=helptext, wraplength=200)
@staticmethod
def control_to_optionsframe(control, frame, var, choices, helptext):
""" Standard non-check buttons sit in the main options frame """
ctl = control(frame, textvariable=var)
ctl.pack(padx=5, pady=5, fill=tk.X, expand=True)
if control == ttk.Combobox:
ctl["values"] = [choice for choice in choices]
Tooltip(ctl, text=helptext, wraplength=200)
def add_browser_buttons(self, frame, sysbrowser, filepath):
""" Add correct file browser button for control """
img = Images().icons[sysbrowser]
action = getattr(self, "ask_" + sysbrowser)
filetypes = self.option.get("filetypes", "default")
fileopn = ttk.Button(frame, image=img,
command=lambda cmd=action: cmd(filepath,
filetypes))
fileopn.pack(padx=(0, 5), side=tk.RIGHT)
@staticmethod
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 = FileHandler("dir", filetypes).retfile
if dirname:
filepath.set(dirname)
@staticmethod
def ask_load(filepath, filetypes):
""" Pop-up to get path to a file """
filename = FileHandler("filename", filetypes).retfile
if filename:
filepath.set(filename)
@staticmethod
def ask_save(filepath, filetypes=None):
""" Pop-up to get path to save a new file """
filename = FileHandler("savefilename", filetypes).retfile
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_context(self, filepath, filetypes):
""" Method to pop the correct dialog depending on context """
selected_action = self.option["action_option"].get()
selected_variable = self.option["dest"]
filename = FileHandler("context",
filetypes,
command=self.command,
action=selected_action,
variable=selected_variable).retfile
if filename:
filepath.set(filename)
class ActionFrame(ttk.Frame):
"""Action Frame - Displays action controls for the command tab """
def __init__(self, parent):
ttk.Frame.__init__(self, parent)
self.pack(fill=tk.BOTH, padx=5, pady=5, side=tk.BOTTOM, anchor=tk.N)
self.command = parent.command
self.title = self.command.title()
self.add_action_button(parent.category,
parent.actionbtns,
parent.tk_vars)
self.add_util_buttons(parent.cli_opts, parent.tk_vars)
def add_action_button(self, category, actionbtns, tk_vars):
""" Add the action buttons for page """
actframe = ttk.Frame(self)
actframe.pack(fill=tk.X, side=tk.LEFT)
var_value = "{},{}".format(category, self.command)
btnact = ttk.Button(actframe,
text=self.title,
width=10,
command=lambda: tk_vars["action"].set(var_value))
btnact.pack(side=tk.LEFT)
Tooltip(btnact,
text="Run the {} script".format(self.title),
wraplength=200)
actionbtns[self.command] = btnact
btngen = ttk.Button(actframe,
text="Generate",
width=10,
command=lambda: tk_vars["generate"].set(var_value))
btngen.pack(side=tk.RIGHT, padx=5)
Tooltip(btngen,
text="Output command line options to the console",
wraplength=200)
def add_util_buttons(self, cli_options, tk_vars):
""" Add the section utility buttons """
utlframe = ttk.Frame(self)
utlframe.pack(side=tk.RIGHT)
config = Config(cli_options, tk_vars)
for utl in ("load", "save", "clear", "reset"):
img = Images().icons[utl]
action_cls = config if utl in (("save", "load")) else cli_options
action = getattr(action_cls, utl)
btnutl = ttk.Button(utlframe,
image=img,
command=lambda cmd=action: cmd(self.command))
btnutl.pack(padx=2, side=tk.LEFT)
Tooltip(btnutl,
text=utl.capitalize() + " " + self.title + " config",
wraplength=200)

92
lib/gui/display.py Normal file
View File

@ -0,0 +1,92 @@
#!/usr/bin python3
""" Display Frame of the Faceswap GUI
What is displayed in the Display Frame varies
depending on what tasked is being run """
import tkinter as tk
from tkinter import ttk
from .display_analysis import Analysis
from .display_command import GraphDisplay, PreviewExtract, PreviewTrain
class DisplayNotebook(ttk.Notebook):
""" The display tabs """
def __init__(self, parent, session, tk_vars, scaling_factor):
ttk.Notebook.__init__(self, parent, width=780)
parent.add(self)
self.wrapper_var = tk_vars["display"]
self.runningtask = tk_vars["runningtask"]
self.session = session
self.set_wrapper_var_trace()
self.add_static_tabs(scaling_factor)
self.static_tabs = [child for child in self.tabs()]
def set_wrapper_var_trace(self):
""" Set the trigger actions for the display vars
when they have been triggered in the Process Wrapper """
self.wrapper_var.trace("w", self.update_displaybook)
def add_static_tabs(self, scaling_factor):
""" Add tabs that are permanently available """
for tab in ("job queue", "analysis"):
if tab == "job queue":
continue # Not yet implemented
if tab == "analysis":
helptext = {"stats":
"Summary statistics for each training session"}
frame = Analysis(self, tab, helptext, scaling_factor)
else:
frame = self.add_frame()
self.add(frame, text=tab.title())
def add_frame(self):
""" Add a single frame for holding tab's contents """
frame = ttk.Frame(self)
frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=5, pady=5)
return frame
def command_display(self, command):
""" Select what to display based on incoming
command """
build_tabs = getattr(self, "{}_tabs".format(command))
build_tabs()
def extract_tabs(self):
""" Build the extract tabs """
helptext = ("Updates preview from output every 5 "
"seconds to limit disk contention")
PreviewExtract(self, "preview", helptext, 5000)
def train_tabs(self):
""" Build the train tabs """
for tab in ("graph", "preview"):
if tab == "graph":
helptext = "Graph showing Loss vs Iterations"
GraphDisplay(self, "graph", helptext, 5000)
elif tab == "preview":
helptext = "Training preview. Updated on every save iteration"
PreviewTrain(self, "preview", helptext, 5000)
def convert_tabs(self):
""" Build the convert tabs
Currently identical to Extract, so just call that """
self.extract_tabs()
def remove_tabs(self):
""" Remove all command specific tabs """
for child in self.tabs():
if child not in self.static_tabs:
self.forget(child)
def update_displaybook(self, *args):
""" Set the display tabs based on executing task """
command = self.wrapper_var.get()
self.remove_tabs()
if not command or command not in ("extract", "train", "convert"):
return
self.command_display(command)

505
lib/gui/display_analysis.py Normal file
View File

@ -0,0 +1,505 @@
#!/usr/bin python3
""" Analysis tab of Display Frame of the Faceswap GUI """
import csv
import tkinter as tk
from tkinter import ttk
from .display_graph import SessionGraph
from .display_page import DisplayPage
from .stats import Calculations, SavedSessions, SessionsSummary, SessionsTotals
from .tooltip import Tooltip
from .utils import Images, FileHandler
class Analysis(DisplayPage):
""" Session analysis tab """
def __init__(self, parent, tabname, helptext, scaling_factor):
DisplayPage.__init__(self, parent, tabname, helptext)
self.summary = None
self.add_options()
self.add_main_frame(scaling_factor)
def set_vars(self):
""" Analysis specific vars """
selected_id = tk.StringVar()
filename = tk.StringVar()
return {"selected_id": selected_id,
"filename": filename}
def add_main_frame(self, scaling_factor):
""" Add the main frame to the subnotebook
to hold stats and session data """
mainframe = self.subnotebook_add_page("stats")
self.stats = StatsData(mainframe,
self.vars["filename"],
self.vars["selected_id"],
self.helptext["stats"],
scaling_factor)
def add_options(self):
""" Add the options bar """
self.reset_session_info()
options = Options(self)
options.add_options()
def reset_session_info(self):
""" Reset the session info status to default """
self.vars["filename"].set(None)
self.set_info("No session data loaded")
def load_session(self):
""" Load previously saved sessions """
self.clear_session()
filename = FileHandler("open", "session").retfile
if not filename:
return
filename = filename.name
loaded_data = SavedSessions(filename).sessions
msg = filename
if len(filename) > 70:
msg = "...{}".format(filename[-70:])
self.set_session_summary(loaded_data, msg)
self.vars["filename"].set(filename)
def reset_session(self):
""" Load previously saved sessions """
self.clear_session()
if self.session.stats["iterations"] == 0:
print("Training not running")
return
loaded_data = self.session.historical.sessions
msg = "Currently running training session"
self.set_session_summary(loaded_data, msg)
self.vars["filename"].set("Currently running training session")
def set_session_summary(self, data, message):
""" Set the summary data and info message """
self.summary = SessionsSummary(data).summary
self.set_info("Session: {}".format(message))
self.stats.loaded_data = data
self.stats.tree_insert_data(self.summary)
def clear_session(self):
""" Clear sessions stats """
self.summary = None
self.stats.loaded_data = None
self.stats.tree_clear()
self.reset_session_info()
def save_session(self):
""" Save sessions stats to csv """
if not self.summary:
print("No summary data loaded. Nothing to save")
return
savefile = FileHandler("save", "csv").retfile
if not savefile:
return
write_dicts = [val for val in self.summary.values()]
fieldnames = sorted(key for key in write_dicts[0].keys())
with savefile as outfile:
csvout = csv.DictWriter(outfile, fieldnames)
csvout.writeheader()
for row in write_dicts:
csvout.writerow(row)
class Options(object):
""" Options bar of Analysis tab """
def __init__(self, parent):
self.optsframe = parent.optsframe
self.parent = parent
def add_options(self):
""" Add the display tab options """
self.add_buttons()
def add_buttons(self):
""" Add the option buttons """
for btntype in ("reset", "clear", "save", "load"):
cmd = getattr(self.parent, "{}_session".format(btntype))
btn = ttk.Button(self.optsframe,
image=Images().icons[btntype],
command=cmd)
btn.pack(padx=2, side=tk.RIGHT)
hlp = self.set_help(btntype)
Tooltip(btn, text=hlp, wraplength=200)
@staticmethod
def set_help(btntype):
""" Set the helptext for option buttons """
hlp = ""
if btntype == "reset":
hlp = "Load/Refresh stats for the currently training session"
elif btntype == "clear":
hlp = "Clear currently displayed session stats"
elif btntype == "save":
hlp = "Save session stats to csv"
elif btntype == "load":
hlp = "Load saved session stats"
return hlp
class StatsData(ttk.Frame):
""" Stats frame of analysis tab """
def __init__(self,
parent,
filename,
selected_id,
helptext,
scaling_factor):
ttk.Frame.__init__(self, parent)
self.pack(side=tk.TOP,
padx=5,
pady=5,
expand=True,
fill=tk.X,
anchor=tk.N)
self.filename = filename
self.loaded_data = None
self.selected_id = selected_id
self.popup_positions = list()
self.scaling_factor = scaling_factor
self.add_label()
self.tree = ttk.Treeview(self, height=1, selectmode=tk.BROWSE)
self.scrollbar = ttk.Scrollbar(self,
orient="vertical",
command=self.tree.yview)
self.columns = self.tree_configure(helptext)
def add_label(self):
""" Add Treeview Title """
lbl = ttk.Label(self, text="Session Stats", anchor=tk.CENTER)
lbl.pack(side=tk.TOP, expand=True, fill=tk.X, padx=5, pady=5)
def tree_configure(self, helptext):
""" Build a treeview widget to hold the sessions stats """
self.tree.configure(yscrollcommand=self.scrollbar.set)
self.tree.tag_configure("total",
background="black",
foreground="white")
self.tree.pack(side=tk.LEFT, expand=True, fill=tk.X)
self.scrollbar.pack(side=tk.RIGHT, fill=tk.Y)
self.tree.bind("<ButtonRelease-1>", self.select_item)
Tooltip(self.tree, text=helptext, wraplength=200)
return self.tree_columns()
def tree_columns(self):
""" Add the columns to the totals treeview """
columns = (("session", 40, "#"),
("start", 130, None),
("end", 130, None),
("elapsed", 90, None),
("batch", 50, None),
("iterations", 90, None),
("rate", 60, "EGs/sec"))
self.tree["columns"] = [column[0] for column in columns]
for column in columns:
text = column[2] if column[2] else column[0].title()
self.tree.heading(column[0], text=text)
self.tree.column(column[0],
width=column[1],
anchor=tk.E,
minwidth=40)
self.tree.column("#0", width=40)
self.tree.heading("#0", text="Graphs")
return [column[0] for column in columns]
def tree_insert_data(self, sessions):
""" Insert the data into the totals treeview """
self.tree.configure(height=len(sessions))
for item in sessions:
values = [item[column] for column in self.columns]
kwargs = {"values": values, "image": Images().icons["graph"]}
if values[0] == "Total":
kwargs["tags"] = "total"
self.tree.insert("", "end", **kwargs)
def tree_clear(self):
""" Clear the totals tree """
self.tree.delete(* self.tree.get_children())
self.tree.configure(height=1)
def select_item(self, event):
""" Update the session summary info with
the selected item or launch graph """
region = self.tree.identify("region", event.x, event.y)
selection = self.tree.focus()
values = self.tree.item(selection, "values")
if values:
self.selected_id.set(values[0])
if region == "tree":
self.data_popup()
def data_popup(self):
""" Pop up a window and control it's position """
toplevel = SessionPopUp(self.loaded_data, self.selected_id.get())
toplevel.title(self.data_popup_title())
position = self.data_popup_get_position()
height = int(720 * self.scaling_factor)
width = int(400 * self.scaling_factor)
toplevel.geometry("{}x{}+{}+{}".format(str(height),
str(width),
str(position[0]),
str(position[1])))
toplevel.update()
def data_popup_title(self):
""" Set the data popup title """
selected_id = self.selected_id.get()
title = "All Sessions"
if selected_id != "Total":
title = "Session #{}".format(selected_id)
return "{} - {}".format(title, self.filename.get())
def data_popup_get_position(self):
""" Get the position of the next window """
init_pos = [120, 120]
pos = init_pos
while True:
if pos not in self.popup_positions:
self.popup_positions.append(pos)
break
pos = [item + 200 for item in pos]
init_pos, pos = self.data_popup_check_boundaries(init_pos, pos)
return pos
def data_popup_check_boundaries(self, initial_position, position):
""" Check that the popup remains within the screen boundaries """
boundary_x = self.winfo_screenwidth() - 120
boundary_y = self.winfo_screenheight() - 120
if position[0] >= boundary_x or position[1] >= boundary_y:
initial_position = [initial_position[0] + 50, initial_position[1]]
position = initial_position
return initial_position, position
class SessionPopUp(tk.Toplevel):
""" Pop up for detailed grap/stats for selected session """
def __init__(self, data, session_id):
tk.Toplevel.__init__(self)
self.is_totals = True if session_id == "Total" else False
self.data = self.set_session_data(data, session_id)
self.graph = None
self.display_data = None
self.vars = dict()
self.graph_initialised = False
self.build()
def set_session_data(self, sessions, session_id):
""" Set the correct list index based on the passed in session is """
if self.is_totals:
data = SessionsTotals(sessions).stats
else:
data = sessions[int(session_id) - 1]
return data
def build(self):
""" Build the popup window """
optsframe, graphframe = self.layout_frames()
self.opts_build(optsframe)
self.compile_display_data()
self.graph_build(graphframe)
def layout_frames(self):
""" Top level container frames """
leftframe = ttk.Frame(self)
leftframe.pack(side=tk.LEFT, expand=False, fill=tk.BOTH, pady=5)
sep = ttk.Frame(self, width=2, relief=tk.RIDGE)
sep.pack(fill=tk.Y, side=tk.LEFT)
rightframe = ttk.Frame(self)
rightframe.pack(side=tk.RIGHT, fill=tk.BOTH, pady=5, expand=True)
return leftframe, rightframe
def opts_build(self, frame):
""" Options in options to the optsframe """
self.opts_combobox(frame)
self.opts_checkbuttons(frame)
self.opts_entry(frame)
self.opts_buttons(frame)
sep = ttk.Frame(frame, height=2, relief=tk.RIDGE)
sep.pack(fill=tk.X, pady=(5, 0), side=tk.BOTTOM)
def opts_combobox(self, frame):
""" Add the options combo boxes """
choices = {"Display": ("Loss", "Rate"),
"Scale": ("Linear", "Log")}
for item in ["Display", "Scale"]:
var = tk.StringVar()
cmd = self.optbtn_reset if item == "Display" else self.graph_scale
var.trace("w", cmd)
cmbframe = ttk.Frame(frame)
cmbframe.pack(fill=tk.X, pady=5, padx=5, side=tk.TOP)
lblcmb = ttk.Label(cmbframe,
text="{}:".format(item),
width=7,
anchor=tk.W)
lblcmb.pack(padx=(0, 2), side=tk.LEFT)
cmb = ttk.Combobox(cmbframe, textvariable=var, width=10)
cmb["values"] = choices[item]
cmb.current(0)
cmb.pack(fill=tk.X, side=tk.RIGHT)
self.vars[item.lower().strip()] = var
hlp = self.set_help(item)
Tooltip(cmbframe, text=hlp, wraplength=200)
def opts_checkbuttons(self, frame):
""" Add the options check buttons """
for item in ("raw", "trend", "avg", "outliers"):
if item == "avg":
text = "Show Rolling Average"
elif item == "outliers":
text = "Flatten Outliers"
else:
text = "Show {}".format(item.title())
var = tk.BooleanVar()
if item == "raw":
var.set(True)
var.trace("w", self.optbtn_reset)
self.vars[item] = var
ctl = ttk.Checkbutton(frame, variable=var, text=text)
ctl.pack(side=tk.TOP, padx=5, pady=5, anchor=tk.W)
hlp = self.set_help(item)
Tooltip(ctl, text=hlp, wraplength=200)
def opts_entry(self, frame):
""" Add the options entry boxes """
for item in ("avgiterations", ):
if item == "avgiterations":
text = "Iterations to Average:"
default = "10"
entframe = ttk.Frame(frame)
entframe.pack(fill=tk.X, pady=5, padx=5, side=tk.TOP)
lbl = ttk.Label(entframe, text=text, anchor=tk.W)
lbl.pack(padx=(0, 2), side=tk.LEFT)
ctl = ttk.Entry(entframe, width=4, justify=tk.RIGHT)
ctl.pack(side=tk.RIGHT, anchor=tk.W)
ctl.insert(0, default)
hlp = self.set_help(item)
Tooltip(entframe, text=hlp, wraplength=200)
self.vars[item] = ctl
def opts_buttons(self, frame):
""" Add the option buttons """
btnframe = ttk.Frame(frame)
btnframe.pack(fill=tk.X, pady=5, padx=5, side=tk.BOTTOM)
for btntype in ("reset", "save"):
cmd = getattr(self, "optbtn_{}".format(btntype))
btn = ttk.Button(btnframe,
image=Images().icons[btntype],
command=cmd)
btn.pack(padx=2, side=tk.RIGHT)
hlp = self.set_help(btntype)
Tooltip(btn, text=hlp, wraplength=200)
def optbtn_save(self):
""" Action for save button press """
savefile = FileHandler("save", "csv").retfile
if not savefile:
return
save_data = self.display_data.stats
fieldnames = sorted(key for key in save_data.keys())
with savefile as outfile:
csvout = csv.writer(outfile, delimiter=",")
csvout.writerow(fieldnames)
csvout.writerows(zip(*[save_data[key] for key in fieldnames]))
def optbtn_reset(self, *args):
""" Action for reset button press and checkbox changes"""
if not self.graph_initialised:
return
self.compile_display_data()
self.graph.refresh(self.display_data,
self.vars["display"].get(),
self.vars["scale"].get())
def graph_scale(self, *args):
""" Action for changing graph scale """
if not self.graph_initialised:
return
self.graph.set_yscale_type(self.vars["scale"].get())
@staticmethod
def set_help(control):
""" Set the helptext for option buttons """
hlp = ""
control = control.lower()
if control == "reset":
hlp = "Refresh graph"
elif control == "save":
hlp = "Save display data to csv"
elif control == "avgiterations":
hlp = "Number of data points to sample for rolling average"
elif control == "outliers":
hlp = "Flatten data points that fall more than 1 standard " \
"deviation from the mean to the mean value."
elif control == "avg":
hlp = "Display rolling average of the data"
elif control == "raw":
hlp = "Display raw data"
elif control == "trend":
hlp = "Display polynormal data trend"
elif control == "display":
hlp = "Set the data to display"
elif control == "scale":
hlp = "Change y-axis scale"
return hlp
def compile_display_data(self):
""" Compile the data to be displayed """
self.display_data = Calculations(self.data,
self.vars["display"].get(),
self.selections_to_list(),
self.vars["avgiterations"].get(),
self.vars["outliers"].get(),
self.is_totals)
def selections_to_list(self):
""" Compile checkbox selections to list """
selections = list()
for key, val in self.vars.items():
if (isinstance(val, tk.BooleanVar)
and key != "outliers"
and val.get()):
selections.append(key)
return selections
def graph_build(self, frame):
""" Build the graph in the top right paned window """
self.graph = SessionGraph(frame,
self.display_data,
self.vars["display"].get(),
self.vars["scale"].get())
self.graph.pack(expand=True, fill=tk.BOTH)
self.graph.build()
self.graph_initialised = True

185
lib/gui/display_command.py Normal file
View File

@ -0,0 +1,185 @@
#!/usr/bin python3
""" Command specific tabs of Display Frame of the Faceswap GUI """
import datetime
import os
import tkinter as tk
from tkinter import ttk
from .display_graph import TrainingGraph
from .display_page import DisplayOptionalPage
from .tooltip import Tooltip
from .stats import Calculations
from .utils import Images, FileHandler
class PreviewExtract(DisplayOptionalPage):
""" Tab to display output preview images for extract and convert """
def display_item_set(self):
""" Load the latest preview if available """
Images().load_latest_preview()
self.display_item = Images().previewoutput
def display_item_process(self):
""" Display the preview """
if not self.subnotebook.children:
self.add_child()
else:
self.update_child()
def add_child(self):
""" Add the preview label child """
preview = self.subnotebook_add_page(self.tabname, widget=None)
lblpreview = ttk.Label(preview, image=Images().previewoutput[1])
lblpreview.pack(side=tk.TOP, anchor=tk.NW)
Tooltip(lblpreview, text=self.helptext, wraplength=200)
def update_child(self):
""" Update the preview image on the label """
for widget in self.subnotebook_get_widgets():
widget.configure(image=Images().previewoutput[1])
def save_items(self):
""" Open save dialogue and save preview """
location = FileHandler("dir").retfile
if not location:
return
filename = "extract_convert_preview"
now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
filename = os.path.join(location,
"{}_{}.{}".format(filename,
now,
"png"))
Images().previewoutput[0].save(filename)
print("Saved preview to {}".format(filename))
class PreviewTrain(DisplayOptionalPage):
""" Training preview image(s) """
def display_item_set(self):
""" Load the latest preview if available """
Images().load_training_preview()
self.display_item = Images().previewtrain
def display_item_process(self):
""" Display the preview(s) resized as appropriate """
sortednames = sorted([name for name in Images().previewtrain.keys()])
existing = self.subnotebook_get_titles_ids()
for name in sortednames:
if name not in existing.keys():
self.add_child(name)
else:
tab_id = existing[name]
self.update_child(tab_id, name)
def add_child(self, name):
""" Add the preview canvas child """
preview = PreviewTrainCanvas(self.subnotebook, name)
preview = self.subnotebook_add_page(name, widget=preview)
Tooltip(preview, text=self.helptext, wraplength=200)
self.vars["modified"].set(Images().previewtrain[name][2])
def update_child(self, tab_id, name):
""" Update the preview canvas """
if self.vars["modified"].get() != Images().previewtrain[name][2]:
self.vars["modified"].set(Images().previewtrain[name][2])
widget = self.subnotebook_page_from_id(tab_id)
widget.reload()
def save_items(self):
""" Open save dialogue and save preview """
location = FileHandler("dir").retfile
if not location:
return
for preview in self.subnotebook.children.values():
preview.save_preview(location)
class PreviewTrainCanvas(ttk.Frame):
""" Canvas to hold a training preview image """
def __init__(self, parent, previewname):
ttk.Frame.__init__(self, parent)
self.name = previewname
Images().resize_image(self.name, None)
self.previewimage = Images().previewtrain[self.name][1]
self.canvas = tk.Canvas(self, bd=0, highlightthickness=0)
self.canvas.pack(side=tk.TOP, fill=tk.BOTH, expand=True)
self.imgcanvas = self.canvas.create_image(0,
0,
image=self.previewimage,
anchor=tk.NW)
self.bind("<Configure>", self.resize)
def resize(self, event):
""" Resize the image to fit the frame, maintaining aspect ratio """
framesize = (event.width, event.height)
# Sometimes image is resized before frame is drawn
framesize = None if framesize == (1, 1) else framesize
Images().resize_image(self.name, framesize)
self.reload()
def reload(self):
""" Reload the preview image """
self.previewimage = Images().previewtrain[self.name][1]
self.canvas.itemconfig(self.imgcanvas, image=self.previewimage)
def save_preview(self, location):
""" Save the figure to file """
filename = self.name
now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
filename = os.path.join(location,
"{}_{}.{}".format(filename,
now,
"png"))
Images().previewtrain[self.name][0].save(filename)
print("Saved preview to {}".format(filename))
class GraphDisplay(DisplayOptionalPage):
""" The Graph Tab of the Display section """
def display_item_set(self):
""" Load the graph(s) if available """
if self.session.stats["iterations"] == 0:
self.display_item = None
else:
self.display_item = self.session.stats
def display_item_process(self):
""" Add a single graph to the graph window """
losskeys = self.display_item["losskeys"]
loss = self.display_item["loss"]
tabcount = int(len(losskeys) / 2)
existing = self.subnotebook_get_titles_ids()
for i in range(tabcount):
selectedkeys = losskeys[i * 2:(i + 1) * 2]
name = " - ".join(selectedkeys).title().replace("_", " ")
if name not in existing.keys():
selectedloss = loss[i * 2:(i + 1) * 2]
selection = {"loss": selectedloss,
"losskeys": selectedkeys}
data = Calculations(session=selection,
display="loss",
selections=["raw", "trend"])
self.add_child(name, data)
def add_child(self, name, data):
""" Add the graph for the selected keys """
graph = TrainingGraph(self.subnotebook, data, "Loss")
graph.build()
graph = self.subnotebook_add_page(name, widget=graph)
Tooltip(graph, text=self.helptext, wraplength=200)
def save_items(self):
""" Open save dialogue and save graphs """
graphlocation = FileHandler("dir").retfile
if not graphlocation:
return
for graph in self.subnotebook.children.values():
graph.save_fig(graphlocation)

333
lib/gui/display_graph.py Executable file
View File

@ -0,0 +1,333 @@
#!/usr/bin python3
""" Graph functions for Display Frame of the Faceswap GUI """
import datetime
import os
import tkinter as tk
from tkinter import ttk
from math import ceil, floor
import matplotlib
matplotlib.use("TkAgg")
import matplotlib.animation as animation
from matplotlib import pyplot as plt, style
from matplotlib.backends.backend_tkagg import (FigureCanvasTkAgg,
NavigationToolbar2Tk)
from .tooltip import Tooltip
from .utils import Images
class NavigationToolbar(NavigationToolbar2Tk):
""" Same as default, but only including buttons we need
with custom icons and layout
From: https://stackoverflow.com/questions/12695678 """
toolitems = [t for t in NavigationToolbar2Tk.toolitems if
t[0] in ("Home", "Pan", "Zoom", "Save")]
@staticmethod
def _Button(frame, text, file, command, extension=".gif"):
""" Map Buttons to their own frame.
Use custom button icons,
Use ttk buttons
pack to the right """
iconmapping = {"home": "reset",
"filesave": "save",
"zoom_to_rect": "zoom"}
icon = iconmapping[file] if iconmapping.get(file, None) else file
img = Images().icons[icon]
btn = ttk.Button(frame, text=text, image=img, command=command)
btn.pack(side=tk.RIGHT, padx=2)
return btn
def _init_toolbar(self):
""" Same as original but ttk widgets and standard
tooltips used. Separator added and message label
packed to the left """
xmin, xmax = self.canvas.figure.bbox.intervalx
height, width = 50, xmax-xmin
ttk.Frame.__init__(self, master=self.window,
width=int(width), height=int(height))
sep = ttk.Frame(self, height=2, relief=tk.RIDGE)
sep.pack(fill=tk.X, pady=(5, 0), side=tk.TOP)
self.update() # Make axes menu
btnframe = ttk.Frame(self)
btnframe.pack(fill=tk.X, padx=5, pady=5, side=tk.RIGHT)
for text, tooltip_text, image_file, callback in self.toolitems:
if text is None:
# Add a spacer; return value is unused.
self._Spacer()
else:
button = self._Button(btnframe, text=text, file=image_file,
command=getattr(self, callback))
if tooltip_text is not None:
Tooltip(button, text=tooltip_text, wraplength=200)
self.message = tk.StringVar(master=self)
self._message_label = ttk.Label(master=self, textvariable=self.message)
self._message_label.pack(side=tk.LEFT, padx=5)
self.pack(side=tk.BOTTOM, fill=tk.X)
class GraphBase(ttk.Frame):
""" Base class for matplotlib line graphs """
def __init__(self, parent, data, ylabel):
ttk.Frame.__init__(self, parent)
style.use("ggplot")
self.calcs = data
self.ylabel = ylabel
self.colourmaps = ["Reds", "Blues", "Greens",
"Purples", "Oranges", "Greys",
"copper", "summer", "bone"]
self.lines = list()
self.toolbar = None
self.fig = plt.figure(figsize=(4, 4), dpi=75)
self.ax1 = self.fig.add_subplot(1, 1, 1)
self.plotcanvas = FigureCanvasTkAgg(self.fig, self)
self.initiate_graph()
self.update_plot(initiate=True)
def initiate_graph(self):
""" Place the graph canvas """
self.plotcanvas.get_tk_widget().pack(side=tk.TOP,
padx=5,
fill=tk.BOTH,
expand=True)
plt.subplots_adjust(left=0.100,
bottom=0.100,
right=0.95,
top=0.95,
wspace=0.2,
hspace=0.2)
def update_plot(self, initiate=True):
""" Update the plot with incoming data """
if initiate:
self.lines = list()
self.ax1.clear()
self.axes_labels_set()
fulldata = [item for item in self.calcs.stats.values()]
self.axes_limits_set(fulldata)
xrng = [x for x in range(self.calcs.iterations)]
keys = list(self.calcs.stats.keys())
for idx, item in enumerate(self.lines_sort(keys)):
if initiate:
self.lines.extend(self.ax1.plot(xrng,
self.calcs.stats[item[0]],
label=item[1],
linewidth=item[2],
color=item[3]))
else:
self.lines[idx].set_data(xrng, self.calcs.stats[item[0]])
if initiate:
self.legend_place()
def axes_labels_set(self):
""" Set the axes label and range """
self.ax1.set_xlabel("Iterations")
self.ax1.set_ylabel(self.ylabel)
def axes_limits_set_default(self):
""" Set default axes limits """
self.ax1.set_ylim(0.00, 100.0)
self.ax1.set_xlim(0, 1)
def axes_limits_set(self, data):
""" Set the axes limits """
xmax = self.calcs.iterations - 1 if self.calcs.iterations > 1 else 1
if data:
ymin, ymax = self.axes_data_get_min_max(data)
self.ax1.set_ylim(ymin, ymax)
self.ax1.set_xlim(0, xmax)
else:
self.axes_limits_set_default()
@staticmethod
def axes_data_get_min_max(data):
""" Return the minimum and maximum values from list of lists """
ymin, ymax = list(), list()
for item in data:
dataset = list(filter(lambda x: x is not None, item))
if not dataset:
continue
ymin.append(min(dataset) * 1000)
ymax.append(max(dataset) * 1000)
ymin = floor(min(ymin)) / 1000
ymax = ceil(max(ymax)) / 1000
return ymin, ymax
def axes_set_yscale(self, scale):
""" Set the Y-Scale to log or linear """
self.ax1.set_yscale(scale)
def lines_sort(self, keys):
""" Sort the data keys into consistent order
and set line colourmap and line width """
raw_lines = list()
sorted_lines = list()
for key in sorted(keys):
title = key.replace("_", " ").title()
if key.startswith(("avg", "trend")):
sorted_lines.append([key, title])
else:
raw_lines.append([key, title])
groupsize = self.lines_groupsize(raw_lines, sorted_lines)
sorted_lines = raw_lines + sorted_lines
lines = self.lines_style(sorted_lines, groupsize)
return lines
@staticmethod
def lines_groupsize(raw_lines, sorted_lines):
""" Get the number of items in each group.
If raw data isn't selected, then check
the length of remaining groups until
something is found """
groupsize = 1
if raw_lines:
groupsize = len(raw_lines)
else:
for check in ("avg", "trend"):
if any(item[0].startswith(check) for item in sorted_lines):
groupsize = len([item for item in sorted_lines
if item[0].startswith(check)])
break
return groupsize
def lines_style(self, lines, groupsize):
""" Set the colourmap and linewidth for each group """
groups = int(len(lines) / groupsize)
colours = self.lines_create_colors(groupsize, groups)
for idx, item in enumerate(lines):
linewidth = ceil((idx + 1) / groupsize)
item.extend((linewidth, colours[idx]))
return lines
def lines_create_colors(self, groupsize, groups):
""" Create the colours """
colours = list()
for i in range(1, groups + 1):
for colour in self.colourmaps[0:groupsize]:
cmap = matplotlib.cm.get_cmap(colour)
cpoint = 1 - (i / 5)
colours.append(cmap(cpoint))
return colours
def legend_place(self):
""" Place and format legend """
self.ax1.legend(loc="upper right", ncol=2)
def toolbar_place(self, parent):
""" Add Graph Navigation toolbar """
self.toolbar = NavigationToolbar(self.plotcanvas, parent)
self.toolbar.pack(side=tk.BOTTOM)
self.toolbar.update()
class TrainingGraph(GraphBase):
""" Live graph to be displayed during training. """
def __init__(self, parent, data, ylabel):
GraphBase.__init__(self, parent, data, ylabel)
self.anim = None
def build(self):
""" Update the plot area with loss values and cycle through to
animate """
self.anim = animation.FuncAnimation(self.fig,
self.animate,
interval=200,
blit=False)
self.plotcanvas.draw()
def animate(self, i):
""" Read loss data and apply to graph """
self.calcs.refresh()
self.update_plot(initiate=False)
def set_animation_rate(self, iterations):
""" Change the animation update interval based on how
many iterations have been
There's no point calculating a graph over thousands of
points of data when the change will be miniscule """
if iterations > 30000:
speed = 60000 # 1 min updates
elif iterations > 20000:
speed = 30000 # 30 sec updates
elif iterations > 10000:
speed = 10000 # 10 sec updates
elif iterations > 5000:
speed = 5000 # 5 sec updates
elif iterations > 1000:
speed = 2000 # 2 sec updates
elif iterations > 500:
speed = 1000 # 1 sec updates
elif iterations > 100:
speed = 500 # 0.5 sec updates
else:
speed = 200 # 200ms updates
if not self.anim.event_source.interval == speed:
self.anim.event_source.interval = speed
def save_fig(self, location):
""" Save the figure to file """
keys = sorted([key.replace("raw_", "")
for key in self.calcs.stats.keys()
if key.startswith("raw_")])
filename = " - ".join(keys)
now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
filename = os.path.join(location,
"{}_{}.{}".format(filename,
now,
"png"))
self.fig.set_size_inches(16, 9)
self.fig.savefig(filename, bbox_inches="tight", dpi=120)
print("Saved graph to {}".format(filename))
self.resize_fig()
def resize_fig(self):
""" Resize the figure back to the canvas """
class Event(object):
""" Event class that needs to be passed to
plotcanvas.resize """
pass
Event.width = self.winfo_width()
Event.height = self.winfo_height()
self.plotcanvas.resize(Event)
class SessionGraph(GraphBase):
""" Session Graph for session pop-up """
def __init__(self, parent, data, ylabel, scale):
GraphBase.__init__(self, parent, data, ylabel)
self.scale = scale
def build(self):
""" Build the session graph """
self.toolbar_place(self)
self.plotcanvas.draw()
def refresh(self, data, ylabel, scale):
""" Refresh graph data """
self.calcs = data
self.ylabel = ylabel
self.set_yscale_type(scale)
def set_yscale_type(self, scale):
""" switch the y-scale and redraw """
self.scale = scale
self.update_plot(initiate=True)
self.axes_set_yscale(self.scale)
self.plotcanvas.draw()

234
lib/gui/display_page.py Normal file
View File

@ -0,0 +1,234 @@
#!/usr/bin python3
""" Display Page parent classes for display section of the Faceswap GUI """
import tkinter as tk
from tkinter import ttk
from .tooltip import Tooltip
from .utils import Images
class DisplayPage(ttk.Frame):
""" Parent frame holder for each tab.
Defines uniform structure for each tab to inherit from """
def __init__(self, parent, tabname, helptext):
ttk.Frame.__init__(self, parent)
self.pack(fill=tk.BOTH, side=tk.TOP, anchor=tk.NW)
self.session = parent.session
self.runningtask = parent.runningtask
self.helptext = helptext
self.tabname = tabname
self.vars = {"info": tk.StringVar()}
self.add_optional_vars(self.set_vars())
self.subnotebook = self.add_subnotebook()
self.optsframe = self.add_options_frame()
self.add_options_info()
self.add_frame_separator()
self.set_mainframe_single_tab_style()
parent.add(self, text=self.tabname.title())
def add_optional_vars(self, varsdict):
""" Add page specific variables """
if isinstance(varsdict, dict):
for key, val in varsdict.items():
self.vars[key] = val
@staticmethod
def set_vars():
""" Overide to return a dict of page specific variables """
return dict()
def add_subnotebook(self):
""" Add the main frame notebook """
notebook = ttk.Notebook(self)
notebook.pack(side=tk.TOP, anchor=tk.NW, fill=tk.BOTH, expand=True)
return notebook
def add_options_frame(self):
""" Add the display tab options """
optsframe = ttk.Frame(self)
optsframe.pack(side=tk.BOTTOM, padx=5, pady=5, fill=tk.X)
return optsframe
def add_options_info(self):
""" Add the info bar """
lblinfo = ttk.Label(self.optsframe,
textvariable=self.vars["info"],
anchor=tk.W,
width=70)
lblinfo.pack(side=tk.LEFT, padx=5, pady=5, anchor=tk.W)
def set_info(self, msg):
""" Set the info message """
self.vars["info"].set(msg)
def add_frame_separator(self):
""" Add a separator between top and bottom frames """
sep = ttk.Frame(self, height=2, relief=tk.RIDGE)
sep.pack(fill=tk.X, pady=(5, 0), side=tk.BOTTOM)
@staticmethod
def set_mainframe_single_tab_style():
""" Configure ttk notebook style to represent a single frame """
nbstyle = ttk.Style()
nbstyle.configure("single.TNotebook", borderwidth=0)
nbstyle.layout("single.TNotebook.Tab", [])
def subnotebook_add_page(self, tabtitle, widget=None):
""" Add a page to the sub notebook """
frame = widget if widget else ttk.Frame(self.subnotebook)
frame.pack(padx=5, pady=5, fill=tk.BOTH, expand=True)
self.subnotebook.add(frame, text=tabtitle)
self.subnotebook_configure()
return frame
def subnotebook_configure(self):
""" Configure notebook to display or hide tabs """
if len(self.subnotebook.children) == 1:
self.subnotebook.configure(style="single.TNotebook")
else:
self.subnotebook.configure(style="TNotebook")
def subnotebook_hide(self):
""" Hide the subnotebook. Used for hiding
Optional displays """
if self.subnotebook.winfo_ismapped():
self.subnotebook.pack_forget()
def subnotebook_show(self):
""" Show subnotebook. Used for displaying
Optional displays """
if not self.subnotebook.winfo_ismapped():
self.subnotebook.pack(side=tk.TOP,
anchor=tk.NW,
fill=tk.BOTH,
expand=True)
def subnotebook_get_widgets(self):
""" Return each widget that sits within each
subnotebook frame """
for child in self.subnotebook.winfo_children():
for widget in child.winfo_children():
yield widget
def subnotebook_get_titles_ids(self):
""" Return tabs ids and titles """
tabs = dict()
for tab_id in range(0, self.subnotebook.index("end")):
tabs[self.subnotebook.tab(tab_id, "text")] = tab_id
return tabs
def subnotebook_page_from_id(self, tab_id):
""" Return subnotebook tab widget from it's ID """
tab_name = self.subnotebook.tabs()[tab_id].split(".")[-1]
return self.subnotebook.children[tab_name]
class DisplayOptionalPage(DisplayPage):
""" Parent Context Sensitive Display Tab """
def __init__(self, parent, tabname, helptext, waittime):
DisplayPage.__init__(self, parent, tabname, helptext)
self.display_item = None
self.set_info_text()
self.add_options()
parent.select(self)
self.update_idletasks()
self.update_page(waittime)
@staticmethod
def set_vars():
""" Analysis specific vars """
enabled = tk.BooleanVar()
enabled.set(True)
ready = tk.BooleanVar()
ready.set(False)
modified = tk.DoubleVar()
modified.set(None)
return {"enabled": enabled,
"ready": ready,
"modified": modified}
# INFO LABEL
def set_info_text(self):
""" Set waiting for display text """
if not self.vars["enabled"].get():
self.set_info("{} disabled".format(self.tabname.title()))
elif self.vars["enabled"].get() and not self.vars["ready"].get():
self.set_info("Waiting for {}...".format(self.tabname))
else:
self.set_info("Displaying {}".format(self.tabname))
# DISPLAY OPTIONS BAR
def add_options(self):
""" Add the additional options """
self.add_option_save()
self.add_option_enable()
def add_option_save(self):
""" Add save button to save page output to file """
btnsave = ttk.Button(self.optsframe,
image=Images().icons["save"],
command=self.save_items)
btnsave.pack(padx=2, side=tk.RIGHT)
Tooltip(btnsave,
text="Save {}(s) to file".format(self.tabname),
wraplength=200)
def add_option_enable(self):
""" Add checkbutton to enable/disable page """
chkenable = ttk.Checkbutton(self.optsframe,
variable=self.vars["enabled"],
text="Enable {}".format(self.tabname),
command=self.on_chkenable_change)
chkenable.pack(side=tk.RIGHT, padx=5, anchor=tk.W)
Tooltip(chkenable,
text="Enable or disable {} display".format(self.tabname),
wraplength=200)
def save_items(self):
""" Save items. Override for display specific saving """
raise NotImplementedError()
def on_chkenable_change(self):
""" Update the display immediately on a checkbutton change """
if self.vars["enabled"].get():
self.subnotebook_show()
else:
self.subnotebook_hide()
self.set_info_text()
def update_page(self, waittime):
""" Update the latest preview item """
if not self.runningtask.get():
return
if self.vars["enabled"].get():
self.display_item_set()
self.load_display()
self.after(waittime, lambda t=waittime: self.update_page(t))
def display_item_set(self):
""" Override for display specific loading """
raise NotImplementedError()
def load_display(self):
""" Load the display """
if not self.display_item:
return
self.display_item_process()
self.vars["ready"].set(True)
self.set_info_text()
def display_item_process(self):
""" Override for display specific loading """
raise NotImplementedError()

259
lib/gui/options.py Normal file
View File

@ -0,0 +1,259 @@
#!/usr/bin python3
""" Cli Options and Config functions for the GUI """
import inspect
from argparse import SUPPRESS
from tkinter import ttk
import lib.cli as cli
from lib.Serializer import JSONSerializer
import tools.cli as ToolsCli
from .utils import FileHandler, Images
class CliOptions(object):
""" Class and methods for the command line options """
def __init__(self):
self.categories = ("faceswap", "tools")
self.commands = dict()
self.opts = dict()
self.build_options()
def build_options(self):
""" Get the commands that belong to each category """
for category in self.categories:
src = ToolsCli if category == "tools" else cli
mod_classes = self.get_cli_classes(src)
self.commands[category] = self.sort_commands(category, mod_classes)
self.opts.update(self.extract_options(src, mod_classes))
@staticmethod
def get_cli_classes(cli_source):
""" Parse the cli scripts for the arg classes """
mod_classes = list()
for name, obj in inspect.getmembers(cli_source):
if inspect.isclass(obj) and name.lower().endswith("args") \
and name.lower() not in (("faceswapargs",
"extractconvertargs",
"guiargs")):
mod_classes.append(name)
return mod_classes
def sort_commands(self, category, classes):
""" Format classes into command names and sort:
Specific workflow order for faceswap.
Alphabetical for all others """
commands = sorted(self.format_command_name(command)
for command in classes)
if category == "faceswap":
ordered = ["extract", "train", "convert"]
commands = ordered + [command for command in commands
if command not in ordered]
return commands
@staticmethod
def format_command_name(classname):
""" Format args class name to command """
return classname.lower()[:-4]
def extract_options(self, cli_source, mod_classes):
""" Extract the existing ArgParse Options
into master options Dictionary """
subopts = dict()
for classname in mod_classes:
command = self.format_command_name(classname)
options = self.get_cli_arguments(cli_source, classname, command)
self.process_options(options)
subopts[command] = options
return subopts
@staticmethod
def get_cli_arguments(cli_source, classname, command):
""" Extract the options from the main and tools cli files """
meth = getattr(cli_source, classname)(None, command)
return meth.argument_list + meth.optional_arguments
def process_options(self, command_options):
""" Process the options for a single command """
for opt in command_options:
if opt.get("help", "") == SUPPRESS:
command_options.remove(opt)
ctl, sysbrowser, filetypes, action_option = 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["action_option"] = action_option
@staticmethod
def set_control_title(opts):
""" Take the option switch and format it nicely """
ctltitle = opts[1] if len(opts) == 2 else opts[0]
ctltitle = ctltitle.replace("-", " ").replace("_", " ").strip().title()
return ctltitle
def set_control(self, option):
""" Set the control and filesystem browser to use for each option """
sysbrowser = None
action = option.get("action", None)
action_option = option.get("action_option", None)
filetypes = option.get("filetypes", None)
ctl = ttk.Entry
if action in (cli.FullPaths,
cli.DirFullPaths,
cli.FileFullPaths,
cli.SaveFileFullPaths,
cli.ContextFullPaths):
sysbrowser, filetypes = self.set_sysbrowser(action,
filetypes,
action_option)
elif option.get("choices", "") != "":
ctl = ttk.Combobox
elif option.get("action", "") == "store_true":
ctl = ttk.Checkbutton
return ctl, sysbrowser, filetypes, action_option
@staticmethod
def set_sysbrowser(action, filetypes, action_option):
""" Set the correct file system browser and filetypes
for the passed in action """
sysbrowser = "folder"
filetypes = "default" if not filetypes else filetypes
if action == cli.FileFullPaths:
sysbrowser = "load"
elif action == cli.SaveFileFullPaths:
sysbrowser = "save"
elif action == cli.ContextFullPaths and action_option:
sysbrowser = "context"
return sysbrowser, filetypes
def set_context_option(self, command):
""" Set the tk_var for the source action option
that dictates the context sensitive file browser. """
actions = {item["opts"][0]: item["value"]
for item in self.gen_command_options(command)}
for opt in self.gen_command_options(command):
if opt["filesystem_browser"] == "context":
opt["action_option"] = actions[opt["action_option"]]
def gen_command_options(self, command):
""" Yield each option for specified command """
for option in self.opts[command]:
yield option
def options_to_process(self, command=None):
""" Return a consistent object for processing
regardless of whether processing all commands
or just one command for reset and clear """
if command is None:
options = [opt for opts in self.opts.values() for opt in opts]
else:
options = [opt for opt in self.gen_command_options(command)]
return options
def reset(self, command=None):
""" Reset the options for all or passed command
back to default value """
for option in self.options_to_process(command):
default = option.get("default", "")
default = "" if default is None else default
option["value"].set(default)
def clear(self, command=None):
""" Clear the options values for all or passed
commands """
for option in self.options_to_process(command):
if isinstance(option["value"].get(), bool):
option["value"].set(False)
elif isinstance(option["value"].get(), int):
option["value"].set(0)
else:
option["value"].set("")
def get_option_values(self, command=None):
""" Return all or single command control titles
with the associated tk_var value """
ctl_dict = dict()
for cmd, opts in self.opts.items():
if command and command != cmd:
continue
cmd_dict = dict()
for opt in opts:
cmd_dict[opt["control_title"]] = opt["value"].get()
ctl_dict[cmd] = cmd_dict
return ctl_dict
def get_one_option_variable(self, command, title):
""" Return a single tk_var for the specified
command and control_title """
for option in self.gen_command_options(command):
if option["control_title"] == title:
return option["value"]
return None
def gen_cli_arguments(self, command):
""" Return the generated cli arguments for
the selected command """
for option in self.gen_command_options(command):
optval = str(option.get("value", "").get())
opt = option["opts"][0]
if command in ("extract", "convert") and opt == "-o":
Images().pathoutput = optval
if optval == "False" or optval == "":
continue
elif optval == "True":
yield (opt, )
else:
if option.get("nargs", None):
optval = optval.split(" ")
opt = [opt] + optval
else:
opt = (opt, optval)
yield opt
class Config(object):
""" Actions for loading and saving Faceswap GUI command configurations """
def __init__(self, cli_opts, tk_vars):
self.cli_opts = cli_opts
self.serializer = JSONSerializer
self.tk_vars = tk_vars
def load(self, command=None):
""" Load a saved config file """
cfgfile = FileHandler("open", "config").retfile
if not cfgfile:
return
cfg = self.serializer.unmarshal(cfgfile.read())
opts = self.get_command_options(cfg, command) if command else cfg
for cmd, opts in opts.items():
self.set_command_args(cmd, opts)
def get_command_options(self, cfg, command):
""" return the saved options for the requested
command, if not loading global options """
opts = cfg.get(command, None)
if not opts:
self.tk_vars["consoleclear"].set(True)
print("No " + command + " section found in file")
return {command: opts}
def set_command_args(self, command, options):
""" Pass the saved config items back to the CliOptions """
if not options:
return
for srcopt, srcval in options.items():
optvar = self.cli_opts.get_one_option_variable(command, srcopt)
if not optvar:
continue
optvar.set(srcval)
def save(self, command=None):
""" Save the current GUI state to a config file in json format """
cfgfile = FileHandler("save", "config").retfile
if not cfgfile:
return
cfg = self.cli_opts.get_option_values(command)
cfgfile.write(self.serializer.marshal(cfg))
cfgfile.close()

353
lib/gui/stats.py Normal file
View File

@ -0,0 +1,353 @@
#!/usr/bin python3
""" Stats functions for the GUI """
import time
import os
import warnings
from math import ceil, sqrt
import numpy as np
from lib.Serializer import PickleSerializer
class SavedSessions(object):
""" Saved Training Session """
def __init__(self, sessions_data):
self.serializer = PickleSerializer
self.sessions = self.load_sessions(sessions_data)
def load_sessions(self, filename):
""" Load previously saved sessions """
stats = list()
if os.path.isfile(filename):
with open(filename, self.serializer.roptions) as sessions:
stats = self.serializer.unmarshal(sessions.read())
return stats
def save_sessions(self, filename):
""" Save the session file """
with open(filename, self.serializer.woptions) as session:
session.write(self.serializer.marshal(self.sessions))
print("Saved session stats to: {}".format(filename))
class CurrentSession(object):
""" The current training session """
def __init__(self):
self.stats = {"iterations": 0,
"batchsize": None, # Set and reset by wrapper
"timestamps": [],
"loss": [],
"losskeys": []}
self.timestats = {"start": None,
"elapsed": None}
self.modeldir = None # Set and reset by wrapper
self.filename = None
self.historical = None
def initialise_session(self, currentloss):
""" Initialise the training session """
self.load_historical()
for item in currentloss:
self.stats["losskeys"].append(item[0])
self.stats["loss"].append(list())
self.timestats["start"] = time.time()
def load_historical(self):
""" Load historical data and add current session to the end """
self.filename = os.path.join(self.modeldir, "trainingstats.fss")
self.historical = SavedSessions(self.filename)
self.historical.sessions.append(self.stats)
def add_loss(self, currentloss):
""" Add a loss item from the training process """
if self.stats["iterations"] == 0:
self.initialise_session(currentloss)
self.stats["iterations"] += 1
self.add_timestats()
for idx, item in enumerate(currentloss):
self.stats["loss"][idx].append(float(item[1]))
def add_timestats(self):
""" Add timestats to loss dict and timestats """
now = time.time()
self.stats["timestamps"].append(now)
elapsed_time = now - self.timestats["start"]
self.timestats["elapsed"] = time.strftime("%H:%M:%S",
time.gmtime(elapsed_time))
def save_session(self):
""" Save the session file to the modeldir """
if self.stats["iterations"] > 0:
print("Saving session stats...")
self.historical.save_sessions(self.filename)
class SessionsTotals(object):
""" The compiled totals of all saved sessions """
def __init__(self, all_sessions):
self.stats = {"split": [],
"iterations": 0,
"batchsize": [],
"timestamps": [],
"loss": [],
"losskeys": []}
self.initiate(all_sessions)
self.compile(all_sessions)
def initiate(self, sessions):
""" Initiate correct losskey titles and number of loss lists """
for losskey in sessions[0]["losskeys"]:
self.stats["losskeys"].append(losskey)
self.stats["loss"].append(list())
def compile(self, sessions):
""" Compile all of the sessions into totals """
current_split = 0
for session in sessions:
iterations = session["iterations"]
current_split += iterations
self.stats["split"].append(current_split)
self.stats["iterations"] += iterations
self.stats["timestamps"].extend(session["timestamps"])
self.stats["batchsize"].append(session["batchsize"])
self.add_loss(session["loss"])
def add_loss(self, session_loss):
""" Add loss vals to each of their respective lists """
for idx, loss in enumerate(session_loss):
self.stats["loss"][idx].extend(loss)
class SessionsSummary(object):
""" Calculations for analysis summary stats """
def __init__(self, raw_data):
self.summary = list()
self.summary_stats_compile(raw_data)
def summary_stats_compile(self, raw_data):
""" Compile summary stats """
raw_summaries = list()
for idx, session in enumerate(raw_data):
raw_summaries.append(self.summarise_session(idx, session))
totals_summary = self.summarise_totals(raw_summaries)
raw_summaries.append(totals_summary)
self.format_summaries(raw_summaries)
# Compile Session Summaries
@staticmethod
def summarise_session(idx, session):
""" Compile stats for session passed in """
starttime = session["timestamps"][0]
endtime = session["timestamps"][-1]
elapsed = endtime - starttime
rate = (session["batchsize"] * session["iterations"]) / elapsed
return {"session": idx + 1,
"start": starttime,
"end": endtime,
"elapsed": elapsed,
"rate": rate,
"batch": session["batchsize"],
"iterations": session["iterations"]}
@staticmethod
def summarise_totals(raw_summaries):
""" Compile the stats for all sessions combined """
elapsed = 0
rate = 0
batchset = set()
iterations = 0
total_summaries = len(raw_summaries)
for idx, summary in enumerate(raw_summaries):
if idx == 0:
starttime = summary["start"]
if idx == total_summaries - 1:
endtime = summary["end"]
elapsed += summary["elapsed"]
rate += summary["rate"]
batchset.add(summary["batch"])
iterations += summary["iterations"]
batch = ",".join(str(bs) for bs in batchset)
return {"session": "Total",
"start": starttime,
"end": endtime,
"elapsed": elapsed,
"rate": rate / total_summaries,
"batch": batch,
"iterations": iterations}
def format_summaries(self, raw_summaries):
""" Format the summaries nicely for display """
for summary in raw_summaries:
summary["start"] = time.strftime("%x %X",
time.gmtime(summary["start"]))
summary["end"] = time.strftime("%x %X",
time.gmtime(summary["end"]))
summary["elapsed"] = time.strftime("%H:%M:%S",
time.gmtime(summary["elapsed"]))
summary["rate"] = "{0:.1f}".format(summary["rate"])
self.summary = raw_summaries
class Calculations(object):
""" Class to hold calculations against raw session data """
def __init__(self,
session,
display="loss",
selections=["raw"],
avg_samples=10,
flatten_outliers=False,
is_totals=False):
warnings.simplefilter("ignore", np.RankWarning)
self.session = session
if display.lower() == "loss":
display = self.session["losskeys"]
else:
display = [display]
self.args = {"display": display,
"selections": selections,
"avg_samples": int(avg_samples),
"flatten_outliers": flatten_outliers,
"is_totals": is_totals}
self.iterations = 0
self.stats = None
self.refresh()
def refresh(self):
""" Refresh the stats """
self.iterations = 0
self.stats = self.get_raw()
self.get_calculations()
self.remove_raw()
def get_raw(self):
""" Add raw data to stats dict """
raw = dict()
for idx, item in enumerate(self.args["display"]):
if item.lower() == "rate":
data = self.calc_rate(self.session)
else:
data = self.session["loss"][idx][:]
if self.args["flatten_outliers"]:
data = self.flatten_outliers(data)
if self.iterations == 0:
self.iterations = len(data)
raw["raw_{}".format(item)] = data
return raw
def remove_raw(self):
""" Remove raw values from stats if not requested """
if "raw" in self.args["selections"]:
return
for key in list(self.stats.keys()):
if key.startswith("raw"):
del self.stats[key]
def calc_rate(self, data):
""" Calculate rate per iteration
NB: For totals, gaps between sessions can be large
so time diffeence has to be reset for each session's
rate calculation """
batchsize = data["batchsize"]
if self.args["is_totals"]:
split = data["split"]
else:
batchsize = [batchsize]
split = [len(data["timestamps"])]
prev_split = 0
rate = list()
for idx, current_split in enumerate(split):
prev_time = data["timestamps"][prev_split]
timestamp_chunk = data["timestamps"][prev_split:current_split]
for item in timestamp_chunk:
current_time = item
timediff = current_time - prev_time
iter_rate = 0 if timediff == 0 else batchsize[idx] / timediff
rate.append(iter_rate)
prev_time = current_time
prev_split = current_split
if self.args["flatten_outliers"]:
rate = self.flatten_outliers(rate)
return rate
@staticmethod
def flatten_outliers(data):
""" Remove the outliers from a provided list """
retdata = list()
samples = len(data)
mean = (sum(data) / samples)
limit = sqrt(sum([(item - mean)**2 for item in data]) / samples)
for item in data:
if (mean - limit) <= item <= (mean + limit):
retdata.append(item)
else:
retdata.append(mean)
return retdata
def get_calculations(self):
""" Perform the required calculations """
for selection in self.get_selections():
if selection[0] == "raw":
continue
method = getattr(self, "calc_{}".format(selection[0]))
key = "{}_{}".format(selection[0], selection[1])
raw = self.stats["raw_{}".format(selection[1])]
self.stats[key] = method(raw)
def get_selections(self):
""" Compile a list of data to be calculated """
for summary in self.args["selections"]:
for item in self.args["display"]:
yield summary, item
def calc_avg(self, data):
""" Calculate rolling average """
avgs = list()
presample = ceil(self.args["avg_samples"] / 2)
postsample = self.args["avg_samples"] - presample
datapoints = len(data)
if datapoints <= (self.args["avg_samples"] * 2):
print("Not enough data to compile rolling average")
return avgs
for idx in range(0, datapoints):
if idx < presample or idx >= datapoints - postsample:
avgs.append(None)
continue
else:
avg = sum(data[idx - presample:idx + postsample]) \
/ self.args["avg_samples"]
avgs.append(avg)
return avgs
@staticmethod
def calc_trend(data):
""" Compile trend data """
points = len(data)
if points < 10:
dummy = [None for i in range(points)]
return dummy
x_range = range(points)
fit = np.polyfit(x_range, data, 3)
poly = np.poly1d(fit)
trend = poly(x_range)
return trend

82
lib/gui/statusbar.py Normal file
View File

@ -0,0 +1,82 @@
#!/usr/bin python3
""" Status bar for the GUI """
import tkinter as tk
from tkinter import ttk
class StatusBar(ttk.Frame):
""" Status Bar for displaying the Status Message and
Progress Bar """
def __init__(self, parent):
ttk.Frame.__init__(self, parent)
self.pack(side=tk.BOTTOM, padx=10, pady=2, fill=tk.X, expand=False)
self.status_message = tk.StringVar()
self.pbar_message = tk.StringVar()
self.pbar_position = tk.IntVar()
self.status_message.set("Ready")
self.status()
self.pbar = self.progress_bar()
def status(self):
""" Place Status into bottom bar """
statusframe = ttk.Frame(self)
statusframe.pack(side=tk.LEFT, anchor=tk.W, fill=tk.X, expand=False)
lbltitle = ttk.Label(statusframe, text="Status:", width=6, anchor=tk.W)
lbltitle.pack(side=tk.LEFT, expand=False)
lblstatus = ttk.Label(statusframe,
width=20,
textvariable=self.status_message,
anchor=tk.W)
lblstatus.pack(side=tk.LEFT, anchor=tk.W, fill=tk.X, expand=True)
def progress_bar(self):
""" Place progress bar into bottom bar """
progressframe = ttk.Frame(self)
progressframe.pack(side=tk.RIGHT, anchor=tk.E, fill=tk.X)
lblmessage = ttk.Label(progressframe, textvariable=self.pbar_message)
lblmessage.pack(side=tk.LEFT, padx=3, fill=tk.X, expand=True)
pbar = ttk.Progressbar(progressframe,
length=200,
variable=self.pbar_position,
maximum=1000,
mode="determinate")
pbar.pack(side=tk.LEFT, padx=2, fill=tk.X, expand=True)
pbar.pack_forget()
return pbar
def progress_start(self, mode):
""" Set progress bar mode and display """
self.progress_set_mode(mode)
self.pbar.pack()
def progress_stop(self):
""" Reset progress bar and hide """
self.pbar_message.set("")
self.pbar_position.set(0)
self.progress_set_mode("determinate")
self.pbar.pack_forget()
def progress_set_mode(self, mode):
""" Set the progress bar mode """
self.pbar.config(mode=mode)
if mode == "indeterminate":
self.pbar.config(maximum=100)
self.pbar.start()
else:
self.pbar.stop()
self.pbar.config(maximum=1000)
def progress_update(self, message, position, update_position=True):
""" Update the GUIs progress bar and position """
self.pbar_message.set(message)
if update_position:
self.pbar_position.set(position)

165
lib/gui/tooltip.py Executable file
View File

@ -0,0 +1,165 @@
#!/usr/bin python3
""" Tooltip. Pops up help messages for the GUI """
import platform
import tkinter as tk
class Tooltip:
"""
Create a tooltip for a given widget as the mouse goes on it.
Adapted from StackOverflow:
http://stackoverflow.com/questions/3221956/
what-is-the-simplest-way-to-make-tooltips-
in-tkinter/36221216#36221216
http://www.daniweb.com/programming/software-development/
code/484591/a-tooltip-class-for-tkinter
- Originally written by vegaseat on 2014.09.09.
- Modified to include a delay time by Victor Zaccardo on 2016.03.25.
- Modified
- to correct extreme right and extreme bottom behavior,
- to stay inside the screen whenever the tooltip might go out on
the top but still the screen is higher than the tooltip,
- to use the more flexible mouse positioning,
- to add customizable background color, padding, waittime and
wraplength on creation
by Alberto Vassena on 2016.11.05.
Tested on Ubuntu 16.04/16.10, running Python 3.5.2
"""
def __init__(self, widget,
*,
background="#FFFFEA",
pad=(5, 3, 5, 3),
text="widget info",
waittime=400,
wraplength=250):
self.waittime = waittime # in miliseconds, originally 500
self.wraplength = wraplength # in pixels, originally 180
self.widget = widget
self.text = text
self.widget.bind("<Enter>", self.on_enter)
self.widget.bind("<Leave>", self.on_leave)
self.widget.bind("<ButtonPress>", self.on_leave)
self.background = background
self.pad = pad
self.ident = None
self.topwidget = None
def on_enter(self, event=None):
""" Schedule on an enter event """
self.schedule()
def on_leave(self, event=None):
""" Unschedule on a leave event """
self.unschedule()
self.hide()
def schedule(self):
""" Show the tooltip after wait period """
self.unschedule()
self.ident = self.widget.after(self.waittime, self.show)
def unschedule(self):
""" Hide the tooltip """
id_ = self.ident
self.ident = None
if id_:
self.widget.after_cancel(id_)
def show(self):
""" Show the tooltip """
def tip_pos_calculator(widget, label,
*,
tip_delta=(10, 5), pad=(5, 3, 5, 3)):
""" Calculate the tooltip position """
s_width, s_height = widget.winfo_screenwidth(), widget.winfo_screenheight()
width, height = (pad[0] + label.winfo_reqwidth() + pad[2],
pad[1] + label.winfo_reqheight() + pad[3])
mouse_x, mouse_y = widget.winfo_pointerxy()
x_1, y_1 = mouse_x + tip_delta[0], mouse_y + tip_delta[1]
x_2, y_2 = x_1 + width, y_1 + height
x_delta = x_2 - s_width
if x_delta < 0:
x_delta = 0
y_delta = y_2 - s_height
if y_delta < 0:
y_delta = 0
offscreen = (x_delta, y_delta) != (0, 0)
if offscreen:
if x_delta:
x_1 = mouse_x - tip_delta[0] - width
if y_delta:
y_1 = mouse_y - tip_delta[1] - height
offscreen_again = y_1 < 0 # out on the top
if offscreen_again:
# No further checks will be done.
# TIP:
# A further mod might automagically augment the
# wraplength when the tooltip is too high to be
# kept inside the screen.
y_1 = 0
return x_1, y_1
background = self.background
pad = self.pad
widget = self.widget
# creates a toplevel window
self.topwidget = tk.Toplevel(widget)
if platform.system() == "Darwin":
# For Mac OS
self.topwidget.tk.call("::tk::unsupported::MacWindowStyle",
"style", self.topwidget._w,
"help", "none")
# Leaves only the label and removes the app window
self.topwidget.wm_overrideredirect(True)
win = tk.Frame(self.topwidget,
background=background,
borderwidth=0)
label = tk.Label(win,
text=self.text,
justify=tk.LEFT,
background=background,
relief=tk.SOLID,
borderwidth=0,
wraplength=self.wraplength)
label.grid(padx=(pad[0], pad[2]),
pady=(pad[1], pad[3]),
sticky=tk.NSEW)
win.grid()
xpos, ypos = tip_pos_calculator(widget, label)
self.topwidget.wm_geometry("+%d+%d" % (xpos, ypos))
def hide(self):
""" Hide the tooltip """
topwidget = self.topwidget
if topwidget:
topwidget.destroy()
self.topwidget = None

351
lib/gui/utils.py Normal file
View File

@ -0,0 +1,351 @@
#!/usr/bin/env python3
""" Utility functions for the GUI """
import os
import sys
import tkinter as tk
from tkinter import filedialog, ttk
from PIL import Image, ImageTk
class Singleton(type):
""" Instigate a singleton.
From: https://stackoverflow.com/questions/6760685
Singletons are often frowned upon.
Feel free to instigate a better solution """
_instances = {}
def __call__(cls, *args, **kwargs):
if cls not in cls._instances:
cls._instances[cls] = super(Singleton,
cls).__call__(*args,
**kwargs)
return cls._instances[cls]
class FileHandler(object):
""" Raise a filedialog box and capture input """
def __init__(self, handletype, filetype, command=None, action=None,
variable=None):
all_files = ("All files", "*.*")
self.filetypes = {"default": (all_files,),
"alignments": (("JSON", "*.json"),
("Pickle", "*.p"),
("YAML", "*.yaml"),
all_files),
"config": (("Faceswap config files", "*.fsw"),
all_files),
"csv": (("Comma separated values", "*.csv"),
all_files),
"image": (("Bitmap", "*.bmp"),
("JPG", "*.jpeg", "*.jpg"),
("PNG", "*.png"),
("TIFF", "*.tif", "*.tiff"),
all_files),
"session": (("Faceswap session files", "*.fss"),
all_files),
"video": (("Audio Video Interleave", "*.avi"),
("Flash Video", "*.flv"),
("Matroska", "*.mkv"),
("MOV", "*.mov"),
("MP4", "*.mp4"),
("MPEG", "*.mpeg"),
("WebM", "*.webm"),
all_files)}
self.contexts = {
"effmpeg": {
"input": {"extract": "open",
"gen-vid": "dir",
"get-fps": "open",
"get-info": "open",
"mux-audio": "open",
"rescale": "open",
"rotate": "open",
"slice": "open"},
"output": {"extract": "dir",
"gen-vid": "save",
"get-fps": "nothing",
"get-info": "nothing",
"mux-audio": "save",
"rescale": "save",
"rotate": "save",
"slice": "save"}
}
}
self.defaults = self.set_defaults()
self.kwargs = self.set_kwargs(handletype, filetype, command, action,
variable)
self.retfile = getattr(self, handletype.lower())()
def set_defaults(self):
""" Set the default filetype to be first in list of filetypes,
or set a custom filetype if the first is not correct """
defaults = {key: val[0][1].replace("*", "")
for key, val in self.filetypes.items()}
defaults["default"] = None
defaults["video"] = ".mp4"
defaults["image"] = ".png"
return defaults
def set_kwargs(self, handletype, filetype, command, action, variable=None):
""" Generate the required kwargs for the requested browser """
kwargs = dict()
if handletype.lower() in ("open", "save", "filename", "savefilename"):
kwargs["filetypes"] = self.filetypes[filetype]
if self.defaults.get(filetype, None):
kwargs['defaultextension'] = self.defaults[filetype]
if handletype.lower() == "save":
kwargs["mode"] = "w"
if handletype.lower() == "open":
kwargs["mode"] = "r"
if handletype.lower() == "context":
kwargs["filetype"] = filetype
kwargs["command"] = command
kwargs["action"] = action
kwargs["variable"] = variable
return kwargs
def open(self):
""" Open a file """
return filedialog.askopenfile(**self.kwargs)
def save(self):
""" Save a file """
return filedialog.asksaveasfile(**self.kwargs)
def dir(self):
""" Get a directory location """
return filedialog.askdirectory(**self.kwargs)
def savedir(self):
""" Get a save dir location """
return filedialog.askdirectory(**self.kwargs)
def filename(self):
""" Get an existing file location """
return filedialog.askopenfilename(**self.kwargs)
def savefilename(self):
""" Get a save file location """
return filedialog.asksaveasfilename(**self.kwargs)
def nothing(self):
""" Method that does nothing, used for disabling open/save pop up """
return
def context(self):
""" Choose the correct file browser action based on context """
command = self.kwargs["command"]
action = self.kwargs["action"]
filetype = self.kwargs["filetype"]
variable = self.kwargs["variable"]
if self.contexts[command].get(variable, None) is not None:
handletype = self.contexts[command][variable][action]
else:
handletype = self.contexts[command][action]
self.kwargs = self.set_kwargs(handletype, filetype, command, action,
variable)
self.retfile = getattr(self, handletype.lower())()
class Images(object, metaclass=Singleton):
""" Holds locations of images and actual images """
def __init__(self, pathcache=None):
self.pathicons = os.path.join(pathcache, "icons")
self.pathpreview = os.path.join(pathcache, "preview")
self.pathoutput = None
self.previewoutput = None
self.previewtrain = dict()
self.errcount = 0
self.icons = dict()
self.icons["folder"] = tk.PhotoImage(file=os.path.join(self.pathicons,
"open_folder.png"))
self.icons["load"] = tk.PhotoImage(file=os.path.join(self.pathicons,
"open_file.png"))
self.icons["context"] = tk.PhotoImage(file=os.path.join(self.pathicons,
"open_file.png"))
self.icons["save"] = tk.PhotoImage(file=os.path.join(self.pathicons,
"save.png"))
self.icons["reset"] = tk.PhotoImage(file=os.path.join(self.pathicons,
"reset.png"))
self.icons["clear"] = tk.PhotoImage(file=os.path.join(self.pathicons,
"clear.png"))
self.icons["graph"] = tk.PhotoImage(file=os.path.join(self.pathicons,
"graph.png"))
self.icons["zoom"] = tk.PhotoImage(file=os.path.join(self.pathicons,
"zoom.png"))
self.icons["move"] = tk.PhotoImage(file=os.path.join(self.pathicons,
"move.png"))
def delete_preview(self):
""" Delete the preview files """
for item in os.listdir(self.pathpreview):
if item.startswith(".gui_preview_") and item.endswith(".jpg"):
fullitem = os.path.join(self.pathpreview, item)
os.remove(fullitem)
self.clear_image_cache()
def clear_image_cache(self):
""" Clear all cached images """
self.pathoutput = None
self.previewoutput = None
self.previewtrain = dict()
@staticmethod
def get_images(imgpath):
""" Get the images stored within the given directory """
if not os.path.isdir(imgpath):
return None
files = [os.path.join(imgpath, f)
for f in os.listdir(imgpath) if f.endswith((".png", ".jpg"))]
return files
def load_latest_preview(self):
""" Load the latest preview image for extract and convert """
imagefiles = self.get_images(self.pathoutput)
if not imagefiles or len(imagefiles) == 1:
self.previewoutput = None
return
# Get penultimate file so we don't accidently
# load a file that is being saved
show_file = sorted(imagefiles, key=os.path.getctime)[-2]
img = Image.open(show_file)
img.thumbnail((768, 432))
self.previewoutput = (img, ImageTk.PhotoImage(img))
def load_training_preview(self):
""" Load the training preview images """
imagefiles = self.get_images(self.pathpreview)
modified = None
if not imagefiles:
self.previewtrain = dict()
return
for img in imagefiles:
modified = os.path.getmtime(img) if modified is None else modified
name = os.path.basename(img)
name = os.path.splitext(name)[0]
name = name[name.rfind("_") + 1:].title()
try:
size = self.get_current_size(name)
self.previewtrain[name] = [Image.open(img), None, modified]
self.resize_image(name, size)
self.errcount = 0
except ValueError:
# This is probably an error reading the file whilst it's
# being saved so ignore it for now and only pick up if
# there have been multiple consecutive fails
if self.errcount < 10:
self.errcount += 1
else:
print("Error reading the preview file for {}".format(name))
self.previewtrain[name] = None
def get_current_size(self, name):
""" Return the size of the currently displayed image """
if not self.previewtrain.get(name, None):
return None
img = self.previewtrain[name][1]
if not img:
return None
return img.width(), img.height()
def resize_image(self, name, framesize):
""" Resize the training preview image
based on the passed in frame size """
displayimg = self.previewtrain[name][0]
if framesize:
frameratio = float(framesize[0]) / float(framesize[1])
imgratio = float(displayimg.size[0]) / float(displayimg.size[1])
if frameratio <= imgratio:
scale = framesize[0] / float(displayimg.size[0])
size = (framesize[0], int(displayimg.size[1] * scale))
else:
scale = framesize[1] / float(displayimg.size[1])
size = (int(displayimg.size[0] * scale), framesize[1])
# Hacky fix to force a reload if it happens to find corrupted
# data, probably due to reading the image whilst it is partially
# saved. If it continues to fail, then eventually raise.
for i in range(0, 1000):
try:
displayimg = displayimg.resize(size, Image.ANTIALIAS)
except OSError:
if i == 999:
raise
else:
continue
break
self.previewtrain[name][1] = ImageTk.PhotoImage(displayimg)
class ConsoleOut(ttk.Frame):
""" The Console out section of the GUI """
def __init__(self, parent, debug, tk_vars):
ttk.Frame.__init__(self, parent)
self.pack(side=tk.TOP, anchor=tk.W, padx=10, pady=(2, 0),
fill=tk.BOTH, expand=True)
self.console = tk.Text(self)
self.console_clear = tk_vars['consoleclear']
self.set_console_clear_var_trace()
self.debug = debug
self.build_console()
def set_console_clear_var_trace(self):
""" Set the trigger actions for the clear console var
when it has been triggered from elsewhere """
self.console_clear.trace("w", self.clear)
def build_console(self):
""" Build and place the console """
self.console.config(width=100, height=6, bg="gray90", fg="black")
self.console.pack(side=tk.LEFT, anchor=tk.N, fill=tk.BOTH, expand=True)
scrollbar = ttk.Scrollbar(self, command=self.console.yview)
scrollbar.pack(side=tk.LEFT, fill="y")
self.console.configure(yscrollcommand=scrollbar.set)
self.redirect_console()
def redirect_console(self):
""" Redirect stdout/stderr to console frame """
if self.debug:
print("Console debug activated. Outputting to main terminal")
else:
sys.stdout = SysOutRouter(console=self.console, out_type="stdout")
sys.stderr = SysOutRouter(console=self.console, out_type="stderr")
def clear(self, *args):
""" Clear the console output screen """
if not self.console_clear.get():
return
self.console.delete(1.0, tk.END)
self.console_clear.set(False)
class SysOutRouter(object):
""" Route stdout/stderr to the console window """
def __init__(self, console=None, out_type=None):
self.console = console
self.out_type = out_type
self.color = ("black" if out_type == "stdout" else "red")
def write(self, string):
""" Capture stdout/stderr """
self.console.insert(tk.END, string, self.out_type)
self.console.tag_config(self.out_type, foreground=self.color)
self.console.see(tk.END)
@staticmethod
def flush():
""" If flush is forced, send it to normal terminal """
sys.__stdout__.flush()

308
lib/gui/wrapper.py Normal file
View File

@ -0,0 +1,308 @@
#!/usr/bin python3
""" Process wrapper for underlying faceswap commands for the GUI """
import os
import re
import signal
import subprocess
from subprocess import PIPE, Popen, TimeoutExpired
import sys
import tkinter as tk
from threading import Thread
from time import time
from .utils import Images
class ProcessWrapper(object):
""" Builds command, launches and terminates the underlying
faceswap process. Updates GUI display depending on state """
def __init__(self, statusbar, session=None, pathscript=None, cliopts=None):
self.tk_vars = self.set_tk_vars()
self.session = session
self.pathscript = pathscript
self.cliopts = cliopts
self.command = None
self.statusbar = statusbar
self.task = FaceswapControl(self)
def set_tk_vars(self):
""" TK Variables to be triggered by ProcessWrapper to indicate
what state various parts of the GUI should be in """
display = tk.StringVar()
display.set(None)
runningtask = tk.BooleanVar()
runningtask.set(False)
actioncommand = tk.StringVar()
actioncommand.set(None)
actioncommand.trace("w", self.action_command)
generatecommand = tk.StringVar()
generatecommand.set(None)
generatecommand.trace("w", self.generate_command)
consoleclear = tk.BooleanVar()
consoleclear.set(False)
return {"display": display,
"runningtask": runningtask,
"action": actioncommand,
"generate": generatecommand,
"consoleclear": consoleclear}
def action_command(self, *args):
""" The action to perform when the action button is pressed """
if not self.tk_vars["action"].get():
return
category, command = self.tk_vars["action"].get().split(",")
if self.tk_vars["runningtask"].get():
self.task.terminate()
else:
self.command = command
args = self.prepare(category)
self.task.execute_script(command, args)
self.tk_vars["action"].set(None)
def generate_command(self, *args):
""" Generate the command line arguments and output """
if not self.tk_vars["generate"].get():
return
category, command = self.tk_vars["generate"].get().split(",")
args = self.build_args(category, command=command, generate=True)
self.tk_vars["consoleclear"].set(True)
print(" ".join(args))
self.tk_vars["generate"].set(None)
def prepare(self, category):
""" Prepare the environment for execution """
self.tk_vars["runningtask"].set(True)
self.tk_vars["consoleclear"].set(True)
print("Loading...")
self.statusbar.status_message.set("Executing - "
+ self.command + ".py")
mode = "indeterminate" if self.command == "train" else "determinate"
self.statusbar.progress_start(mode)
args = self.build_args(category)
self.tk_vars["display"].set(self.command)
return args
def build_args(self, category, command=None, generate=False):
""" Build the faceswap command and arguments list """
command = self.command if not command else command
script = "{}.{}".format(category, "py")
pathexecscript = os.path.join(self.pathscript, script)
args = ["python"] if generate else ["python", "-u"]
args.extend([pathexecscript, command])
for cliopt in self.cliopts.gen_cli_arguments(command):
args.extend(cliopt)
if command == "train" and not generate:
self.set_session_stats(cliopt)
if command == "train" and not generate:
args.append("-gui") # Embed the preview pane
return args
def set_session_stats(self, cliopt):
""" Set the session stats for batchsize and modeldir """
if cliopt[0] == "-bs":
self.session.stats["batchsize"] = int(cliopt[1])
if cliopt[0] == "-m":
self.session.modeldir = cliopt[1]
def terminate(self, message):
""" Finalise wrapper when process has exited """
self.tk_vars["runningtask"].set(False)
self.statusbar.progress_stop()
self.statusbar.status_message.set(message)
self.tk_vars["display"].set(None)
Images().delete_preview()
if self.command == "train":
self.session.save_session()
self.session.__init__()
self.command = None
print("Process exited.")
class FaceswapControl(object):
""" Control the underlying Faceswap tasks """
__group_processes = ["effmpeg"]
def __init__(self, wrapper):
self.wrapper = wrapper
self.statusbar = wrapper.statusbar
self.command = None
self.args = None
self.process = None
self.consoleregex = {"loss": re.compile(r"([a-zA-Z_]+):.*?(\d+\.\d+)"),
"tqdm": re.compile(r"(\d+%|\d+/\d+|\d+:\d+|\d+\.\d+[a-zA-Z/]+)")}
def execute_script(self, command, args):
""" Execute the requested Faceswap Script """
self.command = command
kwargs = {"stdout": PIPE,
"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(args, **kwargs)
self.thread_stdout()
self.thread_stderr()
def read_stdout(self):
""" Read stdout from the subprocess. If training, pass the loss
values to Queue """
while True:
output = self.process.stdout.readline()
if output == "" and self.process.poll() is not None:
break
if output:
if (self.command == "train" and self.capture_loss(output)) or (
self.command != "train" and self.capture_tqdm(output)):
continue
print(output.strip())
returncode = self.process.poll()
message = self.set_final_status(returncode)
self.wrapper.terminate(message)
def read_stderr(self):
""" Read stdout from the subprocess. If training, pass the loss
values to Queue """
while True:
output = self.process.stderr.readline()
if output == "" and self.process.poll() is not None:
break
if output:
if self.command != "train" and self.capture_tqdm(output):
continue
print(output.strip(), file=sys.stderr)
def thread_stdout(self):
""" Put the subprocess stdout so that it can be read without
blocking """
thread = Thread(target=self.read_stdout)
thread.daemon = True
thread.start()
def thread_stderr(self):
""" Put the subprocess stderr so that it can be read without
blocking """
thread = Thread(target=self.read_stderr)
thread.daemon = True
thread.start()
def capture_loss(self, string):
""" Capture loss values from stdout """
if not str.startswith(string, "["):
return False
loss = self.consoleregex["loss"].findall(string)
if len(loss) < 2:
return False
self.wrapper.session.add_loss(loss)
message = ""
for item in loss:
message += "{}: {} ".format(item[0], item[1])
if not message:
return False
elapsed = self.wrapper.session.timestats["elapsed"]
iterations = self.wrapper.session.stats["iterations"]
message = "Elapsed: {} Iteration: {} {}".format(elapsed,
iterations,
message)
self.statusbar.progress_update(message, 0, False)
return True
def capture_tqdm(self, string):
""" Capture tqdm output for progress bar """
tqdm = self.consoleregex["tqdm"].findall(string)
if len(tqdm) != 5:
return False
percent = tqdm[0]
processed = tqdm[1]
processtime = "Elapsed: {} Remaining: {}".format(tqdm[2], tqdm[3])
rate = tqdm[4]
message = "{} | {} | {} | {}".format(processtime,
rate,
processed,
percent)
current, total = processed.split("/")
position = int((float(current) / float(total)) * 1000)
self.statusbar.progress_update(message, position, True)
return True
def terminate(self):
""" Terminate the subprocess """
if self.command == "train":
print("Sending Exit Signal", flush=True)
try:
now = time()
if os.name == "nt":
os.kill(self.process.pid, signal.CTRL_BREAK_EVENT)
else:
self.process.send_signal(signal.SIGINT)
while True:
timeelapsed = time() - now
if self.process.poll() is not None:
break
if timeelapsed > 30:
raise ValueError("Timeout reached sending Exit Signal")
return
except ValueError as err:
print(err)
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 """
if returncode == 0 or returncode == 3221225786:
status = "Ready"
elif returncode == -15:
status = "Terminated - {}.py".format(self.command)
elif returncode == -9:
status = "Killed - {}.py".format(self.command)
elif returncode == -6:
status = "Aborted - {}.py".format(self.command)
else:
status = "Failed - {}.py. Return Code: {}".format(self.command,
returncode)
return status

View File

@ -1,26 +1,31 @@
#!/usr/bin/python3
# Based on the original https://www.reddit.com/r/deepfakes/ code sample + contribs
# Based on https://github.com/iperov/OpenDeepFaceSwap for Decoder multiple res block chain
# Based on the https://github.com/shaoanlu/faceswap-GAN repo
# source : https://github.com/shaoanlu/faceswap-GAN/blob/master/FaceSwap_GAN_v2_sz128_train.ipynbtemp/faceswap_GAN_keras.ipynb
import enum
import os
import sys
import warnings
warnings.filterwarnings("ignore", category=FutureWarning)
from keras.initializers import RandomNormal
from keras.layers import Input, Dense, Flatten, Reshape
from keras.layers import SeparableConv2D, add
from keras.layers.advanced_activations import LeakyReLU
from keras.layers.convolutional import Conv2D
from keras.layers.core import Activation
from keras.models import Model as KerasModel
from keras.optimizers import Adam
from keras.utils import multi_gpu_model
from lib.PixelShuffler import PixelShuffler
import lib.Serializer
from . import __version__
from keras.layers.core import Activation
if isinstance(__version__, (list, tuple)):
version_str = ".".join([str(n) for n in __version__[1:]])
@ -30,51 +35,78 @@ else:
mswindows = sys.platform=="win32"
try:
from lib.utils import backup_file
except ImportError:
pass
class Encoders():
REGULAR = 'v2' # high memory consumption encoder
NEW_SLIM = 'v3' # slightly lighter on resources and taining speed is faster
class EncoderType(enum.Enum):
ORIGINAL = "original"
SHAOANLU = "shaoanlu"
ENCODER = EncoderType.ORIGINAL
if ENCODER==EncoderType.SHAOANLU:
from .instance_normalization import InstanceNormalization
ENCODER = Encoders.NEW_SLIM
def inst_norm():
return InstanceNormalization()
hdf = {'encoderH5': 'encoder_{version_str}{ENCODER}.h5'.format(**vars()),
'decoder_AH5': 'decoder_A_{version_str}{ENCODER}.h5'.format(**vars()),
'decoder_BH5': 'decoder_B_{version_str}{ENCODER}.h5'.format(**vars())}
conv_init = RandomNormal(0, 0.02)
def inst_norm():
return InstanceNormalization()
class EncoderType(enum.Enum):
ORIGINAL = "original"
SHAOANLU = "shaoanlu"
ENCODER = EncoderType.ORIGINAL
hdf = {'encoderH5': 'encoder_{version_str}{ENCODER.value}.h5'.format(**vars()),
'decoder_AH5': 'decoder_A_{version_str}{ENCODER.value}.h5'.format(**vars()),
'decoder_BH5': 'decoder_B_{version_str}{ENCODER.value}.h5'.format(**vars())}
class Model():
# still playing with dims
ENCODER_DIM = 2048
IMAGE_SIZE = 128, 128
IMAGE_DEPTH = len('RGB') # good to let ppl know what these are...
IMAGE_SHAPE = *IMAGE_SIZE, IMAGE_DEPTH
def __init__(self, model_dir, gpus):
ENCODER_DIM = 1024 # dense layer size
IMAGE_SHAPE = 128, 128 # image shape
assert [n for n in IMAGE_SHAPE if n>=16]
IMAGE_WIDTH = max(IMAGE_SHAPE)
IMAGE_WIDTH = (IMAGE_WIDTH//16 + (1 if (IMAGE_WIDTH%16)>=8 else 0))*16
IMAGE_SHAPE = IMAGE_WIDTH, IMAGE_WIDTH, len('BRG') # good to let ppl know what these are...
def __init__(self, model_dir, gpus, encoder_type=ENCODER):
if mswindows:
from ctypes import cdll
mydll = cdll.LoadLibrary("user32.dll")
mydll.SetProcessDPIAware(True)
mydll.SetProcessDPIAware(True)
self._encoder_type = encoder_type
self.model_dir = model_dir
# can't chnage gpu's when the model is initialized no point in making it r/w
self._gpus = gpus
Encoder = getattr(self, "Encoder") if not ENCODER else getattr(self, "Encoder_{}".format(ENCODER))
Encoder = getattr(self, "Encoder_{}".format(self._encoder_type.value))
Decoder = getattr(self, "Decoder_{}".format(self._encoder_type.value))
self.encoder = Encoder()
self.decoder_A = self.Decoder()
self.decoder_B = self.Decoder()
self.decoder_A = Decoder()
self.decoder_B = Decoder()
self.initModel()
@ -96,8 +128,18 @@ class Model():
def load(self, swapped):
from json import JSONDecodeError
face_A, face_B = (hdf['decoder_AH5'], hdf['decoder_BH5']) if not swapped else (hdf['decoder_BH5'], hdf['decoder_AH5'])
state_dir = os.path.join(self.model_dir, 'state_{version_str}_{ENCODER.value}.json'.format(**globals()))
ser = lib.Serializer.get_serializer('json')
try:
with open(state_dir, 'rb') as fp:
state = ser.unmarshal(fp.read())
self._epoch_no = state['epoch_no']
except (JSONDecodeError, IOError) as e:
print('Failed loading training state metadata', e)
self._epoch_no = 0
try:
self.encoder.load_weights(os.path.join(self.model_dir, hdf['encoderH5']))
@ -106,37 +148,39 @@ class Model():
print('loaded model weights')
return True
except Exception as e:
print('Failed loading existing training data.')
print(e)
print('Failed loading existing training data.', e)
return False
def converter(self, swap):
autoencoder = self.autoencoder_B if not swap else self.autoencoder_A
return autoencoder.predict
def conv(self, filters, kernel_size=4, strides=2):
def conv(self, filters, kernel_size=5, strides=2, **kwargs):
def block(x):
x = Conv2D(filters, kernel_size=kernel_size, strides=strides, padding='same')(x)
x = Conv2D(filters, kernel_size=kernel_size, strides=strides, kernel_initializer=conv_init, padding='same', **kwargs)(x)
x = LeakyReLU(0.1)(x)
return x
return block
return block
def conv_sep2(self, filters, kernel_size=5, strides=2, use_instance_norm=True, **kwargs):
def block(x):
x = SeparableConv2D(filters, kernel_size=kernel_size, strides=strides, kernel_initializer=conv_init, padding='same', **kwargs)(x)
x = Activation("relu")(x)
return x
return block
def conv_sep(self, filters):
def conv_sep3(self, filters, kernel_size=3, strides=2, use_instance_norm=True, **kwargs):
def block(x):
x = SeparableConv2D(filters, kernel_size=4, strides=2, padding='same')(x)
x = LeakyReLU(0.1)(x)
return x
x = SeparableConv2D(filters, kernel_size=kernel_size, strides=strides, kernel_initializer=conv_init, padding='same', **kwargs)(x)
if use_instance_norm:
x = inst_norm()(x)
x = Activation("relu")(x)
return x
return block
def conv_sep_v3(self, filters, kernel_size=4, strides=2):
def block(x):
x = SeparableConv2D(filters, kernel_size=kernel_size, strides=strides, padding="same")(x)
x = Activation("relu")(x)
return x
return block
def upscale(self, filters):
def upscale(self, filters, **kwargs):
def block(x):
x = Conv2D(filters * 4, kernel_size=3, padding='same')(x)
x = LeakyReLU(0.1)(x)
@ -144,92 +188,81 @@ class Model():
return x
return block
def upscale_sep(self, filters):
def upscale_sep3(self, filters, use_instance_norm=True, **kwargs):
def block(x):
x = SeparableConv2D(filters * 4, kernel_size=3, padding='same')(x)
x = Conv2D(filters*4, kernel_size=3, use_bias=False, kernel_initializer=RandomNormal(0, 0.02), padding='same', **kwargs)(x)
if use_instance_norm:
x = inst_norm()(x)
x = LeakyReLU(0.1)(x)
x = PixelShuffler()(x)
return x
return block
def res(self, filters, dilation_rate=1):
def block(x):
rb = Conv2D(filters, kernel_size=3, padding="same", dilation_rate=dilation_rate, use_bias=False)(x)
rb = LeakyReLU(alpha=0.2)(rb)
rb = Conv2D(filters, kernel_size=3, padding="same", dilation_rate=dilation_rate, use_bias=False)(rb)
x = add([rb, x])
x = LeakyReLU(alpha=0.2)(x)
return x
return block
def res_block(self, filters, dilation_rate=1):
def block(x):
rb = Conv2D(filters, kernel_size=3, padding="same", dilation_rate=dilation_rate, use_bias=False)(x)
rb = LeakyReLU(alpha=0.2)(rb)
rb = Conv2D(filters, kernel_size=3, padding="same", dilation_rate=dilation_rate, use_bias=False)(rb)
x = add([rb, x])
x = LeakyReLU(alpha=0.2)(x)
return x
return block
def Encoder_v3(self):
"""Lighter on resources encoder with bigger first conv layer"""
retina = Input(shape=self.IMAGE_SHAPE)
x = self.conv_sep_v3(192)(retina)
x = self.conv(256)(x)
x = self.conv(384)(x)
x = self.conv_sep_v3(512)(x)
x = self.conv(768)(x)
x = self.conv_sep_v3(1024)(x)
x = Dense(self.ENCODER_DIM)(Flatten()(x))
x = Dense(4 * 4 * 1024)(x)
x = Reshape((4, 4, 1024))(x)
out = self.upscale(512)(x)
return KerasModel(retina, out)
def Encoder_v2(self):
"""Old algorithm; pretty good but slow"""
retina = Input(shape=self.IMAGE_SHAPE)
x = self.conv(128)(retina)
x = self.conv(144)(x)
x = self.conv_sep(256)(x)
x = self.conv(448)(x)
x = self.conv_sep(512)(x)
x = self.conv(768)(x)
x = self.conv_sep(1024)(x)
x = Dense(self.ENCODER_DIM)(Flatten()(x))
x = Dense(4 * 4 * 1024)(x)
x = Reshape((4, 4, 1024))(x)
out = self.upscale(512)(x)
return KerasModel(retina, out)
def Encoder_original(self, **kwargs):
impt = Input(shape=self.IMAGE_SHAPE)
in_conv_filters = self.IMAGE_SHAPE[0] if self.IMAGE_SHAPE[0] <= 128 else 128 + (self.IMAGE_SHAPE[0]-128)//4
def Decoder(self):
inp = Input(shape=(8, 8, 512))
x = self.upscale(384)(inp)
x = self.res_block(384)(x)
x = self.upscale_sep(192)(x)
x = self.res_block(192)(x)
x = self.upscale(128)(x)
x = self.res_block(128)(x)
x = self.upscale(64)(x)
x = self.res_block(64)(x)
x = self.conv(in_conv_filters)(impt)
x = self.conv_sep2(256)(x)
x = self.conv(512)(x)
x = self.conv_sep2(1024)(x)
# rb = Conv2D(64, kernel_size=3, padding="same", dilation_rate=2)(x)
# rb = LeakyReLU(alpha=0.2)(rb)
# rb = Conv2D(64, kernel_size=3, padding="same", dilation_rate=2)(rb)
# x = add([rb, x])
#
#x = self.upscale(32)(x)
out = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
return KerasModel(inp, out)
dense_shape = self.IMAGE_SHAPE[0] // 16
x = Dense(self.ENCODER_DIM)(Flatten()(x))
x = Dense(dense_shape * dense_shape * 512)(x)
x = Reshape((dense_shape, dense_shape, 512))(x)
x = self.upscale(512)(x)
return KerasModel(impt, x, **kwargs)
def Encoder_shaoanlu(self, **kwargs):
impt = Input(shape=self.IMAGE_SHAPE)
in_conv_filters = self.IMAGE_SHAPE[0] if self.IMAGE_SHAPE[0] <= 128 else 128 + (self.IMAGE_SHAPE[0]-128)//4
x = Conv2D(in_conv_filters, kernel_size=5, kernel_initializer=conv_init, use_bias=False, padding="same")(impt)
x = self.conv_sep3(in_conv_filters+32, use_instance_norm=False)(x)
x = self.conv_sep3(256)(x)
x = self.conv_sep3(512)(x)
x = self.conv_sep3(1024)(x)
dense_shape = self.IMAGE_SHAPE[0] // 16
x = Dense(self.ENCODER_DIM)(Flatten()(x))
x = Dense(dense_shape * dense_shape * 768)(x)
x = Reshape((dense_shape, dense_shape, 768))(x)
x = self.upscale(512)(x)
return KerasModel(impt, x, **kwargs)
def Decoder_original(self):
decoder_shape = self.IMAGE_SHAPE[0]//8
inpt = Input(shape=(decoder_shape, decoder_shape, 512))
x = self.upscale(384, kernel_initializer=RandomNormal(0, 0.02))(inpt)
x = self.upscale(256-32, kernel_initializer=RandomNormal(0, 0.02))(x)
x = self.upscale(self.IMAGE_SHAPE[0], kernel_initializer=RandomNormal(0, 0.02))(x)
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
return KerasModel(inpt, x)
def Decoder_shaoanlu(self):
decoder_shape = self.IMAGE_SHAPE[0]//8
inpt = Input(shape=(decoder_shape, decoder_shape, 512))
x = self.upscale_sep3(512)(inpt)
x = self.upscale_sep3(256)(x)
x = self.upscale_sep3(self.IMAGE_SHAPE[0])(x)
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
return KerasModel(inpt, x)
def save_weights(self):
from threading import Thread
from time import sleep
@ -240,24 +273,30 @@ class Model():
for model in hdf.values():
backup_file(model_dir, model)
except NameError:
print('backup functionality not available\n')
print('backup functionality not available\n')
state_dir = os.path.join(self.model_dir, 'state_{version_str}_{ENCODER.value}.json'.format(**globals()))
ser = lib.Serializer.get_serializer('json')
try:
with open(state_dir, 'wb') as fp:
state_json = ser.marshal({
'epoch_no' : self._epoch_no
})
fp.write(state_json.encode('utf-8'))
except IOError as e:
pass
print('\nsaving model weights', end='', flush=True)
from concurrent.futures import ThreadPoolExecutor, as_completed
# thought maybe I/O bound, sometimes saving in parallel is faster
threads = []
t = Thread(target=self.encoder.save_weights, args=(str(self.model_dir / hdf['encoderH5']),))
threads.append(t)
t = Thread(target=self.decoder_A.save_weights, args=(str(self.model_dir / hdf['decoder_AH5']),))
threads.append(t)
t = Thread(target=self.decoder_B.save_weights, args=(str(self.model_dir / hdf['decoder_BH5']),))
threads.append(t)
for thread in threads:
thread.start()
while any([t.is_alive() for t in threads]):
sleep(0.1)
print('saved model weights')
with ThreadPoolExecutor(max_workers=4) as executor:
futures = [executor.submit(getattr(self, mdl_name.rstrip('H5')).save_weights, str(self.model_dir / mdl_H5_fn)) for mdl_name, mdl_H5_fn in hdf.items()]
for future in as_completed(futures):
future.result()
print('.', end='', flush=True)
print('done', flush=True)
@property
@ -267,15 +306,14 @@ class Model():
@property
def model_name(self):
try:
return self._model_nomen
return self._model_name
except AttributeError:
self._model_nomen = self._model_nomen = os.path.split(os.path.dirname(__file__))[1].replace("Model_", "")
return self._model_nomen
self._model_name = os.path.split(os.path.dirname(__file__))[1].replace("Model_", "")
return self._model_name
def __str__(self):
return "<{}: ver={}, nn_dims={}, img_size={}>".format(self.model_name,
return "<{}: ver={}, dense_dim={}, img_size={}>".format(self.model_name,
version_str,
self.ENCODER_DIM,
"x".join([str(n) for n in self.IMAGE_SHAPE[:2]]))
"x".join([str(n) for n in self.IMAGE_SHAPE[:2]]))

View File

@ -1,17 +1,14 @@
import time
import numpy
from lib.training_data import TrainingDataGenerator, stack_images
TRANSFORM_PRC = 115.
#TRANSFORM_PRC = 150.
class Trainer():
#
_random_transform_args = {
'rotation_range': 10 * (TRANSFORM_PRC * .01),
'zoom_range': 0.05 * (TRANSFORM_PRC * .01),
@ -22,13 +19,10 @@ class Trainer():
def __init__(self, model, fn_A, fn_B, batch_size, *args):
self.batch_size = batch_size
self.model = model
#generator = TrainingDataGenerator(self.random_transform_args, 160)
# make sre to keep zoom=2 or you won't get 128x128 vectors as input
#generator = TrainingDataGenerator(self.random_transform_args, 220, 5, zoom=2)
generator = TrainingDataGenerator(self.random_transform_args, 160, 6, zoom=2)
#generator = TrainingDataGenerator(self.random_transform_args, 180, 7, zoom=2)
from timeit import default_timer as clock
self._clock = clock
generator = TrainingDataGenerator(self.random_transform_args, 160, 5, zoom=2)
self.images_A = generator.minibatchAB(fn_A, self.batch_size)
self.images_B = generator.minibatchAB(fn_B, self.batch_size)
@ -37,19 +31,22 @@ class Trainer():
def train_one_step(self, iter_no, viewer):
when = self._clock()
_, warped_A, target_A = next(self.images_A)
_, warped_B, target_B = next(self.images_B)
loss_A = self.model.autoencoder_A.train_on_batch(warped_A, target_A)
loss_B = self.model.autoencoder_B.train_on_batch(warped_B, target_B)
print("[{0}] [#{1:05d}] loss_A: {2:.5f}, loss_B: {3:.5f}".format(
time.strftime("%H:%M:%S"), iter_no, loss_A, loss_B),
loss_B = self.model.autoencoder_B.train_on_batch(warped_B, target_B)
self.model._epoch_no += 1
print("[{0}] [#{1:05d}] [{2:.3f}s] loss_A: {3:.5f}, loss_B: {4:.5f}".format(
time.strftime("%H:%M:%S"), self.model._epoch_no, self._clock()-when, loss_A, loss_B),
end='\r')
if viewer is not None:
viewer(self.show_sample(target_A[0:24], target_B[0:24]), "training using {}, bs={}".format(self.model, self.batch_size))
viewer(self.show_sample(target_A[0:8], target_B[0:8]), "training using {}, bs={}".format(self.model, self.batch_size))
def show_sample(self, test_A, test_B):

View File

@ -0,0 +1,145 @@
from keras.engine import Layer, InputSpec
from keras import initializers, regularizers, constraints
from keras import backend as K
from keras.utils.generic_utils import get_custom_objects
import numpy as np
class InstanceNormalization(Layer):
"""Instance normalization layer (Lei Ba et al, 2016, Ulyanov et al., 2016).
Normalize the activations of the previous layer at each step,
i.e. applies a transformation that maintains the mean activation
close to 0 and the activation standard deviation close to 1.
# Arguments
axis: Integer, the axis that should be normalized
(typically the features axis).
For instance, after a `Conv2D` layer with
`data_format="channels_first"`,
set `axis=1` in `InstanceNormalization`.
Setting `axis=None` will normalize all values in each instance of the batch.
Axis 0 is the batch dimension. `axis` cannot be set to 0 to avoid errors.
epsilon: Small float added to variance to avoid dividing by zero.
center: If True, add offset of `beta` to normalized tensor.
If False, `beta` is ignored.
scale: If True, multiply by `gamma`.
If False, `gamma` is not used.
When the next layer is linear (also e.g. `nn.relu`),
this can be disabled since the scaling
will be done by the next layer.
beta_initializer: Initializer for the beta weight.
gamma_initializer: Initializer for the gamma weight.
beta_regularizer: Optional regularizer for the beta weight.
gamma_regularizer: Optional regularizer for the gamma weight.
beta_constraint: Optional constraint for the beta weight.
gamma_constraint: Optional constraint for the gamma weight.
# Input shape
Arbitrary. Use the keyword argument `input_shape`
(tuple of integers, does not include the samples axis)
when using this layer as the first layer in a model.
# Output shape
Same shape as input.
# References
- [Layer Normalization](https://arxiv.org/abs/1607.06450)
- [Instance Normalization: The Missing Ingredient for Fast Stylization](https://arxiv.org/abs/1607.08022)
"""
def __init__(self,
axis=None,
epsilon=1e-3,
center=True,
scale=True,
beta_initializer='zeros',
gamma_initializer='ones',
beta_regularizer=None,
gamma_regularizer=None,
beta_constraint=None,
gamma_constraint=None,
**kwargs):
super(InstanceNormalization, self).__init__(**kwargs)
self.supports_masking = True
self.axis = axis
self.epsilon = epsilon
self.center = center
self.scale = scale
self.beta_initializer = initializers.get(beta_initializer)
self.gamma_initializer = initializers.get(gamma_initializer)
self.beta_regularizer = regularizers.get(beta_regularizer)
self.gamma_regularizer = regularizers.get(gamma_regularizer)
self.beta_constraint = constraints.get(beta_constraint)
self.gamma_constraint = constraints.get(gamma_constraint)
def build(self, input_shape):
ndim = len(input_shape)
if self.axis == 0:
raise ValueError('Axis cannot be zero')
if (self.axis is not None) and (ndim == 2):
raise ValueError('Cannot specify axis for rank 1 tensor')
self.input_spec = InputSpec(ndim=ndim)
if self.axis is None:
shape = (1,)
else:
shape = (input_shape[self.axis],)
if self.scale:
self.gamma = self.add_weight(shape=shape,
name='gamma',
initializer=self.gamma_initializer,
regularizer=self.gamma_regularizer,
constraint=self.gamma_constraint)
else:
self.gamma = None
if self.center:
self.beta = self.add_weight(shape=shape,
name='beta',
initializer=self.beta_initializer,
regularizer=self.beta_regularizer,
constraint=self.beta_constraint)
else:
self.beta = None
self.built = True
def call(self, inputs, training=None):
input_shape = K.int_shape(inputs)
reduction_axes = list(range(0, len(input_shape)))
if (self.axis is not None):
del reduction_axes[self.axis]
del reduction_axes[0]
mean = K.mean(inputs, reduction_axes, keepdims=True)
stddev = K.std(inputs, reduction_axes, keepdims=True) + self.epsilon
normed = (inputs - mean) / stddev
broadcast_shape = [1] * len(input_shape)
if self.axis is not None:
broadcast_shape[self.axis] = input_shape[self.axis]
if self.scale:
broadcast_gamma = K.reshape(self.gamma, broadcast_shape)
normed = normed * broadcast_gamma
if self.center:
broadcast_beta = K.reshape(self.beta, broadcast_shape)
normed = normed + broadcast_beta
return normed
def get_config(self):
config = {
'axis': self.axis,
'epsilon': self.epsilon,
'center': self.center,
'scale': self.scale,
'beta_initializer': initializers.serialize(self.beta_initializer),
'gamma_initializer': initializers.serialize(self.gamma_initializer),
'beta_regularizer': regularizers.serialize(self.beta_regularizer),
'gamma_regularizer': regularizers.serialize(self.gamma_regularizer),
'beta_constraint': constraints.serialize(self.beta_constraint),
'gamma_constraint': constraints.serialize(self.gamma_constraint)
}
base_config = super(InstanceNormalization, self).get_config()
return dict(list(base_config.items()) + list(config.items()))
get_custom_objects().update({'InstanceNormalization': InstanceNormalization})

View File

@ -1,18 +1,16 @@
pathlib==1.0.1
scandir==1.6
h5py==2.7.1
Keras==2.1.2
opencv-python==3.3.0.10
scandir==1.7
h5py==2.8.0
Keras==2.2.0
opencv-python==3.4.1.15
scikit-image
face_recognition
cmake
dlib
face-recognition
tqdm
matplotlib
matplotlib==2.2.2
ffmpy==0.2.2
# tensorflow is included within the docker image.
# If you are looking for dependencies for a manual install,
# you may want to install tensorflow-gpu==1.4.0 for CUDA 8.0 or tensorflow-gpu>=1.5.0 for CUDA 9.0
# cmake needs to be installed before compiling dlib.

File diff suppressed because it is too large Load Diff

View File

@ -118,10 +118,10 @@ class Train(object):
def run_training_cycle(self, model, trainer):
""" Perform the training cycle """
for epoch in range(0, self.args.epochs):
save_iteration = epoch % self.args.save_interval == 0
for iteration in range(0, self.args.iterations):
save_iteration = iteration % self.args.save_interval == 0
viewer = self.show if save_iteration or self.save_now else None
trainer.train_one_step(epoch, viewer)
trainer.train_one_step(iteration, viewer)
if self.stop:
break
elif save_iteration:
@ -186,12 +186,11 @@ class Train(object):
img = "_sample_{}.jpg".format(name)
imgfile = os.path.join(scriptpath, img)
cv2.imwrite(imgfile, image)
if self.args.redirect_gui:
img = ".gui_preview.png"
imgfile = os.path.join(scriptpath, img)
img = ".gui_preview_{}.jpg".format(name)
imgfile = os.path.join(scriptpath, "lib", "gui", ".cache", "preview", img)
cv2.imwrite(imgfile, image)
elif self.args.preview:
if self.args.preview:
with self.lock:
self.preview_buffer[name] = image
except Exception as err:

View File

@ -1,9 +1,9 @@
#!/usr/bin/env python3
""" The master tools.py script """
import sys
# Importing the various tools
from tools.sort import SortProcessor
from tools.effmpeg import Effmpeg
import lib.cli as cli
import tools.cli as cli
from lib.cli import FullHelpArgumentParser, GuiArgs
# Python version check
if sys.version_info[0] < 3:
@ -13,6 +13,7 @@ if sys.version_info[0] == 3 and sys.version_info[1] < 2:
def bad_args(args):
""" Print help on bad arguments """
PARSER.print_help()
exit(0)
@ -23,18 +24,17 @@ if __name__ == "__main__":
_tools_warning += "understand how it works."
print(_tools_warning)
PARSER = cli.FullHelpArgumentParser()
PARSER = 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)
EFFMPEG = cli.EffmpegArgs(SUBPARSER,
"effmpeg",
"This command allows you to easily execute common ffmpeg tasks.")
SORT = cli.SortArgs(SUBPARSER,
"sort",
"This command lets you sort images using various methods.")
GUI = GuiArgs(SUBPARSER,
"gui",
"Launch the Faceswap Tools Graphical User Interface.")
PARSER.set_defaults(func=bad_args)
ARGUMENTS = PARSER.parse_args()
ARGUMENTS.func(ARGUMENTS)

347
tools/cli.py Normal file
View File

@ -0,0 +1,347 @@
#!/usr/bin/env python3
""" Command Line Arguments for tools """
from lib.cli import FaceSwapArgs
from lib.cli import ContextFullPaths, DirFullPaths, FileFullPaths, SaveFileFullPaths
from lib.utils import _image_extensions
class EffmpegArgs(FaceSwapArgs):
""" Class to parse the command line arguments for EFFMPEG tool """
@staticmethod
def __parse_transpose(value):
index = 0
opts = ["(0, 90CounterClockwise&VerticalFlip)",
"(1, 90Clockwise)",
"(2, 90CounterClockwise)",
"(3, 90Clockwise&VerticalFlip)"]
if len(value) == 1:
index = int(value)
else:
for i in range(5):
if value in opts[i]:
index = i
break
return opts[index]
def get_argument_list(self):
argument_list = list()
argument_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.\n"
"'slice' cuts a portion of the video "
"into a separate video file.\n"
"'get-fps' returns the chosen video's "
"fps."})
argument_list.append({"opts": ('-i', '--input'),
"action": ContextFullPaths,
"dest": "input",
"default": "input",
"help": "Input file.",
"required": True,
"action_option": "-a",
"filetypes": "video"})
argument_list.append({"opts": ('-o', '--output'),
"action": ContextFullPaths,
"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.\n"
"Note: the chosen output file "
"extension will determine the file "
"encoding.",
"action_option": "-a",
"filetypes": "video"})
argument_list.append({"opts": ('-r', '--reference-video'),
"action": FileFullPaths,
"dest": "ref_vid",
"default": None,
"help": "Path to reference video if 'input' "
"was not a video.",
"filetypes": "video"})
argument_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."})
argument_list.append({"opts": ("-ef", "--extract-filetype"),
"choices": _image_extensions,
"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."})
argument_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.\n"
"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."})
argument_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.\n"
"Default: 00:00:00, in HH:MM:SS."})
argument_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.\n"
"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."})
argument_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."})
argument_list.append({"opts": ('-tr', '--transpose'),
"choices": ("(0, 90CounterClockwise&VerticalFlip)",
"(1, 90Clockwise)",
"(2, 90CounterClockwise)",
"(3, 90Clockwise&VerticalFlip)"),
"type": lambda v: self.__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"})
argument_list.append({"opts": ('-de', '--degrees'),
"type": str,
"dest": "degrees",
"default": None,
"help": "Rotate the video clockwise by the "
"given number of degrees."})
argument_list.append({"opts": ('-sc', '--scale'),
"type": str,
"dest": "scale",
"default": "1920x1080",
"help": "Set the new resolution scale if the "
"chosen action is 'rescale'."})
argument_list.append({"opts": ('-pr', '--preview'),
"action": "store_true",
"dest": "preview",
"default": False,
"help": "Uses ffplay to preview the effects of "
"the intended action. "
"This functionality is not yet fully "
"implemented."})
argument_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."})
argument_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 argument_list
class SortArgs(FaceSwapArgs):
""" Class to parse the command line arguments for sort tool """
@staticmethod
def get_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": ('-i', '--input'),
"action": DirFullPaths,
"dest": "input_dir",
"default": "input_dir",
"help": "Input directory of aligned faces.",
"required": True})
argument_list.append({"opts": ('-o', '--output'),
"action": DirFullPaths,
"dest": "output_dir",
"default": "_output_dir",
"help": "Output directory for sorted aligned "
"faces."})
argument_list.append({"opts": ('-fp', '--final-process'),
"type": str,
"choices": ("folders", "rename"),
"dest": 'final_process',
"default": "rename",
"help": "'folders': files are sorted using the "
"-s/--sort-by method, then they are "
"organized into folders using the "
"-g/--group-by grouping method. "
"'rename': files are sorted using the "
"-s/--sort-by then they are renamed. "
"Default: rename"})
argument_list.append({"opts": ('-k', '--keep'),
"action": 'store_true',
"dest": 'keep_original',
"default": False,
"help": "Keeps the original files in the input "
"directory. Be careful when using this "
"with rename grouping and no specified "
"output directory as this would keep "
"the original and renamed files in the "
"same directory."})
argument_list.append({"opts": ('-s', '--sort-by'),
"type": str,
"choices": ("blur", "face", "face-cnn",
"face-cnn-dissim", "face-dissim",
"face-yaw", "hist",
"hist-dissim"),
"dest": 'sort_method',
"default": "hist",
"help": "Sort by method. "
"Choose how images are sorted. "
"Default: hist"})
argument_list.append({"opts": ('-g', '--group-by'),
"type": str,
"choices": ("blur", "face", "face-cnn",
"face-yaw", "hist"),
"dest": 'group_method',
"default": "hist",
"help": "Group by method. "
"When -fp/--final-processing by "
"folders choose the how the images are "
"grouped after sorting. "
"Default: hist"})
argument_list.append({"opts": ('-t', '--ref_threshold'),
"type": float,
"dest": 'min_threshold',
"default": -1.0,
"help": "Float value. "
"Minimum threshold to use for grouping "
"comparison with 'face' and 'hist' "
"methods. The lower the value the more "
"discriminating the grouping is. "
"Leaving -1.0 will make the program "
"set the default value automatically. "
"For face 0.6 should be enough, with "
"0.5 being very discriminating. "
"For face-cnn 7.2 should be enough, "
"with 4 being very discriminating. "
"For hist 0.3 should be enough, with "
"0.2 being very discriminating. "
"Be careful setting a value that's too "
"low in a directory with many images, "
"as this could result in a lot of "
"directories being created. "
"Defaults: face 0.6, face-cnn 7.2, "
"hist 0.3"})
argument_list.append({"opts": ('-b', '--bins'),
"type": int,
"dest": 'num_bins',
"default": 5,
"help": "Integer value. "
"Number of folders that will be used "
"to group by blur and face-yaw. "
"For blur folder 0 will be the least "
"blurry, while the last folder will be "
"the blurriest. "
"For face-yaw the number of bins is by "
"how much 180 degrees is divided. So "
"if you use 18, then each folder will "
"be a 10 degree increment. Folder 0 "
"will contain faces looking the most "
"to the left whereas the last folder "
"will contain the faces looking the "
"most to the right. "
"If the number of images doesn't "
"divide evenly into the number of "
"bins, the remaining images get put in "
"the last bin."
"Default value: 5"})
argument_list.append({"opts": ('-l', '--log-changes'),
"action": 'store_true',
"dest": 'log_changes',
"default": False,
"help": "Logs file renaming changes if "
"grouping by renaming, or it logs the "
"file copying/movement if grouping by "
"folders. If no log file is specified "
"with '--log-file', then a "
"'sort_log.json' file will be created "
"in the input directory."})
argument_list.append({"opts": ('-lf', '--log-file'),
"action": SaveFileFullPaths,
"filetypes": "alignments",
"dest": 'log_file_path',
"default": 'sort_log.json',
"help": "Specify a log file to use for saving "
"the renaming or grouping information. "
"If specified extension isn't 'json' "
"or 'yaml', then json will be used as "
"the serializer, with the supplied "
"filename. "
"Default: sort_log.json"})
return argument_list

View File

@ -5,7 +5,9 @@ Created on 2018-03-16 15:14
@author: Lev Velykoivanenko (velykoivanenko.lev@gmail.com)
"""
import argparse
# TODO: fix file handlers for effmpeg in gui (changes what needs to be opened)
# TODO: add basic cli preview to effmpeg
# TODO: integrate preview into gui window
import os
import sys
import subprocess
@ -14,9 +16,9 @@ import datetime
from ffmpy import FFprobe, FFmpeg, FFRuntimeError
# faceswap imports
from lib.cli import FileFullPaths, ComboFullPaths
from lib.cli import FullHelpArgumentParser
from lib.utils import _image_extensions, _video_extensions
from . import cli
if sys.version_info[0] < 3:
raise Exception("This program requires at least python3.2")
@ -113,6 +115,8 @@ class Effmpeg(object):
_actions_req_fps = ["extract", "gen_vid"]
_actions_req_ref_video = ["mux_audio"]
_actions_can_preview = ["gen_vid", "mux_audio", "rescale", "rotate",
"slice"]
_actions_can_use_ref_video = ["gen_vid"]
_actions_have_dir_output = ["extract"]
_actions_have_vid_output = ["gen_vid", "mux_audio", "rescale", "rotate",
@ -122,6 +126,9 @@ class Effmpeg(object):
_actions_have_vid_input = ["extract", "get_fps", "get_info", "rescale",
"rotate", "slice"]
# Class variable that stores the target executable (ffmpeg or ffplay)
_executable = 'ffmpeg'
# Class variable that stores the common ffmpeg arguments based on verbosity
__common_ffmpeg_args_dict = {"normal": "-hide_banner ",
"quiet": "-loglevel panic -hide_banner ",
@ -132,10 +139,8 @@ class Effmpeg(object):
# 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
def __init__(self, arguments):
self.args = arguments
self.input = DataItem()
self.output = DataItem()
self.ref_vid = DataItem()
@ -143,257 +148,8 @@ class Effmpeg(object):
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
def process(self):
# Format action to match the method name
self.args.action = self.args.action.replace('-', '_')
@ -410,7 +166,8 @@ class Effmpeg(object):
else:
self.output = DataItem(path=self.__get_default_output())
if self.args.ref_vid.lower() == "none" or self.args.ref_vid == '':
if self.args.ref_vid is None \
or self.args.ref_vid == '':
self.args.ref_vid = None
# Instantiate ref_vid DataItem object
@ -475,13 +232,16 @@ class Effmpeg(object):
self.args.fps = self.input.fps
# Processing transpose
if self.args.transpose.lower() == "none":
if self.args.transpose is None or \
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 == '':
if self.args.degrees is None \
or self.args.degrees.lower() == "none" \
or self.args.degrees == '':
self.args.degrees = None
elif self.args.transpose is None:
try:
@ -491,6 +251,13 @@ class Effmpeg(object):
"{}".format(self.args.degrees), file=sys.stderr)
exit(1)
# Set executable based on whether previewing or not
"""
if self.args.preview and self.args.action in self._actions_can_preview:
Effmpeg._executable = 'ffplay'
self.output = DataItem()
"""
# Set verbosity of output
self.__set_verbosity(self.args.quiet, self.args.verbose)
@ -498,9 +265,9 @@ class Effmpeg(object):
if self.args.action in self._actions_have_print_output:
self.print_ = True
self.process()
self.effmpeg_process()
def process(self):
def effmpeg_process(self):
kwargs = {"input_": self.input,
"output": self.output,
"ref_vid": self.ref_vid,
@ -512,7 +279,8 @@ class Effmpeg(object):
"degrees": self.args.degrees,
"transpose": self.args.transpose,
"scale": self.args.scale,
"print_": self.print_}
"print_": self.print_,
"preview": self.args.preview}
action = getattr(self, self.args.action)
action(**kwargs)
@ -524,13 +292,12 @@ class Effmpeg(object):
_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)
Effmpeg.__run_ffmpeg(inputs=_input, outputs=_output)
@staticmethod
def gen_vid(input_=None, output=None, fps=None, mux_audio=False,
ref_vid=None, **kwargs):
ref_vid=None, preview=None, **kwargs):
filename = Effmpeg.__get_extracted_filename(input_.path)
_input_opts = Effmpeg._common_ffmpeg_args[:]
_input_path = os.path.join(input_.path, filename)
@ -542,8 +309,7 @@ class Effmpeg(object):
else:
_inputs = {_input_path: _input_opts}
_outputs = {output.path: _output_opts}
ff = FFmpeg(inputs=_inputs, outputs=_outputs)
Effmpeg.__run_ffmpeg(ff)
Effmpeg.__run_ffmpeg(inputs=_inputs, outputs=_outputs)
@staticmethod
def get_fps(input_=None, print_=False, **kwargs):
@ -575,17 +341,16 @@ class Effmpeg(object):
return out
@staticmethod
def rescale(input_=None, output=None, scale=None, **kwargs):
def rescale(input_=None, output=None, scale=None, preview=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)
Effmpeg.__run_ffmpeg(inputs=_inputs, outputs=_outputs)
@staticmethod
def rotate(input_=None, output=None, degrees=None, transpose=None,
**kwargs):
preview=None, **kwargs):
if transpose is None and degrees is None:
raise ValueError("You have not supplied a valid transpose or "
"degrees value:\ntranspose: {}\ndegrees: "
@ -604,29 +369,28 @@ class Effmpeg(object):
_inputs = {input_.path: _input_opts}
_outputs = {output.path: _output_opts}
ff = FFmpeg(inputs=_inputs, outputs=_outputs)
Effmpeg.__run_ffmpeg(ff)
Effmpeg.__run_ffmpeg(inputs=_inputs, outputs=_outputs)
@staticmethod
def mux_audio(input_=None, output=None, ref_vid=None, **kwargs):
def mux_audio(input_=None, output=None, ref_vid=None, preview=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)
Effmpeg.__run_ffmpeg(inputs=_inputs, outputs=_outputs)
@staticmethod
def slice(input_=None, output=None, start=None, duration=None, **kwargs):
def slice(input_=None, output=None, start=None, duration=None,
preview=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)
Effmpeg.__run_ffmpeg(inputs=_inputs, outputs=_output)
# Various helper methods
@classmethod
@ -667,8 +431,9 @@ class Effmpeg(object):
return all(getattr(self, i).fps is None for i in items_to_check)
@staticmethod
def __run_ffmpeg(ff):
@classmethod
def __run_ffmpeg(cls, inputs=None, outputs=None):
ff = FFmpeg(executable=cls._executable, inputs=inputs, outputs=outputs)
try:
ff.run(stderr=subprocess.STDOUT)
except FFRuntimeError as ffe:
@ -712,33 +477,23 @@ class Effmpeg(object):
name = '.'.join(filename[:-2])
vid_ext = ''
underscore = ''
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, ''))
zero_pad = zero_pad.replace(ve, '')
if '_' in zero_pad:
zero_pad = len(zero_pad.replace('_', ''))
underscore = '_'
else:
zero_pad = len(zero_pad)
break
zero_pad = str(zero_pad).zfill(2)
filename_list = [name, vid_ext + '%0' + zero_pad + 'd', img_ext]
filename_list = [name, vid_ext + underscore + '%' + 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(':', '')
@ -760,18 +515,19 @@ class Effmpeg(object):
def bad_args(args):
parser.print_help()
""" Print help on bad arguments """
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 = FullHelpArgumentParser()
SUBPARSER = PARSER.add_subparsers()
EFFMPEG = cli.EffmpegArgs(
SUBPARSER, "effmpeg", "Wrapper for various common ffmpeg commands.")
PARSER.set_defaults(func=bad_args)
ARGUMENTS = PARSER.parse_args()
ARGUMENTS.func(ARGUMENTS)
parser.set_defaults(func=bad_args)
arguments = parser.parse_args()
arguments.func(arguments)

View File

@ -2,7 +2,6 @@
"""
A tool that allows for sorting and grouping images in different ways.
"""
import argparse
import os
import sys
import operator
@ -12,22 +11,16 @@ from tqdm import tqdm
from shutil import copyfile
# faceswap imports
from lib.cli import DirFullPaths, FileFullPaths
import face_recognition
from lib.cli import DirFullPaths, FileFullPaths, FullHelpArgumentParser
import lib.Serializer as Serializer
from . import cli
# DLIB is a GPU Memory hog, so the following modules should only be imported
# when required
face_recognition = None
FaceLandmarksExtractor = None
def import_face_recognition():
""" Import the face_recognition module only when it is required """
global face_recognition
if face_recognition is None:
import face_recognition
def import_FaceLandmarksExtractor():
""" Import the FaceLandmarksExtractor module only when it is required """
global FaceLandmarksExtractor
@ -36,197 +29,13 @@ def import_FaceLandmarksExtractor():
FaceLandmarksExtractor = lib.FaceLandmarksExtractor
class SortProcessor(object):
def __init__(self, subparser, command, description='default'):
self.argument_list = self.get_argument_list()
self.optional_arguments = self.get_optional_arguments()
self.args = None
class Sort(object):
def __init__(self, arguments):
self.args = arguments
self.changes = None
self.serializer = None
self.parse_arguments(description, subparser, command)
@staticmethod
def get_argument_list():
log_filetypes = [["Serializers", ['json', 'yaml']],
["JSON", ["json"]],
["YAML", ["yaml"]]]
log_filetypes = FileFullPaths.prep_filetypes(log_filetypes)
arguments_list = list()
arguments_list.append({"opts": ('-i', '--input'),
"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 "
"faces."})
arguments_list.append({"opts": ('-fp', '--final-process'),
"type": str,
"choices": ("folders", "rename"),
"dest": 'final_process',
"default": "rename",
"help": "'folders': files are sorted using the "
"-s/--sort-by method, then they are "
"organized into folders using the "
"-g/--group-by grouping method. "
"'rename': files are sorted using the "
"-s/--sort-by then they are renamed. "
"Default: rename"})
arguments_list.append({"opts": ('-k', '--keep'),
"action": 'store_true',
"dest": 'keep_original',
"default": False,
"help": "Keeps the original files in the input "
"directory. Be careful when using this "
"with rename grouping and no specified "
"output directory as this would keep "
"the original and renamed files in the "
"same directory."})
arguments_list.append({"opts": ('-s', '--sort-by'),
"type": str,
"choices": ("blur", "face", "face-cnn",
"face-cnn-dissim", "face-dissim",
"face-yaw", "hist",
"hist-dissim"),
"dest": 'sort_method',
"default": "hist",
"help": "Sort by method. "
"Choose how images are sorted. "
"Default: hist"})
arguments_list.append({"opts": ('-g', '--group-by'),
"type": str,
"choices": ("blur", "face", "face-cnn",
"face-yaw", "hist"),
"dest": 'group_method',
"default": "hist",
"help": "Group by method. "
"When -fp/--final-processing by "
"folders choose the how the images are "
"grouped after sorting. "
"Default: hist"})
arguments_list.append({"opts": ('-t', '--ref_threshold'),
"type": float,
"dest": 'min_threshold',
"default": -1.0,
"help": "Float value. "
"Minimum threshold to use for grouping "
"comparison with 'face' and 'hist' "
"methods. The lower the value the more "
"discriminating the grouping is. "
"Leaving -1.0 will make the program "
"set the default value automatically. "
"For face 0.6 should be enough, with "
"0.5 being very discriminating. "
"For face-cnn 7.2 should be enough, "
"with 4 being very discriminating. "
"For hist 0.3 should be enough, with "
"0.2 being very discriminating. "
"Be careful setting a value that's too "
"low in a directory with many images, "
"as this could result in a lot of "
"directories being created. "
"Defaults: face 0.6, face-cnn 7.2, "
"hist 0.3"})
arguments_list.append({"opts": ('-b', '--bins'),
"type": int,
"dest": 'num_bins',
"default": 5,
"help": "Integer value. "
"Number of folders that will be used "
"to group by blur and face-yaw. "
"For blur folder 0 will be the least "
"blurry, while the last folder will be "
"the blurriest. "
"For face-yaw the number of bins is by "
"how much 180 degrees is divided. So "
"if you use 18, then each folder will "
"be a 10 degree increment. Folder 0 "
"will contain faces looking the most "
"to the left whereas the last folder "
"will contain the faces looking the "
"most to the right. "
"If the number of images doesn't "
"divide evenly into the number of "
"bins, the remaining images get put in "
"the last bin."
"Default value: 5"})
arguments_list.append({"opts": ('-l', '--log-changes'),
"action": 'store_true',
"dest": 'log_changes',
"default": False,
"help": "Logs file renaming changes if "
"grouping by renaming, or it logs the "
"file copying/movement if grouping by "
"folders. If no log file is specified "
"with '--log-file', then a "
"'sort_log.json' file will be created "
"in the input directory."})
arguments_list.append({"opts": ('-lf', '--log-file'),
"action": FileFullPaths,
"filetypes": log_filetypes,
"dest": 'log_file_path',
"default": 'sort_log.json',
"help": "Specify a log file to use for saving "
"the renaming or grouping information. "
"If specified extension isn't 'json' "
"or 'yaml', then json will be used as "
"the serializer, with the supplied "
"filename. "
"Default: sort_log.json"})
return arguments_list
@staticmethod
def get_optional_arguments():
"""
Put the arguments in a list so that they are accessible from both
argparse and gui.
"""
# Override this for custom arguments
argument_list = []
return argument_list
def parse_arguments(self, description, subparser, command):
parser = subparser.add_parser(
command,
help="This command lets you sort images using various "
"methods.",
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 = self.add_optional_arguments(parser)
parser.set_defaults(func=self.process_arguments)
def add_optional_arguments(self, parser):
for option in self.optional_arguments:
args = option['opts']
kwargs = {key: option[key] for key in option.keys() if key != 'opts'}
parser.add_argument(*args, **kwargs)
return parser
def process_arguments(self, arguments):
self.args = arguments
def process(self):
# Setting default argument values that cannot be set by argparse
@ -266,9 +75,9 @@ class SortProcessor(object):
self.args.group_method = _group.replace('-', '_')
self.args.final_process = _final.replace('-', '_')
self.process()
self.sort_process()
def process(self):
def sort_process(self):
"""
This method dynamically assigns the functions that will be used to run
the core process of sorting, optionally grouping, renaming/moving into
@ -306,8 +115,6 @@ class SortProcessor(object):
return img_list
def sort_face(self):
import_face_recognition()
input_dir = self.args.input_dir
print("Sorting by face similarity...")
@ -337,8 +144,6 @@ class SortProcessor(object):
return img_list
def sort_face_dissim(self):
import_face_recognition()
input_dir = self.args.input_dir
print("Sorting by face dissimilarity...")
@ -753,8 +558,6 @@ class SortProcessor(object):
:return: img_list but with the comparative values that the chosen
grouping method expects.
"""
import_face_recognition()
input_dir = self.args.input_dir
print("Preparing to group...")
if group_method == 'group_blur':
@ -920,7 +723,6 @@ class SortProcessor(object):
@staticmethod
def get_avg_score_faces(f1encs, references):
import_face_recognition()
scores = []
for f2encs in references:
score = face_recognition.face_distance(f1encs, f2encs)[0]
@ -937,7 +739,8 @@ class SortProcessor(object):
def bad_args(args):
parser.print_help()
""" Print help on bad arguments """
PARSER.print_help()
exit(0)
@ -948,11 +751,11 @@ if __name__ == "__main__":
print(__warning_string)
print("Images sort tool.\n")
parser = argparse.ArgumentParser()
subparser = parser.add_subparsers()
sort = SortProcessor(
subparser, "sort", "Sort images using various methods.")
parser.set_defaults(func=bad_args)
arguments = parser.parse_args()
arguments.func(arguments)
PARSER = FullHelpArgumentParser()
SUBPARSER = PARSER.add_subparsers()
SORT = cli.SortArgs(
SUBPARSER, "sort", "Sort images using various methods.")
PARSER.set_defaults(func=bad_args)
ARGUMENTS = PARSER.parse_args()
ARGUMENTS.func(ARGUMENTS)