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
7
.gitignore
vendored
|
|
@ -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
|
||||
|
|
|
|||
21
faceswap.py
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
209
lib/cli.py
|
|
@ -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,
|
||||
|
|
|
|||
|
Before Width: | Height: | Size: 691 B After Width: | Height: | Size: 691 B |
BIN
lib/gui/.cache/icons/graph.png
Executable file
|
After Width: | Height: | Size: 574 B |
BIN
lib/gui/.cache/icons/move.png
Executable file
|
After Width: | Height: | Size: 550 B |
|
Before Width: | Height: | Size: 406 B After Width: | Height: | Size: 406 B |
|
Before Width: | Height: | Size: 263 B After Width: | Height: | Size: 263 B |
|
Before Width: | Height: | Size: 773 B After Width: | Height: | Size: 773 B |
|
Before Width: | Height: | Size: 530 B After Width: | Height: | Size: 530 B |
BIN
lib/gui/.cache/icons/zoom.png
Executable file
|
After Width: | Height: | Size: 642 B |
0
lib/gui/.cache/preview/.keep
Normal file
7
lib/gui/__init__.py
Normal 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
|
|
@ -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
|
|
@ -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
|
|
@ -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
|
|
@ -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
|
|
@ -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
|
|
@ -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
|
|
@ -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
|
|
@ -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
|
|
@ -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
|
|
@ -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
|
|
@ -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
|
|
@ -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
|
||||
|
|
@ -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]]))
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
145
plugins/Model_OriginalHighRes/instance_normalization.py
Normal 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})
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
1346
scripts/gui.py
|
|
@ -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:
|
||||
|
|
|
|||
28
tools.py
|
|
@ -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
|
|
@ -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
|
||||
374
tools/effmpeg.py
|
|
@ -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)
|
||||
|
|
|
|||
235
tools/sort.py
|
|
@ -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)
|
||||
|
|
|
|||