def _rotate_image_by_angle(self, image, angle):
""" Rotate an image by a given angle.
From: https://stackoverflow.com/questions/22041699 """
logger.trace("Rotating image: (image: %s, angle: %s)", image.shape,
angle)
channels_first = image.shape[0] <= 4
if channels_first:
image = np.moveaxis(image, 0, 2)
height, width = image.shape[:2]
image_center = (width / 2, height / 2)
rotation_matrix = cv2.getRotationMatrix2D( # pylint: disable=no-member
image_center, -1. * angle, 1.)
rotation_matrix[0, 2] += self.input_size / 2 - image_center[0]
rotation_matrix[1, 2] += self.input_size / 2 - image_center[1]
logger.trace("Rotated image: (rotation_matrix: %s", rotation_matrix)
image = cv2.warpAffine(
image, # pylint: disable=no-member
rotation_matrix,
(self.input_size, self.input_size))
if channels_first:
image = np.moveaxis(image, 2, 0)
return image, rotation_matrix
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 _get_adjusted_boxes(self, original_boxes):
""" Obtain an array of adjusted bounding boxes based on the number of re-feed iterations
that have been selected and the minimum dimension of the original bounding box.
Parameters
----------
original_boxes: :class:`numpy.ndarray`
The original ('x', 'y', 'w', 'h') detected face boxes corresponding to the incoming
detected face objects
Returns
-------
:class:`numpy.ndarray`
The original boxes (in position 0) and the randomly adjusted bounding boxes
"""
if self._re_feed == 0:
return original_boxes[None, ...]
beta = 0.05
max_shift = np.min(original_boxes[..., 2:], axis=1) * beta
rands = np.random.rand(self._re_feed, *original_boxes.shape) * 2 - 1
new_boxes = np.rint(original_boxes +
(rands * max_shift[None, :, None])).astype("int32")
retval = np.concatenate((original_boxes[None, ...], new_boxes))
logger.trace(retval)
return retval
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)
batch = self.predict(batch)
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 _remove_zero_sized_faces(batch):
""" Remove items from dict where detected face is of zero size
or face falls entirely outside of image """
dims = [img.shape[:2] for img in batch["image"]]
logger.trace("image dims: %s", dims)
batch["detected_faces"] = [[
face for face in faces if face.right > 0 and face.left < dim[1]
and face.bottom > 0 and face.top < dim[0]
] for dim, faces in zip(dims, batch.get("detected_faces", list()))]
def get_batch(self, queue):
""" Get items for inputting to the detector plugin in batches
Items are returned from the ``queue`` in batches of
:attr:`~plugins.extract._base.Extractor.batchsize`
Remember to put ``'EOF'`` to the out queue after processing
the final batch
Outputs items in the following format. All lists are of length
:attr:`~plugins.extract._base.Extractor.batchsize`:
>>> {'filename': [<filenames of source frames>],
>>> 'image': [<source images>],
>>> 'scaled_image': <np.array of images standardized for prediction>,
>>> 'scale': [<scaling factors for each image>],
>>> 'pad': [<padding for each image>],
>>> 'detected_faces': [[<lib.faces_detect.DetectedFace objects]]}
Parameters
----------
queue : queue.Queue()
The ``queue`` that the batch will be fed from. This will be a queue that loads
images.
Returns
-------
exhausted, bool
``True`` if queue is exhausted, ``False`` if not.
batch, dict
A dictionary of lists of :attr:`~plugins.extract._base.Extractor.batchsize`.
"""
exhausted = False
batch = dict()
for _ in range(self.batchsize):
item = self._get_item(queue)
if item == "EOF":
exhausted = True
break
for key, val in item.items():
batch.setdefault(key, []).append(val)
scaled_image, scale, pad = self._compile_detection_image(
item["image"])
batch.setdefault("scaled_image", []).append(scaled_image)
batch.setdefault("scale", []).append(scale)
batch.setdefault("pad", []).append(pad)
if batch:
batch["scaled_image"] = np.array(batch["scaled_image"],
dtype="float32")
logger.trace(
"Returning batch: %s", {
k: v.shape if isinstance(v, np.ndarray) else v
for k, v in batch.items() if k != "image"
})
else:
logger.trace(item)
return exhausted, batch
def finalize(self, batch):
""" Finalize the output from Masker
This should be called as the final task of each `plugin`.
It strips unneeded items from the :attr:`batch` ``dict`` and pairs the detected faces back
up with their original frame before yielding each frame.
Outputs items in the format:
>>> {'image': [<original frame>],
>>> 'filename': [<frame filename>),
>>> 'detected_faces': [<lib.faces_detect.DetectedFace objects>]}
Parameters
----------
batch : dict
The final ``dict`` from the `plugin` process. It must contain the `keys`:
``detected_faces``, ``filename``, ``image``
Yields
------
dict
A ``dict`` for each frame containing the ``image``, ``filename`` and list of
:class:`lib.faces_detect.DetectedFace` objects.
"""
for mask, face in zip(batch["prediction"], batch["detected_faces"]):
face.add_mask(self._storage_name,
mask,
face.feed_matrix,
face.feed_interpolators[1],
storage_size=self._storage_size)
face.feed = dict()
self._remove_invalid_keys(batch,
("detected_faces", "filename", "image"))
logger.trace(
"Item out: %s",
{key: val
for key, val in batch.items() if key != "image"})
for filename, image, face in zip(batch["filename"], batch["image"],
batch["detected_faces"]):
self._output_faces.append(face)
if len(self._output_faces) != self._faces_per_filename[filename]:
continue
retval = dict(filename=filename,
image=image,
detected_faces=self._output_faces)
self._output_faces = []
logger.trace(
"Yielding: (filename: '%s', image: %s, detected_faces: %s)",
retval["filename"], retval["image"].shape,
len(retval["detected_faces"]))
yield retval
def _scale_image(image, image_size, scale):
""" Scale the image and optional pad to given size """
interpln = cv2.INTER_CUBIC if scale > 1.0 else cv2.INTER_AREA
if scale != 1.0:
dims = (int(image_size[1] * scale), int(image_size[0] * scale))
logger.trace("Resizing detection image from %s to %s. Scale=%s",
"x".join(str(i) for i in reversed(image_size)),
"x".join(str(i) for i in dims), scale)
image = cv2.resize(image, dims, interpolation=interpln)
logger.trace("Resized image shape: %s", image.shape)
return image
def _remove_zero_sized_faces(self, batch_faces):
""" Remove items from batch_faces where detected face is of zero size
or face falls entirely outside of image """
logger.trace("Input sizes: %s", [len(face) for face in batch_faces])
retval = [[face
for face in faces
if face.right > 0 and face.left < self.input_size
and face.bottom > 0 and face.top < self.input_size]
for faces in batch_faces]
logger.trace("Output sizes: %s", [len(face) for face in retval])
return retval
def _normalize_faces(self, faces):
""" Normalizes the face for feeding into model
The normalization method is dictated by the command line argument `-nh (--normalization)`
"""
if self.normalize_method is None:
return faces
logger.trace("Normalizing faces")
meth = getattr(self, "_normalize_{}".format(self.normalize_method.lower()))
faces = [meth(face) for face in faces]
logger.trace("Normalized faces")
return faces
def _pad_image(self, image):
""" Pad a resized image to input size """
height, width = image.shape[:2]
if width < self.input_size or height < self.input_size:
pad_l = (self.input_size - width) // 2
pad_r = (self.input_size - width) - pad_l
pad_t = (self.input_size - height) // 2
pad_b = (self.input_size - height) - pad_t
image = cv2.copyMakeBorder( # pylint:disable=no-member
image, pad_t, pad_b, pad_l, pad_r, cv2.BORDER_CONSTANT) # pylint:disable=no-member
logger.trace("Padded image shape: %s", image.shape)
return image
def _compile_detection_image(self, input_image):
""" Compile the detection image for feeding into the model"""
image = self._convert_color(input_image)
image_size = image.shape[:2]
scale = self._set_scale(image_size)
pad = self._set_padding(image_size, scale)
image = self._scale_image(image, image_size, scale)
image = self._pad_image(image)
logger.trace("compiled: (images shape: %s, scale: %s, pad: %s)", image.shape, scale, pad)
return image, scale, pad
def finalize(self, batch):
""" Finalize the output from Masker
This should be called as the final task of each `plugin`.
Pairs the detected faces back up with their original frame before yielding each frame.
Parameters
----------
batch : dict
The final ``dict`` from the `plugin` process. It must contain the `keys`:
``detected_faces``, ``filename``, ``feed_faces``, ``roi_masks``
Yields
------
:class:`~plugins.extract.pipeline.ExtractMedia`
The :attr:`DetectedFaces` list will be populated for this class with the bounding
boxes, landmarks and masks for the detected faces found in the frame.
"""
for mask, face, feed_face, roi_mask in zip(batch["prediction"],
batch["detected_faces"],
batch["feed_faces"],
batch["roi_masks"]):
self._crop_out_of_bounds(mask, roi_mask)
face.add_mask(self._storage_name,
mask,
feed_face.adjusted_matrix,
feed_face.interpolators[1],
storage_size=self._storage_size,
storage_centering=self._storage_centering)
del batch["feed_faces"]
logger.trace(
"Item out: %s", {
key: val.shape if isinstance(val, np.ndarray) else val
for key, val in batch.items()
})
for filename, face in zip(batch["filename"], batch["detected_faces"]):
self._output_faces.append(face)
if len(self._output_faces) != self._faces_per_filename[filename]:
continue
output = self._extract_media.pop(filename)
output.add_detected_faces(self._output_faces)
self._output_faces = []
logger.trace(
"Yielding: (filename: '%s', image: %s, detected_faces: %s)",
output.filename, output.image_shape,
len(output.detected_faces))
yield output
def _collect_item(self, queue):
""" Collect the item from the :attr:`_rollover` dict or from the queue
Add face count per frame to self._faces_per_filename for joining
batches back up in finalize """
if self._rollover is not None:
logger.trace("Getting from _rollover: (filename: `%s`, faces: %s)",
self._rollover.filename, len(self._rollover.detected_faces))
item = self._rollover
self._rollover = None
else:
item = self._get_item(queue)
if item != "EOF":
logger.trace("Getting from queue: (filename: %s, faces: %s)",
item.filename, len(item.detected_faces))
self._faces_per_filename[item.filename] = len(item.detected_faces)
return item
def _compile_detection_image(self, item):
""" Compile the detection image for feeding into the model
Parameters
----------
item: :class:`plugins.extract.pipeline.ExtractMedia`
The input item from the pipeline
"""
image = item.get_image_copy(self.color_format)
scale = self._set_scale(item.image_size)
pad = self._set_padding(item.image_size, scale)
image = self._scale_image(image, item.image_size, scale)
image = self._pad_image(image)
logger.trace("compiled: (images shape: %s, scale: %s, pad: %s)", image.shape, scale, pad)
return image, scale, pad
def _collect_item(self, queue):
""" Collect the item from the _rollover dict or from the queue
Add face count per frame to self._faces_per_filename for joining
batches back up in finalize """
if self._rollover:
logger.trace("Getting from _rollover: (filename: `%s`, faces: %s)",
self._rollover["filename"], len(self._rollover["detected_faces"]))
item = self._rollover
self._rollover = dict()
else:
item = self._get_item(queue)
if item != "EOF":
logger.trace("Getting from queue: (filename: %s, faces: %s)",
item["filename"], len(item["detected_faces"]))
self._faces_per_filename[item["filename"]] = len(item["detected_faces"])
return item
def finalize(self, batch):
""" Finalize the output from Aligner
This should be called as the final task of each `plugin`.
Pairs the detected faces back up with their original frame before yielding each frame.
Parameters
----------
batch : dict
The final ``dict`` from the `plugin` process. It must contain the `keys`:
``detected_faces``, ``landmarks``, ``filename``
Yields
------
:class:`~plugins.extract.pipeline.ExtractMedia`
The :attr:`DetectedFaces` list will be populated for this class with the bounding boxes
and landmarks for the detected faces found in the frame.
"""
for face, landmarks in zip(batch["detected_faces"],
batch["landmarks"]):
if not isinstance(landmarks, np.ndarray):
landmarks = np.array(landmarks)
face.landmarks_xy = landmarks
logger.trace(
"Item out: %s", {
key: val.shape if isinstance(val, np.ndarray) else val
for key, val in batch.items()
})
for filename, face in zip(batch["filename"], batch["detected_faces"]):
self._output_faces.append(face)
if len(self._output_faces) != self._faces_per_filename[filename]:
continue
output = self._extract_media.pop(filename)
output.add_detected_faces(self._output_faces)
self._output_faces = []
logger.trace(
"Final Output: (filename: '%s', image shape: %s, detected_faces: %s, "
"item: %s)", output.filename, output.image_shape,
output.detected_faces, output)
yield output
def _rotate_face(face, rotation_matrix):
""" Rotates the detection bounding box around the given rotation matrix.
Parameters
----------
face: :class:`DetectedFace`
A :class:`DetectedFace` containing the `x`, `w`, `y`, `h` detection bounding box
points.
rotation_matrix: numpy.ndarray
The rotation matrix to rotate the given object by.
Returns
-------
:class:`DetectedFace`
The same class with the detection bounding box points rotated by the given matrix.
"""
logger.trace("Rotating face: (face: %s, rotation_matrix: %s)", face, rotation_matrix)
bounding_box = [[face.left, face.top],
[face.right, face.top],
[face.right, face.bottom],
[face.left, face.bottom]]
rotation_matrix = cv2.invertAffineTransform(rotation_matrix)
points = np.array(bounding_box, "int32")
points = np.expand_dims(points, axis=0)
transformed = cv2.transform(points, rotation_matrix).astype("int32")
rotated = transformed.squeeze()
# Bounding box should follow x, y planes, so get min/max for non-90 degree rotations
pt_x = min([pnt[0] for pnt in rotated])
pt_y = min([pnt[1] for pnt in rotated])
pt_x1 = max([pnt[0] for pnt in rotated])
pt_y1 = max([pnt[1] for pnt in rotated])
width = pt_x1 - pt_x
height = pt_y1 - pt_y
face.x = int(pt_x)
face.y = int(pt_y)
face.w = int(width)
face.h = int(height)
return face
def finalize(self, batch):
""" Finalize the output from Detector
This should be called as the final task of each ``plugin``.
Parameters
----------
batch : dict
The final ``dict`` from the `plugin` process. It must contain the keys ``filename``,
``faces``
Yields
------
:class:`~plugins.extract.pipeline.ExtractMedia`
The :attr:`DetectedFaces` list will be populated for this class with the bounding boxes
for the detected faces found in the frame.
"""
if not isinstance(batch, dict):
logger.trace("Item out: %s", batch)
return batch
logger.trace("Item out: %s", {k: v.shape if isinstance(v, np.ndarray) else v
for k, v in batch.items()})
batch_faces = [[self.to_detected_face(face[0], face[1], face[2], face[3])
for face in faces]
for faces in batch["prediction"]]
# Rotations
if any(m.any() for m in batch["rotmat"]) and any(batch_faces):
batch_faces = [[self._rotate_face(face, rotmat) if rotmat.any() else face
for face in faces]
for faces, rotmat in zip(batch_faces, batch["rotmat"])]
# Remove zero sized faces
batch_faces = self._remove_zero_sized_faces(batch_faces)
# Scale back out to original frame
batch["detected_faces"] = [[self.to_detected_face((face.left - pad[0]) / scale,
(face.top - pad[1]) / scale,
(face.right - pad[0]) / scale,
(face.bottom - pad[1]) / scale)
for face in faces]
for scale, pad, faces in zip(batch["scale"],
batch["pad"],
batch_faces)]
if self.min_size > 0 and batch.get("detected_faces", None):
batch["detected_faces"] = self._filter_small_faces(batch["detected_faces"])
batch = self._dict_lists_to_list_dicts(batch)
for item in batch:
output = self._extract_media.pop(item["filename"])
output.add_detected_faces(item["detected_faces"])
logger.trace("final output: (filename: '%s', image shape: %s, detected_faces: %s, "
"item: %s", output.filename, output.image_shape, output.detected_faces,
output)
yield output
def get_batch(self, queue):
""" Get items for inputting into the masker from the queue in batches
Items are returned from the ``queue`` in batches of
:attr:`~plugins.extract._base.Extractor.batchsize`
To ensure consistent batch sizes for masker the items are split into separate items for
each :class:`lib.faces_detect.DetectedFace` object.
Remember to put ``'EOF'`` to the out queue after processing
the final batch
Outputs items in the following format. All lists are of length
:attr:`~plugins.extract._base.Extractor.batchsize`:
>>> {'filename': [<filenames of source frames>],
>>> 'image': [<source images>],
>>> 'detected_faces': [[<lib.faces_detect.DetectedFace objects]]}
Parameters
----------
queue : queue.Queue()
The ``queue`` that the plugin will be fed from.
Returns
-------
exhausted, bool
``True`` if queue is exhausted, ``False`` if not
batch, dict
A dictionary of lists of :attr:`~plugins.extract._base.Extractor.batchsize`:
"""
exhausted = False
batch = dict()
idx = 0
while idx < self.batchsize:
item = self._collect_item(queue)
if item == "EOF":
logger.trace("EOF received")
exhausted = True
break
# Put frames with no faces into the out queue to keep TQDM consistent
if not item["detected_faces"]:
self._queues["out"].put(item)
continue
for f_idx, face in enumerate(item["detected_faces"]):
face.image = self._convert_color(item["image"])
face.load_feed_face(face.image,
size=self.input_size,
coverage_ratio=1.0,
dtype="float32",
is_aligned_face=self._image_is_aligned)
batch.setdefault("detected_faces", []).append(face)
batch.setdefault("filename", []).append(item["filename"])
batch.setdefault("image", []).append(item["image"])
idx += 1
if idx == self.batchsize:
frame_faces = len(item["detected_faces"])
if f_idx + 1 != frame_faces:
self._rollover = {
k: v[f_idx + 1:] if k == "detected_faces" else v
for k, v in item.items()
}
logger.trace(
"Rolled over %s faces of %s to next batch for '%s'",
len(self._rollover["detected_faces"]), frame_faces,
item["filename"])
break
if batch:
logger.trace(
"Returning batch: %s", {
k: v.shape if isinstance(v, np.ndarray) else v
for k, v in batch.items()
})
else:
logger.trace(item)
return exhausted, batch
def get_batch(self, queue):
""" Get items for inputting into the aligner from the queue in batches
Items are returned from the ``queue`` in batches of
:attr:`~plugins.extract._base.Extractor.batchsize`
Items are received as :class:`~plugins.extract.pipeline.ExtractMedia` objects and converted
to ``dict`` for internal processing.
To ensure consistent batch sizes for aligner the items are split into separate items for
each :class:`~lib.faces_detect.DetectedFace` object.
Remember to put ``'EOF'`` to the out queue after processing
the final batch
Outputs items in the following format. All lists are of length
:attr:`~plugins.extract._base.Extractor.batchsize`:
>>> {'filename': [<filenames of source frames>],
>>> 'image': [<source images>],
>>> 'detected_faces': [[<lib.faces_detect.DetectedFace objects]]}
Parameters
----------
queue : queue.Queue()
The ``queue`` that the plugin will be fed from.
Returns
-------
exhausted, bool
``True`` if queue is exhausted, ``False`` if not
batch, dict
A dictionary of lists of :attr:`~plugins.extract._base.Extractor.batchsize`:
"""
exhausted = False
batch = dict()
idx = 0
while idx < self.batchsize:
item = self._collect_item(queue)
if item == "EOF":
logger.trace("EOF received")
exhausted = True
break
# Put frames with no faces into the out queue to keep TQDM consistent
if not item.detected_faces:
self._queues["out"].put(item)
continue
converted_image = item.get_image_copy(self.color_format)
for f_idx, face in enumerate(item.detected_faces):
batch.setdefault("image", []).append(converted_image)
batch.setdefault("detected_faces", []).append(face)
batch.setdefault("filename", []).append(item.filename)
idx += 1
if idx == self.batchsize:
frame_faces = len(item.detected_faces)
if f_idx + 1 != frame_faces:
self._rollover = ExtractMedia(
item.filename,
item.image,
detected_faces=item.detected_faces[f_idx + 1:])
logger.trace(
"Rolled over %s faces of %s to next batch for '%s'",
len(self._rollover.detected_faces), frame_faces,
item.filename)
break
if batch:
logger.trace(
"Returning batch: %s", {
k: v.shape if isinstance(v, np.ndarray) else v
for k, v in batch.items()
})
else:
logger.trace(item)
return exhausted, batch
def _set_scale(self, image_size):
""" Set the scale factor for incoming image """
scale = self.input_size / max(image_size)
logger.trace("Detector scale: %s", scale)
return scale
def get_batch(self, queue):
""" Get items for inputting into the masker from the queue in batches
Items are returned from the ``queue`` in batches of
:attr:`~plugins.extract._base.Extractor.batchsize`
Items are received as :class:`~plugins.extract.pipeline.ExtractMedia` objects and converted
to ``dict`` for internal processing.
To ensure consistent batch sizes for masker the items are split into separate items for
each :class:`~lib.align.DetectedFace` object.
Remember to put ``'EOF'`` to the out queue after processing
the final batch
Outputs items in the following format. All lists are of length
:attr:`~plugins.extract._base.Extractor.batchsize`:
>>> {'filename': [<filenames of source frames>],
>>> 'detected_faces': [[<lib.align.DetectedFace objects]]}
Parameters
----------
queue : queue.Queue()
The ``queue`` that the plugin will be fed from.
Returns
-------
exhausted, bool
``True`` if queue is exhausted, ``False`` if not
batch, dict
A dictionary of lists of :attr:`~plugins.extract._base.Extractor.batchsize`:
"""
exhausted = False
batch = dict()
idx = 0
while idx < self.batchsize:
item = self._collect_item(queue)
if item == "EOF":
logger.trace("EOF received")
exhausted = True
break
# Put frames with no faces into the out queue to keep TQDM consistent
if not item.detected_faces:
self._queues["out"].put(item)
continue
for f_idx, face in enumerate(item.detected_faces):
image = item.get_image_copy(self.color_format)
roi = np.ones((*item.image_size[:2], 1), dtype="float32")
if not self._image_is_aligned:
# Add the ROI mask to image so we can get the ROI mask with a single warp
image = np.concatenate([image, roi], axis=-1)
feed_face = AlignedFace(face.landmarks_xy,
image=image,
centering=self._storage_centering,
size=self.input_size,
coverage_ratio=self.coverage_ratio,
dtype="float32",
is_aligned=self._image_is_aligned)
if not self._image_is_aligned:
# Split roi mask from feed face alpha channel
roi_mask = feed_face.face[..., 3]
feed_face._face = feed_face.face[..., :3] # pylint:disable=protected-access
else:
# We have to do the warp here as AlignedFace did not perform it
roi_mask = transform_image(roi,
feed_face.matrix,
feed_face.size,
padding=feed_face.padding)
batch.setdefault("roi_masks", []).append(roi_mask)
batch.setdefault("detected_faces", []).append(face)
batch.setdefault("feed_faces", []).append(feed_face)
batch.setdefault("filename", []).append(item.filename)
idx += 1
if idx == self.batchsize:
frame_faces = len(item.detected_faces)
if f_idx + 1 != frame_faces:
self._rollover = ExtractMedia(
item.filename,
item.image,
detected_faces=item.detected_faces[f_idx + 1:])
logger.trace(
"Rolled over %s faces of %s to next batch for '%s'",
len(self._rollover.detected_faces), frame_faces,
item.filename)
break
if batch:
logger.trace(
"Returning batch: %s", {
k: v.shape if isinstance(v, np.ndarray) else v
for k, v in batch.items()
})
else:
logger.trace(item)
return exhausted, batch
def _rotate_rect(bounding_box, rotation_matrix):
""" Rotate a bounding box dict based on the rotation_matrix"""
logger.trace("Rotating bounding box")
bounding_box = rotate_landmarks(bounding_box, rotation_matrix)
return bounding_box
def finalize(self, batch):
""" Finalize the output from Detector
This should be called as the final task of each ``plugin``.
It strips unneeded items from the :attr:`batch` ``dict`` and performs standard final
processing on each item
Outputs items in the format:
>>> {'image': [<original frame>],
>>> 'filename': [<frame filename>),
>>> 'detected_faces': [<lib.faces_detect.DetectedFace objects>]}
Parameters
----------
batch : dict
The final ``dict`` from the `plugin` process. It must contain the keys ``image``,
``filename``, ``faces``
Yields
------
dict
A ``dict`` for each frame containing the ``image``, ``filename`` and ``list`` of
``detected_faces``
"""
if not isinstance(batch, dict):
logger.trace("Item out: %s", batch)
return batch
logger.trace(
"Item out: %s", {
k: v.shape if isinstance(v, np.ndarray) else v
for k, v in batch.items()
})
batch_faces = [[
self.to_detected_face(face[0], face[1], face[2], face[3])
for face in faces
] for faces in batch["prediction"]]
# Rotations
if any(m.any() for m in batch["rotmat"]) and any(batch_faces):
batch_faces = [[
self._rotate_rect(face, rotmat) if rotmat.any() else face
for face in faces
] for faces, rotmat in zip(batch_faces, batch["rotmat"])]
# Scale back out to original frame
batch["detected_faces"] = [[
self.to_detected_face(
(face.left - pad[0]) / scale, (face.top - pad[1]) / scale,
(face.right - pad[0]) / scale, (face.bottom - pad[1]) / scale)
for face in faces
] for scale, pad, faces in zip(batch["scale"], batch["pad"],
batch_faces)]
# Remove zero sized faces
self._remove_zero_sized_faces(batch)
if self.min_size > 0 and batch.get("detected_faces", None):
batch["detected_faces"] = self._filter_small_faces(
batch["detected_faces"])
self._remove_invalid_keys(batch,
("detected_faces", "filename", "image"))
batch = self._dict_lists_to_list_dicts(batch)
for item in batch:
logger.trace(
"final output: %s", {
k: v.shape if isinstance(v, np.ndarray) else v
for k, v in item.items()
})
yield item
