Skip to content

Commit

Permalink
Bugfixes for video file alignments storage:
Browse files Browse the repository at this point in the history
  - 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
torzdf committed Apr 18, 2024
1 parent 3f69d9f commit d75898f
Showing 5 changed files with 19 additions and 8 deletions.
4 changes: 2 additions & 2 deletions scripts/fsmedia.py
Original file line number Diff line number Diff line change
@@ -287,12 +287,12 @@ def _load_video_frames(self) -> Generator[tuple[str, np.ndarray], None, None]:
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()
2 changes: 1 addition & 1 deletion tools/manual/faceviewer/viewport.py
Original file line number Diff line number Diff line change
@@ -481,7 +481,7 @@ def get_mesh(self, face: DetectedFace) -> dict[T.Literal["polygon", "line"], lis
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


4 changes: 2 additions & 2 deletions tools/manual/thumbnails.py
Original file line number Diff line number Diff line change
@@ -209,11 +209,11 @@ def _load_from_video(self,
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
8 changes: 6 additions & 2 deletions tools/mask/mask.py
Original file line number Diff line number Diff line change
@@ -149,6 +149,10 @@ def __init__(self, arguments: Namespace) -> None:

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 @@ def _get_alignments(self, alignments: str | None, input_location: str) -> Alignm
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:
9 changes: 8 additions & 1 deletion tools/preview/preview.py
Original file line number Diff line number Diff line change
@@ -290,9 +290,15 @@ def __init__(self, app: Preview, arguments: Namespace, sample_size: int) -> None
"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 @@ def _get_filelist(self) -> list[str]:
"""
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

0 comments on commit d75898f

Please sign in to comment.