repo_description: Deepfakes Software For All
repo_topics: deep-face-swap deep-learning deep-neural-networks deepface deepfakes deeplearning face-swap faceswap fakeapp machine-learning myfakeapp neural-nets neural-networks openfaceswap
code: #!/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__)
  