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
  """Dataclass objects for holding and serializing alignments data"""
  from __future__ import annotations
  
  from dataclasses import dataclass, field, fields, MISSING
  import types
  import typing as T
  
  import numpy as np
  import numpy.typing as npt
  
  from lib.logger import format_array
  
  from .constants import CenteringType
  
  
  @dataclass
  class DataclassDict:
      """Parent DataClass that has methods for loading to and from a dict for data serialization"""
      def __repr__(self) -> str:
          """Pretty print for logging"""
          params = {}
          for k, v in self.__dict__.items():
              if isinstance(v, np.ndarray):
                  params[k] = format_array(v)
                  continue
              if isinstance(v, bytes):
                  params[k] = f"{len(v)}b"
                  continue
              params[k] = v
          s_params = ", ".join(f"{k}={repr(v)}" for k, v in params.items())
          return f"{self.__class__.__name__}({s_params})"
  
      @classmethod
      def _object_to_serial(cls, obj: T.Any) -> T.Any:
          """Convert object lists or DataclassDicts serializable items
  
          Parameters
          ----------
          obj
              The object to convert
  
          Returns
          -------
          The converted object or original object if not to be converted
          """
          if isinstance(obj, np.ndarray):
              return obj.tolist()
          if isinstance(obj, DataclassDict):
              return obj.to_dict()
          return obj
  
      def to_dict(self) -> dict[str, T.Any]:
          """Obtain the contents of the dataclass object as a python dictionary
  
          Returns
          -------
          The dataclass object as a python dictionary, with numpy arrays converted to lists
          """
          retval: dict[str, T.Any] = {}
          for k, v in self.__dict__.items():
              if isinstance(v, (list, tuple)):
                  retval[k] = [self._object_to_serial(x) for x in v]
              elif isinstance(v, dict):
                  retval[k] = {x: self._object_to_serial(y) for x, y in v.items()}
              else:
                  retval[k] = self._object_to_serial(v)
          return retval
  
      @classmethod
      def _convert_dtype(cls, data_type: T.Any, val: T.Any) -> DataclassDict | np.ndarray | None:
          """Convert a serialized dict to a DataclassDict or list to a numpy array of the correct
          dtype
  
          Parameters
          ----------
          field_type
              The field type for the incoming value
          value
              The list to convert to a numpy array or DataclassDict
  
          Returns
          -------
          The inbound item to a DataclassDict or numpy array. ``None`` if the item does not convert
          """
          if isinstance(data_type, type) and issubclass(data_type, DataclassDict):
              return data_type.from_dict(val)
  
          origin = T.get_origin(data_type)
          if origin is types.UnionType:
              args = tuple(a for a in T.get_args(data_type) if a is not types.NoneType)
              assert len(args) == 1
              if val is None:
                  return val
              data_type = args[0]
              origin = T.get_origin(data_type)
  
          if origin is not np.ndarray:
              return None
  
          args = T.get_args(data_type)
          dtype = T.get_args(args[1])[0]
          return np.array(val, dtype=dtype)
  
      @classmethod
      def _parse_dict(cls, field_type: T.Any, value: dict[str, T.Any]) -> dict[str, T.Any]:
          """Parse incoming serialized dicts into their correct nested objects
  
          Parameters
          ----------
          field_type
              The field type for the incoming value
          value
              The dictionary to parse
  
          Returns
          -------
          The dictionary with its values converted to the correct datatype
          """
          assert T.get_origin(field_type) is dict
          dtype = T.get_args(field_type)[1]
          retval = {}
          for k, v in value.items():
              converted = cls._convert_dtype(dtype, v)
              if converted is not None:
                  retval[k] = converted
                  continue
              retval[k] = v
          return retval
  
      @classmethod
      def _parse_list(cls, field_type: T.Any, value: list[T.Any] | tuple[T.Any]
                      ) -> list[T.Any] | tuple[T.Any]:
          """Parse incoming serialized lists into their correct nested objects
  
          Parameters
          ----------
          field_type
              The field type for the incoming value
          value
              The list to parse
  
          Returns
          -------
          The list with its values converted to the correct datatype
          """
          origin = T.get_origin(field_type)
          assert origin in (list, tuple), (
              f"value: {type(value)} field: {T.get_origin(field_type)}")
          dtype = T.get_args(field_type)[0]
          items = []
          for v in value:
              converted = cls._convert_dtype(dtype, v)
              if converted is not None:
                  items.append(converted)
                  continue
              items.append(v)
          retval = T.cast(list[T.Any] | tuple[T.Any], tuple(items) if origin is tuple else items)
          return retval
  
      @classmethod
      def from_dict(cls, data_dict: dict[str, T.Any]) -> T.Self:
          """Load the contents from a serialized python dict into this dataclass
  
          Parameters
          ----------
          data_dict
              The data to load into the dataclass
          """
          inbound = set(data_dict)
          all_fields = set(f.name for f in fields(cls))
          required = set(f.name for f in fields(cls)
                         if f.default is MISSING and f.default_factory is MISSING)
          if not inbound.issubset(all_fields):
              raise ValueError(f"Dictionary keys {sorted(inbound)} should be a subset of dataclass "
                               f"params {sorted(all_fields)}")
          if not required.issubset(inbound):
              raise ValueError(f"Dataclass params {sorted(required)} should be a subset of "
                               f"dictionary keys {sorted(inbound)}")
          type_hints = T.get_type_hints(cls)
          kwargs: dict[str, T.Any] = {}
          for f in fields(cls):
              if f.name not in data_dict:
                  continue
              field_type = type_hints.get(f.name)
              val = data_dict[f.name]
              converted = cls._convert_dtype(field_type, val)
              if converted is not None:
                  kwargs[f.name] = converted
                  continue
              if isinstance(val, dict):
                  kwargs[f.name] = cls._parse_dict(field_type, val)
                  continue
              if isinstance(val, (list, tuple)):
                  kwargs[f.name] = cls._parse_list(field_type, val)
                  continue
              kwargs[f.name] = val
          return cls(**kwargs)
  
  
  @dataclass(repr=False)
  class MaskAlignmentsFile(DataclassDict):
      """Dataclass for storing Masks in alignments files and PNG Headers"""
      mask: bytes
      """The zlib compressed UINT8 mask of shape (stored_size, stored_size)"""
      affine_matrix: npt.NDArray[np.float32]
      """The affine matrix that takes the mask from stored space to frame space"""
      interpolator: int
      """The interpolator required to take the mask from stored space to frame space"""
      stored_size: int
      """The size the mask is stored at"""
      stored_centering: CenteringType
      """The (legacy, face, head) centering type of the mask"""
  
  
  @dataclass(repr=False)
  class PNGAlignments(DataclassDict):
      """Base Dataclass for storing a single faces' Alignment Information in Alignments files and PNG
      Headers."""
      x: int
      """The left most point of the bounding box"""
      y: int
      """The top most point of the bounding box"""
      w: int
      """The width of the bounding box"""
      h: int
      """The height of the bounding box"""
      landmarks_xy: npt.NDArray[np.float32]
      """The (x, y) landmark points of the face"""
      mask: dict[str, MaskAlignmentsFile] = field(default_factory=dict)
      """The masks stored for the face"""
      identity: dict[str, npt.NDArray[np.float32]] = field(default_factory=dict)
      """The identity vectors stored for the face"""
  
      def __repr__(self) -> str:
          """Pretty print for logging"""
          params: dict[str, T.Any] = {}
          for k, v in self.__dict__.items():
              if k in ("landmarks_xy", "thumb"):
                  params[k] = None if v is None else f"{type(v)}[{len(v)}]"
                  continue
              if k == "identity":
                  params[k] = {n: f"{type(i)}[{len(i)}]" for n, i in v.items()}
                  continue
              params[k] = v
          s_params = ", ".join(f"{k}={v}" for k, v in params.items())
          return f"{self.__class__.__name__}({s_params})"
  
  
  @dataclass(repr=False)
  class PNGSource(DataclassDict):
      """Dataclass for storing additional meta information in PNG headers."""
      alignments_version: float
      """The alignments file version that created the alignments data"""
      original_filename: str
      """The original filename that this face was saved with"""
      face_index: int
      """The index of this face within the frame"""
      source_filename: str
      """The filename of the original frame the face was extracted from"""
      source_is_video: bool
      """``True`` if the face was extracted from a video. ``False`` if from an image"""
      source_frame_dims: tuple[int, int]
      """The (Height, Width) dimensions of the original frame the face was extracted from"""
  
  
  @dataclass(repr=False)
  class PNGHeader(DataclassDict):
      """Dataclass for storing all alignment and meta information in PNG Headers."""
      alignments: PNGAlignments
      """The alignment information for the face"""
      source: PNGSource
      """The frame source information for the face"""
  
  
  @dataclass(repr=False)
  class FileAlignments(PNGAlignments):
      """Dataclass that holds the same information as PNGAlignments as well as a thumbnail for a
      single face"""
      thumb: npt.NDArray[np.uint8] | None = None
      """96px JPEG thumbnail of the aligned face image stored as a list"""
  
  
  @dataclass(repr=False)
  class AlignmentsEntry(DataclassDict):
      """Holds the alignments entry for a single frame in the Alignments data dictionary"""
      faces: list[FileAlignments] = field(default_factory=list)
      """The detected faces in a frame"""
      video_meta: dict[T.Literal["pts_time", "keyframe"], int] = field(default_factory=dict)
      """The keyframe to pts timestamp mapping for video data"""
  