Files
torzdf 0a0d1fabf0 convert bugfixes:
- ffmpeg writer: Exit audio buffer when no audio available
  - convert: correctly apply y-offset
  - Fix color adjust when "none" selected
2026-03-28 14:13:10 +00:00

661 lines
25 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.image import SingleFrameLoader
from lib.infer.objects import FrameFaces
from lib.convert import Converter
from lib.utils import get_module_objects, FaceswapError, handle_deprecated_cli_opts
from lib.queue_manager import queue_manager
from lib.video import check_for_video
from scripts.fs_media import Alignments
from scripts.convert import Predict, ConvertItem
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
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_cli_opts(arguments)
self._config_tools = ConfigTools(arguments.config_file)
self._lock = Lock()
self._dispatcher = Dispatcher(self)
num_faces = 5
self._display = FacesDisplay(self, 256, 64, num_faces)
self._samples = Samples(self, arguments, num_faces)
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":
"""The object responsible for parsing configuration options and updating to/from the GUI"""
return self._config_tools
@property
def dispatcher(self) -> "Dispatcher":
"""Responsible for triggering events and variables and handling global GUI state"""
return self._dispatcher
@property
def display(self) -> FacesDisplay:
"""The object that holds the sample, converted and patched faces"""
return self._display
@property
def lock(self) -> Lock:
"""The threading lock object for the Preview GUI"""
return self._lock
@property
def progress_bar(self) -> 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
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
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._event_needs_patch = Event()
self._is_updating = False
self._stacked_event = False
logger.debug("Initialized %s", self.__class__.__name__)
@property
def needs_patch(self) -> Event:
"""Set by the parent and cleared by the child. Informs the child patching thread that a
run needs to be processed"""
return self._event_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._event_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._event_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
The main tkinter Preview app
arguments
The :mod:`argparse` arguments as passed in from :mod:`tools.py`
sample_size
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]] = []
is_video = check_for_video(arguments.input_dir)
self._alignments = Alignments(arguments.alignments_path,
arguments.input_dir,
is_extract=False,
input_is_video=is_video)
if not self._alignments.have_alignments_file:
logger.error("Alignments file not found at: '%s'", self._alignments.file)
sys.exit(1)
video_meta = self._alignments.video_meta_data
self._images = SingleFrameLoader(arguments.input_dir, video_meta_data=video_meta)
if is_video and video_meta is None:
video_meta = self._images.video_meta_data
assert video_meta is not None
self._alignments.save_video_meta_data(video_meta["pts_time"], video_meta["keyframes"])
if self._images.is_video:
self._alignments.update_legacy_has_source(os.path.basename(arguments.input_dir))
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_offset(self._predictor.centering,
self._predictor.y_offset)
self.generate()
logger.debug("Initialized %s", self.__class__.__name__)
@property
def available_masks(self) -> list[str]:
"""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:
"""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]]:
"""The predicted faces output from the Faceswap model"""
return self._predicted_images
@property
def alignments(self) -> Alignments:
"""The alignments for the preview faces"""
return self._alignments
@property
def predictor(self) -> Predict:
"""The Predictor for the Faceswap model"""
return self._predictor
@property
def _random_choice(self) -> list[int]:
"""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")
retval = [filename for filename in self._images.file_list
if self._alignments.frame_has_faces(os.path.basename(filename))]
logger.debug("Filtered out frames: %s", self._images.count - 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 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 = self._filelist[selection]
basename = os.path.basename(filename)
if self._images.is_video and basename.isdigit():
frame_no = int(basename)
elif self._images.is_video:
frame_no = int(os.path.splitext(basename)[0][filename.rfind("_") + 1:])
logger.trace( # type:ignore[attr-defined]
"Extracted frame_no %s from filename '%s'", frame_no, basename)
else:
frame_no = self._images.file_list.index(filename)
_, image = self._images.image_from_index(frame_no)
# Get first face only
face = self._alignments.get_faces_in_frame(basename)[0]
detected_face = DetectedFace()
detected_face.from_alignment(face, image=image)
inbound = FrameFaces(filename=basename, image=image)
inbound.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
The main tkinter Preview app
arguments
The :mod:`argparse` arguments as passed in from :mod:`tools.py`
Attributes
----------
converter_arguments
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
config_file = arguments.config_file if hasattr(arguments, "config_file") 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),
config_file=config_file)
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:
"""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
The :mod:`argparse` arguments as passed in from :mod:`tools.py`
available_masks: list
The masks that are available for convert
Returns
----------
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
The input queue for the patching process
trigger_event
The event that indicates a patching run needs to be processed
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()
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
The input queue for the patching process
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
The input queue for the patching process
queue_out
The output queue from the patching process
sample_size
The number of samples to be displayed
Returns
-------
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
__all__ = get_module_objects(__name__)