189 lines
9.9 KiB
Python
189 lines
9.9 KiB
Python
# MIT License
|
|
#
|
|
# Copyright (c) 2019-2024 Iván de Paz Centeno
|
|
#
|
|
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
# of this software and associated documentation files (the "Software"), to deal
|
|
# in the Software without restriction, including without limitation the rights
|
|
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
# copies of the Software, and to permit persons to whom the Software is
|
|
# furnished to do so, subject to the following conditions:
|
|
#
|
|
# The above copyright notice and this permission notice shall be included in all
|
|
# copies or substantial portions of the Software.
|
|
#
|
|
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
# SOFTWARE.
|
|
|
|
import tensorflow as tf
|
|
import numpy as np
|
|
|
|
from mtcnn.stages import StagePNet, StageRNet, StageONet
|
|
|
|
from mtcnn.utils.images import load_images_batch, standarize_batch
|
|
from mtcnn.utils.bboxes import fix_bboxes_offsets, limit_bboxes, to_json
|
|
|
|
|
|
COMMON_STAGES = {
|
|
"face_detection_only": [StagePNet, StageRNet],
|
|
"face_and_landmarks_detection": [StagePNet, StageRNet, StageONet],
|
|
}
|
|
|
|
|
|
class MTCNN:
|
|
"""
|
|
MTCNN class for detecting faces and landmarks through configurable stages.
|
|
This structure allows skipping certain stages to optimize performance based on the user's needs.
|
|
|
|
Args:
|
|
stages (str or list, optional): Defines the pipeline stages. It can be a string to choose from predefined
|
|
configurations or a list of stage classes or instances.
|
|
Options: "face_detection_only", "face_and_landmarks_detection".
|
|
Default is "face_and_landmarks_detection".
|
|
device (str, optional): The device where the model will be run. Can be "CPU:0", "GPU:0", "GPU:1", ...
|
|
Default is "CPU:0".
|
|
|
|
"""
|
|
|
|
def __init__(self, stages="face_and_landmarks_detection", device="CPU:0"):
|
|
if isinstance(stages, str):
|
|
if stages not in COMMON_STAGES:
|
|
raise ValueError(f"Invalid stages option: {stages}. Must be one of {list(COMMON_STAGES.keys())}.")
|
|
stages = COMMON_STAGES[stages]
|
|
|
|
# Instantiate stages if necessary (can pass already instantiated stages too)
|
|
self._stages = [stage() if isinstance(stage, type) else stage for stage in stages]
|
|
self._device = device
|
|
|
|
@property
|
|
def device(self):
|
|
"""Returns the device where the algorithm is executed"""
|
|
return self._device
|
|
|
|
@property
|
|
def stages(self):
|
|
"""Returns the list of pipeline stages."""
|
|
return self._stages
|
|
|
|
def get_stage(self, stage_id=None, stage_name=None):
|
|
"""
|
|
Retrieves a stage by its ID or name.
|
|
|
|
Args:
|
|
stage_id (int, optional): The ID of the stage.
|
|
stage_name (str, optional): The name of the stage.
|
|
|
|
Returns:
|
|
The matching stage if found, otherwise None.
|
|
"""
|
|
for stage in self._stages:
|
|
if stage.id == stage_id or stage.name == stage_name:
|
|
return stage
|
|
|
|
return None
|
|
|
|
def predict(self, image, fit_to_image=True, limit_boundaries_landmarks=False, box_format="xywh", output_type="json", postprocess=True,
|
|
**kwargs):
|
|
"""
|
|
Alias for detect_faces().
|
|
"""
|
|
return self.detect_faces(image, fit_to_image=fit_to_image, limit_boundaries_landmarks=limit_boundaries_landmarks,
|
|
box_format=box_format, output_type=output_type, postprocess=postprocess, **kwargs)
|
|
|
|
def detect_faces(self, image, fit_to_image=True, limit_boundaries_landmarks=False, box_format="xywh", output_type="json",
|
|
postprocess=True, batch_stack_justification="center", **kwargs):
|
|
"""
|
|
Runs face detection on a single image or batch of images through the configured stages.
|
|
|
|
Args:
|
|
image (str, bytes, np.ndarray or tf.Tensor or list): The input image or batch of images.
|
|
It can be a file path, a tensor, or raw bytes.
|
|
fit_to_image (bool, optional): Whether to fit bounding boxes and landmarks within image boundaries. Default is True.
|
|
limit_boundaries_landmarks (bool, optional): Whether to ensure landmarks stay within image boundaries. Default is False.
|
|
box_format (str, optional): The format of the bounding box. Can be "xywh" for [X1, Y1, width, height] or "xyxy" for [X1, Y1, X2, Y2].
|
|
Default is "xywh".
|
|
output_type (str, optional): The output format. Can be "json" for dictionary output or "numpy" for numpy array output. Default is "json".
|
|
postprocess (bool, optional): Flag to enable postprocessing. The postprocessing includes functionality affected by `fit_to_image`,
|
|
`limit_boundaries_landmarks` and removing padding effects caused by batching images with different shapes.
|
|
batch_stack_justification (str, optional): The justification of the smaller images w.r.t. the largest images when
|
|
stacking in batch processing, which requires padding smaller images to the size of the
|
|
biggest one.
|
|
**kwargs: Additional parameters passed to the stages. The following parameters are used:
|
|
|
|
- **StagePNet**:
|
|
- min_face_size (int, optional): The minimum size of a face to detect. Default is 20.
|
|
- min_size (int, optional): The minimum size to start the image pyramid. Default is 12.
|
|
- scale_factor (float, optional): The scaling factor for the image pyramid. Default is 0.709.
|
|
- threshold_pnet (float, optional): The confidence threshold for proposals from PNet. Default is 0.6.
|
|
- nms_pnet1 (float, optional): The IoU threshold for the first round of NMS per scale. Default is 0.5.
|
|
- nms_pnet2 (float, optional): The IoU threshold for the second round of NMS across all scales. Default is 0.7.
|
|
|
|
- **StageRNet**:
|
|
- threshold_rnet (float, optional): Confidence threshold for RNet proposals. Default is 0.7.
|
|
- nms_rnet (float, optional): IoU threshold for Non-Maximum Suppression in RNet. Default is 0.7.
|
|
|
|
- **StageONet**:
|
|
- threshold_onet (float, optional): Confidence threshold for ONet proposals. Default is 0.8.
|
|
- nms_onet (float, optional): IoU threshold for Non-Maximum Suppression in ONet. Default is 0.7.
|
|
|
|
Returns:
|
|
list or list of lists: A list of detected faces (in case a single image) or a list of lists of detected faces
|
|
(one per image in the batch). If the stages are `face_and_landmarks_detection`,
|
|
the output will have the detected faces and landmarks in JSON format.
|
|
In case of `face_detection_only`, only the bounding boxes will be provided in
|
|
JSON format.
|
|
"""
|
|
return_tensor = output_type == "numpy"
|
|
as_width_height = box_format == "xywh"
|
|
|
|
is_batch = isinstance(image, list)
|
|
images = image if is_batch else [image]
|
|
|
|
with tf.device(self._device):
|
|
# Load the images into memory and normalize them into a single tensor
|
|
try:
|
|
images_raw = load_images_batch(images)
|
|
images_normalized, images_oshapes, pad_param = standarize_batch(images_raw,
|
|
justification=batch_stack_justification,
|
|
normalize=True)
|
|
|
|
bboxes_batch = None
|
|
|
|
# Process images through each stage (PNet, RNet, ONet)
|
|
for stage in self.stages:
|
|
bboxes_batch = stage(bboxes_batch=bboxes_batch, images_normalized=images_normalized, images_oshapes=images_oshapes, **kwargs)
|
|
|
|
except tf.errors.InvalidArgumentError: # No faces found
|
|
bboxes_batch = np.empty((0, 16))
|
|
pad_param = None
|
|
|
|
if postprocess and pad_param is not None:
|
|
# Adjust bounding boxes and landmarks to account for padding offsets
|
|
bboxes_batch = fix_bboxes_offsets(bboxes_batch, pad_param)
|
|
|
|
# Optionally, limit the bounding boxes and landmarks to stay within image boundaries
|
|
if fit_to_image:
|
|
bboxes_batch = limit_bboxes(bboxes_batch, images_shapes=images_oshapes, limit_landmarks=limit_boundaries_landmarks)
|
|
|
|
# Convert bounding boxes and landmarks to JSON format if required
|
|
if return_tensor:
|
|
result = bboxes_batch
|
|
|
|
if as_width_height:
|
|
result[:, 3] = result[:, 3] - result[:, 1]
|
|
result[:, 4] = result[:, 4] - result[:, 2]
|
|
|
|
else:
|
|
result = to_json(bboxes_batch,
|
|
images_count=len(images),
|
|
output_as_width_height=as_width_height,
|
|
input_as_width_height=False)
|
|
result = result[0] if (not is_batch and len(result) > 0) else result
|
|
|
|
return result
|