simcop2387/faceswap
1
0
Fork 0
mirror of https://github.com/deepfakes/faceswap synced 2025-06-07 10:43:27 -04:00
Code Issues Releases Wiki Activity

convert: Bugfix: Update legacy .png video alignments to include video file extension

Browse source
This commit is contained in:
torzdf 2024-04-17 12:37:47 +01:00
parent 3d5b962d29
commit 957734dfc0
5 changed files with 107 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

82
lib/align/alignments.py
View file

@ -10,14 +10,14 @@ from datetime import datetime
import numpy as np
from lib.serializer import get_serializer, get_serializer_from_filename
from lib.utils import FaceswapError
from lib.utils import FaceswapError, VIDEO_EXTENSIONS
if T.TYPE_CHECKING:
from collections.abc import Generator
from .aligned_face import CenteringType
logger = logging.getLogger(__name__)
_VERSION = 2.3
_VERSION = 2.4
# VERSION TRACKING
# 1.0 - Never really existed. Basically any alignments file prior to version 2.0
# 2.0 - Implementation of full head extract. Any alignments version below this will have used
@ -27,6 +27,7 @@ _VERSION = 2.3
# 2.2 - Add support for differently centered masks (i.e. not all masks stored as face centering)
# 2.3 - Add 'identity' key to alignments file. May or may not be populated, to contain vggface2
# embeddings. Make 'video_meta' key a standard key. Can be unpopulated
# 2.4 - Update video file alignment keys to end in the video extension rather than '.png'
# TODO Convert these to Dataclasses
@ -584,6 +585,21 @@ class Alignments():
frame_name, face_count, frame_fullname)
yield frame_name, val["faces"], face_count, frame_fullname
def update_legacy_has_source(self, filename: str) -> None:
""" Update legacy alignments files when we have the source filename available.
Updates here can only be performed when we have the source filename
Parameters
----------
filename: str:
The filename/folder of the original source images/video for the current alignments
"""
updates = [updater.is_updated for updater in (_VideoExtension(self, filename), )]
if any(updates):
self._io.update_version()
self.save()
class _IO():
""" Class to handle the saving/loading of an alignments file.
@ -719,10 +735,14 @@ class _IO():
_MaskCentering(self._alignments),
_IdentityAndVideoMeta(self._alignments))]
if any(updates):
self._version = _VERSION
logger.info("Updating alignments file to version %s", self._version)
self.update_version()
self.save()
def update_version(self) -> None:
""" Update the version of the alignments file to the latest version """
self._version = _VERSION
logger.info("Updating alignments file to version %s", self._version)
def load(self) -> dict[str, AlignmentDict]:
""" Load the alignments data from the serialized alignments :attr:`file`.
@ -908,6 +928,60 @@ class _Updater():
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 alignments version is < 2.4
Returns
-------
bool
``True`` if the key extensions need updating otherwise ``False``
"""
retval = self._alignments.version < 2.4 and self._extension in VIDEO_EXTENSIONS
logger.debug("Needs update for video extension: %s (version: %s, extension: %s)",
retval, self._alignments.version, self._extension)
return retval
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>}}

26
scripts/convert.py
View file

@ -85,13 +85,7 @@ class Convert():
self._args = handle_deprecated_cliopts(arguments)
self._images = ImagesLoader(self._args.input_dir, fast_count=True)
self._alignments = Alignments(self._args, False, self._images.is_video)
if self._alignments.version == 1.0:
logger.error("The alignments file format has been updated since the given alignments "
"file was generated. You need to update the file to proceed.")
logger.error("To do this run the 'Alignments Tool' > 'Extract' Job.")
sys.exit(1)
self._alignments = self._get_alignments()
self._opts = OptionalActions(self._args, self._images.file_list, self._alignments)
self._add_queues()
@ -133,6 +127,24 @@ class Convert():
logger.debug(retval)
return retval
def _get_alignments(self) -> Alignments:
""" Perform validation checks and legacy updates and return alignemnts object
Returns
-------
:class:`~lib.align.alignments.Alignments`
The alignments file for the extract job
"""
retval = Alignments(self._args, False, self._images.is_video)
if retval.version == 1.0:
logger.error("The alignments file format has been updated since the given alignments "
"file was generated. You need to update the file to proceed.")
logger.error("To do this run the 'Alignments Tool' > 'Extract' Job.")
sys.exit(1)
retval.update_legacy_has_source(os.path.basename(self._args.input_dir))
return retval
def _validate(self) -> None:
""" Validate the Command Line Options.

5
tools/alignments/alignments.py
View file

@ -260,6 +260,11 @@ class _Alignments():
else:
self.alignments = AlignmentData(self._find_alignments())
if (self.alignments is not None and
arguments.frames_dir and
os.path.isfile(arguments.frames_dir)):
self.alignments.update_legacy_has_source(os.path.basename(arguments.frames_dir))
logger.debug("Initialized %s", self.__class__.__name__)
def _find_alignments(self) -> str:

6
tools/alignments/media.py
View file

@ -473,14 +473,14 @@ class Frames(MediaLoader):
The full framename, the filename and the file extension of the frame
"""
logger.info("Loading video frames from %s", self.folder)
vidname = os.path.splitext(os.path.basename(self.folder))[0]
vidname, ext = os.path.splitext(os.path.basename(self.folder))
for i in range(self.count):
idx = i + 1
# Keep filename format for outputted face
filename = f"{vidname}_{idx:06d}"
retval = {"frame_fullname": f"{filename}.png",
retval = {"frame_fullname": f"{filename}{ext}",
"frame_name": filename,
"frame_extension": ".png"}
"frame_extension": ext}
logger.trace(retval) # type: ignore
yield retval

2
tools/manual/detected_faces.py
View file

@ -58,6 +58,8 @@ class DetectedFaces():
self._updated_frame_indices: set[int] = set()
self._alignments: Alignments = self._get_alignments(alignments_path, input_location)
self._alignments.update_legacy_has_source(os.path.basename(input_location))
self._extractor = extractor
self._tk_vars = self._set_tk_vars()