mirror of
https://github.com/deepfakes/faceswap
synced 2025-06-07 19:05:02 -04:00
353 lines
13 KiB
Python
353 lines
13 KiB
Python
#!/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():
|
|
""" Raise a filedialog box and capture input """
|
|
|
|
def __init__(self, handletype, filetype, command=None, action=None,
|
|
variable=None):
|
|
|
|
self.handletype = handletype
|
|
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": "filename",
|
|
"gen-vid": "dir",
|
|
"get-fps": "filename",
|
|
"get-info": "filename",
|
|
"mux-audio": "filename",
|
|
"rescale": "filename",
|
|
"rotate": "filename",
|
|
"slice": "filename"},
|
|
"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(filetype, command, action, variable)
|
|
self.retfile = getattr(self, 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, filetype, command, action, variable=None):
|
|
""" Generate the required kwargs for the requested browser """
|
|
kwargs = dict()
|
|
if self.handletype.lower() == "context":
|
|
self.set_context_handletype(command, action, variable)
|
|
|
|
if self.handletype.lower() in (
|
|
"open", "save", "filename", "savefilename"):
|
|
kwargs["filetypes"] = self.filetypes[filetype]
|
|
if self.defaults.get(filetype, None):
|
|
kwargs['defaultextension'] = self.defaults[filetype]
|
|
if self.handletype.lower() == "save":
|
|
kwargs["mode"] = "w"
|
|
if self.handletype.lower() == "open":
|
|
kwargs["mode"] = "r"
|
|
return kwargs
|
|
|
|
def set_context_handletype(self, command, action, variable):
|
|
""" Choose the correct file browser action based on context """
|
|
if self.contexts[command].get(variable, None) is not None:
|
|
self.handletype = self.contexts[command][variable][action]
|
|
else:
|
|
self.handletype = self.contexts[command][action]
|
|
|
|
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)
|
|
|
|
@staticmethod
|
|
def nothing():
|
|
""" Method that does nothing, used for disabling open/save pop up """
|
|
return
|
|
|
|
|
|
class Images(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): # pylint: disable=too-many-ancestors
|
|
""" 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():
|
|
""" 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()
|