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 __init__(self, arguments):
logger.debug("Initializing %s: (args: %s", self.__class__.__name__,
arguments)
self._args = arguments
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__)
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_opts = dict(blur_kernel=arguments.blur_kernel,
threshold=arguments.threshold)
self._output_type = arguments.output_type
self._output_full_frame = arguments.full_frame
self._output_suffix = self._get_output_suffix()
self._face_count = 0
self._skip_count = 0
self._update_count = 0
self._check_input(arguments.input)
self._saver = self._set_saver(arguments)
self._loader = ImagesLoader(arguments.input,
load_with_hash=self._input_is_faces)
self._alignments = Alignments(os.path.dirname(arguments.alignments),
filename=os.path.basename(
arguments.alignments))
self._extractor = self._get_extractor()
self._extractor_input_thread = self._feed_extractor()
logger.debug("Initialized %s", self.__class__.__name__)
def __init__(self, arguments):
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 = list()
self._verify_output = False
logger.debug("Initialized %s", self.__class__.__name__)
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s)", self.__class__.__name__,
arguments)
self._args = arguments
self._patch_threads = None
self._images = ImagesLoader(self._args.input_dir, fast_count=True)
self._alignments = Alignments(self._args, False, self._images.is_video)
self._opts = OptionalActions(self._args, self._images.file_list,
self._alignments)
self._add_queues()
self._disk_io = DiskIO(self._alignments, self._images, arguments)
self._predictor = Predict(self._disk_io.load_queue, self._queue_size,
arguments)
self._validate()
get_folder(self._args.output_dir)
configfile = self._args.configfile if hasattr(self._args,
"configfile") else None
self._converter = Converter(self._predictor.output_size,
self._predictor.coverage_ratio,
self._disk_io.draw_transparent,
self._disk_io.pre_encode,
arguments,
configfile=configfile)
logger.debug("Initialized %s", self.__class__.__name__)
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 sort_face(self):
""" Sort by identity similarity """
logger.info("Sorting by identity similarity...")
# TODO This should be set in init
self._loader = ImagesLoader(self.args.input_dir)
filenames = []
preds = np.empty((self._loader.count, 512), dtype="float32")
for idx, (filename, image) in enumerate(tqdm(self._loader.load(),
desc="Classifying Faces...",
total=self._loader.count)):
filenames.append(filename)
preds[idx] = self.vgg_face.predict(image)
logger.info("Sorting by ward linkage...")
indices = self.vgg_face.sorted_similarity(preds, method="ward")
img_list = np.array(filenames)[indices]
return img_list
def stream(self, skip_list=None):
""" Load the images in :attr:`folder` in the order they are received from
:class:`lib.image.ImagesLoader` in a background thread.
Parameters
----------
skip_list: list, optional
A list of frame indices that should not be loaded. Pass ``None`` if all images should
be loaded. Default: ``None``
Yields
------
str
The filename of the image that is being returned
numpy.ndarray
The image that has been loaded from disk
"""
loader = ImagesLoader(self.folder, queue_size=32)
if skip_list is not None:
loader.add_skip_list(skip_list)
for filename, image in loader.load():
yield filename, image
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s)", self.__class__.__name__,
arguments)
self._args = arguments
self._patch_threads = None
self._images = ImagesLoader(self._args.input_dir, fast_count=True)
self._alignments = Alignments(self._args, False, self._images.is_video)
if self._alignments.version == 1.0:
logger.error(
"The alignments file format has been updated since the given alignments "
"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)
self._opts = OptionalActions(self._args, self._images.file_list,
self._alignments)
self._add_queues()
self._disk_io = DiskIO(self._alignments, self._images, arguments)
self._predictor = Predict(self._disk_io.load_queue, self._queue_size,
arguments)
self._validate()
get_folder(self._args.output_dir)
configfile = self._args.configfile if hasattr(self._args,
"configfile") else None
self._converter = Converter(self._predictor.output_size,
self._predictor.coverage_ratio,
self._predictor.centering,
self._disk_io.draw_transparent,
self._disk_io.pre_encode,
arguments,
configfile=configfile)
logger.debug("Initialized %s", self.__class__.__name__)
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: 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
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()
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 = ExtractMedia(filename, 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 :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):
""" 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
phase_desc = "Extraction"
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()
if self._args.singleprocess:
phase_desc = self._extractor.phase.title()
desc = "Running pass {} of {}: {}".format(phase + 1,
self._extractor.passes,
phase_desc)
status_bar = tqdm(self._extractor.detected_faces(),
total=self._images.process_count,
file=sys.stdout,
desc=desc)
for idx, extract_media in enumerate(status_bar):
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
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, extract_media, size):
""" Prepare faces for output
Loads the aligned face, 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)
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, extract_media):
""" 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 = list()
filename, extension = os.path.splitext(
os.path.basename(extract_media.filename))
for idx, face in enumerate(extract_media.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(
extract_media.filename)] = dict(faces=final_faces)
del extract_media
class Mask():
""" 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_opts = dict(blur_kernel=arguments.blur_kernel,
threshold=arguments.threshold)
self._output_type = arguments.output_type
self._output_full_frame = arguments.full_frame
self._output_suffix = self._get_output_suffix()
self._face_count = 0
self._skip_count = 0
self._update_count = 0
self._check_input(arguments.input)
self._saver = self._set_saver(arguments)
self._loader = ImagesLoader(arguments.input,
load_with_hash=self._input_is_faces)
self._alignments = Alignments(os.path.dirname(arguments.alignments),
filename=os.path.basename(
arguments.alignments))
self._extractor = self._get_extractor()
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):
""" Obtain a Mask extractor plugin and launch it
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,
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
"""
logger.debug("args: %s", args)
if self._update_type != "output":
queue = args[0]
for filename, image, hsh in tqdm(self._loader.load(),
total=self._loader.count):
if hsh not in self._alignments.hashes_to_frame:
self._skip_count += 1
logger.warning("Skipping face not in alignments file: '%s'",
filename)
continue
frames = self._alignments.hashes_to_frame[hsh]
if len(frames) > 1:
# Filter the output by filename in case of multiple frames with the same face
logger.debug(
"Filtering multiple hashes to current filename: (filename: '%s', "
"frames: %s", filename, frames)
lookup = os.path.splitext(os.path.basename(filename))[0]
frames = {
k: v
for k, v in frames.items()
if lookup.startswith(os.path.splitext(k)[0])
}
logger.debug("Filtered: (filename: '%s', frame: '%s')",
filename, frames)
for frame, idx in frames.items():
self._face_count += 1
alignment = self._alignments.get_faces_in_frame(frame)[idx]
if self._check_for_missing(frame, idx, alignment):
continue
detected_face = self._get_detected_face(alignment)
if self._update_type == "output":
detected_face.image = image
self._save(frame, idx, detected_face)
else:
queue.put(
ExtractMedia(filename,
image,
detected_faces=[detected_face]))
self._update_count += 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._skip_count += 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._face_count += 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._update_count += 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):
""" The filename suffix, based on selected output options
Returns
-------
str:
The suffix to be appended to the output filename
"""
sfx = "{}_mask_preview_".format(self._mask_type)
sfx += "face_" if not self._output_full_frame or self._input_is_faces else "frame_"
sfx += "{}.png".format(self._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":
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._update_count != 0:
self._alignments.backup()
self._alignments.save()
else:
self._extractor_input_thread.join()
self._saver.close()
if self._skip_count != 0:
logger.warning(
"%s face(s) skipped due to not existing in the alignments file",
self._skip_count)
if self._update_type != "output":
if self._update_count == 0:
logger.warning("No masks were updated of the %s faces seen",
self._face_count)
else:
logger.info("Updated masks for %s faces of %s",
self._update_count, self._face_count)
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:
for frame, idx in self._alignments.hashes_to_frame[
face.hash].items():
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 _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:
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 = detected_face.image
else:
detected_face.load_aligned(detected_face.image)
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 = detected_face.image
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 Sort():
""" Sorts folders of faces based on input criteria """
# pylint: disable=no-member
def __init__(self, arguments):
self.args = arguments
self.changes = None
self.serializer = None
self.vgg_face = None
# TODO set this as ImagesLoader in init. Need to move all processes to use it
self._loader = None
def process(self):
""" Main processing function of the sort tool """
# Setting default argument values that cannot be set by argparse
# Set output dir to the same value as input dir
# if the user didn't specify it.
if self.args.output_dir is None:
logger.verbose(
"No output directory provided. Using input dir as output dir.")
self.args.output_dir = self.args.input_dir
# Assigning default threshold values based on grouping method
if (self.args.final_process == "folders"
and self.args.min_threshold < 0.0):
method = self.args.group_method.lower()
if method == 'face-cnn':
self.args.min_threshold = 7.2
elif method == 'hist':
self.args.min_threshold = 0.3
# Load VGG Face if sorting by face
if self.args.sort_method.lower() == "face":
self.vgg_face = VGGFace(backend=self.args.backend,
loglevel=self.args.loglevel)
# If logging is enabled, prepare container
if self.args.log_changes:
self.changes = dict()
# Assign default sort_log.json value if user didn't specify one
if self.args.log_file_path == 'sort_log.json':
self.args.log_file_path = os.path.join(self.args.input_dir,
'sort_log.json')
# Set serializer based on logfile extension
self.serializer = get_serializer_from_filename(
self.args.log_file_path)
# Prepare sort, group and final process method names
_sort = "sort_" + self.args.sort_method.lower()
_group = "group_" + self.args.group_method.lower()
_final = "final_process_" + self.args.final_process.lower()
if _sort.startswith('sort_color-'):
self.args.color_method = _sort.replace('sort_color-', '')
_sort = _sort[:10]
self.args.sort_method = _sort.replace('-', '_')
self.args.group_method = _group.replace('-', '_')
self.args.final_process = _final.replace('-', '_')
self.sort_process()
@staticmethod
def launch_aligner():
""" Load the aligner plugin to retrieve landmarks """
extractor = Extractor(None, "fan", None, normalize_method="hist")
extractor.set_batchsize("align", 1)
extractor.launch()
return extractor
@staticmethod
def alignment_dict(filename, image):
""" Set the image to an ExtractMedia object for alignment """
height, width = image.shape[:2]
face = DetectedFace(x=0, w=width, y=0, h=height)
return ExtractMedia(filename, image, detected_faces=[face])
def _get_landmarks(self):
""" Multi-threaded, parallel and sequentially ordered landmark loader """
extractor = self.launch_aligner()
filename_list, image_list = self._get_images()
feed_list = list(map(Sort.alignment_dict, filename_list, image_list))
landmarks = np.zeros((len(feed_list), 68, 2), dtype='float32')
logger.info("Finding landmarks in images...")
# TODO thread the put to queue so we don't have to put and get at the same time
# Or even better, set up a proper background loader from disk (i.e. use lib.image.ImageIO)
for idx, feed in enumerate(
tqdm(feed_list, desc="Aligning...", file=sys.stdout)):
extractor.input_queue.put(feed)
landmarks[idx] = next(
extractor.detected_faces()).detected_faces[0].landmarks_xy
return filename_list, image_list, landmarks
def _get_images(self):
""" Multi-threaded, parallel and sequentially ordered image loader """
logger.info("Loading images...")
filename_list = self.find_images(self.args.input_dir)
with futures.ThreadPoolExecutor() as executor:
image_list = list(
tqdm(executor.map(read_image, filename_list),
desc="Loading Images...",
file=sys.stdout,
total=len(filename_list)))
return filename_list, image_list
def sort_process(self):
"""
This method dynamically assigns the functions that will be used to run
the core process of sorting, optionally grouping, renaming/moving into
folders. After the functions are assigned they are executed.
"""
sort_method = self.args.sort_method.lower()
group_method = self.args.group_method.lower()
final_method = self.args.final_process.lower()
img_list = getattr(self, sort_method)()
if "folders" in final_method:
# Check if non-dissim sort method and group method are not the same
if group_method.replace('group_', '') not in sort_method:
img_list = self.reload_images(group_method, img_list)
img_list = getattr(self, group_method)(img_list)
else:
img_list = getattr(self, group_method)(img_list)
getattr(self, final_method)(img_list)
logger.info("Done.")
# Methods for sorting
def sort_blur(self):
""" Sort by blur amount """
logger.info("Sorting by estimated image blur...")
filename_list, image_list = self._get_images()
logger.info("Estimating blur...")
blurs = [self.estimate_blur(img) for img in image_list]
logger.info("Sorting...")
matched_list = list(zip(filename_list, blurs))
img_list = sorted(matched_list,
key=operator.itemgetter(1),
reverse=True)
return img_list
def sort_face(self):
""" Sort by identity similarity """
logger.info("Sorting by identity similarity...")
# TODO This should be set in init
self._loader = ImagesLoader(self.args.input_dir)
filenames = []
preds = np.empty((self._loader.count, 512), dtype="float32")
for idx, (filename, image) in enumerate(
tqdm(self._loader.load(),
desc="Classifying Faces...",
total=self._loader.count)):
filenames.append(filename)
preds[idx] = self.vgg_face.predict(image)
logger.info("Sorting by ward linkage...")
indices = self.vgg_face.sorted_similarity(preds, method="ward")
img_list = np.array(filenames)[indices]
return img_list
def sort_face_cnn(self):
""" Sort by landmark similarity """
logger.info("Sorting by landmark similarity...")
filename_list, _, landmarks = self._get_landmarks()
img_list = list(zip(filename_list, landmarks))
logger.info("Comparing landmarks and sorting...")
img_list_len = len(img_list)
for i in tqdm(range(0, img_list_len - 1),
desc="Comparing...",
file=sys.stdout):
min_score = float("inf")
j_min_score = i + 1
for j in range(i + 1, img_list_len):
fl1 = img_list[i][1]
fl2 = img_list[j][1]
score = np.sum(np.absolute((fl2 - fl1).flatten()))
if score < min_score:
min_score = score
j_min_score = j
(img_list[i + 1], img_list[j_min_score]) = (img_list[j_min_score],
img_list[i + 1])
return img_list
def sort_face_cnn_dissim(self):
""" Sort by landmark dissimilarity """
logger.info("Sorting by landmark dissimilarity...")
filename_list, _, landmarks = self._get_landmarks()
scores = np.zeros(len(filename_list), dtype='float32')
img_list = list(
list(items) for items in zip(filename_list, landmarks, scores))
logger.info("Comparing landmarks...")
img_list_len = len(img_list)
for i in tqdm(range(0, img_list_len - 1),
desc="Comparing...",
file=sys.stdout):
score_total = 0
for j in range(i + 1, img_list_len):
if i == j:
continue
fl1 = img_list[i][1]
fl2 = img_list[j][1]
score_total += np.sum(np.absolute((fl2 - fl1).flatten()))
img_list[i][2] = score_total
logger.info("Sorting...")
img_list = sorted(img_list, key=operator.itemgetter(2), reverse=True)
return img_list
def sort_face_yaw(self):
""" Sort by estimated face yaw angle """
logger.info("Sorting by estimated face yaw angle..")
filename_list, _, landmarks = self._get_landmarks()
logger.info("Estimating yaw...")
yaws = [self.calc_landmarks_face_yaw(mark) for mark in landmarks]
logger.info("Sorting...")
matched_list = list(zip(filename_list, yaws))
img_list = sorted(matched_list,
key=operator.itemgetter(1),
reverse=True)
return img_list
def sort_hist(self):
""" Sort by image histogram similarity """
logger.info("Sorting by histogram similarity...")
filename_list, image_list = self._get_images()
distance = cv2.HISTCMP_BHATTACHARYYA
logger.info("Calculating histograms...")
histograms = [
cv2.calcHist([img], [0], None, [256], [0, 256])
for img in image_list
]
img_list = list(zip(filename_list, histograms))
logger.info("Comparing histograms and sorting...")
img_list_len = len(img_list)
for i in tqdm(range(0, img_list_len - 1),
desc="Comparing",
file=sys.stdout):
min_score = float("inf")
j_min_score = i + 1
for j in range(i + 1, img_list_len):
score = cv2.compareHist(img_list[i][1], img_list[j][1],
distance)
if score < min_score:
min_score = score
j_min_score = j
(img_list[i + 1], img_list[j_min_score]) = (img_list[j_min_score],
img_list[i + 1])
return img_list
def sort_hist_dissim(self):
""" Sort by image histogram dissimilarity """
logger.info("Sorting by histogram dissimilarity...")
filename_list, image_list = self._get_images()
scores = np.zeros(len(filename_list), dtype='float32')
distance = cv2.HISTCMP_BHATTACHARYYA
logger.info("Calculating histograms...")
histograms = [
cv2.calcHist([img], [0], None, [256], [0, 256])
for img in image_list
]
img_list = list(
list(items) for items in zip(filename_list, histograms, scores))
logger.info("Comparing histograms...")
img_list_len = len(img_list)
for i in tqdm(range(0, img_list_len),
desc="Comparing",
file=sys.stdout):
score_total = 0
for j in range(0, img_list_len):
if i == j:
continue
score_total += cv2.compareHist(img_list[i][1], img_list[j][1],
distance)
img_list[i][2] = score_total
logger.info("Sorting...")
img_list = sorted(img_list, key=operator.itemgetter(2), reverse=True)
return img_list
def sort_color(self):
""" Score by channel average intensity """
logger.info("Sorting by channel average intensity...")
desired_channel = {'gray': 0, 'luma': 0, 'orange': 1, 'green': 2}
method = self.args.color_method
channel_to_sort = next(v for (k, v) in desired_channel.items()
if method.endswith(k))
filename_list, image_list = self._get_images()
logger.info("Converting to appropriate colorspace...")
same_size = all(img.size == image_list[0].size for img in image_list)
images = np.array(image_list,
dtype='float32')[None,
...] if same_size else image_list
converted_images = self._convert_color(images, same_size, method)
logger.info("Scoring each image...")
if same_size:
scores = np.average(converted_images[0], axis=(1, 2))
else:
progress_bar = tqdm(converted_images,
desc="Scoring",
file=sys.stdout)
scores = np.array(
[np.average(image, axis=(0, 1)) for image in progress_bar])
logger.info("Sorting...")
matched_list = list(zip(filename_list, scores[:, channel_to_sort]))
sorted_file_img_list = sorted(matched_list,
key=operator.itemgetter(1),
reverse=True)
return sorted_file_img_list
# Methods for grouping
def group_blur(self, img_list):
""" Group into bins by blur """
# Starting the binning process
num_bins = self.args.num_bins
# The last bin will get all extra images if it's
# not possible to distribute them evenly
num_per_bin = len(img_list) // num_bins
remainder = len(img_list) % num_bins
logger.info("Grouping by blur...")
bins = [[] for _ in range(num_bins)]
idx = 0
for i in range(num_bins):
for _ in range(num_per_bin):
bins[i].append(img_list[idx][0])
idx += 1
# If remainder is 0, nothing gets added to the last bin.
for i in range(1, remainder + 1):
bins[-1].append(img_list[-i][0])
return bins
def group_face_cnn(self, img_list):
""" Group into bins by CNN face similarity """
logger.info("Grouping by face-cnn similarity...")
# Groups are of the form: group_num -> reference faces
reference_groups = dict()
# Bins array, where index is the group number and value is
# an array containing the file paths to the images in that group.
bins = []
# Comparison threshold used to decide how similar
# faces have to be to be grouped together.
# It is multiplied by 1000 here to allow the cli option to use smaller
# numbers.
min_threshold = self.args.min_threshold * 1000
img_list_len = len(img_list)
for i in tqdm(range(0, img_list_len - 1),
desc="Grouping",
file=sys.stdout):
fl1 = img_list[i][1]
current_best = [-1, float("inf")]
for key, references in reference_groups.items():
try:
score = self.get_avg_score_faces_cnn(fl1, references)
except TypeError:
score = float("inf")
except ZeroDivisionError:
score = float("inf")
if score < current_best[1]:
current_best[0], current_best[1] = key, score
if current_best[1] < min_threshold:
reference_groups[current_best[0]].append(fl1[0])
bins[current_best[0]].append(img_list[i][0])
else:
reference_groups[len(reference_groups)] = [img_list[i][1]]
bins.append([img_list[i][0]])
return bins
def group_face_yaw(self, img_list):
""" Group into bins by yaw of face """
# Starting the binning process
num_bins = self.args.num_bins
# The last bin will get all extra images if it's
# not possible to distribute them evenly
num_per_bin = len(img_list) // num_bins
remainder = len(img_list) % num_bins
logger.info("Grouping by face-yaw...")
bins = [[] for _ in range(num_bins)]
idx = 0
for i in range(num_bins):
for _ in range(num_per_bin):
bins[i].append(img_list[idx][0])
idx += 1
# If remainder is 0, nothing gets added to the last bin.
for i in range(1, remainder + 1):
bins[-1].append(img_list[-i][0])
return bins
def group_hist(self, img_list):
""" Group into bins by histogram """
logger.info("Grouping by histogram...")
# Groups are of the form: group_num -> reference histogram
reference_groups = dict()
# Bins array, where index is the group number and value is
# an array containing the file paths to the images in that group
bins = []
min_threshold = self.args.min_threshold
img_list_len = len(img_list)
reference_groups[0] = [img_list[0][1]]
bins.append([img_list[0][0]])
for i in tqdm(range(1, img_list_len), desc="Grouping",
file=sys.stdout):
current_best = [-1, float("inf")]
for key, value in reference_groups.items():
score = self.get_avg_score_hist(img_list[i][1], value)
if score < current_best[1]:
current_best[0], current_best[1] = key, score
if current_best[1] < min_threshold:
reference_groups[current_best[0]].append(img_list[i][1])
bins[current_best[0]].append(img_list[i][0])
else:
reference_groups[len(reference_groups)] = [img_list[i][1]]
bins.append([img_list[i][0]])
return bins
# Final process methods
def final_process_rename(self, img_list):
""" Rename the files """
output_dir = self.args.output_dir
process_file = self.set_process_file_method(self.args.log_changes,
self.args.keep_original)
# Make sure output directory exists
if not os.path.exists(output_dir):
os.makedirs(output_dir)
description = ("Copying and Renaming"
if self.args.keep_original else "Moving and Renaming")
for i in tqdm(range(0, len(img_list)),
desc=description,
leave=False,
file=sys.stdout):
src = img_list[i] if isinstance(img_list[i],
str) else img_list[i][0]
src_basename = os.path.basename(src)
dst = os.path.join(output_dir, '{:05d}_{}'.format(i, src_basename))
try:
process_file(src, dst, self.changes)
except FileNotFoundError as err:
logger.error(err)
logger.error('fail to rename %s', src)
for i in tqdm(range(0, len(img_list)),
desc=description,
file=sys.stdout):
renaming = self.set_renaming_method(self.args.log_changes)
fname = img_list[i] if isinstance(img_list[i],
str) else img_list[i][0]
src, dst = renaming(fname, output_dir, i, self.changes)
try:
os.rename(src, dst)
except FileNotFoundError as err:
logger.error(err)
logger.error('fail to rename %s', format(src))
if self.args.log_changes:
self.write_to_log(self.changes)
def final_process_folders(self, bins):
""" Move the files to folders """
output_dir = self.args.output_dir
process_file = self.set_process_file_method(self.args.log_changes,
self.args.keep_original)
# First create new directories to avoid checking
# for directory existence in the moving loop
logger.info("Creating group directories.")
for i in range(len(bins)):
directory = os.path.join(output_dir, str(i))
if not os.path.exists(directory):
os.makedirs(directory)
description = ("Copying into Groups"
if self.args.keep_original else "Moving into Groups")
logger.info("Total groups found: %s", len(bins))
for i in tqdm(range(len(bins)), desc=description, file=sys.stdout):
for j in range(len(bins[i])):
src = bins[i][j]
src_basename = os.path.basename(src)
dst = os.path.join(output_dir, str(i), src_basename)
try:
process_file(src, dst, self.changes)
except FileNotFoundError as err:
logger.error(err)
logger.error("Failed to move '%s' to '%s'", src, dst)
if self.args.log_changes:
self.write_to_log(self.changes)
# Various helper methods
def write_to_log(self, changes):
""" Write the changes to log file """
logger.info("Writing sort log to: '%s'", self.args.log_file_path)
self.serializer.save(self.args.log_file_path, changes)
def reload_images(self, group_method, img_list):
"""
Reloads the image list by replacing the comparative values with those
that the chosen grouping method expects.
:param group_method: str name of the grouping method that will be used.
:param img_list: image list that has been sorted by one of the sort
methods.
:return: img_list but with the comparative values that the chosen
grouping method expects.
"""
logger.info("Preparing to group...")
if group_method == 'group_blur':
filename_list, image_list = self._get_images()
blurs = [self.estimate_blur(img) for img in image_list]
temp_list = list(zip(filename_list, blurs))
elif group_method == 'group_face_cnn':
filename_list, image_list, landmarks = self._get_landmarks()
temp_list = list(zip(filename_list, landmarks))
elif group_method == 'group_face_yaw':
filename_list, image_list, landmarks = self._get_landmarks()
yaws = [self.calc_landmarks_face_yaw(mark) for mark in landmarks]
temp_list = list(zip(filename_list, yaws))
elif group_method == 'group_hist':
filename_list, image_list = self._get_images()
histograms = [
cv2.calcHist([img], [0], None, [256], [0, 256])
for img in image_list
]
temp_list = list(zip(filename_list, histograms))
else:
raise ValueError("{} group_method not found.".format(group_method))
return self.splice_lists(img_list, temp_list)
@staticmethod
def _convert_color(imgs, same_size, method):
""" Helper function to convert colorspaces """
if method.endswith('gray'):
conversion = np.array([[0.0722], [0.7152], [0.2126]])
else:
conversion = np.array([[0.25, 0.5, 0.25], [-0.5, 0.0, 0.5],
[-0.25, 0.5, -0.25]])
if same_size:
path = 'greedy'
operation = 'bijk, kl -> bijl' if method.endswith(
'gray') else 'bijl, kl -> bijk'
else:
operation = 'ijk, kl -> ijl' if method.endswith(
'gray') else 'ijl, kl -> ijk'
path = np.einsum_path(operation,
imgs[0][..., :3],
conversion,
optimize='optimal')[0]
progress_bar = tqdm(imgs, desc="Converting", file=sys.stdout)
images = [
np.einsum(operation, img[..., :3], conversion,
optimize=path).astype('float32') for img in progress_bar
]
return images
@staticmethod
def splice_lists(sorted_list, new_vals_list):
"""
This method replaces the value at index 1 in each sub-list in the
sorted_list with the value that is calculated for the same img_path,
but found in new_vals_list.
Format of lists: [[img_path, value], [img_path2, value2], ...]
:param sorted_list: list that has been sorted by one of the sort
methods.
:param new_vals_list: list that has been loaded by a different method
than the sorted_list.
:return: list that is sorted in the same way as the input sorted list
but the values corresponding to each image are from new_vals_list.
"""
new_list = []
# Make new list of just image paths to serve as an index
val_index_list = [i[0] for i in new_vals_list]
for i in tqdm(range(len(sorted_list)),
desc="Splicing",
file=sys.stdout):
current_img = sorted_list[i] if isinstance(
sorted_list[i], str) else sorted_list[i][0]
new_val_index = val_index_list.index(current_img)
new_list.append([current_img, new_vals_list[new_val_index][1]])
return new_list
@staticmethod
def find_images(input_dir):
""" Return list of images at specified location """
result = []
extensions = [".jpg", ".png", ".jpeg"]
for root, _, files in os.walk(input_dir):
for file in files:
if os.path.splitext(file)[1].lower() in extensions:
result.append(os.path.join(root, file))
break
return result
@staticmethod
def estimate_blur(image):
"""
Estimate the amount of blur an image has with the variance of the Laplacian.
Normalize by pixel number to offset the effect of image size on pixel gradients & variance
"""
if image.ndim == 3:
image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
blur_map = cv2.Laplacian(image, cv2.CV_32F)
score = np.var(blur_map) / np.sqrt(image.shape[0] * image.shape[1])
return score
@staticmethod
def calc_landmarks_face_pitch(flm):
""" UNUSED - Calculate the amount of pitch in a face """
var_t = ((flm[6][1] - flm[8][1]) + (flm[10][1] - flm[8][1])) / 2.0
var_b = flm[8][1]
return var_b - var_t
@staticmethod
def calc_landmarks_face_yaw(flm):
""" Calculate the amount of yaw in a face """
var_l = ((flm[27][0] - flm[0][0]) + (flm[28][0] - flm[1][0]) +
(flm[29][0] - flm[2][0])) / 3.0
var_r = ((flm[16][0] - flm[27][0]) + (flm[15][0] - flm[28][0]) +
(flm[14][0] - flm[29][0])) / 3.0
return var_r - var_l
@staticmethod
def set_process_file_method(log_changes, keep_original):
"""
Assigns the final file processing method based on whether changes are
being logged and whether the original files are being kept in the
input directory.
Relevant cli arguments: -k, -l
:return: function reference
"""
if log_changes:
if keep_original:
def process_file(src, dst, changes):
""" Process file method if logging changes
and keeping original """
copyfile(src, dst)
changes[src] = dst
else:
def process_file(src, dst, changes):
""" Process file method if logging changes
and not keeping original """
os.rename(src, dst)
changes[src] = dst
else:
if keep_original:
def process_file(src, dst, changes): # pylint: disable=unused-argument
""" Process file method if not logging changes
and keeping original """
copyfile(src, dst)
else:
def process_file(src, dst, changes): # pylint: disable=unused-argument
""" Process file method if not logging changes
and not keeping original """
os.rename(src, dst)
return process_file
@staticmethod
def set_renaming_method(log_changes):
""" Set the method for renaming files """
if log_changes:
def renaming(src, output_dir, i, changes):
""" Rename files method if logging changes """
src_basename = os.path.basename(src)
__src = os.path.join(output_dir,
'{:05d}_{}'.format(i, src_basename))
dst = os.path.join(
output_dir,
'{:05d}{}'.format(i,
os.path.splitext(src_basename)[1]))
changes[src] = dst
return __src, dst
else:
def renaming(src, output_dir, i, changes): # pylint: disable=unused-argument
""" Rename files method if not logging changes """
src_basename = os.path.basename(src)
src = os.path.join(output_dir,
'{:05d}_{}'.format(i, src_basename))
dst = os.path.join(
output_dir,
'{:05d}{}'.format(i,
os.path.splitext(src_basename)[1]))
return src, dst
return renaming
@staticmethod
def get_avg_score_hist(img1, references):
""" Return the average histogram score between a face and
reference image """
scores = []
for img2 in references:
score = cv2.compareHist(img1, img2, cv2.HISTCMP_BHATTACHARYYA)
scores.append(score)
return sum(scores) / len(scores)
@staticmethod
def get_avg_score_faces_cnn(fl1, references):
""" Return the average CNN similarity score
between a face and reference image """
scores = []
for fl2 in references:
score = np.sum(np.absolute((fl2 - fl1).flatten()))
scores.append(score)
return sum(scores) / len(scores)
def __init__(self, arguments):
logger.debug("Initializing %s: (args: %s)", self.__class__.__name__, arguments)
self._args = arguments
# load faces
faces_alignments = AlignmentsBase(self._args.faces_align_dir)
print()
print(f'Faces alignments: {len(faces_alignments._data.keys())}')
print(faces_alignments._data.keys())
self._faces = {}
faces_loader = ImagesLoader(self._args.faces_dir)
for filename, image in faces_loader.load():
face_name = os.path.basename(filename)
faces = faces_alignments.get_faces_in_frame(face_name)
detected_faces = list()
for rawface in faces:
face = DetectedFace()
face.from_alignment(rawface, image=image)
feed_face = AlignedFace(face.landmarks_xy,
image=image,
centering='face',
size=image.shape[0],
coverage_ratio=1.0,
dtype="float32")
detected_faces.append(feed_face)
self._faces[face_name] = (filename, image, detected_faces)
print('Faces:', len(self._faces))
print(self._faces.keys())
print()
self._patch_threads = None
self._images = ImagesLoader(self._args.input_dir, fast_count=True)
self._alignments = Alignments(self._args, False, self._images.is_video)
if self._alignments.version == 1.0:
logger.error("The alignments file format has been updated since the given alignments "
"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)
self._opts = OptionalActions(self._args, self._images.file_list, self._alignments)
self._add_queues()
self._disk_io = DiskIO(self._alignments, self._images, arguments)
self._predictor = Predict(self._disk_io.load_queue, self._queue_size, self._faces, arguments)
self._validate()
get_folder(self._args.output_dir)
configfile = self._args.configfile if hasattr(self._args, "configfile") else None
self._converter = Converter(self._predictor.output_size,
self._predictor.coverage_ratio,
self._predictor.centering,
self._disk_io.draw_transparent,
self._disk_io.pre_encode,
arguments,
configfile=configfile)
logger.debug("Initialized %s", self.__class__.__name__)