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 _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) extension = ".png" for frame_idx, (filename, image) in enumerate(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), extension) 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, extension, metadata=meta) saver.save(output, b_image) 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