mirror of
https://github.com/deepfakes/faceswap
synced 2025-06-07 19:05:02 -04:00
379 lines
13 KiB
Python
379 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
""" Handles updating of an alignments file from an older version to the current version. """
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
import typing as T
|
|
|
|
import numpy as np
|
|
|
|
from lib.logger import parse_class_init
|
|
from lib.utils import VIDEO_EXTENSIONS
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
if T.TYPE_CHECKING:
|
|
from .alignments import Alignments, AlignmentFileDict
|
|
|
|
|
|
class _Updater():
|
|
""" Base class for inheriting to test for and update of an alignments file property
|
|
|
|
Parameters
|
|
----------
|
|
alignments: :class:`~Alignments`
|
|
The alignments object that is being tested and updated
|
|
"""
|
|
def __init__(self, alignments: Alignments) -> None:
|
|
logger.debug(parse_class_init(locals()))
|
|
self._alignments = alignments
|
|
self._needs_update = self._test()
|
|
if self._needs_update:
|
|
self._update()
|
|
logger.debug("Initialized: %s", self.__class__.__name__)
|
|
|
|
@property
|
|
def is_updated(self) -> bool:
|
|
""" bool. ``True`` if this updater has been run otherwise ``False`` """
|
|
return self._needs_update
|
|
|
|
def _test(self) -> bool:
|
|
""" Calls the child's :func:`test` method and logs output
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
``True`` if the test condition is met otherwise ``False``
|
|
"""
|
|
logger.debug("checking %s", self.__class__.__name__)
|
|
retval = self.test()
|
|
logger.debug("legacy %s: %s", self.__class__.__name__, retval)
|
|
return retval
|
|
|
|
def test(self) -> bool:
|
|
""" Override to set the condition to test for.
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
``True`` if the test condition is met otherwise ``False``
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
def _update(self) -> int:
|
|
""" Calls the child's :func:`update` method, logs output and sets the
|
|
:attr:`is_updated` flag
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The number of items that were updated
|
|
"""
|
|
retval = self.update()
|
|
logger.debug("Updated %s: %s", self.__class__.__name__, retval)
|
|
return retval
|
|
|
|
def update(self) -> int:
|
|
""" Override to set the action to perform on the alignments object if the test has
|
|
passed
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The number of items that were updated
|
|
"""
|
|
raise NotImplementedError()
|
|
|
|
|
|
class VideoExtension(_Updater):
|
|
""" Alignments files from video files used to have a dummy '.png' extension for each of the
|
|
keys. This has been changed to be file extension of the original input video (for better)
|
|
identification of alignments files generated from video files
|
|
|
|
Parameters
|
|
----------
|
|
alignments: :class:`~Alignments`
|
|
The alignments object that is being tested and updated
|
|
video_filename: str
|
|
The video filename that holds these alignments
|
|
"""
|
|
def __init__(self, alignments: Alignments, video_filename: str) -> None:
|
|
self._video_name, self._extension = os.path.splitext(video_filename)
|
|
super().__init__(alignments)
|
|
|
|
def test(self) -> bool:
|
|
""" Requires update if the extension of the key in the alignment file is not the same
|
|
as for the input video file
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
``True`` if the key extensions need updating otherwise ``False``
|
|
"""
|
|
# Note: Don't check on alignments file version. It's possible that the file gets updated to
|
|
# a newer version before this check is run
|
|
if self._extension.lower() not in VIDEO_EXTENSIONS:
|
|
return False
|
|
|
|
exts = set(os.path.splitext(k)[-1] for k in self._alignments.data)
|
|
if len(exts) != 1:
|
|
logger.debug("Alignments file has multiple key extensions. Skipping")
|
|
return False
|
|
|
|
if self._extension in exts:
|
|
logger.debug("Alignments file contains correct key extensions. Skipping")
|
|
return False
|
|
|
|
logger.debug("Needs update for video extension (version: %s, extension: %s)",
|
|
self._alignments.version, self._extension)
|
|
return True
|
|
|
|
def update(self) -> int:
|
|
""" Update alignments files that have been extracted from videos to have the key end in the
|
|
video file extension rather than ',png' (the old way)
|
|
|
|
Parameters
|
|
----------
|
|
video_filename: str
|
|
The filename of the video file that created these alignments
|
|
"""
|
|
updated = 0
|
|
for key in list(self._alignments.data):
|
|
val = self._alignments.data[key]
|
|
fname = os.path.splitext(key)[0]
|
|
if fname.rsplit("_")[0] != self._video_name:
|
|
continue # Key is from a different source
|
|
|
|
new_key = f"{fname}{self._extension}"
|
|
del self._alignments.data[key]
|
|
self._alignments.data[new_key] = val
|
|
updated += 1
|
|
|
|
logger.debug("Updated alignemnt keys for video extension: %s", updated)
|
|
return updated
|
|
|
|
|
|
class FileStructure(_Updater):
|
|
""" Alignments were structured: {frame_name: <list of faces>}. We need to be able to store
|
|
information at the frame level, so new structure is: {frame_name: {faces: <list of faces>}}
|
|
"""
|
|
def test(self) -> bool:
|
|
""" Test whether the alignments file is laid out in the old structure of
|
|
`{frame_name: [faces]}`
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
``True`` if the file has legacy structure otherwise ``False``
|
|
"""
|
|
return any(isinstance(val, list) for val in self._alignments.data.values())
|
|
|
|
def update(self) -> int:
|
|
""" Update legacy alignments files from the format `{frame_name: [faces}` to the
|
|
format `{frame_name: {faces: [faces]}`.
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The number of items that were updated
|
|
"""
|
|
updated = 0
|
|
for key, val in self._alignments.data.items():
|
|
if not isinstance(val, list):
|
|
continue
|
|
self._alignments.data[key] = {"faces": val}
|
|
updated += 1
|
|
return updated
|
|
|
|
|
|
class LandmarkRename(_Updater):
|
|
""" Landmarks renamed from landmarksXY to landmarks_xy for PEP compliance """
|
|
def test(self) -> bool:
|
|
""" check for legacy landmarksXY keys.
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
``True`` if the alignments file contains legacy `landmarksXY` keys otherwise ``False``
|
|
"""
|
|
return (any(key == "landmarksXY"
|
|
for val in self._alignments.data.values()
|
|
for alignment in val["faces"]
|
|
for key in alignment))
|
|
|
|
def update(self) -> int:
|
|
""" Update legacy `landmarksXY` keys to PEP compliant `landmarks_xy` keys.
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The number of landmarks keys that were changed
|
|
"""
|
|
update_count = 0
|
|
for val in self._alignments.data.values():
|
|
for alignment in val["faces"]:
|
|
if "landmarksXY" in alignment:
|
|
alignment["landmarks_xy"] = alignment.pop("landmarksXY") # type:ignore
|
|
update_count += 1
|
|
return update_count
|
|
|
|
|
|
class ListToNumpy(_Updater):
|
|
""" Landmarks stored as list instead of numpy array """
|
|
def test(self) -> bool:
|
|
""" check for legacy landmarks stored as `list` rather than :class:`numpy.ndarray`.
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
``True`` if not all landmarks are :class:`numpy.ndarray` otherwise ``False``
|
|
"""
|
|
return not all(isinstance(face["landmarks_xy"], np.ndarray)
|
|
for val in self._alignments.data.values()
|
|
for face in val["faces"])
|
|
|
|
def update(self) -> int:
|
|
""" Update landmarks stored as `list` to :class:`numpy.ndarray`.
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The number of landmarks keys that were changed
|
|
"""
|
|
update_count = 0
|
|
for val in self._alignments.data.values():
|
|
for alignment in val["faces"]:
|
|
test = alignment["landmarks_xy"]
|
|
if not isinstance(test, np.ndarray):
|
|
alignment["landmarks_xy"] = np.array(test, dtype="float32")
|
|
update_count += 1
|
|
return update_count
|
|
|
|
|
|
class MaskCentering(_Updater):
|
|
""" Masks not containing the stored_centering parameters. Prior to this implementation all
|
|
masks were stored with face centering """
|
|
|
|
def test(self) -> bool:
|
|
""" Mask centering was introduced in alignments version 2.2
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
``True`` mask centering requires updating otherwise ``False``
|
|
"""
|
|
return self._alignments.version < 2.2
|
|
|
|
def update(self) -> int:
|
|
""" Add the mask key to the alignment file and update the centering of existing masks
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The number of masks that were updated
|
|
"""
|
|
update_count = 0
|
|
for val in self._alignments.data.values():
|
|
for alignment in val["faces"]:
|
|
if "mask" not in alignment:
|
|
alignment["mask"] = {}
|
|
for mask in alignment["mask"].values():
|
|
mask["stored_centering"] = "face"
|
|
update_count += 1
|
|
return update_count
|
|
|
|
|
|
class IdentityAndVideoMeta(_Updater):
|
|
""" Prior to version 2.3 the identity key did not exist and the video_meta key was not
|
|
compulsory. These should now both always appear, but do not need to be populated. """
|
|
|
|
def test(self) -> bool:
|
|
""" Identity Key was introduced in alignments version 2.3
|
|
|
|
Returns
|
|
-------
|
|
bool
|
|
``True`` identity key needs inserting otherwise ``False``
|
|
"""
|
|
return self._alignments.version < 2.3
|
|
|
|
# Identity information was not previously stored in the alignments file.
|
|
def update(self) -> int:
|
|
""" Add the video_meta and identity keys to the alignment file and leave empty
|
|
|
|
Returns
|
|
-------
|
|
int
|
|
The number of keys inserted
|
|
"""
|
|
update_count = 0
|
|
for val in self._alignments.data.values():
|
|
this_update = 0
|
|
if "video_meta" not in val:
|
|
val["video_meta"] = {}
|
|
this_update = 1
|
|
for alignment in val["faces"]:
|
|
if "identity" not in alignment:
|
|
alignment["identity"] = {}
|
|
this_update = 1
|
|
update_count += this_update
|
|
return update_count
|
|
|
|
|
|
class Legacy():
|
|
""" Legacy alignments properties that are no longer used, but are still required for backwards
|
|
compatibility/upgrading reasons.
|
|
|
|
Parameters
|
|
----------
|
|
alignments: :class:`~Alignments`
|
|
The alignments object that requires these legacy properties
|
|
"""
|
|
def __init__(self, alignments: Alignments) -> None:
|
|
self._alignments = alignments
|
|
self._hashes_to_frame: dict[str, dict[str, int]] = {}
|
|
self._hashes_to_alignment: dict[str, AlignmentFileDict] = {}
|
|
|
|
@property
|
|
def hashes_to_frame(self) -> dict[str, dict[str, int]]:
|
|
""" dict: The SHA1 hash of the face mapped to the frame(s) and face index within the frame
|
|
that the hash corresponds to. The structure of the dictionary is:
|
|
|
|
{**SHA1_hash** (`str`): {**filename** (`str`): **face_index** (`int`)}}.
|
|
|
|
Notes
|
|
-----
|
|
This method is deprecated and exists purely for updating legacy hash based alignments
|
|
to new png header storage in :class:`lib.align.update_legacy_png_header`.
|
|
|
|
The first time this property is referenced, the dictionary will be created and cached.
|
|
Subsequent references will be made to this cached dictionary.
|
|
"""
|
|
if not self._hashes_to_frame:
|
|
logger.debug("Generating hashes to frame")
|
|
for frame_name, val in self._alignments.data.items():
|
|
for idx, face in enumerate(val["faces"]):
|
|
self._hashes_to_frame.setdefault(
|
|
face["hash"], {})[frame_name] = idx # type:ignore
|
|
return self._hashes_to_frame
|
|
|
|
@property
|
|
def hashes_to_alignment(self) -> dict[str, AlignmentFileDict]:
|
|
""" dict: The SHA1 hash of the face mapped to the alignment for the face that the hash
|
|
corresponds to. The structure of the dictionary is:
|
|
|
|
Notes
|
|
-----
|
|
This method is deprecated and exists purely for updating legacy hash based alignments
|
|
to new png header storage in :class:`lib.align.update_legacy_png_header`.
|
|
|
|
The first time this property is referenced, the dictionary will be created and cached.
|
|
Subsequent references will be made to this cached dictionary.
|
|
"""
|
|
if not self._hashes_to_alignment:
|
|
logger.debug("Generating hashes to alignment")
|
|
self._hashes_to_alignment = {face["hash"]: face # type:ignore
|
|
for val in self._alignments.data.values()
|
|
for face in val["faces"]}
|
|
return self._hashes_to_alignment
|