mirror of
https://github.com/zebrajr/faceswap.git
synced 2026-01-15 12:15:15 +00:00
bugfix: Correctly generate thumbnails for alignments when frames are missing
This commit is contained in:
21
lib/image.py
21
lib/image.py
@@ -1087,21 +1087,22 @@ class ImagesLoader(ImageIO):
|
|||||||
self._fps = self._get_fps()
|
self._fps = self._get_fps()
|
||||||
|
|
||||||
self._count = None
|
self._count = None
|
||||||
self._file_list = None
|
self._file_list: list[str] = []
|
||||||
self._get_count_and_filelist(fast_count, count)
|
self._get_count_and_filelist(fast_count, count)
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def count(self):
|
def count(self) -> int:
|
||||||
""" int: The number of images or video frames in the source location. This count includes
|
""" int: The number of images or video frames in the source location. This count includes
|
||||||
any files that will ultimately be skipped if a :attr:`skip_list` has been provided. See
|
any files that will ultimately be skipped if a :attr:`skip_list` has been provided. See
|
||||||
also: :attr:`process_count`"""
|
also: :attr:`process_count`"""
|
||||||
|
assert self._count is not None
|
||||||
return self._count
|
return self._count
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def process_count(self):
|
def process_count(self) -> int:
|
||||||
""" int: The number of images or video frames to be processed (IE the total count less
|
""" int: The number of images or video frames to be processed (IE the total count less
|
||||||
items that are to be skipped from the :attr:`skip_list`)"""
|
items that are to be skipped from the :attr:`skip_list`)"""
|
||||||
return self._count - len(self._skip_list)
|
return self.count - len(self._skip_list)
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def is_video(self):
|
def is_video(self):
|
||||||
@@ -1115,10 +1116,10 @@ class ImagesLoader(ImageIO):
|
|||||||
return self._fps
|
return self._fps
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def file_list(self):
|
def file_list(self) -> list[str]:
|
||||||
""" list: A full list of files in the source location. This includes any files that will
|
""" list[str]: A full list of files in the source location. This includes any files that
|
||||||
ultimately be skipped if a :attr:`skip_list` has been provided. If the input is a video
|
will ultimately be skipped if a :attr:`skip_list` has been provided. If the input is a
|
||||||
then this is a list of dummy filenames as corresponding to an alignments file """
|
video then this is a list of dummy filenames as corresponding to an alignments file """
|
||||||
return self._file_list
|
return self._file_list
|
||||||
|
|
||||||
def add_skip_list(self, skip_list):
|
def add_skip_list(self, skip_list):
|
||||||
@@ -1436,7 +1437,7 @@ class SingleFrameLoader(ImagesLoader):
|
|||||||
self._video_meta_data = video_meta_data
|
self._video_meta_data = video_meta_data
|
||||||
super()._get_count_and_filelist(fast_count, count)
|
super()._get_count_and_filelist(fast_count, count)
|
||||||
|
|
||||||
def image_from_index(self, index):
|
def image_from_index(self, index: int) -> tuple[str, np.ndarray]:
|
||||||
""" Return a single image from :attr:`file_list` for the given index.
|
""" Return a single image from :attr:`file_list` for the given index.
|
||||||
|
|
||||||
Parameters
|
Parameters
|
||||||
@@ -1468,7 +1469,7 @@ class SingleFrameLoader(ImagesLoader):
|
|||||||
else:
|
else:
|
||||||
file_list = [f for idx, f in enumerate(self._file_list)
|
file_list = [f for idx, f in enumerate(self._file_list)
|
||||||
if idx not in self._skip_list] if self._skip_list else self._file_list
|
if idx not in self._skip_list] if self._skip_list else self._file_list
|
||||||
|
|
||||||
filename = file_list[index]
|
filename = file_list[index]
|
||||||
image = read_image(filename, raise_error=True)
|
image = read_image(filename, raise_error=True)
|
||||||
filename = os.path.basename(filename)
|
filename = os.path.basename(filename)
|
||||||
|
|||||||
@@ -170,14 +170,18 @@ class ThumbsCreator():
|
|||||||
thread for some speed up.
|
thread for some speed up.
|
||||||
"""
|
"""
|
||||||
reader = SingleFrameLoader(self._location)
|
reader = SingleFrameLoader(self._location)
|
||||||
num_threads = min(reader.count, self._num_threads)
|
skip_list = [idx for idx, f in enumerate(reader.file_list)
|
||||||
frame_split = reader.count // self._num_threads
|
if os.path.basename(f) not in self._alignments.data]
|
||||||
|
if skip_list:
|
||||||
|
reader.add_skip_list(skip_list)
|
||||||
|
num_threads = min(reader.process_count, self._num_threads)
|
||||||
|
frame_split = reader.process_count // self._num_threads
|
||||||
logger.debug("total images: %s, num_threads: %s, frames_per_thread: %s",
|
logger.debug("total images: %s, num_threads: %s, frames_per_thread: %s",
|
||||||
reader.count, num_threads, frame_split)
|
reader.process_count, num_threads, frame_split)
|
||||||
for idx in range(num_threads):
|
for idx in range(num_threads):
|
||||||
is_final = idx == num_threads - 1
|
is_final = idx == num_threads - 1
|
||||||
start_idx = idx * frame_split
|
start_idx = idx * frame_split
|
||||||
end_idx = reader.count if is_final else start_idx + frame_split
|
end_idx = reader.process_count if is_final else start_idx + frame_split
|
||||||
thread = MultiThread(self._load_from_folder, reader, start_idx, end_idx)
|
thread = MultiThread(self._load_from_folder, reader, start_idx, end_idx)
|
||||||
thread.start()
|
thread.start()
|
||||||
self._threads.append(thread)
|
self._threads.append(thread)
|
||||||
|
|||||||
Reference in New Issue
Block a user