/
githubmirror
/
faceswap
Обзор
Документация
Войти
/
githubmirror
/
faceswap
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
v3.0.0
lib/training/cache.py
703 строки
29 KB
torzdf
Faceswap 3 (#1516)
21 дек 2025, 05:45
Не верифицирован
21 дек 2025, 05:45
837bc2d
Код
Авторство
О чём код?
#!/usr/bin/env python3 """ Holds the data cache for training data generators """ from __future__ import annotations import logging import os import typing as T from dataclasses import dataclass, field from threading import Lock import cv2 import numpy as np from tqdm import tqdm from lib.align import CenteringType, DetectedFace, LandmarkType from lib.image import read_image_batch, read_image_meta_batch from lib.logger import parse_class_init from lib.utils import FaceswapError, get_module_objects from plugins.train import train_config as cfg if T.TYPE_CHECKING: from lib.align.alignments import PNGHeaderAlignmentsDict, PNGHeaderDict from lib import align logger = logging.getLogger(__name__) _FACE_CACHES: dict[str, Cache] = {} @dataclass class _MaskConfig: """ Holds the constants required for manipulating training masks """ # pylint:disable=unnecessary-lambda penalized: bool = field(default_factory=lambda: cfg.Loss.penalized_mask_loss()) learn: bool = field(default_factory=lambda: cfg.Loss.learn_mask()) mask_type: str | None = field(default_factory=lambda: None if cfg.Loss.mask_type() == "none" else cfg.Loss.mask_type()) dilation: float = field(default_factory=lambda: cfg.Loss.mask_dilation()) kernel: int = field(default_factory=lambda: cfg.Loss.mask_blur_kernel()) threshold: int = field(default_factory=lambda: cfg.Loss.mask_threshold()) multiplier_enabled: bool = field( default_factory=lambda: ((cfg.Loss.eye_multiplier() > 1 or cfg.Loss.mouth_multiplier() > 1) and cfg.Loss.penalized_mask_loss())) @property def mask_enabled(self) -> bool: """ bool : ``True`` if any of :attr:`penalized` or :attr:`learn` are true and :attr:`mask_type` is not ``None`` """ return self.mask_type is not None and (self.learn or self.penalized) class _MaskProcessing: """ Handle the extraction and processing of masks from faceswap PNG headers for caching Parameters ---------- size : int The largest output size of the model coverage_ratio : float The coverage ratio that the model is using. centering : Literal["face", "head", "legacy"] """ def __init__(self, size: int, coverage_ratio: float, centering: CenteringType) -> None: assert isinstance(size, int) assert isinstance(coverage_ratio, float) assert centering in T.get_args(CenteringType) self._size = size self._coverage = coverage_ratio self._centering: CenteringType = centering self._config = _MaskConfig() logger.debug("Initialized %s", self) def __repr__(self) -> str: """ Pretty print for logging """ params = f"coverage_ratio={repr(self._coverage)}, centering={repr(self._centering)}" return f"{self.__class__.__name__}({params})" def _check_mask_exists(self, filename: str, detected_face: DetectedFace) -> None: """ Check that the requested mask exists for the current detected face Parameters ---------- filename : str The file path for the current image detected_face : :class:`~lib.align.detected_face.DetectedFace` The detected face object that holds the masks Raises ------ FaceswapError If the requested mask type is not available an error is returned along with a list of available masks """ if self._config.mask_type in detected_face.mask: return exist_masks = list(detected_face.mask) msg = "No masks exist for this face" if exist_masks: msg = f"The masks that exist for this face are: {exist_masks}" raise FaceswapError( f"You have selected the mask type '{self._config.mask_type}' but at least one " "face does not contain the selected mask.\n" f"The face that failed was: '{filename}'\n{msg}") def _preprocess(self, detected_face: DetectedFace, mask_type: str) -> align.aligned_mask.Mask: """ Apply pre-processing to the mask Parameters ---------- detected_face : :class:`~lib.align.detected_face.DetectedFace` The detected face object that holds the masks mask_type : str The stored mask type to use Returns ------- :class:`~lib.align.aligned_mask.Mask` The pre-processed mask at its stored size and crop """ mask = detected_face.mask[mask_type] mask.set_dilation(self._config.dilation) mask.set_blur_and_threshold(blur_kernel=self._config.kernel, threshold=self._config.threshold) return mask def _crop_and_resize(self, detected_face: DetectedFace, mask: align.aligned_mask.Mask) -> np.ndarray: """ Crop and resize the mask to the correct centering and training size Parameters ---------- detected_face : :class:`~lib.align.detected_face.DetectedFace` The detected face object that holds the masks mask : :class:`~lib.align.aligned_mask.Mask` The pre-processed mask at its stored size and crop Returns ------- :class:`numpy.ndarray` The processed, cropped and resized final mask """ pose = detected_face.aligned.pose mask.set_sub_crop(pose.offset[mask.stored_centering], pose.offset[self._centering], self._centering, self._coverage, detected_face.aligned.y_offset) face_mask = mask.mask if self._size != face_mask.shape[0]: interpolator = cv2.INTER_CUBIC if mask.stored_size < self._size else cv2.INTER_AREA face_mask = cv2.resize(face_mask, (self._size, self._size), interpolation=interpolator)[..., None] return face_mask def _get_face_mask(self, filename: str, detected_face: DetectedFace) -> np.ndarray | None: """ Obtain the training sized face mask from the DetectedFace for the requested mask type. Parameters ---------- filename : str The file path for the current image detected_face : :class:`~lib.align.detected_face.DetectedFace` The detected face object that holds the masks Returns ------- :class:`numpy.ndarray` | None The face mask used for training or ``None`` if masks are disabled """ if not self._config.mask_enabled: return None assert self._config.mask_type is not None self._check_mask_exists(filename, detected_face) mask = self._preprocess(detected_face, self._config.mask_type) retval = self._crop_and_resize(detected_face, mask) logger.trace("Obtained face mask for: %s %s", # type:ignore[attr-defined] filename, retval.shape) return retval def _get_localized_mask(self, filename: str, detected_face: DetectedFace, area: T.Literal["eye", "mouth"]) -> np.ndarray | None: """ Obtain a localized mask for the given area if it is required for training. Parameters ---------- filename : str The file path for the current image detected_face : :class:`~lib.align.detected_face.DetectedFace` The detected face object that holds the masks area : Literal["eye", "mouth"] The area of the face to obtain the mask for Raises ------ :class:`~lib.utils.FaceswapError` If landmark data is not available to generate the localized mask """ if not self._config.multiplier_enabled: return None try: mask = detected_face.get_landmark_mask(area, self._size // 16, 2.5) except FaceswapError as err: logger.error(str(err)) raise FaceswapError("Eye/Mouth multiplier masks could not be generated due to missing " f"landmark data. The file that failed was: '{filename}'") from err logger.trace("Caching localized '%s' mask for: %s %s", # type:ignore[attr-defined] area, filename, mask.shape) return mask def __call__(self, filename: str, detected_face: DetectedFace) -> None: """ Prepare the masks required for training and compile into a single compressed array within the given DetectedFaces object Parameters ---------- filename : str The file path for the image that masks are to be prepared for detected_face : :class:`~lib.align.detected_face.DetectedFace` The detected face object that holds the masks """ masks = [(self._get_face_mask(filename, detected_face))] for area in T.get_args(T.Literal["eye", "mouth"]): masks.append(self._get_localized_mask(filename, detected_face, area)) detected_face.store_training_masks(masks, delete_masks=True) logger.trace("Stored masks for filename: %s)", filename) # type:ignore[attr-defined] def _check_reset(face_cache: "Cache") -> bool: """ Check whether a given cache needs to be reset because a face centering change has been detected in the other cache. Parameters ---------- face_cache : :class:`Cache` The cache object that is checking whether it should reset Returns ------- bool ``True`` if the given object should reset the cache, otherwise ``False`` """ check_cache = next((cache for cache in _FACE_CACHES.values() if cache != face_cache), None) retval = False if check_cache is None else check_cache.check_reset() return retval @dataclass class _CacheConfig: """ Holds the configuration options for the cache """ size: int """ int : The size to load images at """ centering: CenteringType """ Literal["face", "head", "legacy"] : The centering type to train at """ coverage: float """ float : The selected coverage ration for training """ class Cache(): """ A thread safe mechanism for collecting and holding face meta information (masks, alignments data etc.) for multiple :class:`~lib.training.generator.TrainingDataGenerator`. Each side may have up to 3 generators (training, preview and time-lapse). To conserve RAM these need to share access to the same face information for the images they are processing. As the cache is populated at run-time, thread safe writes are required for the first epoch. Following that, the cache is only used for reads, which is thread safe intrinsically. It would probably be quicker to set locks on each individual face, but for code complexity reasons, and the fact that the lock is only taken up during cache population, and it should only be being read multiple times on save iterations, we lock the whole cache during writes. Parameters ---------- filenames : list[str] The filenames of all the images. This can either be the full path or the base name. If the full paths are passed in, they are stripped to base name for use as the cache key. size : int The largest output size of the model coverage_ratio : float The coverage ratio that the model is using. """ def __init__(self, filenames: list[str], size: int, coverage_ratio: float) -> None: logger.debug(parse_class_init(locals())) self._lock = Lock() self._cache_info = {"cache_full": False, "has_reset": False} self._partially_loaded: list[str] = [] self._image_count = len(filenames) self._cache: dict[str, DetectedFace] = {} self._aligned_landmarks: dict[str, np.ndarray] = {} self._extract_version = 0.0 self._config = _CacheConfig(size=size, centering=T.cast(CenteringType, cfg.centering()), coverage=coverage_ratio) self._mask_prepare = _MaskProcessing(size, coverage_ratio, self._config.centering) logger.debug("Initialized: %s", self.__class__.__name__) @property def cache_full(self) -> bool: """ bool : ``True`` if the cache has been fully populated. ``False`` if there are items still to be cached. """ if self._cache_info["cache_full"]: return self._cache_info["cache_full"] with self._lock: return self._cache_info["cache_full"] @property def aligned_landmarks(self) -> dict[str, np.ndarray]: """ dict[str, :class:`numpy.ndarray`] : filename as key, aligned landmarks as value. """ # Note: Aligned landmarks are only used for warp-to-landmarks, so this can safely populate # all of the aligned landmarks for the entire cache. if not self._aligned_landmarks: with self._lock: # For Warp-To-Landmarks a race condition can occur where this is referenced from # the opposite side prior to it being populated, so block on a lock. self._aligned_landmarks = {key: face.aligned.landmarks for key, face in self._cache.items()} return self._aligned_landmarks @property def size(self) -> int: """ int : The pixel size of the cropped aligned face """ return self._config.size def get_items(self, filenames: list[str]) -> list[DetectedFace]: """ Obtain the cached items for a list of filenames. The returned list is in the same order as the provided filenames. Parameters ---------- filenames : list[str] A list of image filenames to obtain the cached data for Returns ------- list[:class:`~lib.align.detected_face.DetectedFace`] List of DetectedFace objects holding the cached metadata. The list returns in the same order as the filenames received """ return [self._cache[os.path.basename(filename)] for filename in filenames] def check_reset(self) -> bool: """ Check whether this cache has been reset due to a face centering change, and reset the flag if it has. Returns ------- bool ``True`` if the cache has been reset because of a face centering change due to legacy alignments, otherwise ``False``. """ retval = self._cache_info["has_reset"] if retval: logger.debug("Resetting 'has_reset' flag") self._cache_info["has_reset"] = False return retval def _reset_cache(self, set_flag: bool) -> None: """ In the event that a legacy extracted face has been seen, and centering is not legacy the cache will need to be reset for legacy centering. Parameters ---------- set_flag: bool ``True`` if the flag should be set to indicate that the cache is being reset because of a legacy face set/centering mismatch. ``False`` if the cache is being reset because it has detected a reset flag from the opposite cache. """ if set_flag: logger.warning("You are using legacy extracted faces but have selected '%s' centering " "which is incompatible. Switching centering to 'legacy'", self._config.centering) cfg.centering.set("legacy") self._config.centering = "legacy" self._cache = {} self._cache_info["cache_full"] = False if set_flag: self._cache_info["has_reset"] = True def _validate_version(self, png_meta: PNGHeaderDict, filename: str) -> None: """ Validate that there are not a mix of v1.0 extracted faces and v2.x faces. Parameters ---------- png_meta : :class:`~lib.align.alignments.PNGHeaderDict` The information held within the Faceswap PNG Header filename: str The full path to the file being validated Raises ------ :class:`~lib.utils.FaceswapError` If a version 1.0 face appears in a 2.x set or vice versa """ alignment_version = png_meta["source"]["alignments_version"] if not self._extract_version: logger.debug("Setting initial extract version: %s", alignment_version) self._extract_version = alignment_version if alignment_version == 1.0 and self._config.centering != "legacy": self._reset_cache(True) return if (self._extract_version == 1.0 and alignment_version > 1.0) or ( alignment_version == 1.0 and self._extract_version > 1.0): raise FaceswapError("Mixing legacy and full head extracted facesets is not supported. " "The following folder contains a mix of extracted face types: " f"'{os.path.dirname(filename)}'") self._extract_version = min(alignment_version, self._extract_version) def _load_detected_face(self, filename: str, alignments: PNGHeaderAlignmentsDict) -> DetectedFace: """ Load a :class:`~lib.align.detected_face.DetectedFace` object and load its associated `aligned` property. Parameters ---------- filename : str The file path for the current image alignments : :class:`~lib.align.alignments.PNGHeaderAlignmentsDict` The alignments for a single face, extracted from a PNG header Returns ------- :class:`~lib.align.detected_face.DetectedFace` The loaded Detected Face object """ y_offset = cfg.vertical_offset() detected_face = DetectedFace() detected_face.from_png_meta(alignments) detected_face.load_aligned(None, size=self._config.size, centering=self._config.centering, coverage_ratio=self._config.coverage, y_offset=y_offset / 100., is_aligned=True, is_legacy=self._extract_version == 1.0) logger.trace("Cached aligned face for: %s", filename) # type:ignore[attr-defined] return detected_face def _populate_cache(self, needs_cache: list[str], metadata: list[PNGHeaderDict], filenames: list[str]) -> None: """ Populate the given items into the cache Parameters ---------- needs_cache : list[str] The full path to files within this batch that require caching metadata : list[:class:`~lib.align.alignments.PNGHeaderDict`] The faceswap metadata loaded from the image png header filenames : list[str] Full path to the filenames that are being loaded in this batch """ for filename in needs_cache: key = os.path.basename(filename) meta = metadata[filenames.index(filename)] # Version Check self._validate_version(meta, filename) if self._partially_loaded: # Faces already loaded for Warp-to-landmarks self._partially_loaded.remove(key) detected_face = self._cache[key] else: detected_face = self._load_detected_face(filename, meta["alignments"]) self._mask_prepare(filename, detected_face) self._cache[key] = detected_face def _get_batch_with_metadata(self, filenames: list[str]) -> tuple[np.ndarray, list[PNGHeaderDict]]: """ Load a batch of images along with their faceswap metadata for loading into the cache Parameters ---------- filenames : list[str] Full path to the images to be loaded Returns ------- batch : :class:`numpy.ndarray` The batch of images in a single array metadata : :class:`~lib.align.alignments.PNGHeaderDict` The faceswap metadata corresponding to each image in the batch """ try: batch, metadata = read_image_batch(filenames, with_metadata=True) except ValueError as err: if "inhomogeneous" in str(err): raise FaceswapError( "There was an error loading a batch of images. This is most likely due to " "non-faceswap extracted faces in your training folder." "\nAll training images should be Faceswap extracted faces." "\nAll training images should be the same size." f"\nThe files that caused this error are: {filenames}") from err raise if len(batch.shape) == 1: folder = os.path.dirname(filenames[0]) keys = [os.path.basename(filename) for filename in filenames] details = [ f"{key} ({f'{img.shape[1]}px' if isinstance(img, np.ndarray) else type(img)})" for key, img in zip(keys, batch)] msg = (f"There are mismatched image sizes in the folder '{folder}'. All training " "images for each side must have the same dimensions.\nThe batch that " f"failed contains the following files:\n{details}.") raise FaceswapError(msg) return batch, metadata def _update_cache_full(self, filenames: list[str]) -> None: """ Check if cache is full and update the "cache_full" flag in :attr:`_cache_info` if so Parameters ---------- filenames : list[str] Full path to the filenames being processed in the current batch """ cache_full = not self._partially_loaded and len(self._cache) == self._image_count if cache_full: logger.verbose("Cache filled: '%s'", # type:ignore[attr-defined] os.path.dirname(filenames[0])) self._cache_info["cache_full"] = cache_full def cache_metadata(self, filenames: list[str]) -> np.ndarray: """ Obtain the batch with metadata for items that need caching and cache DetectedFace objects to :attr:`_cache`. Parameters ---------- filenames : list[str] List of full paths to image file names Returns ------- :class:`numpy.ndarray` The batch of face images loaded from disk """ keys = [os.path.basename(filename) for filename in filenames] with self._lock: if _check_reset(self): self._reset_cache(False) needs_cache = [filename for filename, key in zip(filenames, keys) if key not in self._cache or key in self._partially_loaded] logger.trace("Needs cache: %s", needs_cache) # type:ignore[attr-defined] if not needs_cache: # Metadata already cached. Just get images logger.debug("All metadata already cached for: %s", keys) return read_image_batch(filenames) batch, metadata = self._get_batch_with_metadata(filenames) self._populate_cache(needs_cache, metadata, filenames) self._update_cache_full(filenames) return batch def pre_fill(self, filenames: list[str], side: T.Literal["a", "b"]) -> None: """ When warp to landmarks is enabled, the cache must be pre-filled, as each side needs access to the other side's alignments. Parameters ---------- filenames : list[str] The list of full paths to the images to load the metadata from side : Literal["a", "b"] The side of the model being cached. Used for info output Raises ------ :class:`~lib.utils.FaceSwapError` If unsupported landmark type exists or a non-faceswap image is loaded """ with self._lock: for filename, meta in tqdm(read_image_meta_batch(filenames), desc=f"WTL: Caching Landmarks ({side.upper()})", total=len(filenames), leave=False): if "itxt" not in meta or "alignments" not in meta["itxt"]: raise FaceswapError(f"Invalid face image found. Aborting: '{filename}'") meta = meta["itxt"] key = os.path.basename(filename) self._validate_version(meta, filename) detected_face = self._load_detected_face(filename, meta["alignments"]) aligned = detected_face.aligned assert aligned is not None if aligned.landmark_type != LandmarkType.LM_2D_68: raise FaceswapError("68 Point facial Landmarks are required for Warp-to-" f"landmarks. The face that failed was: '{filename}'") self._cache[key] = detected_face self._partially_loaded.append(key) def get_cache(side: T.Literal["a", "b"], filenames: list[str] | None = None, size: int | None = None, coverage_ratio: float | None = None) -> Cache: """ Obtain a :class:`Cache` object for the given side. If the object does not pre-exist then create it. Parameters ---------- side : Literal["a", "b"] The side of the model to obtain the cache for filenames : list[str] | None, optional The filenames of all the images. This can either be the full path or the base name. If the full paths are passed in, they are stripped to base name for use as the cache key. Must be passed for the first call of this function for each side. For subsequent calls this parameter is ignored. Default: ``None`` size: int | None, optional The largest output size of the model. Must be passed for the first call of this function for each side. For subsequent calls this parameter is ignored. Default: ``None`` coverage_ratio : float | None, optional The coverage ratio that the model is using. Must be passed for the first call of this function for each side. For subsequent calls this parameter is ignored. Default: ``None`` Returns ------- :class:`Cache` The face meta information cache for the requested side """ assert side in ("a", "b") if not _FACE_CACHES.get(side): assert filenames is not None, "filenames must be provided for first call to cache" assert size is not None, "size must be provided for first call to cache" assert coverage_ratio is not None, ("coverage_ratio must be provided for first call to " "cache") logger.debug("Creating cache. side: %s, size: %s, coverage_ratio: %s", side, size, coverage_ratio) _FACE_CACHES[side] = Cache(filenames, size, coverage_ratio) return _FACE_CACHES[side] class RingBuffer(): """ Rolling buffer for holding training/preview batches Parameters ---------- batch_size : int The batch size to create the buffer for image_shape : tuple[int, int, int] The height/width/channels shape of a single image in the batch buffer_size : int, optional The number of arrays to hold in the rolling buffer. Default: `2` dtype : str, optional The datatype to create the buffer as. Default: `"uint8"` """ def __init__(self, batch_size: int, image_shape: tuple[int, int, int], buffer_size: int = 2, dtype: str = "uint8") -> None: logger.debug(parse_class_init(locals())) self._max_index = buffer_size - 1 self._index = 0 self._buffer = [np.empty((batch_size, *image_shape), dtype=dtype) for _ in range(buffer_size)] logger.debug("Initialized: %s", self) def __repr__(self) -> str: """ Pretty string representation for logging """ params = {"batch_size": repr(self._buffer[0].shape[0]), "image_shape": repr(self._buffer[0].shape[1:]), "buffer_size": repr(len(self._buffer)), "dtype": repr(str(self._buffer[0].dtype))} str_params = [f"{k}={v}" for k, v in params.items()] return f"{self.__class__.__name__}({', '.join(str_params)})" def __call__(self) -> np.ndarray: """ Obtain the next array from the ring buffer Returns ------- :class:`np.ndarray` A pre-allocated numpy array from the buffer """ retval = self._buffer[self._index] self._index += 1 if self._index < self._max_index else -self._max_index return retval __all__ = get_module_objects(__name__)