/
githubmirror
/
faceswap
Обзор
Документация
Войти
/
githubmirror
/
faceswap
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
lib/convert.py
528 строк
23 KB
torzdf
Train: Migrate data loading to Torch (#1540)
16 апр 2026, 21:19
Не верифицирован
16 апр 2026, 21:19
83fdd5f
Код
Авторство
О чём код?
#!/usr/bin/env python3 """Converter for Faceswap""" from __future__ import annotations import logging import typing as T from dataclasses import dataclass import cv2 import numpy as np from lib.align.aligned_mask import LandmarksMask from lib.utils import get_module_objects from plugins.plugin_loader import PluginLoader if T.TYPE_CHECKING: from argparse import Namespace from collections.abc import Callable from lib.align.aligned_face import AlignedFace, CenteringType from lib.align.detected_face import DetectedFace from lib.queue_manager import EventQueue from scripts.convert import ConvertItem from plugins.convert.color._base import Adjustment as ColorAdjust from plugins.convert.color.seamless_clone import Color as SeamlessAdjust from plugins.convert.mask.mask_blend import Mask as MaskAdjust from plugins.convert.scaling._base import Adjustment as ScalingAdjust logger = logging.getLogger(__name__) @dataclass class Adjustments: """Dataclass to hold the optional processing plugins Parameters ---------- color The selected color processing plugin. Default: `None` mask The selected mask processing plugin. Default: `None` seamless The selected mask processing plugin. Default: `None` sharpening The selected mask processing plugin. Default: `None` """ color: ColorAdjust | None = None mask: MaskAdjust | None = None seamless: SeamlessAdjust | None = None sharpening: ScalingAdjust | None = None class Converter(): # pylint:disable=too-many-instance-attributes """The converter is responsible for swapping the original face(s) in a frame with the output of a trained Faceswap model. Parameters ---------- output_size The size of the face, in pixels, that is output from the Faceswap model coverage_ratio The ratio of the training image that was used for training the Faceswap model centering The extracted face centering that the model was trained on (`"face"` or "`legacy`") draw_transparent Whether the final output should be drawn onto a transparent layer rather than the original frame. Only available with certain writer plugins. pre_encode Some writer plugins support the pre-encoding of images prior to saving out. As patching is done in multiple threads, but writing is done in a single thread, it can speed up the process to do any pre-encoding as part of the converter process. arguments The arguments that were passed to the convert process as generated from Faceswap's command line arguments config_file Optional location of custom configuration ``ini`` file. If ``None`` then use the default config location. Default: ``None`` """ def __init__(self, output_size: int, coverage_ratio: float, centering: CenteringType, draw_transparent: bool, pre_encode: Callable | None, arguments: Namespace, config_file: str | None = None) -> None: logger.debug("Initializing %s: (output_size: %s, coverage_ratio: %s, centering: %s, " "draw_transparent: %s, pre_encode: %s, arguments: %s, config_file: %s)", self.__class__.__name__, output_size, coverage_ratio, centering, draw_transparent, pre_encode, arguments, config_file) self._output_size = output_size self._coverage_ratio = coverage_ratio self._centering: CenteringType = centering self._draw_transparent = draw_transparent self._writer_pre_encode = pre_encode self._args = arguments self._config_file = config_file self._scale = arguments.output_scale / 100 self._face_scale = 1.0 - arguments.face_scale / 100. self._adjustments = Adjustments() self._full_frame_output: bool = arguments.writer != "patch" self._load_plugins() logger.debug("Initialized %s", self.__class__.__name__) @property def cli_arguments(self) -> Namespace: """The command line arguments passed to the convert process""" return self._args def reinitialize(self) -> None: """Reinitialize this :class:`Converter`. Called as part of the :mod:`~tools.preview` tool. Resets all adjustments then loads the plugins as specified in the current config. """ logger.debug("Reinitializing converter") self._face_scale = 1.0 - self._args.face_scale / 100. self._adjustments = Adjustments() self._load_plugins(disable_logging=True) logger.debug("Reinitialized converter") def _load_plugins(self, disable_logging: bool = False) -> None: """Load the requested adjustment plugins. Loads the :mod:`plugins.converter` plugins that have been requested for this conversion session. Parameters ---------- config Optional pre-loaded :class:`lib.config.FaceswapConfig`. If passed, then this will be used over any configuration on disk. If ``None`` then it is ignored. Default: ``None`` """ logger.debug("Loading plugins. disable_logging: %s", disable_logging) self._adjustments.mask = PluginLoader.get_converter("mask", "mask_blend", disable_logging=disable_logging)( self._args.mask_type, self._output_size, self._coverage_ratio, config_file=self._config_file) if self._args.color_adjustment is not None and self._args.color_adjustment != "none": self._adjustments.color = PluginLoader.get_converter("color", self._args.color_adjustment, disable_logging=disable_logging)( config_file=self._config_file) sharpening = PluginLoader.get_converter("scaling", "sharpen", disable_logging=disable_logging)( config_file=self._config_file) self._adjustments.sharpening = sharpening logger.debug("Loaded plugins: %s", self._adjustments) def process(self, in_queue: EventQueue, out_queue: EventQueue): """Main convert process. Takes items from the in queue, runs the relevant adjustments, patches faces to final frame and outputs patched frame to the out queue. Parameters ---------- in_queue The output from :class:`scripts.convert.Predictor`. Contains detected faces from the Faceswap model as well as the frame to be patched. out_queue The queue to place patched frames into for writing by one of Faceswap's :mod:`plugins.convert.writer` plugins. """ logger.debug("Starting convert process. (in_queue: %s, out_queue: %s)", in_queue, out_queue) logged = False while True: inbound: T.Literal["EOF"] | ConvertItem | list[ConvertItem] = in_queue.get() if inbound == "EOF": logger.debug("EOF Received") logger.debug("Patch queue finished") # Signal EOF to other processes in pool logger.debug("Putting EOF back to in_queue") in_queue.put(inbound) break items = inbound if isinstance(inbound, list) else [inbound] for item in items: logger.trace("Patch queue got: '%s'", # type: ignore[attr-defined] item.inbound.filename) try: image = self._patch_image(item) except Exception as err: # pylint:disable=broad-except # Log error and output original frame logger.error("Failed to convert image: '%s'. Reason: %s", item.inbound.filename, str(err)) image = item.inbound.image lvl = logger.trace if logged else logger.warning # type: ignore[attr-defined] lvl("Convert error traceback:", exc_info=True) logged = True # UNCOMMENT THIS CODE BLOCK TO PRINT TRACEBACK ERRORS # import sys; import traceback # exc_info = sys.exc_info(); traceback.print_exception(*exc_info) logger.trace("Out queue put: %s", # type: ignore[attr-defined] item.inbound.filename) out_queue.put((item.inbound.filename, image)) logger.debug("Completed convert process") def _get_warp_matrix(self, matrix: np.ndarray, size: int, y_offset: float = 0.0) -> np.ndarray: """Obtain the final scaled warp transformation matrix based on face scaling from the original transformation matrix Parameters ---------- matrix The transformation for patching the swapped face back onto the output frame size The size of the face patch, in pixels y_offset The amount of offset to apply on the y-axis. Default: 0.0 (no offset) Returns ------- The final transformation matrix with any scaling and y-offset applied """ mat = matrix.copy() if self._scale != 1.0 or y_offset else matrix if self._face_scale != 1.0: mat = matrix * self._face_scale patch_center = (size / 2, size / 2) mat[..., 2] += (1 - self._face_scale) * np.array(patch_center) if y_offset: mat[1, 2] += (y_offset * size) return mat def _patch_image(self, predicted: ConvertItem) -> np.ndarray | list[bytes]: """Patch a swapped face onto a frame. Run selected adjustments and swap the faces in a frame. Parameters ---------- predicted The output from :class:`scripts.convert.Predictor`. Returns ------- The final frame ready for writing by a :mod:`plugins.convert.writer` plugin. Frame is either an array, or the pre-encoded output from the writer's pre-encode function (if it has one) """ logger.trace("Patching image: '%s'", # type: ignore[attr-defined] predicted.inbound.filename) frame_size = (predicted.inbound.image.shape[1], predicted.inbound.image.shape[0]) new_image, background = self._get_new_image(predicted, frame_size) if self._full_frame_output: patched_face = self._post_warp_adjustments(background, new_image) patched_face = self._scale_image(patched_face) patched_face *= 255.0 patched_face = np.rint(patched_face, out=np.empty(patched_face.shape, dtype="uint8"), casting='unsafe') else: patched_face = new_image if self._writer_pre_encode is None: retval: np.ndarray | list[bytes] = patched_face else: kwargs: dict[str, T.Any] = {} if self.cli_arguments.writer == "patch": kwargs["canvas_size"] = (background.shape[1], background.shape[0]) kwargs["matrices"] = np.array([self._get_warp_matrix(face.adjusted_matrix, patched_face.shape[1], face.y_offset) for face in predicted.reference_faces], dtype="float32") retval = self._writer_pre_encode(patched_face, **kwargs) logger.trace("Patched image: '%s'", # type: ignore[attr-defined] predicted.inbound.filename) return retval def _warp_to_frame(self, reference: AlignedFace, face: np.ndarray, frame: np.ndarray) -> None: """Perform affine transformation to place a face patch onto the given frame. Affine is done in place on the `frame` array, so this function does not return a value Parameters ---------- reference The object holding the original aligned face face The swapped face patch frame The frame to affine the face onto """ # Warp face with the mask mat = self._get_warp_matrix(reference.adjusted_matrix, face.shape[0], reference.y_offset) frame_face = np.zeros_like(frame) cv2.warpAffine(face, mat, (frame.shape[1], frame.shape[0]), frame_face, flags=cv2.WARP_INVERSE_MAP | reference.interpolators[1], borderMode=cv2.BORDER_CONSTANT) background = frame[..., :3] alpha = frame[..., 3:4] foreground, mask = np.split(frame_face, # pylint:disable=unbalanced-tuple-unpacking (3, ), axis=-1) background *= (1.0 - mask) background += (foreground * mask) alpha += mask * (1.0 - alpha) # Merge masks def _get_new_image(self, predicted: ConvertItem, frame_size: tuple[int, int]) -> tuple[np.ndarray, np.ndarray]: """Get the new face from the predictor and apply pre-warp manipulations. Applies any requested adjustments to the raw output of the Faceswap model before transforming the image into the target frame. Parameters ---------- predicted The output from :class:`scripts.convert.Predictor`. frame_size The (`width`, `height`) of the final frame in pixels Returns ------- placeholder The original frame with the swapped faces patched onto it background The original frame """ logger.trace("Getting: (filename: '%s', faces: %s)", # type: ignore[attr-defined] predicted.inbound.filename, len(predicted.swapped_faces)) placeholder: np.ndarray = np.zeros((frame_size[1], frame_size[0], 4), dtype="float32") faces: list[np.ndarray] | None = None if self._full_frame_output: background = predicted.inbound.image / np.array(255.0, dtype="float32") placeholder[:, :, :3] = background else: faces = [] # Collect the faces into final array background = placeholder # Used for obtaining original frame dimensions for new_face, detected_face, reference_face in zip(predicted.swapped_faces, predicted.inbound.detected_faces, predicted.reference_faces): predicted_mask = new_face[:, :, -1] if new_face.shape[2] == 4 else None new_face = new_face[:, :, :3] new_face = self._pre_warp_adjustments(new_face, detected_face, reference_face, predicted_mask) if self._full_frame_output: self._warp_to_frame(reference_face, new_face, placeholder,) else: assert faces is not None faces.append(new_face) if not self._full_frame_output: placeholder = np.array(faces, dtype="float32") logger.trace("Got filename: '%s'. (placeholders: %s)", # type: ignore[attr-defined] predicted.inbound.filename, placeholder.shape) return placeholder, background def _pre_warp_adjustments(self, new_face: np.ndarray, detected_face: DetectedFace, reference_face: AlignedFace, predicted_mask: np.ndarray | None) -> np.ndarray: """Run any requested adjustments that can be performed on the raw output from the Faceswap model. Any adjustments that can be performed before warping the face into the final frame are performed here. Parameters ---------- new_face The swapped face received from the faceswap model. detected_face The detected_face object as defined in :class:`scripts.convert.Predictor` reference_face The aligned face object sized to the model output of the original face for reference predicted_mask The predicted mask output from the Faceswap model. ``None`` if the model did not learn a mask Returns ------- The face output from the Faceswap Model with any requested pre-warp adjustments performed. """ logger.trace("new_face shape: %s, predicted_mask shape: %s", # type: ignore[attr-defined] new_face.shape, predicted_mask.shape if predicted_mask is not None else None) old_face = T.cast(np.ndarray, reference_face.face)[..., :3] / 255.0 new_face, raw_mask = self._get_image_mask(new_face, detected_face, predicted_mask, reference_face) if self._adjustments.color is not None: new_face = self._adjustments.color.run(old_face, new_face, raw_mask) if self._adjustments.seamless is not None: new_face = self._adjustments.seamless.run(old_face, new_face, raw_mask) logger.trace("returning: new_face shape %s", new_face.shape) # type: ignore[attr-defined] return new_face def _get_image_mask(self, new_face: np.ndarray, detected_face: DetectedFace, predicted_mask: np.ndarray | None, reference_face: AlignedFace) -> tuple[np.ndarray, np.ndarray]: """Return any selected image mask Places the requested mask into the new face's Alpha channel. Parameters ---------- new_face The swapped face received from the faceswap model. detected_face The detected_face object as defined in :class:`scripts.convert.Predictor` predicted_mask The predicted mask output from the Faceswap model. ``None`` if the model did not learn a mask reference_face The aligned face object sized to the model output of the original face for reference Returns ------- swapped_face The swapped face with the requested mask added to the Alpha channel raw_mask The raw mask with no erosion or blurring applied """ logger.trace("Getting mask. Image shape: %s", new_face.shape) # type: ignore[attr-defined] mask_centering: CenteringType lm_mask = None if self._args.mask_type in ("components", "extended"): mask_centering = reference_face.centering m_type: T.Literal["face", "face_extended"] = ( "face" if self._args.mask_type == "components" else "face_extended" ) lm_mask = LandmarksMask(m_type, reference_face.landmark_type, reference_face.landmarks, reference_face.size) elif self._args.mask_type not in ("none", "predicted"): mask_centering = detected_face.mask[self._args.mask_type].stored_centering else: mask_centering = "face" # Unused but requires a valid value assert self._adjustments.mask is not None mask, raw_mask = self._adjustments.mask.run(detected_face, reference_face.pose.offset[mask_centering], reference_face.pose.offset[self._centering], self._centering, landmarks_mask=lm_mask, predicted_mask=predicted_mask) logger.trace("Adding mask to alpha channel") # type: ignore[attr-defined] new_face = np.concatenate((new_face, mask), -1) logger.trace("Got mask. Image shape: %s", new_face.shape) # type: ignore[attr-defined] return new_face, raw_mask def _post_warp_adjustments(self, background: np.ndarray, new_image: np.ndarray) -> np.ndarray: """Perform any requested adjustments to the swapped faces after they have been transformed into the final frame. Parameters ---------- background The original frame new_image A blank frame of original frame size with the faces warped onto it Returns ------- The final merged and swapped frame with any requested post-warp adjustments applied """ if self._adjustments.sharpening is not None: new_image = self._adjustments.sharpening.run(new_image) if self._draw_transparent: frame = new_image else: # This next code is kinda redundant, but if sharpening is performed it is needed foreground, mask = np.split(new_image, # pylint:disable=unbalanced-tuple-unpacking (3, ), axis=-1) foreground *= mask background *= (1.0 - mask) background += foreground frame = background np.clip(frame, 0.0, 1.0, out=frame) return frame def _scale_image(self, frame: np.ndarray) -> np.ndarray: """Scale the final image if requested. If output scale has been requested in command line arguments, scale the output otherwise return the final frame. Parameters ---------- frame The final frame with faces swapped Returns ------- The final frame scaled by the requested scaling factor """ if self._scale == 1: return frame logger.trace("source frame: %s", frame.shape) # type: ignore[attr-defined] interpolation = cv2.INTER_CUBIC if self._scale > 1 else cv2.INTER_AREA dims = (round((frame.shape[1] / 2 * self._scale) * 2), round((frame.shape[0] / 2 * self._scale) * 2)) frame = cv2.resize(frame, dims, interpolation=interpolation) logger.trace("resized frame: %s", frame.shape) # type: ignore[attr-defined] np.clip(frame, 0.0, 1.0, out=frame) return frame __all__ = get_module_objects(__name__)