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/env/bin/python3
  """Handles face detection plugins and runners """
  from __future__ import annotations
  
  import logging
  import typing as T
  
  import cv2
  import numpy as np
  
  
  from lib.align.aligned_utils import batch_create_matrices
  from lib.logger import format_array, parse_class_init
  from lib.utils import get_module_objects
  
  from .objects import ExtractBatch
  from .handler import ExtractHandler
  
  if T.TYPE_CHECKING:
      import numpy.typing as npt
  
  logger = logging.getLogger(__name__)
  
  
  class Detect(ExtractHandler):
      """Responsible for handling Detection plugins within the extract pipeline
  
      Parameters
      ----------
      plugin
          The plugin that this runner is to use
      rotation | None
          The rotation arguments. Either a list of angles between 0 and 360 to rotate at or a single
          step size. Default: ``None``, no rotations
      min_size
          Minimum percentage of the frame's shortest edge to accept as a successful detection along
          the detection's longest edge Default: `0` (accept all detections)
      max_size
          Maximum percentage of the frame's shortest edge to accept as a successful detection along
          the detection's longest edge Default: `0` (accept all detections)
      compile_model
          ``True`` to compile any PyTorch models
      config_file
          Full path to a custom config file to load. ``None`` for default config
      """
      def __init__(self,
                   plugin: str,
                   rotation: str | None = None,
                   min_size: int = 0,
                   max_size: int = 0,
                   compile_model: bool = False,
                   config_file: str | None = None) -> None:
          logger.debug(parse_class_init(locals()))
          super().__init__(plugin, compile_model=compile_model, config_file=config_file)
          self._rotation = rotation
          self._rotator = Rotator(rotation, self.plugin.input_size)
          """Responsible for rotating feed images for the model"""
          self._empty_bbox = np.empty((0, 4), dtype="float32")
          """An empty detection result, that will never be used so only needs to be created once"""
          self._min_size = min_size / 100.
          """The user selected shortest frame dim multiplier to accept for minimum size"""
          self._max_size = max_size / 100.
          """The user selected shortest frame dim multiplier to accept for maximum size"""
          self._filter_counts = 0
  
      def __repr__(self) -> str:
          """Pretty print for logging"""
          retval = super().__repr__()[:-1]
          retval += (f", rotation={repr(self._rotation)}, min_size={int(self._min_size * 100)}, "
                     f"max_size={int(self._max_size * 100)})")
          return retval
  
      # Pre-processing
      def _get_matrices(self,
                        images: list[npt.NDArray[np.uint8]],
                        filenames: list[str]) -> npt.NDArray[np.float32]:
          """Calculate the scales and padding required to take each image in this batch to model
          input size and store the matrices in the batch object
  
          Parameters
          ----------
          images
              The images to obtain the matrices for
          filenames : list[str]
              The corresponding file names of the images
  
          Returns
          -------
          The transformation matrices for taking the images to model input size
          """
          orig_wh = np.array([x.shape[:2] for x in images])[:, ::-1]
          scales = self.plugin.input_size / orig_wh.max(axis=1)
          new_wh = np.rint(orig_wh * scales[:, None]).astype(np.int32)
          pad_xy = (self.plugin.input_size - new_wh) // 2
  
          retval = np.zeros((len(scales), 3, 3), dtype="float32")
          retval[:, 0, 0] = scales
          retval[:, 1, 1] = scales
          retval[:, 0, 2] = pad_xy[:, 0]
          retval[:, 1, 2] = pad_xy[:, 1]
          retval[:, 2, 2] = 1.
  
          logger.trace(  # type:ignore[attr-defined]
              "[%s_pre_process] filenames: %s, matrices: %s",
              self.plugin.name, filenames, format_array(retval))
          return retval
  
      def _scale_images(self,
                        images: list[npt.NDArray[np.uint8]],
                        matrices: npt.NDArray[np.float32]) -> npt.NDArray[np.uint8]:
          """Scale the image and pad to given size
  
          Parameters
          ----------
          images
              The images to scale
          matrices
              The corresponding warp matrices for scaling the images
  
          Returns
          -------
          The scaled images
          """
          retval = np.zeros((len(images), self.plugin.input_size, self.plugin.input_size, 3),
                            dtype=images[0].dtype)
          interpolators = np.where(matrices[:, 0, 0] < 1.0, cv2.INTER_AREA, cv2.INTER_CUBIC)
          dims = (self.plugin.input_size, self.plugin.input_size)
          warp_mats = matrices[:, :2]
          for idx, (image, mat, interpolator) in enumerate(zip(images, warp_mats, interpolators)):
              image = image[..., 2::-1] if self.plugin.is_rgb else image
              cv2.warpAffine(image, mat, dims, dst=retval[idx], flags=interpolator)
          logger.trace("Resized batch shape: %s", retval.shape)  # type:ignore[attr-defined]
          return retval
  
      def pre_process(self, batch: ExtractBatch) -> None:
          """Perform pre-processing for detection plugins.
  
          - Gets the scale and padding to take the batch of images to model input size
          - Formats the image to the correct color order, dtype and scale for the plugin
          - Performs any plugin specific pre-processing
  
          Parameters
          ----------
          batch
              The incoming ExtractBatch to use for pre-processing
          """
          batch.matrices = self._get_matrices(batch.images, batch.filenames)
          images = self._scale_images(batch.images, batch.matrices)
          images = self._format_images(images)
          batch.data = self.plugin.pre_process(images)
  
      # Processing
      def _process_rotations(self,
                             predictions: npt.NDArray[np.float32],
                             mask_requires: npt.NDArray[np.bool_],
                             indices_angle: npt.NDArray[np.int32],
                             box_list: list[npt.NDArray[np.float32] | None],
                             rotation_index: int) -> None:
          """Process the output after a rotation, and store the discovered boxes and the angle index
          they were discovered at
  
          Parameters
          ----------
          predictions
              The predictions from the model
          mask_requires
              The mask indicating which frames can still be allocated bounding boxes
          indices_angle
              The array that stores the angle index that each frame's faces was found at
          box_list
              The list of final bounding boxes to be output
          rotation_index
              The current angle index we are iterating
          """
          bboxes = (self.plugin.post_process(predictions) if self._overridden["post_process"]
                    else predictions)
          mask_found = np.array([np.any(n) for n in bboxes], dtype="bool")
          indices_requires = np.flatnonzero(mask_requires)
          indices_angle[indices_requires[mask_found]] = rotation_index
          mask_requires[indices_requires[mask_found]] = False
          for i, box in zip(indices_requires[mask_found], bboxes):
              box_list[i] = box
  
      def process(self, batch: ExtractBatch) -> None:
          """Obtain the output from the plugin's model.
  
          Executes the plugin's predict function and stores the output prior to post-processing.
  
          If rotations have been selected, plugin post-processing is done as part of this process as
          the computed bounding boxes are required for re-feeding the model future rotations
  
          Parameters
          ----------
          batch
              The incoming ExtractBatch to use for processing
          """
          process = "process"
          input_images = batch.data
          batch_size = input_images.shape[0]
          box_list: list[None | np.ndarray] = [None for _ in range(batch_size)]
          boxes: np.ndarray | None = None
          indices_angle = np.zeros((batch_size, ), dtype="int32")
  
          idx = 0
          mask_requires = np.array([True for _ in range(batch_size)])
          while True:
              feed = self._rotator.rotate(idx, input_images[mask_requires])
              if feed is None:
                  logger.trace(  # type:ignore[attr-defined]
                      "[%s.%s] No faces found in %s image(s) of %s after %s rotations: %s",
                      self.plugin.name,
                      process,
                      mask_requires.sum(),
                      batch_size,
                      idx,
                      batch.filenames)
  
                  break
              result = self._predict(feed)
              if not self._rotator.enabled:
                  # Not rotating. Do post-processing in next thread
                  boxes = result
                  break
  
              # We are rotating, so we have to do post-processing here, to re-feed model
              self._process_rotations(result, mask_requires, indices_angle, box_list, idx)
              if not np.any(mask_requires):
                  logger.trace(  # type:ignore[attr-defined]
                      "[%s.%s] Found faces for all %s images after %s rotations: %s",
                      self.plugin.name,
                      process,
                      batch_size,
                      idx + 1,
                      batch.filenames)
                  break
              idx += 1
  
          boxes = (np.array([self._empty_bbox if b is None else b for b in box_list],
                            dtype="object")
                   if boxes is None else boxes)
          batch.data = np.empty(2, dtype="object")
          batch.data[0] = indices_angle
          batch.data[1] = boxes
  
      # Post-Processing
      def _stack_boxes(self,
                       batch: ExtractBatch,
                       predictions: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
          """Stack the detected boxes into a single array, remove any zero sized boxes, collate the
          indexing information and add to batch
  
          Parameters
          ----------
          batch
              The detector batch being processed
          predictions
              The face detection bounding boxes received from the plugin
  
          Returns
          -------
          The stacked detection boxes from all frames in the batch.
          """
          valid = np.fromiter((i for i, p in enumerate(predictions) if np.any(p)), dtype=np.int32)
          if not valid.size:
              batch.frame_ids = valid
              return self._empty_bbox
  
          result = [predictions[i] for i in valid]
          lengths = np.fromiter((a.shape[0] for a in result), dtype=np.int32)
          batch.frame_ids = np.repeat(valid, lengths)
          return np.vstack(result).astype(np.float32)
  
      def _scale_boxes(self, batch: ExtractBatch, predictions: npt.NDArray[np.float32]) -> None:
          """Scale the detected faces back out to original image size, round to int and add to the
          batch object
  
          Parameters
          ----------
          batch
              The detector batch being processed
          predictions
              The stacked face detection predictions at model input size
          """
          if not batch.frame_ids.size:
              return
          mats = batch.matrices[batch.frame_ids]
  
          predictions[:, [0, 2]] -= mats[:, 0, 2][:, None]
          predictions[:, [1, 3]] -= mats[:, 1, 2][:, None]
          predictions /= mats[:, 0, 0][:, None]
          np.rint(predictions, out=predictions)
          batch.bboxes = predictions.astype("int32")
          logger.trace("[%s.out] Finalized batch: %s",  # type:ignore[attr-defined]
                       self.plugin.name,
                       batch)
  
      def _filter_boxes(self, batch: ExtractBatch) -> None:
          """Filter out any detections that are smaller or larger than :attr:`_min_size` and
          :attr:`_max_size` along their longest edge
  
          Parameters
          ----------
          batch
              The detector batch being processed with fully scaled bounding boxes
          """
          if not self._min_size and not self._max_size:
              return
  
          frames = np.array([i.shape[:2] for i in batch.images]).min(axis=1)
          sizes = np.maximum(batch.bboxes[:, 2] - batch.bboxes[:, 0],
                             batch.bboxes[:, 3] - batch.bboxes[:, 1])
  
          mins = (frames * self._min_size).astype("int32")[batch.frame_ids]
          if self._max_size:
              maxes = (frames * self._max_size).astype("int32")[batch.frame_ids]
          else:
              maxes = sizes
  
          keep = np.nonzero(np.logical_and(mins <= sizes, maxes >= sizes))[0]
          if len(keep) == len(sizes):
              return
  
          logger.debug(
              "[%s.out] Removing %s face(s) from %s detections as outside size thresholds (min: %s, "
              "max: %s): %s",
              self.plugin.name,
              len(sizes) - len(keep),
              len(sizes),
              int(self._min_size * 100),
              int(self._max_size * 100),
              batch.filenames)
  
          batch.bboxes = batch.bboxes[keep]
          batch.frame_ids = batch.frame_ids[keep]
          self._filter_counts += len(sizes) - len(keep)
  
      def post_process(self, batch: ExtractBatch) -> None:
          """Perform detection post processing.
  
          If no rotations were requested, any plugin post-processing will be done here.
  
          Detection boxes are:
            - stacked into a single array
            - scaled back to frame dimensions,
            - filtered for faces which fall outside min/max thresholds
            - Added to the batch object along with frame to face mapping information.
  
          Parameters
          ----------
          batch
              The incoming ExtractBatch to use for post-processing
          """
          indices_angle, result = batch.data
          if self._overridden["post_process"] and not self._rotator.enabled:
              result = self.plugin.post_process(result)
          else:
              self._rotator.un_rotate(indices_angle, result)
          result = self._stack_boxes(batch, result)
          self._scale_boxes(batch, result)
          self._filter_boxes(batch)
  
      def output_info(self) -> None:
          """Output the counts of filtered items """
          if not self._filter_counts:
              return
          logger.info("[Detect filter] Scale (min: %s, max: %s): %s",
                      f"{int(self._min_size * 100)}%",
                      f"{int(self._max_size * 100)}%",
                      self._filter_counts)
  
  
  class Rotator:
      """Handles pre-calculation of rotation matrices when rotation angles are requested and
      rotating images for feeding the detector. Handles reversing the rotation for any found
      detection bounding boxes.
  
      Parameters
      ----------
      rotation
          List of requested rotation angles in degrees provided in command line arguments
      image_size
          The size of the square image to obtain rotation matrices for
      """
      def __init__(self, angles: str | None, image_size: int) -> None:
          logger.debug(parse_class_init(locals()))
          self._size = image_size
          self._angles = self._get_angles(angles)
          self._matrices = batch_create_matrices(self._size, rotation=self._angles)
          self._matrices_inverse = self._pre_compute_inverse_matrices()
          self._channels_first: bool | None = None
          self.enabled = len(self._angles) > 1
          """``True`` if rotations are to be performed """
  
      @classmethod
      def _angles_from_step(cls, step_size: int) -> npt.NDArray[np.float32]:
          """Obtain the required rotation angles when the cli argument has been passed in as a step
          size
  
          Parameters
          ----------
          step_size
              The requested step size
  
          Returns
          -------
          The rotation angles between 0 and 360 for the given step size
          """
          retval = np.arange(0, 360, step_size, dtype="float32")
          logger.debug("Setting rotation angles to %s from step size: %s", retval, step_size)
          return retval
  
      def _get_angles(self, rotation: str | None) -> npt.NDArray[np.float32]:
          """Set the rotation angles.
  
          Parameters
          ----------
          rotation
              List of requested rotation angles in degrees provided in command line arguments
  
          Returns
          -------
          The complete list of rotation angles to apply in degrees
          """
          if not rotation:
              logger.debug("Not setting rotation angles")
              return np.array([0], dtype=np.float32)
  
          passed_angles = [int(angle) for angle in rotation.split(",") if int(angle) != 0]
          if len(passed_angles) == 1:
              return self._angles_from_step(passed_angles[0])
  
          retval = np.array([0] + passed_angles, dtype=np.float32)
          logger.debug("Setting rotation angles to %s from given: %s", retval, rotation)
          return retval
  
      def _pre_compute_inverse_matrices(self) -> npt.NDArray[np.float32]:
          """Pre-compute the inverse rotation matrices required to perform translation from rotated
          bounding boxes back to original frame
  
          Returns
          -------
          The rotation matrices for the requested rotation angles
          """
          rot = self._matrices[:, :, :2]
          trans = self._matrices[:, :, 2]
          rot_inv = np.transpose(rot, (0, 2, 1))
          trans_inv = -np.einsum('nij, nj->ni', rot_inv, trans)
          retval = np.concatenate([rot_inv, trans_inv[..., None]], axis=2)
          logger.debug("Precomputed inverse rotation matrices: %s", retval.tolist())
          return retval
  
      def rotate(self, rotation_index: int, images: np.ndarray) -> np.ndarray | None:
          """Rotate a batch of images by the matrix provided by the given rotation index. Attempts
          to detect and handle channels first images as well as channels last
  
          Parameters
          ----------
          rotation_index
              The matrix to use. This will be an incrementing index from an enumerated loop that
              selects through the matrices stored for each angle
          images
              The original, correctly orientated, batch of images to rotate
  
          Returns
          -------
          The batch of image rotated by the angle identified by the given rotation index.
          ``None`` if the given rotation index is invalid
          """
          if rotation_index == 0:
              return images
          if rotation_index >= len(self._angles):
              return None
  
          if self._channels_first is None:
              self._channels_first = images.shape[1] in (1, 3, 4)
              logger.debug("Set channels_first to %s", self._channels_first)
  
          if self._channels_first:
              images = images.transpose(0, 2, 3, 1)
  
          retval = np.empty(images.shape, images.dtype)
          mat = self._matrices[rotation_index]
          size = (self._size, self._size)
  
          for i, img in enumerate(images):
              cv2.warpAffine(img,
                             mat,
                             size,
                             dst=retval[i],
                             borderMode=cv2.BORDER_REPLICATE)
  
          if self._channels_first:
              retval = retval.transpose(0, 3, 1, 2)
  
          return retval
  
      def un_rotate(self,
                    indices_angle: npt.NDArray[np.int32],
                    roi: npt.NDArray[np.float32]) -> None:
          """Un-rotate the given bounding boxes for the given angle indices and update in place
  
          Parameters
          ----------
          indices_angle
              The angle indices that correlate to the angle each roi was rotated to to obtain the
              result
          roi
              Ragged array of (B, N, 4) detected bounding discovered at the corresponding angle
              index
          """
          mask_needs_rotate = indices_angle > 0
          if not np.any(mask_needs_rotate):
              return
  
          indices_needs_rotate = np.flatnonzero(mask_needs_rotate)
          matrices = self._matrices_inverse[indices_angle[mask_needs_rotate]]
  
          for pred_idx, mat in zip(indices_needs_rotate, matrices):
              bboxes = roi[pred_idx]
              pts = np.empty((bboxes.shape[0], 4, 2), dtype="float32")
              pts[:, 0] = bboxes[:, [0, 1]]  # lt
              pts[:, 1] = bboxes[:, [2, 1]]  # rt
              pts[:, 2] = bboxes[:, [2, 3]]  # rb
              pts[:, 3] = bboxes[:, [0, 3]]  # lb
  
              pts = pts @ mat[:, :2].T + mat[:, 2]
  
              # boxes must align on (x, y) planes
              bboxes[:, 0] = pts[..., 0].min(axis=1)
              bboxes[:, 1] = pts[..., 1].min(axis=1)
              bboxes[:, 2] = pts[..., 0].max(axis=1)
              bboxes[:, 3] = pts[..., 1].max(axis=1)
  
  
  __all__ = get_module_objects(__name__)
  