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
