def process(self):
""" The entry point for the Mask tool from :file:`lib.tools.cli`. Runs the Mask process """
logger.debug("Starting masker process")
updater = getattr(
self,
"_update_{}".format("faces" if self._input_is_faces else "frames"))
if self._update_type != "output":
if self._input_is_faces:
self._faces_saver = ImagesSaver(self._loader.location,
as_bytes=True)
for extractor_output in self._extractor.detected_faces():
self._extractor_input_thread.check_and_raise_error()
updater(extractor_output)
self._extractor_input_thread.join()
if self._counts["update"] != 0:
self._alignments.backup()
self._alignments.save()
if self._input_is_faces:
self._faces_saver.close()
else:
self._extractor_input_thread.join()
self._saver.close()
if self._counts["skip"] != 0:
logger.warning(
"%s face(s) skipped due to not existing in the alignments file",
self._counts["skip"])
if self._update_type != "output":
if self._counts["update"] == 0:
logger.warning("No masks were updated of the %s faces seen",
self._counts["face"])
else:
logger.info("Updated masks for %s faces of %s",
self._counts["update"], self._counts["face"])
logger.debug("Completed masker process")
def process(self):
""" Run the re-extraction from Alignments file process"""
logger.info("[EXTRACT FACES]") # Tidy up cli output
self._check_folder()
if self._is_legacy:
self._legacy_check()
self._saver = ImagesSaver(self._faces_dir, as_bytes=True)
self._export_faces()
def _set_saver(self, arguments):
""" set the saver in a background thread
Parameters
----------
arguments: :class:`argparse.Namespace`
The :mod:`argparse` arguments as passed in from :mod:`tools.py`
Returns
-------
``None`` or :class:`lib.image.ImagesSaver`:
If output is requested, returns a :class:`lib.image.ImagesSaver` otherwise
returns ``None``
"""
if not hasattr(
arguments,
"output") or arguments.output is None or not arguments.output:
if self._update_type == "output":
logger.error(
"Processing set as 'output' but no output folder provided."
)
sys.exit(0)
logger.debug("No output provided. Not creating saver")
return None
output_dir = str(get_folder(arguments.output, make_folder=True))
logger.info("Saving preview masks to: '%s'", output_dir)
saver = ImagesSaver(output_dir)
logger.debug(saver)
return saver
def _background_extract(self, output_folder, progress_queue):
""" Perform the background extraction in a thread so GUI doesn't become unresponsive.
Parameters
----------
output_folder: str
The location to save the output faces to
progress_queue: :class:`queue.Queue`
The queue to place incremental counts to for updating the GUI's progress bar
"""
_io = dict(saver=ImagesSaver(str(get_folder(output_folder)), as_bytes=True),
loader=ImagesLoader(self._input_location, count=self._alignments.frames_count))
for frame_idx, (filename, image) in enumerate(_io["loader"].load()):
logger.trace("Outputting frame: %s: %s", frame_idx, filename)
src_filename = os.path.basename(filename)
frame_name = os.path.splitext(src_filename)[0]
progress_queue.put(1)
for face_idx, face in enumerate(self._frame_faces[frame_idx]):
output = "{}_{}{}".format(frame_name, str(face_idx), ".png")
aligned = AlignedFace(face.landmarks_xy,
image=image,
centering="head",
size=512) # TODO user selectable size
meta = dict(alignments=face.to_png_meta(),
source=dict(alignments_version=self._alignments.version,
original_filename=output,
face_index=face_idx,
source_filename=src_filename,
source_is_video=self._globals.is_video))
b_image = encode_image(aligned.face, ".png", metadata=meta)
_io["saver"].save(output, b_image)
_io["saver"].close()
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"]
# cache detected faces for next run
detected_faces[faces["filename"]] = faces
status_bar.update(1)
if not is_final:
logger.debug("Reloading images")
self._threaded_redirector("reload", detected_faces)
saver.close()
def _background_extract(self, output_folder, progress_queue):
""" Perform the background extraction in a thread so GUI doesn't become unresponsive.
Parameters
----------
output_folder: str
The location to save the output faces to
progress_queue: :class:`queue.Queue`
The queue to place incremental counts to for updating the GUI's progress bar
"""
saver = ImagesSaver(str(get_folder(output_folder)), as_bytes=True)
loader = ImagesLoader(self._input_location,
count=self._alignments.frames_count)
for frame_idx, (filename, image) in enumerate(loader.load()):
logger.trace("Outputting frame: %s: %s", frame_idx, filename)
basename = os.path.basename(filename)
frame_name, extension = os.path.splitext(basename)
final_faces = []
progress_queue.put(1)
for face_idx, face in enumerate(self._frame_faces[frame_idx]):
output = "{}_{}{}".format(frame_name, str(face_idx), extension)
aligned = AlignedFace(face.landmarks_xy,
image=image,
centering="head",
size=512) # TODO user selectable size
face.hash, b_image = encode_image_with_hash(
aligned.face, extension)
saver.save(output, b_image)
final_faces.append(face.to_alignment())
self._alignments.data[basename]["faces"] = final_faces
saver.close()
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
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()
class Mask(): # pylint:disable=too-few-public-methods
""" This tool is part of the Faceswap Tools suite and should be called from
``python tools.py mask`` command.
Faceswap Masks tool. Generate masks from existing alignments files, and output masks
for preview.
Parameters
----------
arguments: :class:`argparse.Namespace`
The :mod:`argparse` arguments as passed in from :mod:`tools.py`
"""
def __init__(self, arguments):
logger.debug("Initializing %s: (arguments: %s",
self.__class__.__name__, arguments)
self._update_type = arguments.processing
self._input_is_faces = arguments.input_type == "faces"
self._mask_type = arguments.masker
self._output = dict(opts=dict(blur_kernel=arguments.blur_kernel,
threshold=arguments.threshold),
type=arguments.output_type,
full_frame=arguments.full_frame,
suffix=self._get_output_suffix(arguments))
self._counts = dict(face=0, skip=0, update=0)
self._check_input(arguments.input)
self._saver = self._set_saver(arguments)
loader = FacesLoader if self._input_is_faces else ImagesLoader
self._loader = loader(arguments.input)
self._faces_saver = None
self._alignments = Alignments(os.path.dirname(arguments.alignments),
filename=os.path.basename(
arguments.alignments))
self._extractor = self._get_extractor(arguments.exclude_gpus)
self._extractor_input_thread = self._feed_extractor()
logger.debug("Initialized %s", self.__class__.__name__)
def _check_input(self, mask_input):
""" Check the input is valid. If it isn't exit with a logged error
Parameters
----------
mask_input: str
Path to the input folder/video
"""
if not os.path.exists(mask_input):
logger.error("Location cannot be found: '%s'", mask_input)
sys.exit(0)
if os.path.isfile(mask_input) and self._input_is_faces:
logger.error(
"Input type 'faces' was selected but input is not a folder: '%s'",
mask_input)
sys.exit(0)
logger.debug("input '%s' is valid", mask_input)
def _set_saver(self, arguments):
""" set the saver in a background thread
Parameters
----------
arguments: :class:`argparse.Namespace`
The :mod:`argparse` arguments as passed in from :mod:`tools.py`
Returns
-------
``None`` or :class:`lib.image.ImagesSaver`:
If output is requested, returns a :class:`lib.image.ImagesSaver` otherwise
returns ``None``
"""
if not hasattr(
arguments,
"output") or arguments.output is None or not arguments.output:
if self._update_type == "output":
logger.error(
"Processing set as 'output' but no output folder provided."
)
sys.exit(0)
logger.debug("No output provided. Not creating saver")
return None
output_dir = str(get_folder(arguments.output, make_folder=True))
logger.info("Saving preview masks to: '%s'", output_dir)
saver = ImagesSaver(output_dir)
logger.debug(saver)
return saver
def _get_extractor(self, exclude_gpus):
""" Obtain a Mask extractor plugin and launch it
Parameters
----------
exclude_gpus: list or ``None``
A list of indices correlating to connected GPUs that Tensorflow should not use. Pass
``None`` to not exclude any GPUs.
Returns
-------
:class:`plugins.extract.pipeline.Extractor`:
The launched Extractor
"""
if self._update_type == "output":
logger.debug(
"Update type `output` selected. Not launching extractor")
return None
logger.debug("masker: %s", self._mask_type)
extractor = Extractor(None,
None,
self._mask_type,
exclude_gpus=exclude_gpus,
image_is_aligned=self._input_is_faces)
extractor.launch()
logger.debug(extractor)
return extractor
def _feed_extractor(self):
""" Feed the input queue to the Extractor from a faces folder or from source frames in a
background thread
Returns
-------
:class:`lib.multithreading.Multithread`:
The thread that is feeding the extractor.
"""
masker_input = getattr(
self,
"_input_{}".format("faces" if self._input_is_faces else "frames"))
logger.debug("masker_input: %s", masker_input)
args = tuple() if self._update_type == "output" else (
self._extractor.input_queue, )
input_thread = MultiThread(masker_input, *args, thread_count=1)
input_thread.start()
logger.debug(input_thread)
return input_thread
def _input_faces(self, *args):
""" Input pre-aligned faces to the Extractor plugin inside a thread
Parameters
----------
args: tuple
The arguments that are to be loaded inside this thread. Contains the queue that the
faces should be put to
"""
log_once = False
logger.debug("args: %s", args)
if self._update_type != "output":
queue = args[0]
for filename, image, metadata in tqdm(self._loader.load(),
total=self._loader.count):
if not metadata: # Legacy faces. Update the headers
if not log_once:
logger.warning(
"Legacy faces discovered. These faces will be updated")
log_once = True
metadata = update_legacy_png_header(filename, self._alignments)
if not metadata:
# Face not found
self._counts["skip"] += 1
logger.warning(
"Legacy face not found in alignments file. This face has not "
"been updated: '%s'", filename)
continue
frame_name = metadata["source"]["source_filename"]
face_index = metadata["source"]["face_index"]
alignment = self._alignments.get_faces_in_frame(frame_name)
if not alignment or face_index > len(alignment) - 1:
self._counts["skip"] += 1
logger.warning(
"Skipping Face not found in alignments file. skipping: '%s'",
filename)
continue
alignment = alignment[face_index]
self._counts["face"] += 1
if self._check_for_missing(frame_name, face_index, alignment):
continue
detected_face = self._get_detected_face(alignment)
if self._update_type == "output":
detected_face.image = image
self._save(frame_name, face_index, detected_face)
else:
media = ExtractMedia(filename,
image,
detected_faces=[detected_face])
setattr(media, "mask_tool_face_info",
metadata["source"]) # TODO formalize
queue.put(media)
self._counts["update"] += 1
if self._update_type != "output":
queue.put("EOF")
def _input_frames(self, *args):
""" Input frames to the Extractor plugin inside a thread
Parameters
----------
args: tuple
The arguments that are to be loaded inside this thread. Contains the queue that the
faces should be put to
"""
logger.debug("args: %s", args)
if self._update_type != "output":
queue = args[0]
for filename, image in tqdm(self._loader.load(),
total=self._loader.count):
frame = os.path.basename(filename)
if not self._alignments.frame_exists(frame):
self._counts["skip"] += 1
logger.warning("Skipping frame not in alignments file: '%s'",
frame)
continue
if not self._alignments.frame_has_faces(frame):
logger.debug("Skipping frame with no faces: '%s'", frame)
continue
faces_in_frame = self._alignments.get_faces_in_frame(frame)
self._counts["face"] += len(faces_in_frame)
# To keep face indexes correct/cover off where only one face in an image is missing a
# mask where there are multiple faces we process all faces again for any frames which
# have missing masks.
if all(
self._check_for_missing(frame, idx, alignment)
for idx, alignment in enumerate(faces_in_frame)):
continue
detected_faces = [
self._get_detected_face(alignment)
for alignment in faces_in_frame
]
if self._update_type == "output":
for idx, detected_face in enumerate(detected_faces):
detected_face.image = image
self._save(frame, idx, detected_face)
else:
self._counts["update"] += len(detected_faces)
queue.put(
ExtractMedia(filename,
image,
detected_faces=detected_faces))
if self._update_type != "output":
queue.put("EOF")
def _check_for_missing(self, frame, idx, alignment):
""" Check if the alignment is missing the requested mask_type
Parameters
----------
frame: str
The frame name in the alignments file
idx: int
The index of the face for this frame in the alignments file
alignment: dict
The alignment for a face
Returns
-------
bool:
``True`` if the update_type is "missing" and the mask does not exist in the alignments
file otherwise ``False``
"""
retval = (self._update_type == "missing"
and alignment.get("mask", None) is not None
and alignment["mask"].get(self._mask_type, None) is not None)
if retval:
logger.debug("Mask pre-exists for face: '%s' - %s", frame, idx)
return retval
def _get_output_suffix(self, arguments):
""" The filename suffix, based on selected output options.
Parameters
----------
arguments: :class:`argparse.Namespace`
The command line arguments for the mask tool
Returns
-------
str:
The suffix to be appended to the output filename
"""
sfx = "{}_mask_preview_".format(self._mask_type)
sfx += "face_" if not arguments.full_frame or self._input_is_faces else "frame_"
sfx += "{}.png".format(arguments.output_type)
return sfx
@staticmethod
def _get_detected_face(alignment):
""" Convert an alignment dict item to a detected_face object
Parameters
----------
alignment: dict
The alignment dict for a face
Returns
-------
:class:`lib.FacesDetect.detected_face`:
The corresponding detected_face object for the alignment
"""
detected_face = DetectedFace()
detected_face.from_alignment(alignment)
return detected_face
def process(self):
""" The entry point for the Mask tool from :file:`lib.tools.cli`. Runs the Mask process """
logger.debug("Starting masker process")
updater = getattr(
self,
"_update_{}".format("faces" if self._input_is_faces else "frames"))
if self._update_type != "output":
if self._input_is_faces:
self._faces_saver = ImagesSaver(self._loader.location,
as_bytes=True)
for extractor_output in self._extractor.detected_faces():
self._extractor_input_thread.check_and_raise_error()
updater(extractor_output)
self._extractor_input_thread.join()
if self._counts["update"] != 0:
self._alignments.backup()
self._alignments.save()
if self._input_is_faces:
self._faces_saver.close()
else:
self._extractor_input_thread.join()
self._saver.close()
if self._counts["skip"] != 0:
logger.warning(
"%s face(s) skipped due to not existing in the alignments file",
self._counts["skip"])
if self._update_type != "output":
if self._counts["update"] == 0:
logger.warning("No masks were updated of the %s faces seen",
self._counts["face"])
else:
logger.info("Updated masks for %s faces of %s",
self._counts["update"], self._counts["face"])
logger.debug("Completed masker process")
def _update_faces(self, extractor_output):
""" Update alignments for the mask if the input type is a faces folder
If an output location has been indicated, then puts the mask preview to the save queue
Parameters
----------
extractor_output: dict
The output from the :class:`plugins.extract.pipeline.Extractor` object
"""
for face in extractor_output.detected_faces:
frame_name = extractor_output.mask_tool_face_info[
"source_filename"]
face_index = extractor_output.mask_tool_face_info["face_index"]
logger.trace("Saving face: (frame: %s, face index: %s)",
frame_name, face_index)
self._alignments.update_face(frame_name, face_index,
face.to_alignment())
metadata = dict(alignments=face.to_png_meta(),
source=extractor_output.mask_tool_face_info)
self._faces_saver.save(
extractor_output.filename,
encode_image(extractor_output.image, ".png",
metadata=metadata))
if self._saver is not None:
face.image = extractor_output.image
self._save(frame_name, face_index, face)
def _update_frames(self, extractor_output):
""" Update alignments for the mask if the input type is a frames folder or video
If an output location has been indicated, then puts the mask preview to the save queue
Parameters
----------
extractor_output: dict
The output from the :class:`plugins.extract.pipeline.Extractor` object
"""
frame = os.path.basename(extractor_output.filename)
for idx, face in enumerate(extractor_output.detected_faces):
self._alignments.update_face(frame, idx, face.to_alignment())
if self._saver is not None:
face.image = extractor_output.image
self._save(frame, idx, face)
def _save(self, frame, idx, detected_face):
""" Build the mask preview image and save
Parameters
----------
frame: str
The frame name in the alignments file
idx: int
The index of the face for this frame in the alignments file
detected_face: `lib.FacesDetect.detected_face`
A detected_face object for a face
"""
filename = os.path.join(
self._saver.location, "{}_{}_{}".format(
os.path.splitext(frame)[0], idx, self._output["suffix"]))
if detected_face.mask is None or detected_face.mask.get(
self._mask_type, None) is None:
logger.warning(
"Mask type '%s' does not exist for frame '%s' index %s. Skipping",
self._mask_type, frame, idx)
return
image = self._create_image(detected_face)
logger.trace("filename: '%s', image_shape: %s", filename, image.shape)
self._saver.save(filename, image)
def _create_image(self, detected_face):
""" Create a mask preview image for saving out to disk
Parameters
----------
detected_face: `lib.FacesDetect.detected_face`
A detected_face object for a face
Returns
numpy.ndarray:
A preview image depending on the output type in one of the following forms:
- Containing 3 sub images: The original face, the masked face and the mask
- The mask only
- The masked face
"""
mask = detected_face.mask[self._mask_type]
mask.set_blur_and_threshold(**self._output["opts"])
if not self._output["full_frame"] or self._input_is_faces:
if self._input_is_faces:
face = AlignedFace(detected_face.landmarks_xy,
image=detected_face.image,
centering="face",
size=detected_face.image.shape[0],
is_aligned=True).face
else:
centering = "legacy" if self._alignments.version == 1.0 else "face"
detected_face.load_aligned(detected_face.image,
centering=centering)
face = detected_face.aligned.face
mask = cv2.resize(detected_face.mask[self._mask_type].mask,
(face.shape[1], face.shape[0]),
interpolation=cv2.INTER_CUBIC)[..., None]
else:
face = np.array(detected_face.image
) # cv2 fails if this comes as imageio.core.Array
mask = mask.get_full_frame_mask(face.shape[1], face.shape[0])
mask = np.expand_dims(mask, -1)
height, width = face.shape[:2]
if self._output["type"] == "combined":
masked = (face.astype("float32") * mask.astype("float32") /
255.).astype("uint8")
mask = np.tile(mask, 3)
for img in (face, masked, mask):
cv2.rectangle(img, (0, 0), (width - 1, height - 1),
(255, 255, 255), 1)
out_image = np.concatenate((face, masked, mask), axis=1)
elif self._output["type"] == "mask":
out_image = mask
elif self._output["type"] == "masked":
out_image = np.concatenate([face, mask], axis=-1)
return out_image
class Extract(): # pylint:disable=too-few-public-methods
""" Re-extract faces from source frames based on Alignment data
Parameters
----------
alignments: :class:`tools.lib_alignments.media.AlignmentData`
The alignments data loaded from an alignments file for this rename job
arguments: :class:`argparse.Namespace`
The :mod:`argparse` arguments as passed in from :mod:`tools.py`
"""
def __init__(self, alignments, arguments):
logger.debug("Initializing %s: (arguments: %s)",
self.__class__.__name__, arguments)
self._arguments = arguments
self._alignments = alignments
self._is_legacy = self._alignments.version == 1.0 # pylint:disable=protected-access
self._mask_pipeline = None
self._faces_dir = arguments.faces_dir
self._frames = Frames(arguments.frames_dir)
self._extracted_faces = ExtractedFaces(self._frames,
self._alignments,
size=arguments.size)
self._saver = None
logger.debug("Initialized %s", self.__class__.__name__)
def process(self):
""" Run the re-extraction from Alignments file process"""
logger.info("[EXTRACT FACES]") # Tidy up cli output
self._check_folder()
if self._is_legacy:
self._legacy_check()
self._saver = ImagesSaver(self._faces_dir, as_bytes=True)
self._export_faces()
def _check_folder(self):
""" Check that the faces folder doesn't pre-exist and create. """
err = None
if not self._faces_dir:
err = "ERROR: Output faces folder not provided."
elif not os.path.isdir(self._faces_dir):
logger.debug("Creating folder: '%s'", self._faces_dir)
os.makedirs(self._faces_dir)
elif os.listdir(self._faces_dir):
err = "ERROR: Output faces folder should be empty: '{}'".format(
self._faces_dir)
if err:
logger.error(err)
sys.exit(0)
logger.verbose("Creating output folder at '%s'", self._faces_dir)
def _legacy_check(self):
""" Check whether the alignments file was created with the legacy extraction method.
If so, force user to re-extract all faces if any options have been specified, otherwise
raise the appropriate warnings and set the legacy options.
"""
if self._arguments.large or self._arguments.extract_every_n != 1:
logger.warning(
"This alignments file was generated with the legacy extraction method."
)
logger.warning(
"You should run this extraction job, but with 'large' deselected and "
"'extract-every-n' set to 1 to update the alignments file.")
logger.warning(
"You can then re-run this extraction job with your chosen options."
)
sys.exit(0)
maskers = ["components", "extended"]
nn_masks = [
mask for mask in list(self._alignments.mask_summary)
if mask not in maskers
]
logtype = logger.warning if nn_masks else logger.info
logtype(
"This alignments file was created with the legacy extraction method and will be "
"updated.")
logtype(
"Faces will be extracted using the new method and landmarks based masks will be "
"regenerated.")
if nn_masks:
logtype(
"However, the NN based masks '%s' will be cropped to the legacy extraction "
"method, so you may want to run the mask tool to regenerate these "
"masks.", "', '".join(nn_masks))
self._mask_pipeline = Extractor(None, None, maskers, multiprocess=True)
self._mask_pipeline.launch()
# Update alignments versioning
self._alignments._version = _VERSION # pylint:disable=protected-access
def _export_faces(self):
""" Export the faces to the output folder. """
extracted_faces = 0
skip_list = self._set_skip_list()
count = self._frames.count if skip_list is None else self._frames.count - len(
skip_list)
for filename, image in tqdm(self._frames.stream(skip_list=skip_list),
total=count,
desc="Saving extracted faces"):
frame_name = os.path.basename(filename)
if not self._alignments.frame_exists(frame_name):
logger.verbose("Skipping '%s' - Alignments not found",
frame_name)
continue
extracted_faces += self._output_faces(frame_name, image)
if self._is_legacy and extracted_faces != 0 and not self._arguments.large:
self._alignments.save()
logger.info("%s face(s) extracted", extracted_faces)
def _set_skip_list(self):
""" Set the indices for frames that should be skipped based on the `extract_every_n`
command line option.
Returns
-------
list or ``None``
A list of indices to be skipped if extract_every_n is not `1` otherwise
returns ``None``
"""
skip_num = self._arguments.extract_every_n
if skip_num == 1:
logger.debug("Not skipping any frames")
return None
skip_list = []
for idx, item in enumerate(self._frames.file_list_sorted):
if idx % skip_num != 0:
logger.trace(
"Adding image '%s' to skip list due to extract_every_n = %s",
item["frame_fullname"], skip_num)
skip_list.append(idx)
logger.debug("Adding skip list: %s", skip_list)
return skip_list
def _output_faces(self, filename, image):
""" For each frame save out the faces
Parameters
----------
filename: str
The filename (without the full path) of the current frame
image: :class:`numpy.ndarray`
The full frame that faces are to be extracted from
Returns
-------
int
The total number of faces that have been extracted
"""
logger.trace("Outputting frame: %s", filename)
face_count = 0
frame_name = os.path.splitext(filename)[0]
faces = self._select_valid_faces(filename, image)
if not faces:
return face_count
if self._is_legacy:
faces = self._process_legacy(filename, image, faces)
for idx, face in enumerate(faces):
output = "{}_{}.png".format(frame_name, str(idx))
meta = dict(alignments=face.to_png_meta(),
source=dict(
alignments_version=self._alignments.version,
original_filename=output,
face_index=idx,
source_filename=filename,
source_is_video=self._frames.is_video))
self._saver.save(
output, encode_image(face.aligned.face, ".png", metadata=meta))
if not self._arguments.large and self._is_legacy:
face.thumbnail = generate_thumbnail(face.aligned.face,
size=96,
quality=60)
self._alignments.data[filename]["faces"][
idx] = face.to_alignment()
face_count += 1
self._saver.close()
return face_count
def _select_valid_faces(self, frame, image):
""" Return the aligned faces from a frame that meet the selection criteria,
Parameters
----------
frame: str
The filename (without the full path) of the current frame
image: :class:`numpy.ndarray`
The full frame that faces are to be extracted from
Returns
-------
list:
List of valid :class:`lib,align.DetectedFace` objects
"""
faces = self._extracted_faces.get_faces_in_frame(frame, image=image)
if not self._arguments.large:
valid_faces = faces
else:
sizes = self._extracted_faces.get_roi_size_for_frame(frame)
valid_faces = [
faces[idx] for idx, size in enumerate(sizes)
if size >= self._extracted_faces.size
]
logger.trace("frame: '%s', total_faces: %s, valid_faces: %s", frame,
len(faces), len(valid_faces))
return valid_faces
def _process_legacy(self, filename, image, detected_faces):
""" Process legacy face extractions to new extraction method.
Updates stored masks to new extract size
Parameters
----------
filename: str
The current frame filename
image: :class:`numpy.ndarray`
The current image the contains the faces
detected_faces: list
list of :class:`lib.align.DetectedFace` objects for the current frame
"""
# Update landmarks based masks for face centering
mask_item = ExtractMedia(filename,
image,
detected_faces=detected_faces)
self._mask_pipeline.input_queue.put(mask_item)
faces = next(self._mask_pipeline.detected_faces()).detected_faces
# Pad and shift Neural Network based masks to face centering
for face in faces:
self._pad_legacy_masks(face)
return faces
@classmethod
def _pad_legacy_masks(cls, detected_face):
""" Recenter legacy Neural Network based masks from legacy centering to face centering
and pad accordingly.
Update the masks back into the detected face objects.
Parameters
----------
detected_face: :class:`lib.align.DetectedFace`
The detected face to update the masks for
"""
offset = detected_face.aligned.pose.offset["face"]
for name, mask in detected_face.mask.items(
): # Re-center mask and pad to face size
if name in ("components", "extended"):
continue
old_mask = mask.mask.astype("float32") / 255.0
size = old_mask.shape[0]
new_size = int(size + (size * _EXTRACT_RATIOS["face"]) / 2)
shift = np.rint(offset *
(size -
(size * _EXTRACT_RATIOS["face"]))).astype("int32")
pos = np.array([(new_size // 2 - size // 2) - shift[1],
(new_size // 2) + (size // 2) - shift[1],
(new_size // 2 - size // 2) - shift[0],
(new_size // 2) + (size // 2) - shift[0]])
bounds = np.array([
max(0, pos[0]),
min(new_size, pos[1]),
max(0, pos[2]),
min(new_size, pos[3])
])
slice_in = [
slice(0 - (pos[0] - bounds[0]), size - (pos[1] - bounds[1])),
slice(0 - (pos[2] - bounds[2]), size - (pos[3] - bounds[3]))
]
slice_out = [
slice(bounds[0], bounds[1]),
slice(bounds[2], bounds[3])
]
new_mask = np.zeros((new_size, new_size, 1), dtype="float32")
new_mask[slice_out[0], slice_out[1], :] = old_mask[slice_in[0],
slice_in[1], :]
mask.replace_mask(new_mask)
# Get the affine matrix from recently generated components mask
# pylint:disable=protected-access
mask._affine_matrix = detected_face.mask[
"components"].affine_matrix
