From d75898f718cc5719630cddbcde174ab11c01338b Mon Sep 17 00:00:00 2001
From: torzdf <36920800+torzdf@users.noreply.github.com>
Date: Thu, 18 Apr 2024 12:45:33 +0100
Subject: [PATCH] Bugfixes for video file alignments storage:   -
 extract/convert: Load images with correct video extension   - Manual tool:
 Cache thumbnails with correct extension   - Mask tool + Preview tool:: Update
 legacy alignment keys for pre-video extension storage

---
 scripts/fsmedia.py                  | 4 ++--
 tools/manual/faceviewer/viewport.py | 2 +-
 tools/manual/thumbnails.py          | 4 ++--
 tools/mask/mask.py                  | 8 ++++++--
 tools/preview/preview.py            | 9 ++++++++-
 5 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/scripts/fsmedia.py b/scripts/fsmedia.py
index 68f503a..36932a4 100644
--- a/scripts/fsmedia.py
+++ b/scripts/fsmedia.py
@@ -287,12 +287,12 @@ class Images():
             A single frame
         """
         logger.debug("Input is video. Capturing frames")
-        vidname = os.path.splitext(os.path.basename(self._args.input_dir))[0]
+        vidname, ext = os.path.splitext(os.path.basename(self._args.input_dir))
         reader = imageio.get_reader(self._args.input_dir, "ffmpeg")  # type:ignore[arg-type]
         for i, frame in enumerate(T.cast(Iterator[np.ndarray], reader)):
             # Convert to BGR for cv2 compatibility
             frame = frame[:, :, ::-1]
-            filename = f"{vidname}_{i + 1:06d}.png"
+            filename = f"{vidname}_{i + 1:06d}{ext}"
             logger.trace("Loading video frame: '%s'", filename)  # type:ignore[attr-defined]
             yield filename, frame
         reader.close()
diff --git a/tools/manual/faceviewer/viewport.py b/tools/manual/faceviewer/viewport.py
index ace7305..94486a8 100644
--- a/tools/manual/faceviewer/viewport.py
+++ b/tools/manual/faceviewer/viewport.py
@@ -481,7 +481,7 @@ class Recycler:
                              asset_type, asset_id)
 
             retval.setdefault(asset_type, []).append(asset_id)
-        logger.info("Got mesh: %s", retval)  # type:ignore[attr-defined]
+        logger.trace("Got mesh: %s", retval)  # type:ignore[attr-defined]
         return retval
 
 
diff --git a/tools/manual/thumbnails.py b/tools/manual/thumbnails.py
index b877346..617d37b 100644
--- a/tools/manual/thumbnails.py
+++ b/tools/manual/thumbnails.py
@@ -209,11 +209,11 @@ class ThumbsCreator():
                      pts_start, pts_end, start_index, segment_count)
         reader = self._get_reader(pts_start, pts_end)
         idx = 0
-        sample_filename = next(fname for fname in self._alignments.data)
+        sample_filename, ext = os.path.splitext(next(fname for fname in self._alignments.data))
         vidname = sample_filename[:sample_filename.rfind("_")]
         for idx, frame in enumerate(reader):
             frame_idx = idx + start_index
-            filename = f"{vidname}_{frame_idx + 1:06d}.png"
+            filename = f"{vidname}_{frame_idx + 1:06d}{ext}"
             self._set_thumbail(filename, frame[..., ::-1], frame_idx)
             if idx == segment_count - 1:
                 # Sometimes extra frames are picked up at the end of a segment, so stop
diff --git a/tools/mask/mask.py b/tools/mask/mask.py
index c3619e0..25e129b 100644
--- a/tools/mask/mask.py
+++ b/tools/mask/mask.py
@@ -149,6 +149,10 @@ class _Mask:
 
         self._loader = Loader(arguments.input, self._input_is_faces)
         self._alignments = self._get_alignments(arguments.alignments, arguments.input)
+
+        if self._loader.is_video and self._alignments is not None:
+            self._alignments.update_legacy_has_source(os.path.basename(self._loader.location))
+
         self._loader.add_alignments(self._alignments)
 
         self._output = Output(arguments, self._alignments, self._loader.file_list)
@@ -206,8 +210,8 @@ class _Mask:
 
         Returns
         -------
-        ``None`` or :class:`lib.align.alignments.Alignments`:
-            If output is requested, returns a :class:`lib.image.ImagesSaver` otherwise
+        ``None`` or :class:`~lib.align.alignments.Alignments`:
+            If output is requested, returns a :class:`~lib.align.alignments.Alignments` otherwise
             returns ``None``
         """
         if alignments:
diff --git a/tools/preview/preview.py b/tools/preview/preview.py
index 4b60ba9..4ac81b3 100644
--- a/tools/preview/preview.py
+++ b/tools/preview/preview.py
@@ -290,9 +290,15 @@ class Samples():
                          "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)
+
         if not self._alignments.have_alignments_file:
             logger.error("Alignments file not found at: '%s'", self._alignments.file)
             sys.exit(1)
+
+        if self._images.is_video:
+            assert isinstance(self._images.input_images, str)
+            self._alignments.update_legacy_has_source(os.path.basename(self._images.input_images))
+
         self._filelist = self._get_filelist()
         self._indices = self._get_indices()
 
@@ -349,7 +355,8 @@ class Samples():
         """
         logger.debug("Filtering file list to frames with faces")
         if isinstance(self._images.input_images, str):
-            filelist = [f"{os.path.splitext(self._images.input_images)[0]}_{frame_no:06d}.png"
+            vid_name, ext = os.path.splitext(self._images.input_images)
+            filelist = [f"{vid_name}_{frame_no:06d}{ext}"
                         for frame_no in range(1, self._images.images_found + 1)]
         else:
             filelist = self._images.input_images