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
This commit is contained in:
torzdf 2024-04-18 12:45:33 +01:00
parent 3f69d9feab
commit d75898f718
5 changed files with 19 additions and 8 deletions

View File

@ -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()

View File

@ -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

View File

@ -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

View File

@ -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:

View File

@ -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