From fd5c004967e46d3ca2df36f5a8cc7d10beeeacd9 Mon Sep 17 00:00:00 2001 From: torzdf <36920800+torzdf@users.noreply.github.com> Date: Thu, 24 Jan 2019 13:07:29 +0000 Subject: [PATCH] Add update_hashes to alignments tool --- tools/alignments.py | 4 ++- tools/cli.py | 12 ++++--- tools/lib_alignments/__init__.py | 2 +- tools/lib_alignments/jobs.py | 55 ++++++++++++++++++++++++++++++++ 4 files changed, 67 insertions(+), 6 deletions(-) diff --git a/tools/alignments.py b/tools/alignments.py index 634df55f74..2d5b65c86d 100644 --- a/tools/alignments.py +++ b/tools/alignments.py @@ -4,7 +4,7 @@ from lib.utils import set_system_verbosity from .lib_alignments import (AlignmentData, Check, Draw, # noqa pylint: disable=unused-import Extract, Legacy, Manual, Merge, Reformat, Rename, - RemoveAlignments, Sort, Spatial) + RemoveAlignments, Sort, Spatial, UpdateHashes) class Alignments(): @@ -27,6 +27,8 @@ def process(self): """ Main processing function of the Align tool """ if self.args.job.startswith("extract"): job = Extract + elif self.args.job == "update-hashes": + job = UpdateHashes elif self.args.job.startswith("remove-"): job = RemoveAlignments elif self.args.job.startswith("sort-"): diff --git a/tools/cli.py b/tools/cli.py index fa1e986009..a67dcbf9c0 100644 --- a/tools/cli.py +++ b/tools/cli.py @@ -70,14 +70,18 @@ def get_argument_list(self): "\n'rename' - Rename faces to correspond with their parent frame and position" "\n\tindex in the alignments file (i.e. how they are named after running" "\n\textract)." + faces_dir + - "\n'sort-x' - Re-index the alignments from left to right. For alignments with" + "\n'sort-x': Re-index the alignments from left to right. For alignments with" "\n\tmultiple faces this will ensure that the left-most face is at index 0" "\n\tOptionally pass in a faces folder (-fc) to also rename extracted faces." - "\n'sort-y' - Re-index the alignments from top to bottom. For alignments with" + "\n'sort-y': Re-index the alignments from top to bottom. For alignments with" "\n\tmultiple faces this will ensure that the top-most face is at index 0" "\n\tOptionally pass in a faces folder (-fc) to also rename extracted faces." - "\n'spatial' - Perform spatial and temporal filtering to smooth alignments" - "\n\t(EXPERIMENTAL!)"}) + "\n'spatial': Perform spatial and temporal filtering to smooth alignments" + "\n\t(EXPERIMENTAL!)" + "\n'update-hashes': Recalculate the face hashes. Only use this if you have " + "\n\taltered the extracted faces (e.g. colour adjust). The files MUST be " + "\n\tnamed '_face index' (i.e. how they are named after running" + "\n\textract)." + faces_dir}) argument_list.append({"opts": ("-a", "--alignments_file"), "action": FileFullPaths, "dest": "alignments_file", diff --git a/tools/lib_alignments/__init__.py b/tools/lib_alignments/__init__.py index 3eb0103082..95a74e602b 100644 --- a/tools/lib_alignments/__init__.py +++ b/tools/lib_alignments/__init__.py @@ -1,4 +1,4 @@ from tools.lib_alignments.media import AlignmentData, ExtractedFaces, Faces, Frames from tools.lib_alignments.annotate import Annotate -from tools.lib_alignments.jobs import Check, Draw, Extract, Legacy, Merge, Reformat, RemoveAlignments, Rename, Sort, Spatial +from tools.lib_alignments.jobs import Check, Draw, Extract, Legacy, Merge, Reformat, RemoveAlignments, Rename, Sort, Spatial, UpdateHashes from tools.lib_alignments.jobs_manual import Manual diff --git a/tools/lib_alignments/jobs.py b/tools/lib_alignments/jobs.py index 1334485b11..27ed0e2492 100644 --- a/tools/lib_alignments/jobs.py +++ b/tools/lib_alignments/jobs.py @@ -962,3 +962,58 @@ def update_alignments(self, landmarks): self.alignments.data[frame][0]["landmarksXY"] = landmarks_xy logger.trace("Updated: (frame: '%s', landmarks: %s)", frame, landmarks_xy) logger.debug("Updated alignments") + + +class UpdateHashes(): + """ Update hashes in an alignments file """ + def __init__(self, alignments, arguments): + logger.debug("Initializing %s: (arguments: %s)", self.__class__.__name__, arguments) + self.alignments = alignments + self.faces = Faces(arguments.faces_dir).file_list_sorted + self.face_hashes = dict() + logger.debug("Initialized %s", self.__class__.__name__) + + def process(self): + """ Update Face Hashes to the alignments file """ + logger.info("[UPDATE FACE HASHES]") # Tidy up cli output + self.get_hashes() + updated = self.update_hashes() + if updated == 0: + logger.info("No hashes were updated. Exiting") + return + self.alignments.save() + logger.info("%s frame(s) had their face hashes updated.", updated) + + def get_hashes(self): + """ Read the face hashes from the faces """ + logger.info("Getting original filenames, indexes and hashes...") + for face in self.faces: + filename = face["face_name"] + extension = face["face_extension"] + if "_" not in face["face_name"]: + logger.warning("Unable to determine index of file. Skipping: '%s'", filename) + continue + index = filename[filename.rfind("_") + 1:] + if not index.isdigit(): + logger.warning("Unable to determine index of file. Skipping: '%s'", filename) + continue + orig_frame = filename[:filename.rfind("_")] + extension + self.face_hashes.setdefault(orig_frame, dict())[int(index)] = face["face_hash"] + + def update_hashes(self): + """ Update hashes to alignments """ + logger.info("Updating hashes to alignments...") + updated = 0 + for frame, hashes in self.face_hashes.items(): + if not self.alignments.frame_exists(frame): + logger.warning("Frame not found in alignments file. Skipping: '%s'", frame) + continue + if not self.alignments.frame_has_faces(frame): + logger.warning("Frame does not have faces. Skipping: '%s'", frame) + continue + existing = [face.get("hash", None) + for face in self.alignments.get_faces_in_frame(frame)] + if any(hsh not in existing for hsh in list(hashes.values())): + self.alignments.add_face_hashes(frame, hashes) + updated += 1 + return updated