Add update_hashes to alignments tool

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-"):

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 '<frame_name>_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",

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

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