1
0
Fork 0
mirror of https://github.com/deepfakes/faceswap synced 2025-06-07 10:43:27 -04:00
faceswap/tools/preview/preview.py
2024-04-05 13:51:57 +01:00

646 lines
26 KiB
Python

#!/usr/bin/env python3
""" Tool to preview swaps and tweak configuration prior to running a convert """
from __future__ import annotations
import gettext
import logging
import random
import tkinter as tk
import typing as T
from tkinter import ttk
import os
import sys
from threading import Event, Lock, Thread
import numpy as np
from lib.align import DetectedFace
from lib.cli.args_extract_convert import ConvertArgs
from lib.gui.utils import get_images, get_config, initialize_config, initialize_images
from lib.convert import Converter
from lib.utils import FaceswapError, handle_deprecated_cliopts
from lib.queue_manager import queue_manager
from scripts.fsmedia import Alignments, Images
from scripts.convert import Predict, ConvertItem
from plugins.extract.pipeline import ExtractMedia
from .control_panels import ActionFrame, ConfigTools, OptionsBook
from .viewer import FacesDisplay, ImagesCanvas
if T.TYPE_CHECKING:
from argparse import Namespace
from lib.queue_manager import EventQueue
from .control_panels import BusyProgressBar
logger = logging.getLogger(__name__)
# LOCALES
_LANG = gettext.translation("tools.preview", localedir="locales", fallback=True)
_ = _LANG.gettext
class Preview(tk.Tk):
""" This tool is part of the Faceswap Tools suite and should be called from
``python tools.py preview`` command.
Loads up 5 semi-random face swaps and displays them, cropped, in place in the final frame.
Allows user to live tweak settings, before saving the final config to
:file:`./config/convert.ini`
Parameters
----------
arguments: :class:`argparse.Namespace`
The :mod:`argparse` arguments as passed in from :mod:`tools.py`
"""
_w: str
def __init__(self, arguments: Namespace) -> None:
logger.debug("Initializing %s: (arguments: '%s'", self.__class__.__name__, arguments)
super().__init__()
arguments = handle_deprecated_cliopts(arguments)
self._config_tools = ConfigTools()
self._lock = Lock()
self._dispatcher = Dispatcher(self)
self._display = FacesDisplay(self, 256, 64)
self._samples = Samples(self, arguments, 5)
self._patch = Patch(self, arguments)
self._initialize_tkinter()
self._image_canvas: ImagesCanvas | None = None
self._opts_book: OptionsBook | None = None
self._cli_frame: ActionFrame | None = None # cli frame holds cli options
logger.debug("Initialized %s", self.__class__.__name__)
@property
def config_tools(self) -> "ConfigTools":
""" :class:`ConfigTools`: The object responsible for parsing configuration options and
updating to/from the GUI """
return self._config_tools
@property
def dispatcher(self) -> "Dispatcher":
""" :class:`Dispatcher`: The object responsible for triggering events and variables and
handling global GUI state """
return self._dispatcher
@property
def display(self) -> FacesDisplay:
""" :class:`~tools.preview.viewer.FacesDisplay`: The object that holds the sample,
converted and patched faces """
return self._display
@property
def lock(self) -> Lock:
""" :class:`threading.Lock`: The threading lock object for the Preview GUI """
return self._lock
@property
def progress_bar(self) -> BusyProgressBar:
""" :class:`~tools.preview.control_panels.BusyProgressBar`: The progress bar that indicates
a swap/patch thread is running """
assert self._cli_frame is not None
return self._cli_frame.busy_progress_bar
def update_display(self):
""" Update the images in the canvas and redraw """
if not hasattr(self, "_image_canvas"): # On first call object not yet created
return
assert self._image_canvas is not None
self._image_canvas.reload()
def _initialize_tkinter(self) -> None:
""" Initialize a standalone tkinter instance. """
logger.debug("Initializing tkinter")
initialize_config(self, None, None)
initialize_images()
get_config().set_geometry(940, 600, fullscreen=False)
self.title("Faceswap.py - Convert Settings")
self.tk.call(
"wm",
"iconphoto",
self._w,
get_images().icons["favicon"]) # pylint:disable=protected-access
logger.debug("Initialized tkinter")
def process(self) -> None:
""" The entry point for the Preview tool from :file:`lib.tools.cli`.
Launch the tkinter preview Window and run main loop.
"""
self._build_ui()
self.mainloop()
def _refresh(self, *args) -> None:
""" Patch faces with current convert settings.
Parameters
----------
*args: tuple
Unused, but required for tkinter callback.
"""
logger.debug("Patching swapped faces. args: %s", args)
self._dispatcher.set_busy()
self._config_tools.update_config()
with self._lock:
assert self._cli_frame is not None
self._patch.converter_arguments = self._cli_frame.convert_args
self._dispatcher.set_needs_patch()
logger.debug("Patched swapped faces")
def _build_ui(self) -> None:
""" Build the elements for displaying preview images and options panels. """
container = ttk.PanedWindow(self,
orient=tk.VERTICAL)
container.pack(fill=tk.BOTH, expand=True)
setattr(container, "preview_display", self._display) # TODO subclass not setattr
self._image_canvas = ImagesCanvas(self, container)
container.add(self._image_canvas, weight=3)
options_frame = ttk.Frame(container)
self._cli_frame = ActionFrame(self, options_frame)
self._opts_book = OptionsBook(options_frame,
self._config_tools,
self._refresh)
container.add(options_frame, weight=1)
self.update_idletasks()
container.sashpos(0, int(400 * get_config().scaling_factor))
class Dispatcher():
""" Handles the app level tk.Variables and the threading events. Dispatches events to the
correct location and handles GUI state whilst events are handled
Parameters
----------
app: :class:`Preview`
The main tkinter Preview app
"""
def __init__(self, app: Preview):
logger.debug("Initializing %s: (app: %s)", self.__class__.__name__, app)
self._app = app
self._tk_busy = tk.BooleanVar(value=False)
self._evnt_needs_patch = Event()
self._is_updating = False
self._stacked_event = False
logger.debug("Initialized %s", self.__class__.__name__)
@property
def needs_patch(self) -> Event:
""":class:`threading.Event`. Set by the parent and cleared by the child. Informs the child
patching thread that a run needs to be processed """
return self._evnt_needs_patch
# TKInter Variables
def set_busy(self) -> None:
""" Set the tkinter busy variable to ``True`` and display the busy progress bar """
if self._tk_busy.get():
logger.debug("Busy event is already set. Doing nothing")
return
if not hasattr(self._app, "progress_bar"):
logger.debug("Not setting busy during initial startup")
return
logger.debug("Setting busy event to True")
self._tk_busy.set(True)
self._app.progress_bar.start()
self._app.update_idletasks()
def _unset_busy(self) -> None:
""" Set the tkinter busy variable to ``False`` and hide the busy progress bar """
self._is_updating = False
if not self._tk_busy.get():
logger.debug("busy unset when already unset. Doing nothing")
return
logger.debug("Setting busy event to False")
self._tk_busy.set(False)
self._app.progress_bar.stop()
self._app.update_idletasks()
# Threading Events
def _wait_for_patch(self) -> None:
""" Wait for a patch thread to complete before triggering a display refresh and unsetting
the busy indicators """
logger.debug("Checking for patch completion...")
if self._evnt_needs_patch.is_set():
logger.debug("Samples not patched. Waiting...")
self._app.after(1000, self._wait_for_patch)
return
logger.debug("Patch completion detected")
self._app.update_display()
self._unset_busy()
if self._stacked_event:
logger.debug("Processing last stacked event")
self.set_busy()
self._stacked_event = False
self.set_needs_patch()
return
def set_needs_patch(self) -> None:
""" Sends a trigger to the patching thread that it needs to be run. Waits for the patching
to complete prior to triggering a display refresh and unsetting the busy indicators """
if self._is_updating:
logger.debug("Request to run patch when it is already running. Adding stacked event.")
self._stacked_event = True
return
self._is_updating = True
logger.debug("Triggering patch")
self._evnt_needs_patch.set()
self._wait_for_patch()
class Samples():
""" The display samples.
Obtains and holds :attr:`sample_size` semi random test faces for displaying in the
preview GUI.
The file list is split into evenly sized groups of :attr:`sample_size`. When a display set is
generated, a random image from each of the groups is selected to provide an array of images
across the length of the video.
Parameters
----------
app: :class:`Preview`
The main tkinter Preview app
arguments: :class:`argparse.Namespace`
The :mod:`argparse` arguments as passed in from :mod:`tools.py`
sample_size: int
The number of samples to take from the input video/images
"""
def __init__(self, app: Preview, arguments: Namespace, sample_size: int) -> None:
logger.debug("Initializing %s: (app: %s, arguments: '%s', sample_size: %s)",
self.__class__.__name__, app, arguments, sample_size)
self._sample_size = sample_size
self._app = app
self._input_images: list[ConvertItem] = []
self._predicted_images: list[tuple[ConvertItem, np.ndarray]] = []
self._images = Images(arguments)
self._alignments = Alignments(arguments,
is_extract=False,
input_is_video=self._images.is_video)
if self._alignments.version == 1.0:
logger.error("The alignments file format has been updated since the given alignments "
"file was generated. You need to update the file to proceed.")
logger.error("To do this run the 'Alignments Tool' > 'Extract' Job.")
sys.exit(1)
if not self._alignments.have_alignments_file:
logger.error("Alignments file not found at: '%s'", self._alignments.file)
sys.exit(1)
self._filelist = self._get_filelist()
self._indices = self._get_indices()
self._predictor = Predict(self._sample_size, arguments)
self._predictor.launch(queue_manager.get_queue("preview_predict_in"))
self._app._display.set_centering(self._predictor.centering)
self.generate()
logger.debug("Initialized %s", self.__class__.__name__)
@property
def available_masks(self) -> list[str]:
""" list: The mask names that are available for every face in the alignments file """
retval = [key
for key, val in self.alignments.mask_summary.items()
if val == self.alignments.faces_count]
return retval
@property
def sample_size(self) -> int:
""" int: The number of samples to take from the input video/images """
return self._sample_size
@property
def predicted_images(self) -> list[tuple[ConvertItem, np.ndarray]]:
""" list: The predicted faces output from the Faceswap model """
return self._predicted_images
@property
def alignments(self) -> Alignments:
""" :class:`~lib.align.Alignments`: The alignments for the preview faces """
return self._alignments
@property
def predictor(self) -> Predict:
""" :class:`~scripts.convert.Predict`: The Predictor for the Faceswap model """
return self._predictor
@property
def _random_choice(self) -> list[int]:
""" list: Random indices from the :attr:`_indices` group """
retval = [random.choice(indices) for indices in self._indices]
logger.debug(retval)
return retval
def _get_filelist(self) -> list[str]:
""" Get a list of files for the input, filtering out those frames which do
not contain faces.
Returns
-------
list
A list of filenames of frames that contain faces.
"""
logger.debug("Filtering file list to frames with faces")
if isinstance(self._images.input_images, str):
filelist = [f"{os.path.splitext(self._images.input_images)[0]}_{frame_no:06d}.png"
for frame_no in range(1, self._images.images_found + 1)]
else:
filelist = self._images.input_images
retval = [filename for filename in filelist
if self._alignments.frame_has_faces(os.path.basename(filename))]
logger.debug("Filtered out frames: %s", self._images.images_found - len(retval))
try:
assert retval
except AssertionError as err:
msg = ("No faces were found in any of the frames passed in. Make sure you are passing "
"in a frames source rather than extracted faces, and that you have provided "
"the correct alignments file.")
raise FaceswapError(msg) from err
return retval
def _get_indices(self) -> list[list[int]]:
""" Get indices for each sample group.
Obtain :attr:`self.sample_size` evenly sized groups of indices
pertaining to the filtered :attr:`self._file_list`
Returns
-------
list
list of indices relating to the filtered file list, split into groups
"""
# Remove start and end values to get a list divisible by self.sample_size
no_files = len(self._filelist)
self._sample_size = min(self._sample_size, no_files)
crop = no_files % self._sample_size
top_tail = list(range(no_files))[
crop // 2:no_files - (crop - (crop // 2))]
# Partition the indices
size = len(top_tail)
retval = [top_tail[start:start + size // self._sample_size]
for start in range(0, size, size // self._sample_size)]
logger.debug("Indices pools: %s", [f"{idx}: (start: {min(pool)}, "
f"end: {max(pool)}, size: {len(pool)})"
for idx, pool in enumerate(retval)])
return retval
def generate(self) -> None:
""" Generate a sample set.
Selects :attr:`sample_size` random faces. Runs them through prediction to obtain the
swap, then trigger the patch event to run the faces through patching.
"""
logger.debug("Generating new random samples")
self._app.dispatcher.set_busy()
self._load_frames()
self._predict()
self._app.dispatcher.set_needs_patch()
logger.debug("Generated new random samples")
def _load_frames(self) -> None:
""" Load a sample of random frames.
* Picks a random face from each indices group.
* Takes the first face from the image (if there are multiple faces). Adds the images to \
:attr:`self._input_images`.
* Sets :attr:`_display.source` to the input images and flags that the display should be \
updated
"""
self._input_images = []
for selection in self._random_choice:
filename = os.path.basename(self._filelist[selection])
image = self._images.load_one_image(self._filelist[selection])
# Get first face only
face = self._alignments.get_faces_in_frame(filename)[0]
detected_face = DetectedFace()
detected_face.from_alignment(face, image=image)
inbound = ExtractMedia(filename=filename, image=image, detected_faces=[detected_face])
self._input_images.append(ConvertItem(inbound=inbound))
self._app.display.source = self._input_images
self._app.display.update_source = True
logger.debug("Selected frames: %s",
[frame.inbound.filename for frame in self._input_images])
def _predict(self) -> None:
""" Predict from the loaded frames.
With a threading lock (to prevent stacking), run the selected faces through the Faceswap
model predict function and add the output to :attr:`predicted`
"""
with self._app.lock:
self._predicted_images = []
for frame in self._input_images:
self._predictor.in_queue.put(frame)
idx = 0
while idx < self._sample_size:
logger.debug("Predicting face %s of %s", idx + 1, self._sample_size)
items: (T.Literal["EOF"] |
list[tuple[ConvertItem, np.ndarray]]) = self._predictor.out_queue.get()
if items == "EOF":
logger.debug("Received EOF")
break
for item in items:
self._predicted_images.append(item)
logger.debug("Predicted face %s of %s", idx + 1, self._sample_size)
idx += 1
logger.debug("Predicted faces")
class Patch():
""" The Patch pipeline
Runs in it's own thread. Takes the output from the Faceswap model predictor and runs the faces
through the convert pipeline using the currently selected options.
Parameters
----------
app: :class:`Preview`
The main tkinter Preview app
arguments: :class:`argparse.Namespace`
The :mod:`argparse` arguments as passed in from :mod:`tools.py`
Attributes
----------
converter_arguments: dict
The currently selected converter command line arguments for the patch queue
"""
def __init__(self, app: Preview, arguments: Namespace) -> None:
logger.debug("Initializing %s: (app: %s, arguments: '%s')",
self.__class__.__name__, app, arguments)
self._app = app
self._queue_patch_in = queue_manager.get_queue("preview_patch_in")
self.converter_arguments: dict[str, T.Any] | None = None # Updated converter args
configfile = arguments.configfile if hasattr(arguments, "configfile") else None
self._converter = Converter(output_size=app._samples.predictor.output_size,
coverage_ratio=app._samples.predictor.coverage_ratio,
centering=app._samples.predictor.centering,
draw_transparent=False,
pre_encode=None,
arguments=self._generate_converter_arguments(
arguments,
app._samples.available_masks),
configfile=configfile)
self._thread = Thread(target=self._process,
name="patch_thread",
args=(self._queue_patch_in,
self._app.dispatcher.needs_patch,
app._samples),
daemon=True)
self._thread.start()
logger.debug("Initializing %s", self.__class__.__name__)
@property
def converter(self) -> Converter:
""" :class:`lib.convert.Converter`: The converter to use for patching the images. """
return self._converter
@staticmethod
def _generate_converter_arguments(arguments: Namespace,
available_masks: list[str]) -> Namespace:
""" Add the default converter arguments to the initial arguments. Ensure the mask selection
is available.
Parameters
----------
arguments: :class:`argparse.Namespace`
The :mod:`argparse` arguments as passed in from :mod:`tools.py`
available_masks: list
The masks that are available for convert
Returns
----------
arguments: :class:`argparse.Namespace`
The :mod:`argparse` arguments as passed in with converter default
arguments added
"""
valid_masks = available_masks + ["none"]
converter_arguments = ConvertArgs(None, "convert").get_optional_arguments() # type: ignore
for item in converter_arguments:
value = item.get("default", None)
# Skip options without a default value
if value is None:
continue
option = item.get("dest", item["opts"][1].replace("--", ""))
if option == "mask_type" and value not in valid_masks:
logger.debug("Amending default mask from '%s' to '%s'", value, valid_masks[0])
value = valid_masks[0]
# Skip options already in arguments
if hasattr(arguments, option):
continue
# Add option to arguments
setattr(arguments, option, value)
logger.debug(arguments)
return arguments
def _process(self,
patch_queue_in: EventQueue,
trigger_event: Event,
samples: Samples) -> None:
""" The face patching process.
Runs in a thread, and waits for an event to be set. Once triggered, runs a patching
cycle and sets the :class:`Display` destination images.
Parameters
----------
patch_queue_in: :class:`~lib.queue_manager.EventQueue`
The input queue for the patching process
trigger_event: :class:`threading.Event`
The event that indicates a patching run needs to be processed
samples: :class:`Samples`
The Samples for display.
"""
logger.debug("Launching patch process thread: (patch_queue_in: %s, trigger_event: %s, "
"samples: %s)", patch_queue_in, trigger_event, samples)
patch_queue_out = queue_manager.get_queue("preview_patch_out")
while True:
trigger = trigger_event.wait(1)
if not trigger:
continue
logger.debug("Patch Triggered")
queue_manager.flush_queue("preview_patch_in")
self._feed_swapped_faces(patch_queue_in, samples)
with self._app.lock:
self._update_converter_arguments()
self._converter.reinitialize(config=self._app.config_tools.config)
swapped = self._patch_faces(patch_queue_in, patch_queue_out, samples.sample_size)
with self._app.lock:
self._app.display.destination = swapped
logger.debug("Patch complete")
trigger_event.clear()
logger.debug("Closed patch process thread")
def _update_converter_arguments(self) -> None:
""" Update the converter arguments to the currently selected values. """
logger.debug("Updating Converter cli arguments")
if self.converter_arguments is None:
logger.debug("No arguments to update")
return
for key, val in self.converter_arguments.items():
logger.debug("Updating %s to %s", key, val)
setattr(self._converter.cli_arguments, key, val)
logger.debug("Updated Converter cli arguments")
@staticmethod
def _feed_swapped_faces(patch_queue_in: EventQueue, samples: Samples) -> None:
""" Feed swapped faces to the converter's in-queue.
Parameters
----------
patch_queue_in: :class:`~lib.queue_manager.EventQueue`
The input queue for the patching process
samples: :class:`Samples`
The Samples for display.
"""
logger.debug("feeding swapped faces to converter")
for item in samples.predicted_images:
patch_queue_in.put(item)
logger.debug("fed %s swapped faces to converter",
len(samples.predicted_images))
logger.debug("Putting EOF to converter")
patch_queue_in.put("EOF")
def _patch_faces(self,
queue_in: EventQueue,
queue_out: EventQueue,
sample_size: int) -> list[np.ndarray]:
""" Patch faces.
Run the convert process on the swapped faces and return the patched faces.
patch_queue_in: :class:`~lib.queue_manager.EventQueue`
The input queue for the patching process
queue_out: :class:`~lib.queue_manager.EventQueue`
The output queue from the patching process
sample_size: int
The number of samples to be displayed
Returns
-------
list
The swapped faces patched with the selected convert settings
"""
logger.debug("Patching faces")
self._converter.process(queue_in, queue_out)
swapped = []
idx = 0
while idx < sample_size:
logger.debug("Patching image %s of %s", idx + 1, sample_size)
item = queue_out.get()
swapped.append(item[1])
logger.debug("Patched image %s of %s", idx + 1, sample_size)
idx += 1
logger.debug("Patched faces")
return swapped