mirror of
https://github.com/deepfakes/faceswap
synced 2025-06-09 04:36:50 -04:00
eceee07241
- Remove legacy update hashes - Remove legacy job from alignment-tools - Remove legacy landmark rotation - Add rotate face method to plugins/extract/detect - Update travis test for new alignments extension - Alignments format to .fsa - Remove serializer option from alignments-tool - Auto update legacy format alignment files to new format
282 lines
11 KiB
Python
282 lines
11 KiB
Python
#!/usr/bin/env python3
|
|
""" Alignments file functions for reading, writing and manipulating
|
|
a serialized alignments file """
|
|
|
|
import logging
|
|
import os
|
|
from datetime import datetime
|
|
|
|
from lib.serializer import get_serializer, get_serializer_from_filename
|
|
from lib.utils import FaceswapError
|
|
|
|
logger = logging.getLogger(__name__) # pylint: disable=invalid-name
|
|
|
|
|
|
class Alignments():
|
|
""" Holds processes pertaining to the alignments file.
|
|
|
|
folder: folder alignments file is stored in
|
|
filename: Filename of alignments file. If a
|
|
valid extension is provided, then it will be used to
|
|
decide the serializer otherwise compressed pickle is used.
|
|
"""
|
|
# pylint: disable=too-many-public-methods
|
|
def __init__(self, folder, filename="alignments"):
|
|
logger.debug("Initializing %s: (folder: '%s', filename: '%s')",
|
|
self.__class__.__name__, folder, filename)
|
|
self.serializer = get_serializer("compressed")
|
|
self.file = self.get_location(folder, filename)
|
|
|
|
self.data = self.load()
|
|
self.update_legacy()
|
|
logger.debug("Initialized %s", self.__class__.__name__)
|
|
|
|
# << PROPERTIES >> #
|
|
|
|
@property
|
|
def frames_count(self):
|
|
""" Return current frames count """
|
|
retval = len(self.data)
|
|
logger.trace(retval)
|
|
return retval
|
|
|
|
@property
|
|
def faces_count(self):
|
|
""" Return current faces count """
|
|
retval = sum(len(faces) for faces in self.data.values())
|
|
logger.trace(retval)
|
|
return retval
|
|
|
|
@property
|
|
def have_alignments_file(self):
|
|
""" Return whether an alignments file exists """
|
|
retval = os.path.exists(self.file)
|
|
logger.trace(retval)
|
|
return retval
|
|
|
|
@property
|
|
def hashes_to_frame(self):
|
|
""" Return a dict of each face_hash with their parent
|
|
frame name(s) and their index in the frame
|
|
"""
|
|
hash_faces = dict()
|
|
for frame_name, faces in self.data.items():
|
|
for idx, face in enumerate(faces):
|
|
hash_faces.setdefault(face["hash"], dict())[frame_name] = idx
|
|
return hash_faces
|
|
|
|
# << INIT FUNCTIONS >> #
|
|
|
|
def get_location(self, folder, filename):
|
|
""" Return the path to alignments file """
|
|
logger.debug("Getting location: (folder: '%s', filename: '%s')", folder, filename)
|
|
extension = os.path.splitext(filename)[1]
|
|
if extension in (".json", ".p", ".yaml", ".yml"):
|
|
# Reformat legacy alignments file
|
|
filename = self.update_file_format(folder, filename)
|
|
logger.debug("Updated legacy alignments. New filename: '%s'", filename)
|
|
elif not extension:
|
|
filename = "{}.{}".format(filename, self.serializer.file_extension)
|
|
logger.debug("File extension set from serializer: '%s'",
|
|
self.serializer.file_extension)
|
|
elif extension != ".fsa":
|
|
raise FaceswapError("{} is not a valid alignments file".format(filename))
|
|
location = os.path.join(str(folder), filename)
|
|
logger.verbose("Alignments filepath: '%s'", location)
|
|
return location
|
|
|
|
# << I/O >> #
|
|
|
|
def load(self):
|
|
""" Load the alignments data
|
|
Override for custom loading logic """
|
|
logger.debug("Loading alignments")
|
|
if not self.have_alignments_file:
|
|
raise FaceswapError("Error: Alignments file not found at "
|
|
"{}".format(self.file))
|
|
|
|
logger.info("Reading alignments from: '%s'", self.file)
|
|
data = self.serializer.load(self.file)
|
|
logger.debug("Loaded alignments")
|
|
return data
|
|
|
|
def reload(self):
|
|
""" Read the alignments data from the correct format """
|
|
logger.debug("Re-loading alignments")
|
|
self.data = self.load()
|
|
logger.debug("Re-loaded alignments")
|
|
|
|
def save(self):
|
|
""" Write the serialized alignments file """
|
|
logger.debug("Saving alignments")
|
|
logger.info("Writing alignments to: '%s'", self.file)
|
|
self.serializer.save(self.file, self.data)
|
|
logger.debug("Saved alignments")
|
|
|
|
def backup(self):
|
|
""" Backup copy of old alignments """
|
|
logger.debug("Backing up alignments")
|
|
if not os.path.isfile(self.file):
|
|
logger.debug("No alignments to back up")
|
|
return
|
|
now = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
src = self.file
|
|
split = os.path.splitext(src)
|
|
dst = split[0] + "_" + now + split[1]
|
|
logger.info("Backing up original alignments to '%s'", dst)
|
|
os.rename(src, dst)
|
|
logger.debug("Backed up alignments")
|
|
|
|
# << VALIDATION >> #
|
|
|
|
def frame_exists(self, frame):
|
|
""" return path of images that have faces """
|
|
retval = frame in self.data.keys()
|
|
logger.trace("'%s': %s", frame, retval)
|
|
return retval
|
|
|
|
def frame_has_faces(self, frame):
|
|
""" Return true if frame exists and has faces """
|
|
retval = bool(self.data.get(frame, list()))
|
|
logger.trace("'%s': %s", frame, retval)
|
|
return retval
|
|
|
|
def frame_has_multiple_faces(self, frame):
|
|
""" Return true if frame exists and has faces """
|
|
if not frame:
|
|
retval = False
|
|
else:
|
|
retval = bool(len(self.data.get(frame, list())) > 1)
|
|
logger.trace("'%s': %s", frame, retval)
|
|
return retval
|
|
|
|
# << DATA >> #
|
|
|
|
def get_faces_in_frame(self, frame):
|
|
""" Return the alignments for the selected frame """
|
|
logger.trace("Getting faces for frame: '%s'", frame)
|
|
return self.data.get(frame, list())
|
|
|
|
def get_full_frame_name(self, frame):
|
|
""" Return a frame with extension for when the extension is
|
|
not known """
|
|
retval = next(key for key in self.data.keys()
|
|
if key.startswith(frame))
|
|
logger.trace("Requested: '%s', Returning: '%s'", frame, retval)
|
|
return retval
|
|
|
|
def count_faces_in_frame(self, frame):
|
|
""" Return number of alignments within frame """
|
|
retval = len(self.data.get(frame, list()))
|
|
logger.trace(retval)
|
|
return retval
|
|
|
|
# << MANIPULATION >> #
|
|
|
|
def delete_face_at_index(self, frame, idx):
|
|
""" Delete the face alignment for given frame at given index """
|
|
logger.debug("Deleting face %s for frame '%s'", idx, frame)
|
|
idx = int(idx)
|
|
if idx + 1 > self.count_faces_in_frame(frame):
|
|
logger.debug("No face to delete: (frame: '%s', idx %s)", frame, idx)
|
|
return False
|
|
del self.data[frame][idx]
|
|
logger.debug("Deleted face: (frame: '%s', idx %s)", frame, idx)
|
|
return True
|
|
|
|
def add_face(self, frame, alignment):
|
|
""" Add a new face for a frame and return it's index """
|
|
logger.debug("Adding face to frame: '%s'", frame)
|
|
if frame not in self.data:
|
|
self.data[frame] = []
|
|
self.data[frame].append(alignment)
|
|
retval = self.count_faces_in_frame(frame) - 1
|
|
logger.debug("Returning new face index: %s", retval)
|
|
return retval
|
|
|
|
def update_face(self, frame, idx, alignment):
|
|
""" Replace a face for given frame and index """
|
|
logger.debug("Updating face %s for frame '%s'", idx, frame)
|
|
self.data[frame][idx] = alignment
|
|
|
|
def filter_hashes(self, hashlist, filter_out=False):
|
|
""" Filter in or out faces that match the hashlist
|
|
|
|
filter_out=True: Remove faces that match in the hashlist
|
|
filter_out=False: Remove faces that are not in the hashlist
|
|
"""
|
|
hashset = set(hashlist)
|
|
for filename, frame in self.data.items():
|
|
for idx, face in reversed(list(enumerate(frame))):
|
|
if ((filter_out and face.get("hash", None) in hashset) or
|
|
(not filter_out and face.get("hash", None) not in hashset)):
|
|
logger.verbose("Filtering out face: (filename: %s, index: %s)", filename, idx)
|
|
del frame[idx]
|
|
else:
|
|
logger.trace("Not filtering out face: (filename: %s, index: %s)",
|
|
filename, idx)
|
|
|
|
# << GENERATORS >> #
|
|
|
|
def yield_faces(self):
|
|
""" Yield face alignments for one image """
|
|
for frame_fullname, alignments in self.data.items():
|
|
frame_name = os.path.splitext(frame_fullname)[0]
|
|
face_count = len(alignments)
|
|
logger.trace("Yielding: (frame: '%s', faces: %s, frame_fullname: '%s')",
|
|
frame_name, face_count, frame_fullname)
|
|
yield frame_name, alignments, face_count, frame_fullname
|
|
|
|
@staticmethod
|
|
def yield_original_index_reverse(image_alignments, number_alignments):
|
|
""" Return the correct original index for
|
|
alignment in reverse order """
|
|
for idx, _ in enumerate(reversed(image_alignments)):
|
|
original_idx = number_alignments - 1 - idx
|
|
logger.trace("Yielding: face index %s", original_idx)
|
|
yield original_idx
|
|
|
|
# << LEGACY FUNCTIONS >> #
|
|
|
|
def update_legacy(self):
|
|
""" Update legacy alignments """
|
|
if self.has_legacy_landmarksxy():
|
|
logger.info("Updating legacy alignments")
|
|
self.update_legacy_landmarksxy()
|
|
|
|
# <File Format> #
|
|
# Serializer is now a compressed pickle .fsa format. This used to be any number of serializers
|
|
def update_file_format(self, folder, filename):
|
|
""" Convert old style alignments format to new style format """
|
|
logger.info("Reformatting legacy alignments file...")
|
|
old_location = os.path.join(str(folder), filename)
|
|
new_location = "{}.{}".format(os.path.splitext(old_location)[0],
|
|
self.serializer.file_extension)
|
|
logger.info("Old location: '%s', New location: '%s'", old_location, new_location)
|
|
|
|
load_serializer = get_serializer_from_filename(old_location)
|
|
data = load_serializer.load(old_location)
|
|
self.serializer.save(new_location, data)
|
|
return os.path.basename(new_location)
|
|
|
|
# <landmarks> #
|
|
# Landmarks renamed from landmarksXY to landmarks_xy for PEP compliance
|
|
def has_legacy_landmarksxy(self):
|
|
""" check for legacy landmarksXY keys """
|
|
logger.debug("checking legacy landmarksXY")
|
|
retval = (any(key == "landmarksXY"
|
|
for alignments in self.data.values()
|
|
for alignment in alignments
|
|
for key in alignment))
|
|
logger.debug("legacy landmarksXY: %s", retval)
|
|
return retval
|
|
|
|
def update_legacy_landmarksxy(self):
|
|
""" Update landmarksXY to landmarks_xy and save alignments """
|
|
update_count = 0
|
|
for alignments in self.data.values():
|
|
for alignment in alignments:
|
|
alignment["landmarks_xy"] = alignment.pop("landmarksXY")
|
|
update_count += 1
|
|
logger.debug("Updated landmarks_xy: %s", update_count)
|
|
self.save()
|