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
  """ Common multi-backend Keras utilities """
  from __future__ import annotations
  import typing as T
  
  import numpy as np
  
  from keras import ops, Variable
  
  from lib.utils import get_module_objects
  
  if T.TYPE_CHECKING:
      from keras import KerasTensor
  
  # TODO these can probably be switched to pure pytorch
  
  
  def frobenius_norm(matrix: KerasTensor,
                     axis: int = -1,
                     keep_dims: bool = True,
                     epsilon: float = 1e-15) -> KerasTensor:
      """ Frobenius normalization for Keras Tensor
  
      Parameters
      ----------
      matrix: :class:`keras.KerasTensor`
          The matrix to normalize
      axis: int, optional
          The axis to normalize. Default: `-1`
      keep_dims: bool, Optional
          Whether to retain the original matrix shape or not. Default:``True``
      epsilon: flot, optional
          Epsilon to apply to the normalization to preven NaN errors on zero values
  
      Returns
      -------
      :class:`keras.KerasTensor`
          The normalized output
      """
      return ops.sqrt(ops.sum(ops.power(matrix, 2), axis=axis, keepdims=keep_dims) + epsilon)
  
  
  def replicate_pad(image: KerasTensor, padding: int) -> KerasTensor:
      """ Apply replication padding to an input batch of images. Expects 4D tensor in BHWC format.
  
      Notes
      -----
      At the time of writing Keras does not have a native replication padding method.
      The implementation here is probably not the most efficient, but it is a pure keras method
      which should work ok.
  
      Parameters
      ----------
      image: :class:`keras.KerasTensor`
          Image tensor to pad
      pad: int
          The amount of padding to apply to each side of the input image
  
      Returns
      -------
      :class:`keras.KerasTensor`
          The input image with replication padding applied
      """
      top_pad = ops.tile(image[:, :1, ...], (1, padding, 1, 1))
      bottom_pad = ops.tile(image[:, -1:, ...], (1, padding, 1, 1))
      pad_top_bottom = ops.concatenate([top_pad, image, bottom_pad], axis=1)
      left_pad = ops.tile(pad_top_bottom[..., :1, :], (1, 1, padding, 1))
      right_pad = ops.tile(pad_top_bottom[..., -1:, :],  (1, 1, padding, 1))
      padded = ops.concatenate([left_pad, pad_top_bottom, right_pad], axis=2)
      return padded
  
  
  class ColorSpaceConvert():
      """ Transforms inputs between different color spaces on the GPU
  
      Notes
      -----
      The following color space transformations are implemented:
          - rgb to lab
          - rgb to xyz
          - srgb to _rgb
          - srgb to ycxcz
          - xyz to ycxcz
          - xyz to lab
          - xyz to rgb
          - ycxcz to rgb
          - ycxcz to xyz
  
      Parameters
      ----------
      from_space: str
          One of `"srgb"`, `"rgb"`, `"xyz"`
      to_space: str
          One of `"lab"`, `"rgb"`, `"ycxcz"`, `"xyz"`
  
      Raises
      ------
      ValueError
          If the requested color space conversion is not defined
      """
      def __init__(self, from_space: str, to_space: str) -> None:
          functions = {"rgb_lab": self._rgb_to_lab,
                       "rgb_xyz": self._rgb_to_xyz,
                       "srgb_rgb": self._srgb_to_rgb,
                       "srgb_ycxcz": self._srgb_to_ycxcz,
                       "xyz_ycxcz": self._xyz_to_ycxcz,
                       "xyz_lab": self._xyz_to_lab,
                       "xyz_rgb": self._xyz_to_rgb,
                       "ycxcz_rgb": self._ycxcz_to_rgb,
                       "ycxcz_xyz": self._ycxcz_to_xyz}
          func_name = f"{from_space.lower()}_{to_space.lower()}"
          if func_name not in functions:
              raise ValueError(f"The color transform {from_space} to {to_space} is not defined.")
  
          self._func = functions[func_name]
          self._ref_illuminant = Variable(np.array([[[0.950428545, 1.000000000, 1.[REDACTED_SSN]]]]),
                                          dtype="float32",
                                          trainable=False)
          self._inv_ref_illuminant = 1. / self._ref_illuminant
  
          self._rgb_xyz_map = self._get_rgb_xyz_map()
          self._xyz_multipliers = Variable([116, 500, 200], dtype="float32", trainable=False)
  
      @classmethod
      def _get_rgb_xyz_map(cls) -> tuple[KerasTensor, KerasTensor]:
          """ Obtain the mapping and inverse mapping for rgb to xyz color space conversion.
  
          Returns
          -------
          tuple
              The mapping and inverse Tensors for rgb to xyz color space conversion
          """
          mapping = np.array([[10135552 / 24577794,  8788810 / 24577794, 4435075 / 24577794],
                              [2613072 / 12288897, 8788810 / 12288897, 887015 / 12288897],
                              [1425312 / 73733382, 8788810 / 73733382, 70074185 / 73733382]])
          inverse = np.linalg.inv(mapping)
          return (Variable(mapping, dtype="float32", trainable=False),
                  Variable(inverse, dtype="float32", trainable=False))
  
      def __call__(self, image: KerasTensor) -> KerasTensor:
          """ Call the colorspace conversion function.
  
          Parameters
          ----------
          image: :class:`keras.KerasTensor`
              The image tensor in the colorspace defined by :attr:`from_space`
  
          Returns
          -------
          :class:`keras.KerasTensor`
              The image tensor in the colorspace defined by :attr:`to_space`
          """
          return self._func(image)
  
      def _rgb_to_lab(self, image: KerasTensor) -> KerasTensor:
          """ RGB to LAB conversion.
  
          Parameters
          ----------
          image: :class:`keras.KerasTensor`
              The image tensor in RGB format
  
          Returns
          -------
          :class:`keras.KerasTensor`
              The image tensor in LAB format
          """
          converted = self._rgb_to_xyz(image)
          return self._xyz_to_lab(converted)
  
      def _rgb_xyz_rgb(self, image: KerasTensor, mapping: KerasTensor) -> KerasTensor:
          """ RGB to XYZ or XYZ to RGB conversion.
  
          Notes
          -----
          The conversion in both directions is the same, but the mappping matrix for XYZ to RGB is
          the inverse of RGB to XYZ.
  
          References
          ----------
          https://www.image-engineering.de/library/technotes/958-how-to-convert-between-srgb-and-ciexyz
  
          Parameters
          ----------
          mapping: :class:`keras.KerasTensor`
              The mapping matrix to perform either the XYZ to RGB or RGB to XYZ color space
              conversion
  
          image: :class:`keras.KerasTensor`
              The image tensor in RGB format
  
          Returns
          -------
          :class:`keras.KerasTensor`
              The image tensor in XYZ format
          """
          dim = image.shape
          image = ops.transpose(image, (0, 3, 1, 2))
          image = ops.reshape(image, (dim[0], dim[3], dim[1] * dim[2]))
          converted = ops.transpose(ops.dot(mapping, image), (0, 2, 1))
          return ops.reshape(converted, dim)
  
      def _rgb_to_xyz(self, image: KerasTensor) -> KerasTensor:
          """ RGB to XYZ conversion.
  
          Parameters
          ----------
          image: :class:`keras.KerasTensor`
              The image tensor in RGB format
  
          Returns
          -------
          :class:`keras.KerasTensor`
              The image tensor in XYZ format
          """
          return self._rgb_xyz_rgb(image, self._rgb_xyz_map[0])
  
      @classmethod
      def _srgb_to_rgb(cls, image: KerasTensor) -> KerasTensor:
          """ SRGB to RGB conversion.
  
          Notes
          -----
          RGB Image is clipped to a small epsilon to stabalize training
  
          Parameters
          ----------
          image: :class:`keras.KerasTensor`
              The image tensor in SRGB format
  
          Returns
          -------
          :class:`keras.KerasTensor`
              The image tensor in RGB format
          """
          limit = np.float32(0.04045)
          return ops.where(image > limit,
                           ops.power((ops.clip(image, limit, np.inf) + 0.055) / 1.055, 2.4),
                           image / 12.92)
  
      def _srgb_to_ycxcz(self, image: KerasTensor) -> KerasTensor:
          """ SRGB to YcXcZ conversion.
  
          Parameters
          ----------
          image: :class:`keras.KerasTensor`
              The image tensor in SRGB format
  
          Returns
          -------
          :class:`keras.KerasTensor`
              The image tensor in YcXcZ format
          """
          converted = self._srgb_to_rgb(image)
          converted = self._rgb_to_xyz(converted)
          return self._xyz_to_ycxcz(converted)
  
      def _xyz_to_lab(self, image: KerasTensor) -> KerasTensor:
          """ XYZ to LAB conversion.
  
          Parameters
          ----------
          image: :class:`keras.KerasTensor`
              The image tensor in XYZ format
  
          Returns
          -------
          :class:`keras.KerasTensor`
              The image tensor in LAB format
          """
          image = image * self._inv_ref_illuminant
          delta = 6 / 29
          delta_cube = delta ** 3
          factor = 1 / (3 * (delta ** 2))
  
          clamped_term = ops.power(ops.clip(image, delta_cube, np.inf), 1.0 / 3.0)
          div = factor * image + (4 / 29)
  
          image = ops.where(image > delta_cube, clamped_term, div)
  
          return ops.concatenate([self._xyz_multipliers[0] * image[..., 1:2] - 16.,
                                  self._xyz_multipliers[1:] * (image[..., :2] - image[..., 1:3])],
                                 axis=-1)
  
      def _xyz_to_rgb(self, image: KerasTensor) -> KerasTensor:
          """ XYZ to YcXcZ conversion.
  
          Parameters
          ----------
          image: :class:`keras.KerasTensor`
              The image tensor in XYZ format
  
          Returns
          -------
          :class:`keras.KerasTensor`
              The image tensor in RGB format
          """
          return self._rgb_xyz_rgb(image, self._rgb_xyz_map[1])
  
      def _xyz_to_ycxcz(self, image: KerasTensor) -> KerasTensor:
          """ XYZ to YcXcZ conversion.
  
          Parameters
          ----------
          image: :class:`keras.KerasTensor`
              The image tensor in XYZ format
  
          Returns
          -------
          :class:`keras.KerasTensor`
              The image tensor in YcXcZ format
          """
          image = image * self._inv_ref_illuminant
          return ops.concatenate([self._xyz_multipliers[0] * image[..., 1:2] - 16.,
                                  self._xyz_multipliers[1:] * (image[..., :2] - image[..., 1:3])],
                                 axis=-1)
  
      def _ycxcz_to_rgb(self, image: KerasTensor) -> KerasTensor:
          """ YcXcZ to RGB conversion.
  
          Parameters
          ----------
          image: :class:`keras.KerasTensor`
              The image tensor in YcXcZ format
  
          Returns
          -------
          :class:`keras.KerasTensor`
              The image tensor in RGB format
          """
          converted = self._ycxcz_to_xyz(image)
          return self._xyz_to_rgb(converted)
  
      def _ycxcz_to_xyz(self, image: KerasTensor) -> KerasTensor:
          """ YcXcZ to XYZ conversion.
  
          Parameters
          ----------
          image: :class:`keras.KerasTensor`
              The image tensor in YcXcZ format
  
          Returns
          -------
          :class:`keras.KerasTensor`
              The image tensor in XYZ format
          """
          ch_y = (image[..., 0:1] + 16.) / self._xyz_multipliers[0]
          return ops.concatenate([ch_y + (image[..., 1:2] / self._xyz_multipliers[1]),
                                  ch_y,
                                  ch_y - (image[..., 2:3] / self._xyz_multipliers[2])],
                                 axis=-1) * self._ref_illuminant
  
  
  __all__ = get_module_objects(__name__)
  