tools/preview/preview.py - master - githubmirror/faceswap

tools/preview/preview.py
660 строк25 KB
torzdf
convert bugfixes:
28 мар 2026, 17:13
28 мар 2026, 17:130a0d1fa
#!/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__)