def _predict(self, batch): """ Wrap models predict function in rotations """ batch["rotmat"] = [np.array([]) for _ in range(len(batch["feed"]))] found_faces = [np.array([]) for _ in range(len(batch["feed"]))] for angle in self.rotation: # Rotate the batch and insert placeholders for already found faces self._rotate_batch(batch, angle) try: batch = self.predict(batch) except tf_errors.ResourceExhaustedError as err: msg = ( "You do not have enough GPU memory available to run detection at the " "selected batch size. You can try a number of things:" "\n1) Close any other application that is using your GPU (web browsers are " "particularly bad for this)." "\n2) Lower the batchsize (the amount of images fed into the model) by " "editing the plugin settings (GUI: Settings > Configure extract settings, " "CLI: Edit the file faceswap/config/extract.ini)." "\n3) Enable 'Single Process' mode.") raise FaceswapError(msg) from err except Exception as err: if get_backend() == "amd": # pylint:disable=import-outside-toplevel from lib.plaidml_utils import is_plaidml_error if (is_plaidml_error(err) and ("CL_MEM_OBJECT_ALLOCATION_FAILURE" in str(err).upper() or "enough memory for the current schedule" in str(err).lower())): msg = ( "You do not have enough GPU memory available to run detection at " "the selected batch size. You can try a number of things:" "\n1) Close any other application that is using your GPU (web " "browsers are particularly bad for this)." "\n2) Lower the batchsize (the amount of images fed into the " "model) by editing the plugin settings (GUI: Settings > Configure " "extract settings, CLI: Edit the file " "faceswap/config/extract.ini).") raise FaceswapError(msg) from err raise if angle != 0 and any([face.any() for face in batch["prediction"]]): logger.verbose("found face(s) by rotating image %s degrees", angle) found_faces = [ face if not found.any() else found for face, found in zip(batch["prediction"], found_faces) ] if all([face.any() for face in found_faces]): logger.trace("Faces found for all images") break batch["prediction"] = found_faces logger.trace( "detect_prediction output: (filenames: %s, prediction: %s, rotmat: %s)", batch["filename"], batch["prediction"], batch["rotmat"]) return batch
def _initialize(self) -> None: """ Initialize PyNVML for Nvidia GPUs. If :attr:`_is_initialized` is ``True`` then this function just returns performing no action. Otherwise :attr:`is_initialized` is set to ``True`` after successfully initializing NVML. Raises ------ FaceswapError If the NVML library could not be successfully loaded """ if self._is_initialized: return try: self._log("debug", "Initializing PyNVML for Nvidia GPU.") pynvml.nvmlInit() except ( pynvml.NVMLError_LibraryNotFound, # pylint:disable=no-member pynvml.NVMLError_DriverNotLoaded, # pylint:disable=no-member pynvml.NVMLError_NoPermission) as err: # pylint:disable=no-member msg = ( "There was an error reading from the Nvidia Machine Learning Library. The most " "likely cause is incorrectly installed drivers. If this is the case, Please " "remove and reinstall your Nvidia drivers before reporting. Original " f"Error: {str(err)}") raise FaceswapError(msg) from err except Exception as err: # pylint: disable=broad-except msg = ( "An unhandled exception occured reading from the Nvidia Machine Learning " f"Library. Original error: {str(err)}") raise FaceswapError(msg) from err super()._initialize()
def _test_for_tf_version(): """ Check that the required Tensorflow version is installed. Raises ------ FaceswapError If Tensorflow is not found, or is not between versions 2.2 and 2.2 """ min_ver = 2.2 max_ver = 2.2 try: # Ensure tensorflow doesn't pin all threads to one core when using Math Kernel Library os.environ["TF_MIN_GPU_MULTIPROCESSOR_COUNT"] = "4" os.environ["KMP_AFFINITY"] = "disabled" import tensorflow as tf # pylint:disable=import-outside-toplevel except ImportError as err: raise FaceswapError("There was an error importing Tensorflow. This is most likely " "because you do not have TensorFlow installed, or you are trying " "to run tensorflow-gpu on a system without an Nvidia graphics " "card. Original import error: {}".format(str(err))) tf_ver = float(".".join(tf.__version__.split(".")[:2])) # pylint:disable=no-member if tf_ver < min_ver: raise FaceswapError("The minimum supported Tensorflow is version {} but you have " "version {} installed. Please upgrade Tensorflow.".format( min_ver, tf_ver)) if tf_ver > max_ver: raise FaceswapError("The maximumum supported Tensorflow is version {} but you have " "version {} installed. Please downgrade Tensorflow.".format( max_ver, tf_ver)) logger.debug("Installed Tensorflow Version: %s", tf_ver)
def get_trainer(self, model_dir): """ Return the trainer name if provided, or read from state file """ if hasattr(self.args, "trainer") and self.args.trainer: logger.debug("Trainer name provided: '%s'", self.args.trainer) return self.args.trainer statefile = [ fname for fname in os.listdir(str(model_dir)) if fname.endswith("_state.json") ] if len(statefile) != 1: raise FaceswapError( "There should be 1 state file in your model folder. {} were " "found. Specify a trainer with the '-t', '--trainer' " "option.".format(len(statefile))) statefile = os.path.join(str(model_dir), statefile[0]) state = self.serializer.load(statefile) trainer = state.get("name", None) if not trainer: raise FaceswapError( "Trainer name could not be read from state file. " "Specify a trainer with the '-t', '--trainer' option.") logger.debug("Trainer from state file: '%s'", trainer) return trainer
def _get_weights_model(self) -> List[keras.models.Model]: """ Obtain a list of all sub-models contained within the weights model. Returns ------- list List of all models contained within the .h5 file Raises ------ FaceswapError In the event of a failure to load the weights, or the weights belonging to a different model """ retval = get_all_sub_models( load_model(self._weights_file, compile=False)) if not retval: raise FaceswapError( f"Error loading weights file {self._weights_file}.") if retval[0].name != self._name: raise FaceswapError( f"You are attempting to load weights from a '{retval[0].name}' " f"model into a '{self._name}' model. This is not supported.") return retval
def _set_timelapse(self): """ Set time-lapse paths if requested. Returns ------- dict The time-lapse keyword arguments for passing to the trainer """ if (not self._args.timelapse_input_a and not self._args.timelapse_input_b and not self._args.timelapse_output): return None if (not self._args.timelapse_input_a or not self._args.timelapse_input_b or not self._args.timelapse_output): raise FaceswapError("To enable the timelapse, you have to supply all the parameters " "(--timelapse-input-A, --timelapse-input-B and " "--timelapse-output).") timelapse_output = str(get_folder(self._args.timelapse_output)) for folder in (self._args.timelapse_input_a, self._args.timelapse_input_b): if folder is not None and not os.path.isdir(folder): raise FaceswapError("The Timelapse path '{}' does not exist".format(folder)) exts = [os.path.splitext(fname)[-1] for fname in os.listdir(folder)] if not any(ext in _image_extensions for ext in exts): raise FaceswapError("The Timelapse path '{}' does not contain any valid " "images".format(folder)) kwargs = {"input_a": self._args.timelapse_input_a, "input_b": self._args.timelapse_input_b, "output": timelapse_output} logger.debug("Timelapse enabled: %s", kwargs) return kwargs
def _validate_encoder_architecture(self): """ Validate that the requested architecture is a valid choice for the running system configuration. If the selection is not valid, an error is logged and system exits. """ arch = self.config["enc_architecture"].lower() model = _MODEL_MAPPING.get(arch) if not model: raise FaceswapError( f"'{arch}' is not a valid choice for encoder architecture. Choose " f"one of {list(_MODEL_MAPPING.keys())}.") if get_backend() == "amd" and model.get("no_amd"): valid = [ x for x in _MODEL_MAPPING if not _MODEL_MAPPING[x].get('no_amd') ] raise FaceswapError( f"'{arch}' is not compatible with the AMD backend. Choose one of " f"{valid}.") tf_ver = float(".".join(tf.__version__.split(".")[:2])) # pylint:disable=no-member tf_min = model.get("tf_min", 2.0) if get_backend() != "amd" and tf_ver < tf_min: raise FaceswapError( f"{arch}' is not compatible with your version of Tensorflow. The " f"minimum version required is {tf_min} whilst you have version " f"{tf_ver} installed.")
def _predict(self, batch): """ Just return the masker's predict function """ try: return self.predict(batch) except tf_errors.ResourceExhaustedError as err: msg = ( "You do not have enough GPU memory available to run detection at the " "selected batch size. You can try a number of things:" "\n1) Close any other application that is using your GPU (web browsers are " "particularly bad for this)." "\n2) Lower the batchsize (the amount of images fed into the model) by " "editing the plugin settings (GUI: Settings > Configure extract settings, " "CLI: Edit the file faceswap/config/extract.ini)." "\n3) Enable 'Single Process' mode.") raise FaceswapError(msg) from err except Exception as err: if get_backend() == "amd": # pylint:disable=import-outside-toplevel from lib.plaidml_utils import is_plaidml_error if (is_plaidml_error(err) and ("CL_MEM_OBJECT_ALLOCATION_FAILURE" in str(err).upper() or "enough memory for the current schedule" in str(err).lower())): msg = ( "You do not have enough GPU memory available to run detection at " "the selected batch size. You can try a number of things:" "\n1) Close any other application that is using your GPU (web " "browsers are particularly bad for this)." "\n2) Lower the batchsize (the amount of images fed into the " "model) by editing the plugin settings (GUI: Settings > Configure " "extract settings, CLI: Edit the file " "faceswap/config/extract.ini).") raise FaceswapError(msg) from err raise
def get_frame_ranges(self): """ split out the frame ranges and parse out 'min' and 'max' values """ if not self.args.frame_ranges: logger.debug("No frame range set") return None minframe, maxframe = None, None if self.images.is_video: minframe, maxframe = 1, self.images.images_found else: indices = [ int(self.imageidxre.findall(os.path.basename(filename))[0]) for filename in self.images.input_images ] if indices: minframe, maxframe = min(indices), max(indices) logger.debug("minframe: %s, maxframe: %s", minframe, maxframe) if minframe is None or maxframe is None: raise FaceswapError( "Frame Ranges specified, but could not determine frame numbering " "from filenames") retval = list() for rng in self.args.frame_ranges: if "-" not in rng: raise FaceswapError( "Frame Ranges not specified in the correct format") start, end = rng.split("-") retval.append((max(int(start), minframe), min(int(end), maxframe))) logger.debug("frame ranges: %s", retval) return retval
def _validate(self): """ Validate the Command Line Options. Ensure that certain cli selections are valid and won't result in an error. Checks: * If frames have been passed in with video output, ensure user supplies reference video. * If "on-the-fly" and an NN mask is selected, output warning and switch to 'extended' * If a mask-type is selected, ensure it exists in the alignments file. * If a predicted mask-type is selected, ensure model has been trained with a mask otherwise attempt to select first available masks, otherwise raise error. Raises ------ FaceswapError If an invalid selection has been found. """ if (self._args.writer == "ffmpeg" and not self._images.is_video and self._args.reference_video is None): raise FaceswapError( "Output as video selected, but using frames as input. You must " "provide a reference video ('-ref', '--reference-video').") if (self._args.on_the_fly and self._args.mask_type not in ("none", "extended", "components")): logger.warning( "You have selected an incompatible mask type ('%s') for On-The-Fly " "conversion. Switching to 'extended'", self._args.mask_type) self._args.mask_type = "extended" if (not self._args.on_the_fly and self._args.mask_type not in ("none", "predicted") and not self._alignments.mask_is_valid(self._args.mask_type)): msg = ( "You have selected the Mask Type `{}` but at least one face does not have this " "mask stored in the Alignments File.\nYou should generate the required masks " "with the Mask Tool or set the Mask Type option to an existing Mask Type.\nA " "summary of existing masks is as follows:\nTotal faces: {}, Masks: " "{}".format(self._args.mask_type, self._alignments.faces_count, self._alignments.mask_summary)) raise FaceswapError(msg) if self._args.mask_type == "predicted" and not self._predictor.has_predicted_mask: available_masks = [ k for k, v in self._alignments.mask_summary.items() if k != "none" and v == self._alignments.faces_count ] if not available_masks: msg = ( "Predicted Mask selected, but the model was not trained with a mask and no " "masks are stored in the Alignments File.\nYou should generate the " "required masks with the Mask Tool or set the Mask Type to `none`." ) raise FaceswapError(msg) mask_type = available_masks[0] logger.warning( "Predicted Mask selected, but the model was not trained with a " "mask. Selecting first available mask: '%s'", mask_type) self._args.mask_type = mask_type
def _set_timelapse(self): """ Set time-lapse paths if requested. Returns ------- dict The time-lapse keyword arguments for passing to the trainer """ if (not self._args.timelapse_input_a and not self._args.timelapse_input_b and not self._args.timelapse_output): return None if (not self._args.timelapse_input_a or not self._args.timelapse_input_b or not self._args.timelapse_output): raise FaceswapError( "To enable the timelapse, you have to supply all the parameters " "(--timelapse-input-A, --timelapse-input-B and " "--timelapse-output).") timelapse_output = get_folder(self._args.timelapse_output) for side in ("a", "b"): folder = getattr(self._args, f"timelapse_input_{side}") if folder is not None and not os.path.isdir(folder): raise FaceswapError( f"The Timelapse path '{folder}' does not exist") training_folder = getattr(self._args, f"input_{side}") if folder == training_folder: continue # Time-lapse folder is training folder filenames = [ fname for fname in os.listdir(folder) if os.path.splitext(fname)[-1].lower() in _image_extensions ] if not filenames: raise FaceswapError( f"The Timelapse path '{folder}' does not contain any valid " "images") # Time-lapse images must appear in the training set, as we need access to alignment and # mask info. Check filenames are there to save failing much later in the process. training_images = [ os.path.basename(img) for img in self._images[side] ] if not all(img in training_images for img in filenames): raise FaceswapError( f"All images in the Timelapse folder '{folder}' must exist in " f"the training folder '{training_folder}'") kwargs = { "input_a": self._args.timelapse_input_a, "input_b": self._args.timelapse_input_b, "output": timelapse_output } logger.debug("Timelapse enabled: %s", kwargs) return kwargs
def _get_face_metadata(self): """ Check for the existence of an aligned directory for identifying which faces in the target frames should be swapped. If it exists, scan the folder for face's metadata Returns ------- dict Dictionary of source frame names with a list of associated face indices to be skipped """ retval = dict() input_aligned_dir = self._args.input_aligned_dir if input_aligned_dir is None: logger.verbose( "Aligned directory not specified. All faces listed in the " "alignments file will be converted") return retval if not os.path.isdir(input_aligned_dir): logger.warning( "Aligned directory not found. All faces listed in the " "alignments file will be converted") return retval log_once = False filelist = get_image_paths(input_aligned_dir) for fullpath, metadata in tqdm(read_image_meta_batch(filelist), total=len(filelist), desc="Reading Face Data", leave=False): if "itxt" not in metadata or "source" not in metadata["itxt"]: # UPDATE LEGACY FACES FROM ALIGNMENTS FILE if not log_once: logger.warning( "Legacy faces discovered in '%s'. These faces will be updated", input_aligned_dir) log_once = True data = update_legacy_png_header(fullpath, self._alignments) if not data: raise FaceswapError( "Some of the faces being passed in from '{}' could not be matched to the " "alignments file '{}'\nPlease double check your sources and try " "again.".format(input_aligned_dir, self._alignments.file)) meta = data["source"] else: meta = metadata["itxt"]["source"] retval.setdefault(meta["source_filename"], list()).append(meta["face_index"]) if not retval: raise FaceswapError( "Aligned directory is empty, no faces will be converted!") if len(retval) <= len(self._input_images) / 3: logger.warning( "Aligned directory contains far fewer images than the input " "directory, are you sure this is the right folder?") return retval
def process_folder(self): """ Iterate through the faces folder pulling out various information for each face. Yields ------ dict A dictionary for each face found containing the keys returned from :class:`lib.image.read_image_meta_batch` """ logger.info("Loading file list from %s", self.folder) if self._alignments is not None: # Legacy updating filelist = [ os.path.join(self.folder, face) for face in os.listdir(self.folder) if self.valid_extension(face) ] else: filelist = [ os.path.join(self.folder, face) for face in os.listdir(self.folder) if os.path.splitext(face)[-1] == ".png" ] log_once = False for fullpath, metadata in tqdm(read_image_meta_batch(filelist), total=len(filelist), desc="Reading Face Data"): if "itxt" not in metadata or "source" not in metadata["itxt"]: if self._alignments is None: # Can't update legacy raise FaceswapError( f"The folder '{self.folder}' contains images that do not include Faceswap " "metadata.\nAll images in the provided folder should contain faces " "generated from Faceswap's extraction process.\nPlease double check the " "source and try again.") if not log_once: logger.warning( "Legacy faces discovered. These faces will be updated") log_once = True data = update_legacy_png_header(fullpath, self._alignments) if not data: raise FaceswapError( "Some of the faces being passed in from '{}' could not be matched to the " "alignments file '{}'\nPlease double check your sources and try " "again.".format(self.folder, self._alignments.file)) retval = data["source"] else: retval = metadata["itxt"]["source"] retval["current_filename"] = os.path.basename(fullpath) yield retval
def _check_location_exists(self): """ Check whether the input location exists. Raises ------ FaceswapError If the given location does not exist """ if isinstance(self.location, str) and not os.path.exists(self.location): raise FaceswapError("The location '{}' does not exist".format(self.location)) if isinstance(self.location, (list, tuple)) and not all(os.path.exists(location) for location in self.location): raise FaceswapError("Not all locations in the input list exist")
def _check_location_exists(self): """ Check whether the output location exists and is a folder Raises ------ FaceswapError If the given location does not exist or the location is not a folder """ if not isinstance(self.location, str): raise FaceswapError("The output location must be a string not a " "{}".format(type(self.location))) super()._check_location_exists() if not os.path.isdir(self.location): raise FaceswapError("The output location '{}' is not a folder".format(self.location))
def get_face_hashes(self): """ Check for the existence of an aligned directory for identifying which faces in the target frames should be swapped. If it exists, obtain the hashes of the faces in the folder """ face_hashes = list() input_aligned_dir = self.args.input_aligned_dir if input_aligned_dir is None: logger.verbose( "Aligned directory not specified. All faces listed in the " "alignments file will be converted") elif not os.path.isdir(input_aligned_dir): logger.warning( "Aligned directory not found. All faces listed in the " "alignments file will be converted") else: file_list = [path for path in get_image_paths(input_aligned_dir)] logger.info("Getting Face Hashes for selected Aligned Images") for face in tqdm(file_list, desc="Hashing Faces"): face_hashes.append(read_image_hash(face)) logger.debug("Face Hashes: %s", (len(face_hashes))) if not face_hashes: raise FaceswapError( "Aligned directory is empty, no faces will be converted!") if len(face_hashes) <= len(self.input_images) / 3: logger.warning( "Aligned directory contains far fewer images than the input " "directory, are you sure this is the right folder?") return face_hashes
def _check_multiple_models(self) -> None: """ Check whether multiple models exist in the model folder, and that no models exist that were trained with a different plugin than the requested plugin. Raises ------ FaceswapError If multiple model files, or models for a different plugin from that requested exists within the model folder """ multiple_models = self._io.multiple_models_in_folder if multiple_models is None: logger.debug("Contents of model folder are valid") return if len(multiple_models) == 1: msg = (f"You have requested to train with the '{self.name}' plugin, but a model file " f"for the '{multiple_models[0]}' plugin already exists in the folder " f"'{self.model_dir}'.\nPlease select a different model folder.") else: ptypes = "', '".join(multiple_models) msg = (f"There are multiple plugin types ('{ptypes}') stored in the model folder '" f"{self.model_dir}'. This is not supported.\nPlease split the model files into " "their own folders before proceeding") raise FaceswapError(msg)
def __call__(self, y_true, y_pred): """ Return the Gradient Magnitude Similarity Deviation Loss. Parameters ---------- y_true: tensor or variable The ground truth value y_pred: tensor or variable The predicted value Returns ------- tensor The loss value """ raise FaceswapError("GMSD Loss is not currently compatible with PlaidML. Please select a " "different Loss method.") true_edge = self._scharr_edges(y_true, True) pred_edge = self._scharr_edges(y_pred, True) ephsilon = 0.0025 upper = 2.0 * true_edge * pred_edge lower = K.square(true_edge) + K.square(pred_edge) gms = (upper + ephsilon) / (lower + ephsilon) gmsd = K.std(gms, axis=(1, 2, 3), keepdims=True) gmsd = K.squeeze(gmsd, axis=-1) return gmsd
def marshal(self, data): """ Serialize an object Parameters ---------- data: varies The data that is to be serialized Returns ------- data: varies The data in a the serialized data format Example ------ >>> serializer = get_serializer('json') >>> data ['foo', 'bar'] >>> json_data = serializer.marshal(data) """ logger.debug("data type: %s", type(data)) try: retval = self._marshal(data) except Exception as err: msg = "Error serializing data for type {}: {}".format( type(data), str(err)) raise FaceswapError(msg) from err logger.debug("returned data type: %s", type(retval)) return retval
def sort_size(self): """ Sort the faces by largest face (in original frame) to smallest """ logger.info("Sorting by original face size...") img_list = [] for filename, image, metadata in tqdm(self._loader.load(), desc="Calculating face sizes", total=self._loader.count, leave=False): if not metadata: msg = ("The images to be sorted do not contain alignment data. Images must have " "been generated by Faceswap's Extract process.\nIf you are sorting an " "older faceset, then you should re-extract the faces from your source " "alignments file to generate this data.") raise FaceswapError(msg) alignments = metadata["alignments"] aligned_face = AlignedFace(np.array(alignments["landmarks_xy"], dtype="float32"), image=image, centering="legacy", is_aligned=True) roi = aligned_face.original_roi size = ((roi[1][0] - roi[0][0]) ** 2 + (roi[1][1] - roi[0][1]) ** 2) ** 0.5 img_list.append((filename, size)) logger.info("Sorting...") return sorted(img_list, key=lambda x: x[1], reverse=True)
def sort_face_yaw(self): """ Sort by estimated face yaw angle """ logger.info("Sorting by estimated face yaw angle..") filenames = [] yaws = [] for filename, image, metadata in tqdm(self._loader.load(), desc="Classifying Faces", total=self._loader.count, leave=False): if not metadata: msg = ("The images to be sorted do not contain alignment data. Images must have " "been generated by Faceswap's Extract process.\nIf you are sorting an " "older faceset, then you should re-extract the faces from your source " "alignments file to generate this data.") raise FaceswapError(msg) alignments = metadata["alignments"] aligned_face = AlignedFace(np.array(alignments["landmarks_xy"], dtype="float32"), image=image, centering="legacy", is_aligned=True) filenames.append(filename) yaws.append(aligned_face.pose.yaw) logger.info("Sorting...") matched_list = list(zip(filenames, yaws)) 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...") filenames = [] preds = [] for filename, image, metadata in tqdm(self._loader.load(), desc="Classifying Faces", total=self._loader.count, leave=False): if not metadata: msg = ("The images to be sorted do not contain alignment data. Images must have " "been generated by Faceswap's Extract process.\nIf you are sorting an " "older faceset, then you should re-extract the faces from your source " "alignments file to generate this data.") raise FaceswapError(msg) alignments = metadata["alignments"] face = AlignedFace(np.array(alignments["landmarks_xy"], dtype="float32"), image=image, centering="legacy", size=self._vgg_face.input_size, is_aligned=True).face filenames.append(filename) preds.append(self._vgg_face.predict(face)) logger.info("Sorting by ward linkage...") indices = self._vgg_face.sorted_similarity(np.array(preds), method="ward") img_list = np.array(filenames)[indices] return img_list
def _validate_version(self, png_meta, filename): """ Validate that there are not a mix of v1.0 extracted faces and v2.x faces. Parameters ---------- png_meta: dict The information held within the Faceswap PNG Header filename: str The full path to the file being validated Raises ------ FaceswapError If a version 1.0 face appears in a 2.x set or vice versa """ alignment_version = png_meta["source"]["alignments_version"] if not self._extract_version: logger.debug("Setting initial extract version: %s", alignment_version) self._extract_version = alignment_version if alignment_version == 1.0 and self._centering != "legacy": self._reset_cache(True) return if (self._extract_version == 1.0 and alignment_version > 1.0) or ( alignment_version == 1.0 and self._extract_version > 1.0): raise FaceswapError( "Mixing legacy and full head extracted facesets is not supported. " "The following folder contains a mix of extracted face types: " "{}".format(os.path.dirname(filename))) self._extract_version = min(alignment_version, self._extract_version)
def unmarshal(self, serialized_data): """ Unserialize data to its original object type Parameters ---------- serialized_data: varies Data in serializer format that is to be unmarshalled to its original object Returns ------- data: varies The data in a python object format Example ------ >>> serializer = get_serializer('json') >>> json_data = <json object> >>> data = serializer.unmarshal(json_data) """ logger.debug("data type: %s", type(serialized_data)) try: retval = self._unmarshal(serialized_data) except Exception as err: msg = "Error unserializing data for type {}: {}".format( type(serialized_data), str(err)) raise FaceswapError(msg) from err logger.debug("returned data type: %s", type(retval)) return retval
def initialize(self, *args, **kwargs): """ Initialize the extractor plugin Should be called from :mod:`~plugins.extract.pipeline` """ logger.debug("initialize %s: (args: %s, kwargs: %s)", self.__class__.__name__, args, kwargs) logger.info("Initializing %s (%s)...", self.name, self._plugin_type.title()) self.queue_size = 1 name = self.name.replace(" ", "_").lower() self._add_queues(kwargs["in_queue"], kwargs["out_queue"], [f"predict_{name}", f"post_{name}"]) self._compile_threads() try: self.init_model() except tf_errors.UnknownError as err: if "failed to get convolution algorithm" in str(err).lower(): msg = ("Tensorflow raised an unknown error. This is most likely caused by a " "failure to launch cuDNN which can occur for some GPU/Tensorflow " "combinations. You should enable `allow_growth` to attempt to resolve this " "issue:" "\nGUI: Go to Settings > Extract Plugins > Global and enable the " "`allow_growth` option." "\nCLI: Go to `faceswap/config/extract.ini` and change the `allow_growth " "option to `True`.") raise FaceswapError(msg) from err raise err logger.info("Initialized %s (%s) with batchsize of %s", self.name, self._plugin_type.title(), self.batchsize)
def _get_landmarks(self, filenames, batch, side): """ Obtains the 68 Point Landmarks for the images in this batch. This is only called if config item ``warp_to_landmarks`` is ``True`` or if :attr:`mask_type` is not ``None``. If the landmarks for an image cannot be found, then an error is raised. """ logger.trace("Retrieving landmarks: (filenames: %s, side: '%s')", filenames, side) src_points = [ self._landmarks[side].get(sha1(face).hexdigest(), None) for face in batch ] # Raise error on missing alignments if not all(isinstance(pts, np.ndarray) for pts in src_points): indices = [ idx for idx, hsh in enumerate(src_points) if hsh is None ] missing = [filenames[idx] for idx in indices] msg = ( "Files missing alignments for this batch: {}" "\nAt least one of your images does not have a matching entry in your " "alignments file." "\nIf you are training with a mask or using 'warp to landmarks' then every " "face you intend to train on must exist within the alignments file." "\nThe specific files that caused this failure are listed above." "\nMost likely there will be more than just these files missing from the " "alignments file. You can use the Alignments Tool to help identify missing " "alignments".format(missing)) raise FaceswapError(msg) logger.trace("Returning: (src_points: %s)", [str(src) for src in src_points]) return np.array(src_points)
def save(self, filename, data): """ Serialize data and save to a file Parameters ---------- filename: str The path to where the serialized file should be saved data: varies The data that is to be serialized to file Example ------ >>> serializer = get_serializer('json') >>> data ['foo', 'bar'] >>> json_file = '/path/to/json/file.json' >>> serializer.save(json_file, data) """ logger.debug("filename: %s, data type: %s", filename, type(data)) filename = self._check_extension(filename) try: with open(filename, self._write_option) as s_file: s_file.write(self.marshal(data)) except IOError as err: msg = "Error writing to '{}': {}".format(filename, err.strerror) raise FaceswapError(msg) from err
def pre_fill(self, filenames, side): """ When warp to landmarks is enabled, the cache must be pre-filled, as each side needs access to the other side's alignments. Parameters ---------- filenames: list The list of full paths to the images to load the metadata from side: str `"a"` or `"b"`. The side of the model being cached. Used for info output """ with self._lock: for filename, meta in tqdm( read_image_meta_batch(filenames), desc="WTL: Caching Landmarks ({})".format(side.upper()), total=len(filenames), leave=False): if "itxt" not in meta or "alignments" not in meta["itxt"]: raise FaceswapError( f"Invalid face image found. Aborting: '{filename}'") size = meta["width"] meta = meta["itxt"] # Version Check self._validate_version(meta, filename) detected_face = self._add_aligned_face(filename, meta["alignments"], size) self._cache[os.path.basename( filename)]["detected_face"] = detected_face self._partial_load = True
def random_warp(self, image): """ get pair of random warped images from aligned face image """ logger.trace("Randomly warping image") height, width = image.shape[0:2] coverage = self.get_coverage(image) try: assert height == width and height % 2 == 0 except AssertionError as err: msg = ( "Training images should be square with an even number of pixels across each " "side. An image was found with width: {}, height: {}." "\nMost likely this is a frame rather than a face within your training set. " "\nMake sure that the only images within your training set are faces generated " "from the Extract process.".format(width, height)) raise FaceswapError(msg) from err range_ = np.linspace(height // 2 - coverage // 2, height // 2 + coverage // 2, 5, dtype='float32') mapx = np.broadcast_to(range_, (5, 5)).copy() mapy = mapx.T # mapx, mapy = np.float32(np.meshgrid(range_,range_)) # instead of broadcast pad = int(1.25 * self.input_size) slices = slice(pad // 10, -pad // 10) dst_slices = [ slice(0, (size + 1), (size // 4)) for size in self.output_sizes ] interp = np.empty((2, self.input_size, self.input_size), dtype='float32') for i, map_ in enumerate([mapx, mapy]): map_ = map_ + np.random.normal(size=(5, 5), scale=self.scale) interp[i] = cv2.resize(map_, (pad, pad))[slices, slices] # pylint:disable=no-member warped_image = cv2.remap( # pylint:disable=no-member image, interp[0], interp[1], cv2.INTER_LINEAR) # pylint:disable=no-member logger.trace("Warped image shape: %s", warped_image.shape) src_points = np.stack([mapx.ravel(), mapy.ravel()], axis=-1) dst_points = [ np.mgrid[dst_slice, dst_slice] for dst_slice in dst_slices ] mats = [ umeyama(src_points, True, dst_pts.T.reshape(-1, 2))[0:2] for dst_pts in dst_points ] target_images = [ cv2.warpAffine( image, # pylint:disable=no-member mat, (self.output_sizes[idx], self.output_sizes[idx])) for idx, mat in enumerate(mats) ] logger.trace("Target image shapes: %s", [tgt.shape for tgt in target_images]) return self.compile_images(warped_image, target_images)
def load(self, filename): """ Load data from an existing serialized file Parameters ---------- filename: str The path to the serialized file Returns ---------- data: varies The data in a python object format Example ------ >>> serializer = get_serializer('json') >>> json_file = '/path/to/json/file.json' >>> data = serializer.load(json_file) """ logger.debug("filename: %s", filename) try: with open(filename, self._read_option) as s_file: data = s_file.read() logger.debug("stored data type: %s", type(data)) retval = self.unmarshal(data) except IOError as err: msg = "Error reading from '{}': {}".format(filename, err.strerror) raise FaceswapError(msg) from err logger.debug("data type: %s", type(retval)) return retval