#!/usr/bin/env python3 """ Handles loading of faces/frames from source locations and pairing with alignments information """ from __future__ import annotations import logging import os import typing as T import numpy as np from tqdm import tqdm from lib.align import alignments, DetectedFace, update_legacy_png_header from lib.image import FacesLoader, ImagesLoader from lib.utils import get_module_objects from plugins.extract import ExtractMedia if T.TYPE_CHECKING: from lib.align.alignments import PNGHeaderDict logger = logging.getLogger(__name__) class Loader: """ Loader for reading source data from disk, and yielding the output paired with alignment information Parameters ---------- location: str Full path to the source files location is_faces: bool ``True`` if the source is a folder of faceswap extracted faces """ def __init__(self, location: str, is_faces: bool) -> None: logger.debug("Initializing %s (location: %s, is_faces: %s)", self.__class__.__name__, location, is_faces) self._is_faces = is_faces self._loader = FacesLoader(location) if is_faces else ImagesLoader(location) self._alignments: alignments.Alignments | None = None self._skip_count = 0 logger.debug("Initialized %s", self.__class__.__name__) @property def file_list(self) -> list[str]: """list[str]: Full file list of source files to be loaded """ return self._loader.file_list @property def is_video(self) -> bool: """bool: ``True`` if the source is a video file otherwise ``False`` """ return self._loader.is_video @property def location(self) -> str: """str: Full path to the source folder/video file location """ return self._loader.location @property def skip_count(self) -> int: """int: The number of faces/frames that have been skipped due to no match in alignments file """ return self._skip_count def add_alignments(self, alignments_object: alignments.Alignments | None) -> None: """ Add the loaded alignments to :attr:`_alignments` for content matching Parameters ---------- alignments_object: :class:`~lib.align.Alignments` | None The alignments file object or ``None`` if not provided """ logger.debug("Adding alignments to loader: %s", alignments_object) self._alignments = alignments_object @classmethod def _get_detected_face(cls, alignment: alignments.AlignmentFileDict) -> DetectedFace: """ Convert an alignment dict item to a detected_face object Parameters ---------- alignment: :class:`lib.align.alignments.AlignmentFileDict` The alignment dict for a face Returns ------- :class:`~lib.align.detected_face.DetectedFace`: The corresponding detected_face object for the alignment """ detected_face = DetectedFace() detected_face.from_alignment(alignment) return detected_face def _process_face(self, filename: str, image: np.ndarray, metadata: PNGHeaderDict) -> ExtractMedia | None: """ Process a single face when masking from face images Parameters ---------- filename: str the filename currently being processed image: :class:`numpy.ndarray` The current face being processed metadata: dict The source frame metadata from the PNG header Returns ------- :class:`plugins.pipeline.ExtractMedia` | None the extract media object for the processed face or ``None`` if alignment information could not be found """ frame_name = metadata["source"]["source_filename"] face_index = metadata["source"]["face_index"] if self._alignments is None: # mask from PNG header lookup_index = 0 aligns = [T.cast(alignments.AlignmentFileDict, metadata["alignments"])] else: # mask from Alignments file lookup_index = face_index aligns = self._alignments.get_faces_in_frame(frame_name) if not aligns or face_index > len(aligns) - 1: self._skip_count += 1 logger.warning("Skipping Face not found in alignments file: '%s'", filename) return None alignment = aligns[lookup_index] detected_face = self._get_detected_face(alignment) retval = ExtractMedia(filename, image, detected_faces=[detected_face], is_aligned=True) retval.add_frame_metadata(metadata["source"]) return retval def _from_faces(self) -> T.Generator[ExtractMedia, None, None]: """ Load content from pre-aligned faces and pair with corresponding metadata Yields ------ :class:`plugins.pipeline.ExtractMedia` the extract media object for the processed face """ log_once = False for filename, image, metadata in tqdm(self._loader.load(), total=self._loader.count): if not metadata: # Legacy faces. Update the headers if self._alignments is None: logger.error("Legacy faces have been discovered, but no alignments file " "provided. You must provide an alignments file for this face set") break if not log_once: logger.warning("Legacy faces discovered. These faces will be updated") log_once = True metadata = update_legacy_png_header(filename, self._alignments) if not metadata: # Face not found self._skip_count += 1 logger.warning("Legacy face not found in alignments file. This face has not " "been updated: '%s'", filename) continue if "source_frame_dims" not in metadata.get("source", {}): logger.error("The faces need to be re-extracted as at least some of them do not " "contain information required to correctly generate masks.") logger.error("You can re-extract the face-set by using the Alignments Tool's " "Extract job.") break retval = self._process_face(filename, image, metadata) if retval is None: continue yield retval def _from_frames(self) -> T.Generator[ExtractMedia, None, None]: """ Load content from frames and and pair with corresponding metadata Yields ------ :class:`plugins.pipeline.ExtractMedia` the extract media object for the processed face """ assert self._alignments is not None for filename, image in tqdm(self._loader.load(), total=self._loader.count): frame = os.path.basename(filename) if not self._alignments.frame_exists(frame): self._skip_count += 1 logger.warning("Skipping frame not in alignments file: '%s'", frame) continue if not self._alignments.frame_has_faces(frame): logger.debug("Skipping frame with no faces: '%s'", frame) continue faces_in_frame = self._alignments.get_faces_in_frame(frame) detected_faces = [self._get_detected_face(alignment) for alignment in faces_in_frame] retval = ExtractMedia(filename, image, detected_faces=detected_faces) yield retval def load(self) -> T.Generator[ExtractMedia, None, None]: """ Load content from source and pair with corresponding alignment data Yields ------ :class:`plugins.pipeline.ExtractMedia` the extract media object for the processed face """ if self._is_faces: iterator = self._from_faces else: iterator = self._from_frames yield from iterator() if self._skip_count > 0: logger.warning("%s face(s) skipped due to not existing in the alignments file", self._skip_count) __all__ = get_module_objects(__name__)