class Extract():
""" The extract process. """
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s", self.__class__.__name__,
arguments)
self.args = arguments
self.output_dir = get_folder(self.args.output_dir)
logger.info("Output Directory: %s", self.args.output_dir)
self.images = Images(self.args)
self.alignments = Alignments(self.args, True, self.images.is_video)
self.plugins = Plugins(self.args)
self.post_process = PostProcess(arguments)
self.verify_output = False
self.save_interval = None
if hasattr(self.args, "save_interval"):
self.save_interval = self.args.save_interval
logger.debug("Initialized %s", self.__class__.__name__)
@property
def skip_num(self):
""" Number of frames to skip if extract_every_n is passed """
return self.args.extract_every_n if hasattr(self.args,
"extract_every_n") else 1
def process(self):
""" Perform the extraction process """
logger.info('Starting, this may take a while...')
Utils.set_verbosity(self.args.loglevel)
# queue_manager.debug_monitor(1)
self.threaded_io("load")
save_thread = self.threaded_io("save")
self.run_extraction()
save_thread.join()
self.alignments.save()
Utils.finalize(self.images.images_found // self.skip_num,
self.alignments.faces_count, self.verify_output)
def threaded_io(self, task, io_args=None):
""" Perform I/O task in a background thread """
logger.debug("Threading task: (Task: '%s')", task)
io_args = tuple() if io_args is None else (io_args, )
if task == "load":
func = self.load_images
elif task == "save":
func = self.save_faces
elif task == "reload":
func = self.reload_images
io_thread = MultiThread(func, *io_args, thread_count=1)
io_thread.start()
return io_thread
def load_images(self):
""" Load the images """
logger.debug("Load Images: Start")
load_queue = queue_manager.get_queue("extract_in")
idx = 0
for filename, image in self.images.load():
idx += 1
if load_queue.shutdown.is_set():
logger.debug("Load Queue: Stop signal received. Terminating")
break
if idx % self.skip_num != 0:
logger.trace("Skipping image '%s' due to extract_every_n = %s",
filename, self.skip_num)
continue
if image is None or not image.any():
logger.warning("Unable to open image. Skipping: '%s'",
filename)
continue
imagename = os.path.basename(filename)
if imagename in self.alignments.data.keys():
logger.trace("Skipping image: '%s'", filename)
continue
item = {"filename": filename, "image": image}
load_queue.put(item)
load_queue.put("EOF")
logger.debug("Load Images: Complete")
def reload_images(self, detected_faces):
""" Reload the images and pair to detected face """
logger.debug("Reload Images: Start. Detected Faces Count: %s",
len(detected_faces))
load_queue = queue_manager.get_queue("detect")
for filename, image in self.images.load():
if load_queue.shutdown.is_set():
logger.debug("Reload Queue: Stop signal received. Terminating")
break
logger.trace("Reloading image: '%s'", filename)
detect_item = detected_faces.pop(filename, None)
if not detect_item:
logger.warning("Couldn't find faces for: %s", filename)
continue
detect_item["image"] = image
load_queue.put(detect_item)
load_queue.put("EOF")
logger.debug("Reload Images: Complete")
@staticmethod
def save_faces():
""" Save the generated faces """
logger.debug("Save Faces: Start")
save_queue = queue_manager.get_queue("extract_out")
while True:
if save_queue.shutdown.is_set():
logger.debug("Save Queue: Stop signal received. Terminating")
break
item = save_queue.get()
if item == "EOF":
break
filename, face = item
logger.trace("Saving face: '%s'", filename)
try:
with open(filename, "wb") as out_file:
out_file.write(face)
except Exception as err: # pylint: disable=broad-except
logger.error("Failed to save image '%s'. Original Error: %s",
filename, err)
continue
logger.debug("Save Faces: Complete")
def run_extraction(self):
""" Run Face Detection """
save_queue = queue_manager.get_queue("extract_out")
to_process = self.process_item_count()
frame_no = 0
size = self.args.size if hasattr(self.args, "size") else 256
align_eyes = self.args.align_eyes if hasattr(self.args,
"align_eyes") else False
if self.plugins.is_parallel:
logger.debug("Using parallel processing")
self.plugins.launch_aligner()
self.plugins.launch_detector()
if not self.plugins.is_parallel:
logger.debug("Using serial processing")
self.run_detection(to_process)
self.plugins.launch_aligner()
for faces in tqdm(self.plugins.detect_faces(extract_pass="******"),
total=to_process,
file=sys.stdout,
desc="Extracting faces"):
filename = faces["filename"]
self.align_face(faces, align_eyes, size, filename)
self.post_process.do_actions(faces)
faces_count = len(faces["detected_faces"])
if faces_count == 0:
logger.verbose("No faces were detected in image: %s",
os.path.basename(filename))
if not self.verify_output and faces_count > 1:
self.verify_output = True
self.output_faces(filename, faces, save_queue)
frame_no += 1
if frame_no == self.save_interval:
self.alignments.save()
frame_no = 0
save_queue.put("EOF")
def process_item_count(self):
""" Return the number of items to be processedd """
processed = sum(
os.path.basename(frame) in self.alignments.data.keys()
for frame in self.images.input_images)
logger.debug("Items already processed: %s", processed)
if processed != 0 and self.args.skip_existing:
logger.info("Skipping previously extracted frames: %s", processed)
if processed != 0 and self.args.skip_faces:
logger.info("Skipping frames with detected faces: %s", processed)
to_process = (self.images.images_found - processed) // self.skip_num
logger.debug("Items to be Processed: %s", to_process)
if to_process == 0:
logger.error("No frames to process. Exiting")
queue_manager.terminate_queues()
exit(0)
return to_process
def run_detection(self, to_process):
""" Run detection only """
self.plugins.launch_detector()
detected_faces = dict()
for detected in tqdm(self.plugins.detect_faces(extract_pass="******"),
total=to_process,
file=sys.stdout,
desc="Detecting faces"):
exception = detected.get("exception", False)
if exception:
break
del detected["image"]
filename = detected["filename"]
detected_faces[filename] = detected
self.threaded_io("reload", detected_faces)
def align_face(self, faces, align_eyes, size, filename):
""" Align the detected face and add the destination file path """
final_faces = list()
image = faces["image"]
landmarks = faces["landmarks"]
detected_faces = faces["detected_faces"]
for idx, face in enumerate(detected_faces):
detected_face = DetectedFace()
detected_face.from_dlib_rect(face, image)
detected_face.landmarksXY = landmarks[idx]
detected_face.load_aligned(image, size=size, align_eyes=align_eyes)
final_faces.append({
"file_location":
self.output_dir / Path(filename).stem,
"face":
detected_face
})
faces["detected_faces"] = final_faces
def output_faces(self, filename, faces, save_queue):
""" Output faces to save thread """
final_faces = list()
for idx, detected_face in enumerate(faces["detected_faces"]):
output_file = detected_face["file_location"]
extension = Path(filename).suffix
out_filename = "{}_{}{}".format(str(output_file), str(idx),
extension)
face = detected_face["face"]
resized_face = face.aligned_face
face.hash, img = hash_encode_image(resized_face, extension)
save_queue.put((out_filename, img))
final_faces.append(face.to_alignment())
self.alignments.data[os.path.basename(filename)] = final_faces
class Extract():
""" The Faceswap Face Extraction Process.
The extraction process is responsible for detecting faces in a series of images/video, aligning
these faces and then generating a mask.
It leverages a series of user selected plugins, chained together using
:mod:`plugins.extract.pipeline`.
The extract process is self contained and should not be referenced by any other scripts, so it
contains no public properties.
Parameters
----------
arguments: argparse.Namespace
The arguments to be passed to the extraction process as generated from Faceswap's command
line arguments
"""
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s", self.__class__.__name__,
arguments)
self._args = arguments
Utils.set_verbosity(self._args.loglevel)
self._output_dir = str(get_folder(self._args.output_dir))
logger.info("Output Directory: %s", self._args.output_dir)
self._images = ImagesLoader(self._args.input_dir,
load_with_hash=False,
fast_count=True)
self._alignments = Alignments(self._args, True, self._images.is_video)
self._existing_count = 0
self._set_skip_list()
self._post_process = PostProcess(arguments)
configfile = self._args.configfile if hasattr(self._args,
"configfile") else None
normalization = None if self._args.normalization == "none" else self._args.normalization
self._extractor = Extractor(self._args.detector,
self._args.aligner,
self._args.masker,
configfile=configfile,
multiprocess=not self._args.singleprocess,
rotate_images=self._args.rotate_images,
min_size=self._args.min_size,
normalize_method=normalization)
self._threads = list()
self._verify_output = False
logger.debug("Initialized %s", self.__class__.__name__)
@property
def _save_interval(self):
""" int: The number of frames to be processed between each saving of the alignments file if
it has been provided, otherwise ``None`` """
if hasattr(self._args, "save_interval"):
return self._args.save_interval
return None
@property
def _skip_num(self):
""" int: Number of frames to skip if extract_every_n has been provided """
return self._args.extract_every_n if hasattr(self._args,
"extract_every_n") else 1
def _set_skip_list(self):
""" Add the skip list to the image loader
Checks against `extract_every_n` and the existence of alignments data (can exist if
`skip_existing` or `skip_existing_faces` has been provided) and compiles a list of frame
indices that should not be processed, providing these to :class:`lib.image.ImagesLoader`.
"""
if self._skip_num == 1 and not self._alignments.data:
logger.debug("No frames to be skipped")
return
skip_list = []
for idx, filename in enumerate(self._images.file_list):
if idx % self._skip_num != 0:
logger.trace(
"Adding image '%s' to skip list due to extract_every_n = %s",
filename, self._skip_num)
skip_list.append(idx)
# Items may be in the alignments file if skip-existing[-faces] is selected
elif os.path.basename(filename) in self._alignments.data:
self._existing_count += 1
logger.trace("Removing image: '%s' due to previously existing",
filename)
skip_list.append(idx)
if self._existing_count != 0:
logger.info(
"Skipping %s frames due to skip_existing/skip_existing_faces.",
self._existing_count)
logger.debug("Adding skip list: %s", skip_list)
self._images.add_skip_list(skip_list)
def process(self):
""" The entry point for triggering the Extraction Process.
Should only be called from :class:`lib.cli.ScriptExecutor`
"""
logger.info('Starting, this may take a while...')
# from lib.queue_manager import queue_manager
# queue_manager.debug_monitor(3)
self._threaded_redirector("load")
self._run_extraction()
for thread in self._threads:
thread.join()
self._alignments.save()
Utils.finalize(self._images.process_count + self._existing_count,
self._alignments.faces_count, self._verify_output)
def _threaded_redirector(self, task, io_args=None):
""" Redirect image input/output tasks to relevant queues in background thread
Parameters
----------
task: str
The name of the task to be put into a background thread
io_args: tuple, optional
Any arguments that need to be provided to the background function
"""
logger.debug("Threading task: (Task: '%s')", task)
io_args = tuple() if io_args is None else (io_args, )
func = getattr(self, "_{}".format(task))
io_thread = MultiThread(func, *io_args, thread_count=1)
io_thread.start()
self._threads.append(io_thread)
def _load(self):
""" Load the images
Loads images from :class:`lib.image.ImagesLoader`, formats them into a dict compatible
with :class:`plugins.extract.Pipeline.Extractor` and passes them into the extraction queue.
"""
logger.debug("Load Images: Start")
load_queue = self._extractor.input_queue
for filename, image in self._images.load():
if load_queue.shutdown.is_set():
logger.debug("Load Queue: Stop signal received. Terminating")
break
item = {"filename": filename, "image": image[..., :3]}
load_queue.put(item)
load_queue.put("EOF")
logger.debug("Load Images: Complete")
def _reload(self, detected_faces):
""" Reload the images and pair to detected face
When the extraction pipeline is running in serial mode, images are reloaded from disk,
paired with their extraction data and passed back into the extraction queue
Parameters
----------
detected_faces: dict
Dictionary of detected_faces with the filename as its key and a list of
:class:`lib.faces_detect.DetectedFace` as the values for pairing with reloaded images.
"""
logger.debug("Reload Images: Start. Detected Faces Count: %s",
len(detected_faces))
load_queue = self._extractor.input_queue
for filename, image in self._images.load():
if load_queue.shutdown.is_set():
logger.debug("Reload Queue: Stop signal received. Terminating")
break
logger.trace("Reloading image: '%s'", filename)
detect_item = detected_faces.pop(filename, None)
if not detect_item:
logger.warning("Couldn't find faces for: %s", filename)
continue
detect_item["image"] = image
load_queue.put(detect_item)
load_queue.put("EOF")
logger.debug("Reload Images: Complete")
def _run_extraction(self):
""" The main Faceswap Extraction process
Receives items from :class:`plugins.extract.Pipeline.Extractor` and either saves out the
faces and data (if on the final pass) or reprocesses data through the pipeline for serial
processing.
"""
size = self._args.size if hasattr(self._args, "size") else 256
saver = ImagesSaver(self._output_dir, as_bytes=True)
exception = False
for phase in range(self._extractor.passes):
if exception:
break
is_final = self._extractor.final_pass
detected_faces = dict()
self._extractor.launch()
self._check_thread_error()
desc = "Running pass {} of {}: {}".format(
phase + 1, self._extractor.passes,
self._extractor.phase.title())
status_bar = tqdm(self._extractor.detected_faces(),
total=self._images.process_count,
file=sys.stdout,
desc=desc)
for idx, faces in enumerate(status_bar):
self._check_thread_error()
exception = faces.get("exception", False)
if exception:
break
if self._extractor.final_pass:
self._output_processing(faces, size)
self._output_faces(saver, faces)
if self._save_interval and (idx +
1) % self._save_interval == 0:
self._alignments.save()
else:
del faces["image"]
status_bar.update(1)
if not is_final:
logger.debug("Reloading images")
self._threaded_redirector("reload", detected_faces)
saver.close()
def _check_thread_error(self):
""" Check if any errors have occurred in the running threads and their errors """
for thread in self._threads:
thread.check_and_raise_error()
def _output_processing(self, faces, size):
""" Prepare faces for output
Loads the aligned face, perform any processing actions and verify the output.
Parameters:
faces: dict
Dictionary output from :class:`plugins.extract.Pipeline.Extractor`
size: int
The size that the aligned face should be created at
"""
for face in faces["detected_faces"]:
face.load_aligned(faces["image"], size=size)
self._post_process.do_actions(faces)
faces_count = len(faces["detected_faces"])
if faces_count == 0:
logger.verbose("No faces were detected in image: %s",
os.path.basename(faces["filename"]))
if not self._verify_output and faces_count > 1:
self._verify_output = True
def _output_faces(self, saver, faces):
""" Output faces to save thread
Set the face filename based on the frame name and put the face to the
:class:`lib.image.ImagesSaver` save queue and add the face information to the alignments
data.
Parameters
----------
saver: lib.images.ImagesSaver
The background saver for saving the image
faces: dict
The output dictionary from :class:`plugins.extract.Pipeline.Extractor`
"""
logger.debug("Save Faces: Start")
final_faces = list()
filename, extension = os.path.splitext(
os.path.basename(faces["filename"]))
for idx, face in enumerate(faces["detected_faces"]):
output_filename = "{}_{}{}".format(filename, str(idx), extension)
face.hash, image = encode_image_with_hash(face.aligned_face,
extension)
saver.save(output_filename, image)
final_faces.append(face.to_alignment())
self._alignments.data[os.path.basename(
faces["filename"])] = final_faces
class Extract():
""" The extract process. """
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s", self.__class__.__name__, arguments)
self.args = arguments
Utils.set_verbosity(self.args.loglevel)
self.output_dir = get_folder(self.args.output_dir)
logger.info("Output Directory: %s", self.args.output_dir)
self.images = Images(self.args)
self.alignments = Alignments(self.args, True, self.images.is_video)
self.post_process = PostProcess(arguments)
self.extractor = Extractor(self.args.detector,
self.args.aligner,
self.args.loglevel,
self.args.multiprocess,
self.args.rotate_images,
self.args.min_size)
self.save_queue = queue_manager.get_queue("extract_save")
self.verify_output = False
self.save_interval = None
if hasattr(self.args, "save_interval"):
self.save_interval = self.args.save_interval
logger.debug("Initialized %s", self.__class__.__name__)
@property
def skip_num(self):
""" Number of frames to skip if extract_every_n is passed """
return self.args.extract_every_n if hasattr(self.args, "extract_every_n") else 1
def process(self):
""" Perform the extraction process """
logger.info('Starting, this may take a while...')
# queue_manager.debug_monitor(3)
self.threaded_io("load")
save_thread = self.threaded_io("save")
self.run_extraction()
save_thread.join()
self.alignments.save()
Utils.finalize(self.images.images_found // self.skip_num,
self.alignments.faces_count,
self.verify_output)
def threaded_io(self, task, io_args=None):
""" Perform I/O task in a background thread """
logger.debug("Threading task: (Task: '%s')", task)
io_args = tuple() if io_args is None else (io_args, )
if task == "load":
func = self.load_images
elif task == "save":
func = self.save_faces
elif task == "reload":
func = self.reload_images
io_thread = MultiThread(func, *io_args, thread_count=1)
io_thread.start()
return io_thread
def load_images(self):
""" Load the images """
logger.debug("Load Images: Start")
load_queue = self.extractor.input_queue
idx = 0
for filename, image in self.images.load():
idx += 1
if load_queue.shutdown.is_set():
logger.debug("Load Queue: Stop signal received. Terminating")
break
if idx % self.skip_num != 0:
logger.trace("Skipping image '%s' due to extract_every_n = %s",
filename, self.skip_num)
continue
if image is None or not image.any():
logger.warning("Unable to open image. Skipping: '%s'", filename)
continue
imagename = os.path.basename(filename)
if imagename in self.alignments.data.keys():
logger.trace("Skipping image: '%s'", filename)
continue
item = {"filename": filename,
"image": image}
load_queue.put(item)
load_queue.put("EOF")
logger.debug("Load Images: Complete")
def reload_images(self, detected_faces):
""" Reload the images and pair to detected face """
logger.debug("Reload Images: Start. Detected Faces Count: %s", len(detected_faces))
load_queue = self.extractor.input_queue
for filename, image in self.images.load():
if load_queue.shutdown.is_set():
logger.debug("Reload Queue: Stop signal received. Terminating")
break
logger.trace("Reloading image: '%s'", filename)
detect_item = detected_faces.pop(filename, None)
if not detect_item:
logger.warning("Couldn't find faces for: %s", filename)
continue
detect_item["image"] = image
load_queue.put(detect_item)
load_queue.put("EOF")
logger.debug("Reload Images: Complete")
def save_faces(self):
""" Save the generated faces """
logger.debug("Save Faces: Start")
while True:
if self.save_queue.shutdown.is_set():
logger.debug("Save Queue: Stop signal received. Terminating")
break
item = self.save_queue.get()
logger.trace(item)
if item == "EOF":
break
filename, face = item
logger.trace("Saving face: '%s'", filename)
try:
with open(filename, "wb") as out_file:
out_file.write(face)
except Exception as err: # pylint: disable=broad-except
logger.error("Failed to save image '%s'. Original Error: %s", filename, err)
continue
logger.debug("Save Faces: Complete")
def process_item_count(self):
""" Return the number of items to be processedd """
processed = sum(os.path.basename(frame) in self.alignments.data.keys()
for frame in self.images.input_images)
logger.debug("Items already processed: %s", processed)
if processed != 0 and self.args.skip_existing:
logger.info("Skipping previously extracted frames: %s", processed)
if processed != 0 and self.args.skip_faces:
logger.info("Skipping frames with detected faces: %s", processed)
to_process = (self.images.images_found - processed) // self.skip_num
logger.debug("Items to be Processed: %s", to_process)
if to_process == 0:
logger.error("No frames to process. Exiting")
queue_manager.terminate_queues()
exit(0)
return to_process
def run_extraction(self):
""" Run Face Detection """
to_process = self.process_item_count()
size = self.args.size if hasattr(self.args, "size") else 256
align_eyes = self.args.align_eyes if hasattr(self.args, "align_eyes") else False
exception = False
for phase in range(self.extractor.passes):
if exception:
break
is_final = self.extractor.final_pass
detected_faces = dict()
self.extractor.launch()
for idx, faces in enumerate(tqdm(self.extractor.detected_faces(),
total=to_process,
file=sys.stdout,
desc="Running pass {} of {}: {}".format(
phase + 1,
self.extractor.passes,
self.extractor.phase.title()))):
exception = faces.get("exception", False)
if exception:
break
filename = faces["filename"]
if self.extractor.final_pass:
self.output_processing(faces, align_eyes, size, filename)
self.output_faces(filename, faces)
if self.save_interval and idx + 1 % self.save_interval == 0:
self.alignments.save()
else:
del faces["image"]
detected_faces[filename] = faces
if is_final:
logger.debug("Putting EOF to save")
self.save_queue.put("EOF")
else:
logger.debug("Reloading images")
self.threaded_io("reload", detected_faces)
def output_processing(self, faces, align_eyes, size, filename):
""" Prepare faces for output """
self.align_face(faces, align_eyes, size, filename)
self.post_process.do_actions(faces)
faces_count = len(faces["detected_faces"])
if faces_count == 0:
logger.verbose("No faces were detected in image: %s",
os.path.basename(filename))
if not self.verify_output and faces_count > 1:
self.verify_output = True
def align_face(self, faces, align_eyes, size, filename):
""" Align the detected face and add the destination file path """
final_faces = list()
image = faces["image"]
landmarks = faces["landmarks"]
detected_faces = faces["detected_faces"]
for idx, face in enumerate(detected_faces):
detected_face = DetectedFace()
detected_face.from_bounding_box(face, image)
detected_face.landmarksXY = landmarks[idx]
detected_face.load_aligned(image, size=size, align_eyes=align_eyes)
final_faces.append({"file_location": self.output_dir / Path(filename).stem,
"face": detected_face})
faces["detected_faces"] = final_faces
def output_faces(self, filename, faces):
""" Output faces to save thread """
final_faces = list()
for idx, detected_face in enumerate(faces["detected_faces"]):
output_file = detected_face["file_location"]
extension = Path(filename).suffix
out_filename = "{}_{}{}".format(str(output_file), str(idx), extension)
face = detected_face["face"]
resized_face = face.aligned_face
face.hash, img = hash_encode_image(resized_face, extension)
self.save_queue.put((out_filename, img))
final_faces.append(face.to_alignment())
self.alignments.data[os.path.basename(filename)] = final_faces
class Extract():
""" The extract process. """
def __init__(self, arguments):
self.args = arguments
self.output_dir = get_folder(self.args.output_dir)
print("Output Directory: {}".format(self.args.output_dir))
self.images = Images(self.args)
self.alignments = Alignments(self.args, True)
self.plugins = Plugins(self.args)
self.post_process = PostProcess(arguments)
self.export_face = True
self.verify_output = False
self.save_interval = None
if hasattr(self.args, "save_interval"):
self.save_interval = self.args.save_interval
def process(self):
""" Perform the extraction process """
print('Starting, this may take a while...')
Utils.set_verbosity(self.args.verbose)
# queue_manager.debug_monitor(1)
self.threaded_io("load")
save_thread = self.threaded_io("save")
self.run_extraction(save_thread)
self.alignments.save()
Utils.finalize(self.images.images_found, self.alignments.faces_count(),
self.verify_output)
def threaded_io(self, task, io_args=None):
""" Load images in a background thread """
io_args = tuple() if io_args is None else (io_args, )
if task == "load":
func = self.load_images
elif task == "save":
func = self.save_faces
elif task == "reload":
func = self.reload_images
io_thread = MultiThread(thread_count=1)
io_thread.in_thread(func, *io_args)
return io_thread
def load_images(self):
""" Load the images """
load_queue = queue_manager.get_queue("load")
for filename, image in self.images.load():
imagename = os.path.basename(filename)
if imagename in self.alignments.data.keys():
continue
load_queue.put((filename, image))
load_queue.put("EOF")
def reload_images(self, detected_faces):
""" Reload the images and pair to detected face """
load_queue = queue_manager.get_queue("detect")
for filename, image in self.images.load():
detect_item = detected_faces.pop(filename, None)
if not detect_item:
continue
detect_item["image"] = image
load_queue.put(detect_item)
load_queue.put("EOF")
def save_faces(self):
""" Save the generated faces """
if not self.export_face:
return
save_queue = queue_manager.get_queue("save")
while True:
item = save_queue.get()
if item == "EOF":
break
filename, output_file, resized_face, idx = item
out_filename = "{}_{}{}".format(str(output_file), str(idx),
Path(filename).suffix)
# pylint: disable=no-member
cv2.imwrite(out_filename, resized_face)
def run_extraction(self, save_thread):
""" Run Face Detection """
to_process = self.process_item_count()
frame_no = 0
if self.plugins.is_parallel:
self.plugins.launch_aligner()
self.plugins.launch_detector()
if not self.plugins.is_parallel:
self.run_detection(to_process)
self.plugins.launch_aligner()
for faces in tqdm(self.plugins.detect_faces(extract_pass="******"),
total=to_process,
file=sys.stdout,
desc="Extracting faces"):
exception = faces.get("exception", False)
if exception:
exit(1)
filename = faces["filename"]
faces["output_file"] = self.output_dir / Path(filename).stem
self.post_process.do_actions(faces)
faces_count = len(faces["detected_faces"])
if self.args.verbose and faces_count == 0:
print("Warning: No faces were detected in image: "
"{}".format(os.path.basename(filename)))
if not self.verify_output and faces_count > 1:
self.verify_output = True
self.process_faces(filename, faces)
frame_no += 1
if frame_no == self.save_interval:
self.alignments.save()
frame_no = 0
if self.export_face:
queue_manager.get_queue("save").put("EOF")
save_thread.join_threads()
def process_item_count(self):
""" Return the number of items to be processedd """
processed = sum(
os.path.basename(frame) in self.alignments.data.keys()
for frame in self.images.input_images)
if processed != 0 and self.args.skip_existing:
print("Skipping {} previously extracted frames".format(processed))
if processed != 0 and self.args.skip_faces:
print("Skipping {} frames with detected faces".format(processed))
to_process = self.images.images_found - processed
if to_process == 0:
print("No frames to process. Exiting")
queue_manager.terminate_queues()
exit(0)
return to_process
def run_detection(self, to_process):
""" Run detection only """
self.plugins.launch_detector()
detected_faces = dict()
for detected in tqdm(self.plugins.detect_faces(extract_pass="******"),
total=to_process,
file=sys.stdout,
desc="Detecting faces"):
exception = detected.get("exception", False)
if exception:
break
del detected["image"]
filename = detected["filename"]
detected_faces[filename] = detected
self.threaded_io("reload", detected_faces)
def process_faces(self, filename, faces):
""" Perform processing on found faces """
final_faces = list()
save_queue = queue_manager.get_queue("save")
filename = faces["filename"]
output_file = faces["output_file"]
resized_faces = faces["resized_faces"]
for idx, face in enumerate(faces["detected_faces"]):
if self.export_face:
save_queue.put(
(filename, output_file, resized_faces[idx], idx))
final_faces.append(face.to_alignment())
self.alignments.data[os.path.basename(filename)] = final_faces
class Extract(): # pylint:disable=too-few-public-methods
""" The Faceswap Face Extraction Process.
The extraction process is responsible for detecting faces in a series of images/video, aligning
these faces and then generating a mask.
It leverages a series of user selected plugins, chained together using
:mod:`plugins.extract.pipeline`.
The extract process is self contained and should not be referenced by any other scripts, so it
contains no public properties.
Parameters
----------
arguments: :class:`argparse.Namespace`
The arguments to be passed to the extraction process as generated from Faceswap's command
line arguments
"""
def __init__(self, arguments: argparse.Namespace) -> None:
logger.debug("Initializing %s: (args: %s", self.__class__.__name__,
arguments)
self._args = arguments
self._output_dir = None if self._args.skip_saving_faces else get_folder(
self._args.output_dir)
logger.info("Output Directory: %s", self._args.output_dir)
self._images = ImagesLoader(self._args.input_dir, fast_count=True)
self._alignments = Alignments(self._args, True, self._images.is_video)
self._existing_count = 0
self._set_skip_list()
self._post_process = PostProcess(arguments)
configfile = self._args.configfile if hasattr(self._args,
"configfile") else None
normalization = None if self._args.normalization == "none" else self._args.normalization
maskers = ["components", "extended"]
maskers += self._args.masker if self._args.masker else []
self._extractor = Extractor(self._args.detector,
self._args.aligner,
maskers,
configfile=configfile,
multiprocess=not self._args.singleprocess,
exclude_gpus=self._args.exclude_gpus,
rotate_images=self._args.rotate_images,
min_size=self._args.min_size,
normalize_method=normalization,
re_feed=self._args.re_feed)
self._threads = []
self._verify_output = False
logger.debug("Initialized %s", self.__class__.__name__)
@property
def _save_interval(self) -> Optional[int]:
""" int: The number of frames to be processed between each saving of the alignments file if
it has been provided, otherwise ``None`` """
if hasattr(self._args, "save_interval"):
return self._args.save_interval
return None
@property
def _skip_num(self) -> int:
""" int: Number of frames to skip if extract_every_n has been provided """
return self._args.extract_every_n if hasattr(self._args,
"extract_every_n") else 1
def _set_skip_list(self) -> None:
""" Add the skip list to the image loader
Checks against `extract_every_n` and the existence of alignments data (can exist if
`skip_existing` or `skip_existing_faces` has been provided) and compiles a list of frame
indices that should not be processed, providing these to :class:`lib.image.ImagesLoader`.
"""
if self._skip_num == 1 and not self._alignments.data:
logger.debug("No frames to be skipped")
return
skip_list = []
for idx, filename in enumerate(self._images.file_list):
if idx % self._skip_num != 0:
logger.trace(
"Adding image '%s' to skip list due to extract_every_n = %s",
filename, self._skip_num)
skip_list.append(idx)
# Items may be in the alignments file if skip-existing[-faces] is selected
elif os.path.basename(filename) in self._alignments.data:
self._existing_count += 1
logger.trace("Removing image: '%s' due to previously existing",
filename)
skip_list.append(idx)
if self._existing_count != 0:
logger.info(
"Skipping %s frames due to skip_existing/skip_existing_faces.",
self._existing_count)
logger.debug("Adding skip list: %s", skip_list)
self._images.add_skip_list(skip_list)
def process(self) -> None:
""" The entry point for triggering the Extraction Process.
Should only be called from :class:`lib.cli.launcher.ScriptExecutor`
"""
logger.info('Starting, this may take a while...')
# from lib.queue_manager import queue_manager ; queue_manager.debug_monitor(3)
self._threaded_redirector("load")
self._run_extraction()
for thread in self._threads:
thread.join()
self._alignments.save()
finalize(self._images.process_count + self._existing_count,
self._alignments.faces_count, self._verify_output)
def _threaded_redirector(self,
task: str,
io_args: Optional[tuple] = None) -> None:
""" Redirect image input/output tasks to relevant queues in background thread
Parameters
----------
task: str
The name of the task to be put into a background thread
io_args: tuple, optional
Any arguments that need to be provided to the background function
"""
logger.debug("Threading task: (Task: '%s')", task)
io_args = tuple() if io_args is None else (io_args, )
func = getattr(self, f"_{task}")
io_thread = MultiThread(func, *io_args, thread_count=1)
io_thread.start()
self._threads.append(io_thread)
def _load(self) -> None:
""" Load the images
Loads images from :class:`lib.image.ImagesLoader`, formats them into a dict compatible
with :class:`plugins.extract.Pipeline.Extractor` and passes them into the extraction queue.
"""
logger.debug("Load Images: Start")
load_queue = self._extractor.input_queue
for filename, image in self._images.load():
if load_queue.shutdown.is_set():
logger.debug("Load Queue: Stop signal received. Terminating")
break
item = ExtractMedia(filename, image[..., :3])
load_queue.put(item)
load_queue.put("EOF")
logger.debug("Load Images: Complete")
def _reload(self, detected_faces: dict[str, ExtractMedia]) -> None:
""" Reload the images and pair to detected face
When the extraction pipeline is running in serial mode, images are reloaded from disk,
paired with their extraction data and passed back into the extraction queue
Parameters
----------
detected_faces: dict
Dictionary of :class:`plugins.extract.pipeline.ExtractMedia` with the filename as the
key for repopulating the image attribute.
"""
logger.debug("Reload Images: Start. Detected Faces Count: %s",
len(detected_faces))
load_queue = self._extractor.input_queue
for filename, image in self._images.load():
if load_queue.shutdown.is_set():
logger.debug("Reload Queue: Stop signal received. Terminating")
break
logger.trace("Reloading image: '%s'", filename)
extract_media = detected_faces.pop(filename, None)
if not extract_media:
logger.warning("Couldn't find faces for: %s", filename)
continue
extract_media.set_image(image)
load_queue.put(extract_media)
load_queue.put("EOF")
logger.debug("Reload Images: Complete")
def _run_extraction(self) -> None:
""" The main Faceswap Extraction process
Receives items from :class:`plugins.extract.Pipeline.Extractor` and either saves out the
faces and data (if on the final pass) or reprocesses data through the pipeline for serial
processing.
"""
size = self._args.size if hasattr(self._args, "size") else 256
saver = None if self._args.skip_saving_faces else ImagesSaver(
self._output_dir, as_bytes=True)
exception = False
for phase in range(self._extractor.passes):
if exception:
break
is_final = self._extractor.final_pass
detected_faces = {}
self._extractor.launch()
self._check_thread_error()
ph_desc = "Extraction" if self._extractor.passes == 1 else self._extractor.phase_text
desc = f"Running pass {phase + 1} of {self._extractor.passes}: {ph_desc}"
for idx, extract_media in enumerate(
tqdm(self._extractor.detected_faces(),
total=self._images.process_count,
file=sys.stdout,
desc=desc)):
self._check_thread_error()
if is_final:
self._output_processing(extract_media, size)
self._output_faces(saver, extract_media)
if self._save_interval and (idx +
1) % self._save_interval == 0:
self._alignments.save()
else:
extract_media.remove_image()
# cache extract_media for next run
detected_faces[extract_media.filename] = extract_media
if not is_final:
logger.debug("Reloading images")
self._threaded_redirector("reload", detected_faces)
if not self._args.skip_saving_faces:
saver.close()
def _check_thread_error(self) -> None:
""" Check if any errors have occurred in the running threads and their errors """
for thread in self._threads:
thread.check_and_raise_error()
def _output_processing(self, extract_media: ExtractMedia,
size: int) -> None:
""" Prepare faces for output
Loads the aligned face, generate the thumbnail, perform any processing actions and verify
the output.
Parameters
----------
extract_media: :class:`plugins.extract.pipeline.ExtractMedia`
Output from :class:`plugins.extract.pipeline.Extractor`
size: int
The size that the aligned face should be created at
"""
for face in extract_media.detected_faces:
face.load_aligned(extract_media.image, size=size, centering="head")
face.thumbnail = generate_thumbnail(face.aligned.face,
size=96,
quality=60)
self._post_process.do_actions(extract_media)
extract_media.remove_image()
faces_count = len(extract_media.detected_faces)
if faces_count == 0:
logger.verbose("No faces were detected in image: %s",
os.path.basename(extract_media.filename))
if not self._verify_output and faces_count > 1:
self._verify_output = True
def _output_faces(self, saver: ImagesSaver,
extract_media: ExtractMedia) -> None:
""" Output faces to save thread
Set the face filename based on the frame name and put the face to the
:class:`~lib.image.ImagesSaver` save queue and add the face information to the alignments
data.
Parameters
----------
saver: lib.images.ImagesSaver
The background saver for saving the image
extract_media: :class:`~plugins.extract.pipeline.ExtractMedia`
The output from :class:`~plugins.extract.Pipeline.Extractor`
"""
logger.trace("Outputting faces for %s", extract_media.filename)
final_faces = []
filename = os.path.splitext(os.path.basename(
extract_media.filename))[0]
extension = ".png"
for idx, face in enumerate(extract_media.detected_faces):
output_filename = f"{filename}_{idx}{extension}"
meta = dict(alignments=face.to_png_meta(),
source=dict(
alignments_version=self._alignments.version,
original_filename=output_filename,
face_index=idx,
source_filename=os.path.basename(
extract_media.filename),
source_is_video=self._images.is_video,
source_frame_dims=extract_media.image_size))
image = encode_image(face.aligned.face, extension, metadata=meta)
if not self._args.skip_saving_faces:
saver.save(output_filename, image)
final_faces.append(face.to_alignment())
self._alignments.data[os.path.basename(
extract_media.filename)] = dict(faces=final_faces)
del extract_media
