def cache_metadata(self, filenames):
""" Obtain the batch with metadata for items that need caching and cache them to
:attr:`_cache`.
Parameters
----------
filenames: list
List of full paths to image file names
Returns
-------
:class:`numpy.ndarray`
The batch of face images loaded from disk
"""
keys = [os.path.basename(filename) for filename in filenames]
with self._lock:
if _check_reset(self):
self._reset_cache(False)
needs_cache = [
filename for filename, key in zip(filenames, keys)
if not self._cache[key]["cached"]
]
logger.trace("Needs cache: %s", needs_cache)
if not needs_cache:
# Don't bother reading the metadata if no images in this batch need caching
logger.debug("All metadata already cached for: %s", keys)
return read_image_batch(filenames)
batch, metadata = read_image_batch(filenames, with_metadata=True)
# Populate items into cache
for filename in needs_cache:
key = os.path.basename(filename)
meta = metadata[filenames.index(filename)]
# Version Check
self._validate_version(meta, filename)
if self._partial_load: # Faces already loaded for Warp-to-landmarks
detected_face = self._cache[key]["detected_face"]
else:
detected_face = self._add_aligned_face(
filename, meta["alignments"], batch.shape[1])
self._add_mask(filename, detected_face)
for area in ("eye", "mouth"):
self._add_localized_mask(filename, detected_face, area)
self._cache[key]["cached"] = True
# Update the :attr:`cache_full` attribute
cache_full = all(item["cached"] for item in self._cache.values())
if cache_full:
logger.verbose("Cache filled: '%s'",
os.path.dirname(filenames[0]))
self._cache_full = cache_full
return batch
def _process_batch(self, filenames, side):
""" Performs the augmentation and compiles target images and samples. See
:func:`minibatch_ab` for more details on the output. """
logger.trace("Process batch: (filenames: '%s', side: '%s')", filenames,
side)
batch = read_image_batch(filenames)
batch = self._apply_mask(filenames, batch, side)
processed = dict()
# Initialize processing training size on first image
if not self._processing.initialized:
self._processing.initialize(batch.shape[1])
# Get Landmarks prior to manipulating the image
if self._warp_to_landmarks:
batch_src_pts = self._get_landmarks(filenames, side)
batch_dst_pts = self._get_closest_match(filenames, side,
batch_src_pts)
warp_kwargs = dict(batch_src_points=batch_src_pts,
batch_dst_points=batch_dst_pts)
else:
warp_kwargs = dict()
# Color Augmentation of the image only
if self._augment_color:
batch[..., :3] = self._processing.color_adjust(batch[..., :3])
# Random Transform and flip
batch = self._processing.transform(batch)
if not self._no_flip:
batch = self._processing.random_flip(batch)
# Add samples to output if this is for display
if self._processing.is_display:
processed["samples"] = batch[..., :3].astype("float32") / 255.0
# Get Targets
processed.update(self._processing.get_targets(batch))
# Random Warp # TODO change masks to have a input mask and a warped target mask
if not self._config["disable_warp"]:
processed["feed"] = [
self._processing.warp(batch[..., :3], self._warp_to_landmarks,
**warp_kwargs)
]
else:
size = (self._model_input_size, self._model_input_size)
processed["feed"] = [
np.array([cv2.resize(img, size)
for img in batch[..., :3]]).astype("float32") / 255.0
]
logger.trace(
"Processed batch: (filenames: %s, side: '%s', processed: %s)",
filenames, side, {
k:
v.shape if isinstance(v, np.ndarray) else [i.shape for i in v]
for k, v in processed.items()
})
return processed
def _process_batch(self, filenames, side):
""" Performs the augmentation and compiles target images and samples. See
:func:`minibatch_ab` for more details on the output. """
logger.trace("Process batch: (filenames: '%s', side: '%s')", filenames,
side)
batch = read_image_batch(filenames)
processed = dict()
to_landmarks = self._training_opts["warp_to_landmarks"]
# Initialize processing training size on first image
if not self._processing.initialized:
self._processing.initialize(batch.shape[1])
# Get Landmarks prior to manipulating the image
if self._mask_class or to_landmarks:
batch_src_pts = self._get_landmarks(filenames, batch, side)
# Color augmentation before mask is added
if self._training_opts["augment_color"]:
batch = self._processing.color_adjust(batch)
# Add mask to batch prior to transforms and warps
if self._mask_class:
batch = np.array([
self._mask_class(src_pts, image, channels=4).mask
for src_pts, image in zip(batch_src_pts, batch)
])
# Random Transform and flip
batch = self._processing.transform(batch)
if not self._training_opts["no_flip"]:
batch = self._processing.random_flip(batch)
# Add samples to output if this is for display
if self._processing.is_display:
processed["samples"] = batch[..., :3].astype("float32") / 255.0
# Get Targets
processed.update(self._processing.get_targets(batch))
# Random Warp
if to_landmarks:
warp_kwargs = dict(batch_src_points=batch_src_pts,
batch_dst_points=self._get_closest_match(
filenames, side, batch_src_pts))
else:
warp_kwargs = dict()
processed["feed"] = self._processing.warp(batch[..., :3], to_landmarks,
**warp_kwargs)
logger.trace(
"Processed batch: (filenames: %s, side: '%s', processed: %s)",
filenames, side, {
k:
v.shape if isinstance(v, np.ndarray) else [i.shape for i in v]
for k, v in processed.items()
})
return processed
def _process_batch(self, filenames, side):
""" Performs the augmentation and compiles target images and samples.
If this is the first time a face has been loaded, then it's meta data is extracted from the
png header and added to :attr:`_face_cache`
See
:func:`minibatch_ab` for more details on the output.
Parameters
----------
filenames: list
List of full paths to image file names
side: str
The side of the model being trained on (`a` or `b`)
"""
logger.trace("Process batch: (filenames: '%s', side: '%s')", filenames,
side)
if not self._face_cache.cache_full:
batch = self._face_cache.cache_metadata(filenames)
else:
batch = read_image_batch(filenames)
cache = self._face_cache.get_items(filenames)
batch, landmarks = self._crop_to_center(filenames, cache, batch, side)
batch = self._apply_mask(filenames, cache, batch, side)
processed = dict()
# Initialize processing training size on first image
if not self._processing.initialized:
self._processing.initialize(batch.shape[1])
# Get Landmarks prior to manipulating the image
if self._warp_to_landmarks:
batch_dst_pts = self._get_closest_match(filenames, side, landmarks)
warp_kwargs = dict(batch_src_points=landmarks,
batch_dst_points=batch_dst_pts)
else:
warp_kwargs = dict()
# Color Augmentation of the image only
if self._augment_color:
batch[..., :3] = self._processing.color_adjust(batch[..., :3])
# Random Transform and flip
batch = self._processing.transform(batch)
if not self._no_flip:
batch = self._processing.random_flip(batch)
# Switch color order for RGB models
if self._color_order == "rgb":
batch = batch[..., [2, 1, 0, 3]]
# Add samples to output if this is for display
if self._processing.is_display:
processed["samples"] = batch[..., :3].astype("float32") / 255.0
# Get Targets
processed.update(self._processing.get_targets(batch))
# Random Warp # TODO change masks to have a input mask and a warped target mask
if self._no_warp:
processed["feed"] = [self._processing.skip_warp(batch[..., :3])]
else:
processed["feed"] = [
self._processing.warp(batch[..., :3], self._warp_to_landmarks,
**warp_kwargs)
]
logger.trace(
"Processed batch: (filenames: %s, side: '%s', processed: %s)",
filenames, side, {
k:
v.shape if isinstance(v, np.ndarray) else [i.shape for i in v]
for k, v in processed.items()
})
return processed
def cache_metadata(self, filenames):
""" Obtain the batch with metadata for items that need caching and cache them to
:attr:`_cache`.
Parameters
----------
filenames: list
List of full paths to image file names
Returns
-------
:class:`numpy.ndarray`
The batch of face images loaded from disk
"""
keys = [os.path.basename(filename) for filename in filenames]
with self._lock:
if _check_reset(self):
self._reset_cache(False)
needs_cache = [
filename for filename, key in zip(filenames, keys)
if not self._cache[key]["cached"]
]
logger.trace("Needs cache: %s", needs_cache)
if not needs_cache:
# Don't bother reading the metadata if no images in this batch need caching
logger.debug("All metadata already cached for: %s", keys)
return read_image_batch(filenames)
batch, metadata = read_image_batch(filenames, with_metadata=True)
if len(batch.shape) == 1:
folder = os.path.dirname(filenames[0])
details = [
"{0} ({1})".format(
key, f"{img.shape[1]}px" if isinstance(
img, np.ndarray) else type(img))
for key, img in zip(keys, batch)
]
msg = (
f"There are mismatched image sizes in the folder '{folder}'. All training "
"images for each side must have the same dimensions.\nThe batch that "
f"failed contains the following files:\n{details}.")
raise FaceswapError(msg)
# Populate items into cache
for filename in needs_cache:
key = os.path.basename(filename)
meta = metadata[filenames.index(filename)]
# Version Check
self._validate_version(meta, filename)
if self._partial_load: # Faces already loaded for Warp-to-landmarks
detected_face = self._cache[key]["detected_face"]
else:
detected_face = self._add_aligned_face(
filename, meta["alignments"], batch.shape[1])
self._add_mask(filename, detected_face)
for area in ("eye", "mouth"):
self._add_localized_mask(filename, detected_face, area)
self._cache[key]["cached"] = True
# Update the :attr:`cache_full` attribute
cache_full = all(item["cached"] for item in self._cache.values())
if cache_full:
logger.verbose("Cache filled: '%s'",
os.path.dirname(filenames[0]))
self._cache_full = cache_full
return batch
