/
githubmirror
/
faceswap
Обзор
Документация
Войти
/
githubmirror
/
faceswap
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
lib/gui/utils/image.py
664 строки
25 KB
torzdf
Train: Migrate data loading to Torch (#1540)
16 апр 2026, 21:19
Не верифицирован
16 апр 2026, 21:19
83fdd5f
Код
Авторство
О чём код?
#!/usr/bin python3 """ Utilities for handling images in the Faceswap GUI """ from __future__ import annotations import logging import os import typing as T import cv2 import numpy as np from PIL import Image, ImageDraw, ImageTk from lib.gui import gui_config as cfg from lib.training.preview_cv import PreviewBuffer from lib.utils import get_module_objects from .config import get_config, PATH_CACHE if T.TYPE_CHECKING: from collections.abc import Sequence logger = logging.getLogger(__name__) _IMAGES: Images | None = None _PREVIEW_TRIGGER: PreviewTrigger | None = None TRAINING_PREVIEW = ".gui_training_preview.png" def initialize_images() -> None: """ Initialize the :class:`Images` handler and add to global constant. This should only be called once on first GUI startup. Future access to :class:`Images` handler should only be executed through :func:`get_images`. """ global _IMAGES # pylint:disable=global-statement if _IMAGES is not None: return logger.debug("Initializing images") _IMAGES = Images() def get_images() -> "Images": """ Get the Master GUI Images handler. Returns ------- :class:`Images` The Master GUI Images handler """ assert _IMAGES is not None return _IMAGES def _get_previews(image_path: str) -> list[str]: """ Get the images stored within the given directory. Parameters ---------- image_path: str The folder containing images to be scanned Returns ------- list: The image filenames stored within the given folder """ logger.debug("Getting images: '%s'", image_path) if not os.path.isdir(image_path): logger.debug("Folder does not exist") return [] files = [os.path.join(image_path, f) for f in os.listdir(image_path) if f.lower().endswith((".png", ".jpg"))] logger.debug("Image files: %s", files) return files class PreviewTrain(): """ Handles the loading of the training preview image(s) and adding to the display buffer Parameters ---------- cache_path: str Full path to the cache folder that contains the preview images """ def __init__(self, cache_path: str) -> None: logger.debug("Initializing %s: (cache_path: '%s')", self.__class__.__name__, cache_path) self._buffer = PreviewBuffer() self._cache_path = cache_path self._modified: float = 0.0 self._error_count: int = 0 logger.debug("Initialized %s", self.__class__.__name__) @property def buffer(self) -> PreviewBuffer: """ :class:`~lib.training.PreviewBuffer` The preview buffer for the training preview image. """ return self._buffer def load(self) -> bool: """ Load the latest training preview image(s) from disk and add to :attr:`buffer` """ logger.trace("Loading Training preview images") # type:ignore image_files = _get_previews(self._cache_path) filename = next((fname for fname in image_files if os.path.basename(fname) == TRAINING_PREVIEW), "") img: np.ndarray | None = None if not filename: logger.trace("No preview to display") # type:ignore return False try: modified = os.path.getmtime(filename) if modified <= self._modified: logger.trace("preview '%s' not updated. Current timestamp: %s, " # type:ignore "existing timestamp: %s", filename, modified, self._modified) return False logger.debug("Loading preview: '%s'", filename) img = cv2.imread(filename, cv2.IMREAD_UNCHANGED) assert img is not None self._modified = modified self._buffer.add_image(os.path.basename(filename), img) self._error_count = 0 except (ValueError, AssertionError): # 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 logger.debug("Unable to display preview: (image: '%s', attempt: %s)", img, self._error_count) if self._error_count < 10: self._error_count += 1 else: logger.error("Error reading the preview file for '%s'", filename) return False logger.debug("Loaded preview: '%s' (%s)", filename, img.shape) return True def reset(self) -> None: """ Reset the preview buffer when the display page has been disabled. Notes ----- The buffer requires resetting, otherwise the re-enabled preview window hangs waiting for a training image that has already been marked as processed """ logger.debug("Resetting training preview") del self._buffer self._buffer = PreviewBuffer() self._modified = 0.0 self._error_count = 0 class PreviewExtract(): """ Handles the loading of preview images for extract and convert Parameters ---------- cache_path: str Full path to the cache folder that contains the preview images """ def __init__(self, cache_path: str) -> None: logger.debug("Initializing %s: (cache_path: '%s')", self.__class__.__name__, cache_path) self._cache_path = cache_path self._batch_mode = False self._output_path = "" self._modified: float = 0.0 self._filenames: list[str] = [] self._images: np.ndarray | None = None self._placeholder: np.ndarray | None = None self._preview_image: Image.Image | None = None self._preview_image_tk: ImageTk.PhotoImage | None = None logger.debug("Initialized %s", self.__class__.__name__) @property def image(self) -> ImageTk.PhotoImage: """:class:`PIL.ImageTk.PhotoImage` The preview image for displaying in a tkinter canvas """ assert self._preview_image_tk is not None return self._preview_image_tk def save(self, filename: str) -> None: """ Save the currently displaying preview image to the given location Parameters ---------- filename: str The full path to the filename to save the preview image to """ logger.debug("Saving preview to %s", filename) assert self._preview_image is not None self._preview_image.save(filename) def set_faceswap_output_path(self, location: str, batch_mode: bool = False) -> None: """ Set the path that will contain the output from an Extract or Convert task. Required so that the GUI can fetch output images to display for return in :attr:`preview_image`. Parameters ---------- location: str The output location that has been specified for an Extract or Convert task batch_mode: bool ``True`` if extracting in batch mode otherwise False """ self._output_path = location self._batch_mode = batch_mode def _get_newest_folder(self) -> str: """ Obtain the most recent folder created in the extraction output folder when processing in batch mode. Returns ------- str The most recently modified folder within the parent output folder. If no folders have been created, returns the parent output folder """ folders = [] if not os.path.exists(self._output_path) else [ os.path.join(self._output_path, folder) for folder in os.listdir(self._output_path) if os.path.isdir(os.path.join(self._output_path, folder))] folders.sort(key=os.path.getmtime) retval = folders[-1] if folders else self._output_path logger.debug("sorted folders: %s, return value: %s", folders, retval) return retval def _get_newest_filenames(self, image_files: list[str]) -> list[str]: """ Return image filenames that have been modified since the last check. Parameters ---------- image_files: list The list of image files to check the modification date for Returns ------- list: A list of images that have been modified since the last check """ if not self._modified: retval = image_files else: retval = [fname for fname in image_files if os.path.getmtime(fname) > self._modified] if not retval: logger.debug("No new images in output folder") else: self._modified = max(os.path.getmtime(img) for img in retval) logger.debug("Number new images: %s, Last Modified: %s", len(retval), self._modified) return retval def _pad_and_border(self, image: Image.Image, size: int) -> np.ndarray: """ Pad rectangle images to a square and draw borders Parameters ---------- image: :class:`PIL.Image` The image to process size: int The size of the image as it should be displayed Returns ------- :class:`numpy.ndarray`: The processed image """ if image.size[0] != image.size[1]: # Pad to square new_img = Image.new("RGB", (size, size)) new_img.paste(image, ((size - image.size[0]) // 2, (size - image.size[1]) // 2)) image = new_img draw = ImageDraw.Draw(image) draw.rectangle(((0, 0), (size, size)), outline="#E5E5E5", width=1) retval = np.array(image) logger.trace("image shape: %s", retval.shape) # type: ignore return retval def _process_samples(self, samples: list[np.ndarray], filenames: list[str], num_images: int) -> bool: """ Process the latest sample images into a displayable image. Parameters ---------- samples: list The list of extract/convert preview images to display filenames: list The full path to the filenames corresponding to the images num_images: int The number of images that should be displayed Returns ------- bool ``True`` if samples successfully compiled otherwise ``False`` """ a_samples = np.array(samples) if not np.any(a_samples): logger.debug("No preview images collected.") return False self._filenames = (self._filenames + filenames)[-num_images:] cache = self._images if cache is None: logger.debug("Creating new cache") cache = a_samples[-num_images:] else: logger.debug("Appending to existing cache") cache = np.concatenate((cache, a_samples))[-num_images:] self._images = cache assert self._images is not None logger.debug("Cache shape: %s", self._images.shape) return True def _load_images_to_cache(self, # pylint:disable=too-many-locals image_files: list[str], frame_dims: tuple[int, int], thumbnail_size: int) -> bool: """ Load preview images to the image cache. Load new images and append to cache, filtering the cache to the number of thumbnails that will fit inside the display panel. Parameters ---------- image_files: list A list of new image files that have been modified since the last check frame_dims: tuple The (width (`int`), height (`int`)) of the display panel that will display the preview thumbnail_size: int The size of each thumbnail that should be created Returns ------- bool ``True`` if images were successfully loaded to cache otherwise ``False`` """ logger.debug("Number image_files: %s, frame_dims: %s, thumbnail_size: %s", len(image_files), frame_dims, thumbnail_size) num_images = (frame_dims[0] // thumbnail_size) * (frame_dims[1] // thumbnail_size) logger.debug("num_images: %s", num_images) if num_images == 0: return False samples: list[np.ndarray] = [] start_idx = len(image_files) - num_images if len(image_files) > num_images else 0 show_files = sorted(image_files, key=os.path.getctime)[start_idx:] dropped_files = [] for fname in show_files: try: img_file = Image.open(fname) except PermissionError as err: logger.debug("Permission error opening preview file: '%s'. Original error: %s", fname, str(err)) dropped_files.append(fname) continue except Exception as err: # pylint:disable=broad-except # Swallow any issues with opening an image rather than spamming console # Can happen when trying to read partially saved images logger.debug("Error opening preview file: '%s'. Original error: %s", fname, str(err)) dropped_files.append(fname) continue width, height = img_file.size scaling = thumbnail_size / max(width, height) logger.debug("image width: %s, height: %s, scaling: %s", width, height, scaling) try: img = img_file.resize((int(width * scaling), int(height * scaling))) except OSError as err: # Image only gets loaded when we call a method, so may error on partial loads logger.debug("OS Error resizing preview image: '%s'. Original error: %s", fname, err) dropped_files.append(fname) continue samples.append(self._pad_and_border(img, thumbnail_size)) return self._process_samples(samples, [fname for fname in show_files if fname not in dropped_files], num_images) def _create_placeholder(self, thumbnail_size: int) -> None: """ Create a placeholder image for when there are fewer thumbnails available than columns to display them. Parameters ---------- thumbnail_size: int The size of the thumbnail that the placeholder should replicate """ logger.debug("Creating placeholder. thumbnail_size: %s", thumbnail_size) placeholder = Image.new("RGB", (thumbnail_size, thumbnail_size)) draw = ImageDraw.Draw(placeholder) draw.rectangle(((0, 0), (thumbnail_size, thumbnail_size)), outline="#E5E5E5", width=1) n_placeholder = np.array(placeholder) self._placeholder = n_placeholder logger.debug("Created placeholder. shape: %s", n_placeholder.shape) def _place_previews(self, frame_dims: tuple[int, int]) -> Image.Image | None: """ Format the preview thumbnails stored in the cache into a grid fitting the display panel. Parameters ---------- frame_dims: tuple The (width (`int`), height (`int`)) of the display panel that will display the preview Returns ------- :class:`PIL.Image`: | None The final preview display image """ if self._images is None: logger.debug("No images in cache. Returning None") return None samples = self._images.copy() num_images, thumbnail_size = samples.shape[:2] if self._placeholder is None: self._create_placeholder(thumbnail_size) logger.debug("num_images: %s, thumbnail_size: %s", num_images, thumbnail_size) cols, rows = frame_dims[0] // thumbnail_size, frame_dims[1] // thumbnail_size logger.debug("cols: %s, rows: %s", cols, rows) if cols == 0 or rows == 0: logger.debug("Cols or Rows is zero. No items to display") return None remainder = (cols * rows) - num_images if remainder != 0: logger.debug("Padding sample display. Remainder: %s", remainder) assert self._placeholder is not None placeholder = np.concatenate([np.expand_dims(self._placeholder, 0)] * remainder) samples = np.concatenate((samples, placeholder)) display = np.vstack([np.hstack(T.cast("Sequence", samples[row * cols: (row + 1) * cols])) for row in range(rows)]) logger.debug("display shape: %s", display.shape) return Image.fromarray(display) def load_latest_preview(self, thumbnail_size: int, frame_dims: tuple[int, int]) -> bool: """ Load the latest preview image for extract and convert. Retrieves the latest preview images from the faceswap output folder, resizes to thumbnails and lays out for display. Places the images into :attr:`preview_image` for loading into the display panel. Parameters ---------- thumbnail_size: int The size of each thumbnail that should be created frame_dims: tuple The (width (`int`), height (`int`)) of the display panel that will display the preview Returns ------- bool ``True`` if a preview was successfully loaded otherwise ``False`` """ logger.debug("Loading preview image: (thumbnail_size: %s, frame_dims: %s)", thumbnail_size, frame_dims) image_path = self._get_newest_folder() if self._batch_mode else self._output_path image_files = _get_previews(image_path) gui_preview = os.path.join(self._output_path, ".gui_preview.jpg") if not image_files or (len(image_files) == 1 and gui_preview not in image_files): logger.debug("No preview to display") return False # Filter to just the gui_preview if it exists in folder output image_files = [gui_preview] if gui_preview in image_files else image_files logger.debug("Image Files: %s", len(image_files)) image_files = self._get_newest_filenames(image_files) if not image_files: return False if not self._load_images_to_cache(image_files, frame_dims, thumbnail_size): logger.debug("Failed to load any preview images") if gui_preview in image_files: # Reset last modified for failed loading of a gui preview image so it is picked # up next time self._modified = 0.0 return False if image_files == [gui_preview]: # Delete the preview image so that the main scripts know to output another logger.debug("Deleting preview image") os.remove(image_files[0]) show_image = self._place_previews(frame_dims) if not show_image: self._preview_image = None self._preview_image_tk = None return False logger.debug("Displaying preview: %s", self._filenames) self._preview_image = show_image self._preview_image_tk = ImageTk.PhotoImage(show_image) return True def delete_previews(self) -> None: """ Remove any image preview files """ for fname in self._filenames: if os.path.basename(fname) == ".gui_preview.jpg": logger.debug("Deleting: '%s'", fname) try: os.remove(fname) except FileNotFoundError: logger.debug("File does not exist: %s", fname) class Images(): """ The centralized image repository for holding all icons and images required by the GUI. This class should be initialized on GUI startup through :func:`initialize_images`. Any further access to this class should be through :func:`get_images`. """ def __init__(self) -> None: logger.debug("Initializing %s", self.__class__.__name__) self._path_preview = os.path.join(PATH_CACHE, "preview") self._batch_mode = False self._preview_train = PreviewTrain(self._path_preview) self._preview_extract = PreviewExtract(self._path_preview) self._icons = self._load_icons() logger.debug("Initialized %s", self.__class__.__name__) @property def preview_train(self) -> PreviewTrain: """ :class:`PreviewTrain` The object handling the training preview images """ return self._preview_train @property def preview_extract(self) -> PreviewExtract: """ :class:`PreviewTrain` The object handling the training preview images """ return self._preview_extract @property def icons(self) -> dict[str, ImageTk.PhotoImage]: """ dict: The faceswap icons for all parts of the GUI. The dictionary key is the icon name (`str`) the value is the icon sized and formatted for display (:class:`PIL.ImageTK.PhotoImage`). Example ------- >>> icons = get_images().icons >>> save = icons["save"] >>> button = ttk.Button(parent, image=save) >>> button.pack() """ return self._icons @staticmethod def _load_icons() -> dict[str, ImageTk.PhotoImage]: """ Scan the icons cache folder and load the icons into :attr:`icons` for retrieval throughout the GUI. Returns ------- dict: The icons formatted as described in :attr:`icons` """ size = cfg.icon_size() size = int(round(size * get_config().scaling_factor)) icons: dict[str, ImageTk.PhotoImage] = {} path_icons = os.path.join(PATH_CACHE, "icons") for fname in os.listdir(path_icons): name, ext = os.path.splitext(fname) if ext != ".png": continue img = Image.open(os.path.join(path_icons, fname)) p_img = ImageTk.PhotoImage(img.resize((size, size), resample=Image.Resampling.HAMMING)) icons[name] = p_img logger.debug(icons) return icons def delete_preview(self) -> None: """ Delete the preview files in the cache folder and reset the image cache. Should be called when terminating tasks, or when Faceswap starts up or shuts down. """ logger.debug("Deleting previews") for item in os.listdir(self._path_preview): if item.startswith(os.path.splitext(TRAINING_PREVIEW)[0]) and item.endswith((".jpg", ".png")): full_item = os.path.join(self._path_preview, item) logger.debug("Deleting: '%s'", full_item) os.remove(full_item) self._preview_extract.delete_previews() del self._preview_train del self._preview_extract self._preview_train = PreviewTrain(self._path_preview) self._preview_extract = PreviewExtract(self._path_preview) class PreviewTrigger(): """ Triggers to indicate to underlying Faceswap process that the preview image should be updated. Writes a file to the cache folder that is picked up by the main process. """ def __init__(self) -> None: logger.debug("Initializing: %s", self.__class__.__name__) self._trigger_files = {"update": os.path.join(PATH_CACHE, ".preview_trigger"), "mask_toggle": os.path.join(PATH_CACHE, ".preview_mask_toggle")} logger.debug("Initialized: %s (trigger_files: %s)", self.__class__.__name__, self._trigger_files) def set(self, trigger_type: T.Literal["update", "mask_toggle"]): """ Place the trigger file into the cache folder Parameters ---------- trigger_type: ["update", "mask_toggle"] The type of action to trigger. 'update': Full preview update. 'mask_toggle': toggle mask on and off """ trigger = self._trigger_files[trigger_type] if not os.path.isfile(trigger): with open(trigger, "w", encoding="utf8"): pass logger.debug("Set preview trigger: %s", trigger) def clear(self, trigger_type: T.Literal["update", "mask_toggle"] | None = None) -> None: """ Remove the trigger file from the cache folder. Parameters ---------- trigger_type: ["update", "mask_toggle", ``None``], optional The trigger to clear. 'update': Full preview update. 'mask_toggle': toggle mask on and off. ``None`` - clear all triggers. Default: ``None`` """ if trigger_type is None: triggers = list(self._trigger_files.values()) else: triggers = [self._trigger_files[trigger_type]] for trigger in triggers: if os.path.isfile(trigger): os.remove(trigger) logger.debug("Removed preview trigger: %s", trigger) def preview_trigger() -> PreviewTrigger: """ Set the global preview trigger if it has not already been set and return. Returns ------- :class:`PreviewTrigger` The trigger to indicate to the main faceswap process that it should perform a training preview update """ global _PREVIEW_TRIGGER # pylint:disable=global-statement if _PREVIEW_TRIGGER is None: _PREVIEW_TRIGGER = PreviewTrigger() return _PREVIEW_TRIGGER __all__ = get_module_objects(__name__)