mirror of
https://github.com/deepfakes/faceswap.git
synced 2026-10-18 22:45:59 +08:00
torzdf
291 lines
10 KiB
Python
291 lines
10 KiB
Python
#! /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"""
|