TKK_E32232028/.venv/lib/python3.10/site-packages/mtcnn/utils/bboxes.py

566 lines
26 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 numpy as np
from mtcnn.utils.landmarks import parse_landmarks
def generate_bounding_box(bbox_reg, bbox_class, threshold_face, strides=2, cell_size=12):
"""
Generates bounding boxes for detected objects (e.g., faces) based on the class and regression outputs of a model,
supporting batch input.
Args:
bbox_reg (tf.Tensor): Bounding box regression predictions with shape (batch_size, height, width, 4).
This contains adjustments to apply to the initial bounding box positions for each image in the batch.
bbox_class (tf.Tensor): Class predictions (e.g., face/non-face) of shape (batch_size, height, width, 2),
where the second channel corresponds to the probability of a face being present.
threshold_face (float): A threshold between 0 and 1 that determines if a detection is considered a face or not.
Bounding boxes are only generated for detections with probabilities greater than this value.
strides (int, optional): The step size (in pixels) used to slide the detection window over the image. Default is 2.
cell_size (int, optional): The size of the sliding window (in pixels) used to detect faces. Default is 12.
Returns:
np.ndarray: An array of bounding boxes for the entire batch, where each box is represented as
[batch_index, x1, y1, x2, y2, confidence].
The `batch_index` indicates which image in the batch the bounding box belongs to.
"""
bbox_reg = bbox_reg.numpy()
bbox_class = bbox_class.numpy()
# Create a mask for detected faces based on the threshold for face probability
confidence_score = bbox_class[:,:,:,1]
# Find the indices where the detection mask is true (i.e., face detected)
index_bboxes = np.stack(np.where(confidence_score > threshold_face)) # batch_size, y, x
filtered_bbox_reg = np.transpose(bbox_reg[index_bboxes[0], index_bboxes[1], index_bboxes[2]], (1,0))
# Extract the regression values
reg_x1, reg_y1, reg_x2, reg_y2 = filtered_bbox_reg
# Convert strides and cell size into arrays for easy broadcasting
strides = np.asarray([[1], [strides], [strides]])
cellsize = [np.asarray([[0], [1], [1]]), np.asarray([[0], [cell_size], [cell_size]])]
# Calculate the top-left and bottom-right corners of the bounding boxes
bbox_up_left = index_bboxes * strides + cellsize[0]
bbox_bottom_right = index_bboxes * strides + cellsize[1]
# Calculate width and height for the bounding boxes
reg_w = bbox_bottom_right[2] - bbox_up_left[2] # width of bounding box
reg_h = bbox_bottom_right[1] - bbox_up_left[1] # height of bounding box
# Apply the regression to adjust the bounding box coordinates
x1 = bbox_up_left[2] + reg_x1 * reg_w # Adjusted x1
y1 = bbox_up_left[1] + reg_y1 * reg_h # Adjusted y1
x2 = bbox_bottom_right[2] + reg_x2 * reg_w # Adjusted x2
y2 = bbox_bottom_right[1] + reg_y2 * reg_h # Adjusted y2
# Concatenate the bounding box coordinates and detection information, keeping batch index
bboxes_result = np.stack([
index_bboxes[0], x1, y1, x2, y2, confidence_score[index_bboxes[0], index_bboxes[1], index_bboxes[2]]
], axis=0).T
# Sort bounding boxes by score in descending order
bboxes_result = sort_by_scores(bboxes_result, scores=bboxes_result[:, -1], ascending=False)
return bboxes_result
def upscale_bboxes(bboxes_result, scales):
"""
Upscales bounding boxes to their original size based on the scaling factors applied during image resizing,
supporting batch input.
Args:
bboxes_result (np.ndarray): Array of bounding boxes, where each box is represented as
[batch_index, x1, y1, x2, y2, confidence, reg_x1, reg_y1, reg_x2, reg_y2].
scales (np.ndarray): Array of scaling factors used during image resizing, typically one scale per image or detection.
The shape of `scales` should be (batch_size,), where each entry corresponds to the scale applied to an
image in the batch.
Returns:
np.ndarray: The input bounding boxes, but with the coordinates scaled back to the original image dimensions,
adjusted for each image in the batch according to its respective scale.
"""
# Broadcast the scales to match the shape of the bounding boxes, ensuring the correct scale is applied to each batch entry
scales_bcast = np.expand_dims(scales[bboxes_result[:,0].astype(int)], axis=-1)
# Scale the bounding box coordinates (x1, y1, x2, y2) back to the original image size
bboxes_result[:,1:5] = bboxes_result[:,1:5] / scales_bcast
return bboxes_result
def iou(bboxes, method="union"):
"""
Computes the Intersection over Union (IoU) for a set of bounding boxes based on the specified method ("union" or "min").
Args:
bboxes (list or np.ndarray): List or array of bounding boxes, where each bounding box is represented as
[row1, col1, row2, col2] (coordinates of the top-left and bottom-right corners).
method (str, optional): Method to compute the IoU. Options are:
- "union": Computes IoU based on the union of the bounding boxes.
- "min": Computes IoU based on the minimum area of the bounding boxes.
Default is "union".
Returns:
np.ndarray: A matrix of shape (N, N) where each element [i, j] represents the IoU between the i-th and j-th bounding box.
The matrix is symmetric, with diagonal elements equal to 1 (IoU of a box with itself).
"""
# Convert the list of bounding boxes to a NumPy array
bboxes = np.stack(bboxes, axis=0)
# Calculate the area of each bounding box
area_bboxes = (bboxes[:, 2] - bboxes[:, 0]) * (bboxes[:, 3] - bboxes[:, 1])
# Expand dimensions to compute pairwise IoU (N x N matrix)
bboxes_a = np.expand_dims(bboxes, axis=0)
bboxes_b = np.expand_dims(bboxes, axis=1)
# Calculate the intersection coordinates
row_inter_top = np.maximum(bboxes_a[:, :, 0], bboxes_b[:, :, 0])
col_inter_left = np.maximum(bboxes_a[:, :, 1], bboxes_b[:, :, 1])
row_inter_bottom = np.minimum(bboxes_a[:, :, 2], bboxes_b[:, :, 2])
col_inter_right = np.minimum(bboxes_a[:, :, 3], bboxes_b[:, :, 3])
# Calculate the intersection area
height_inter = np.maximum(0, row_inter_bottom - row_inter_top)
width_inter = np.maximum(0, col_inter_right - col_inter_left)
area_inter = height_inter * width_inter
# Compute IoU based on the specified method
if method == "union":
# Union: Area of A + Area of B - Intersection
area_union = area_bboxes[:, None] + area_bboxes[None, :] - area_inter
iou_matrix = area_inter / area_union
elif method == "min":
# Minimum: Area of the smaller box between A and B
area_min = np.minimum(area_bboxes[:, None], area_bboxes[None, :])
iou_matrix = area_inter / area_min
else:
raise ValueError("Method should be either 'union' or 'min'.")
return iou_matrix
def sort_by_scores(tensor, scores, ascending=True):
"""
Sorts a tensor based on an array of scores, either in ascending or descending order.
Args:
tensor (np.ndarray): Tensor of shape (N, ...) where N is the number of elements to sort.
scores (np.ndarray): Array of shape (N,) containing scores associated with each element in the tensor.
ascending (bool, optional): Whether to sort in ascending order. Default is True (ascending).
Returns:
np.ndarray: The tensor sorted according to the scores.
"""
# Get the sorted indices based on the scores
sorted_indices = np.argsort(scores)
# Sort the tensor using the sorted indices, reversing if descending
sorted_tensor = tensor[sorted_indices[::(-2 * int(not ascending) + 1)]]
return sorted_tensor
def nms(target_iou, threshold):
"""
Performs Non-Maximum Suppression (NMS) to filter out overlapping bounding boxes based on the IoU threshold.
Args:
target_iou (np.ndarray): A square IoU matrix of shape (N, N) where each element [i, j] represents the IoU
between the i-th and j-th bounding box.
threshold (float): IoU threshold above which boxes are considered to overlap too much and will be suppressed.
Returns:
np.ndarray: Array of indices of bounding boxes that are kept after NMS.
"""
# Step 1: Create a mask for allowed comparisons (upper triangular part of the IoU matrix, excluding the diagonal)
allowed_mask = np.triu(np.ones((target_iou.shape[0], target_iou.shape[0])), k=1)
# Step 2: Create a mask for failed comparisons (IoU above the threshold)
failed_mask = (target_iou > threshold).astype(int)
# Step 3: Combine the masks and get the indices of the remaining boxes
result_indexes = np.where((failed_mask * allowed_mask).sum(axis=0) == 0)[0]
return result_indexes
def smart_nms_from_bboxes(bboxes, threshold, column_image_id=0, columns_bbox=slice(1, 5, None), column_confidence=5,
method="union", initial_sort=True):
"""
Applies Non-Maximum Suppression (NMS) to a set of bounding boxes grouped by image ID.
Args:
bboxes (np.ndarray): Array of bounding boxes, where each box is represented as
[image_id, row1, col1, row2, col2, score].
threshold (float): IoU threshold for NMS. Bounding boxes with IoU higher than this value will be suppressed.
column_image_id (int): Column position in the array indicating the image id.
columns_bbox (slice): Slice of columns containing the BBox coords.
method (str, optional): Method for IoU calculation. Can be "union" or "min". Default is "union".
column_confidence (int): Column containing the value of confidence for each bbox.
initial_sort (bool): True to sort bboxes by confidence value. False otherwise.
Returns:
A np.ndarray containing the filtered bboxes, image-wise.
dict: A dictionary where keys are `image_id` and values are arrays of indices of bounding boxes that are
kept after NMS for each image.
"""
# Step 0: Sort if required
if initial_sort:
bboxes = sort_by_scores(bboxes, scores=bboxes[:, column_confidence], ascending=False)
# Step 1: Get unique image IDs
image_ids = np.unique(bboxes[:, 0])
result = []
# Step 2: Apply NMS per image
for image_id in image_ids:
# Filter bounding boxes for the current image
target_bboxes = bboxes[bboxes[:, column_image_id] == image_id]
# Compute the IoU matrix for the bounding boxes
target_iou = iou(target_bboxes[:, columns_bbox], method=method)
# Perform NMS and get the indices of the boxes to keep
target_indexes = nms(target_iou, threshold)
# Filter the boxes for the image
target_filtered_bboxes = target_bboxes[target_indexes.astype(int)]
# Store the result
result.append(target_filtered_bboxes)
result = np.concatenate(result, axis=0) if len(result) > 0 else np.empty((0, 6))
return result
def resize_to_square(bboxes):
"""
Adjusts bounding boxes to be square by resizing them based on their largest dimension
(width or height). The bounding boxes are resized by expanding the smaller dimension
to match the larger one while keeping the center of the box intact.
Args:
bboxes (np.ndarray): An array of bounding boxes of shape (n, 5), where each row
represents [batch_index, x1, y1, x2, y2].
Returns:
np.ndarray: An array of bounding boxes adjusted to be square, maintaining their center positions.
"""
bboxes = bboxes.copy()
h = bboxes[:, 4] - bboxes[:, 2] # Height of each bounding box
w = bboxes[:, 3] - bboxes[:, 1] # Width of each bounding box
largest_size = np.maximum(w, h) # Largest dimension (width or height)
# Adjust x1 and y1 to center the bounding box and resize to square
bboxes[:, 1] = bboxes[:, 1] + w * 0.5 - largest_size * 0.5
bboxes[:, 2] = bboxes[:, 2] + h * 0.5 - largest_size * 0.5
bboxes[:, 3:5] = bboxes[:, 1:3] + np.tile(largest_size, (2, 1)).T # Resize x2, y2
return bboxes
def replace_confidence(bboxes_batch, new_scores):
"""
Replaces the confidence scores of bounding boxes with new scores provided.
Args:
bboxes_batch (np.ndarray): An array of bounding boxes of shape (n, m), where each row
contains the bounding box coordinates and the confidence score.
The confidence score is expected to be in the last column.
new_scores (np.ndarray): An array of new confidence scores of shape (n, m), where the
confidence score is also expected to be in the last column.
Returns:
np.ndarray: The bounding boxes array with updated confidence scores from `new_scores`.
"""
bboxes_batch[:, -1] = new_scores[:, -1]
return bboxes_batch
def adjust_bboxes(bboxes_batch, bboxes_offsets):
"""
Adjusts the bounding box coordinates by applying the provided offsets.
The offsets are applied to resize and shift the bounding boxes based on their width and height.
Args:
bboxes_batch (np.ndarray): An array of bounding boxes of shape (n, m), where each row contains
the batch index, bounding box coordinates [x1, y1, x2, y2], and
potentially additional data such as scores.
bboxes_offsets (np.ndarray): An array of offsets for adjusting the bounding boxes. The shape should be
(n, 4), where each row contains offsets for [dx1, dy1, dx2, dy2].
Returns:
np.ndarray: The adjusted bounding boxes with updated coordinates, maintaining any additional columns
beyond the bounding box coordinates (such as scores).
"""
bboxes_batch = bboxes_batch.copy()
w = bboxes_batch[:, 3] - bboxes_batch[:, 1] + 1 # Calculate width of each bounding box
h = bboxes_batch[:, 4] - bboxes_batch[:, 2] + 1 # Calculate height of each bounding box
sizes = np.stack([w, h, w, h], axis=-1) # Stack width and height to match bbox_offsets
bboxes_batch[:, 1:5] += bboxes_offsets * sizes # Apply offsets to the coordinates
return bboxes_batch
def pick_matches(bboxes_batch, scores_column=-1, score_threshold=0.7):
"""
Filters bounding boxes based on the confidence score threshold.
Only bounding boxes with a confidence score higher than the specified threshold are returned.
Args:
bboxes_batch (np.ndarray): An array of bounding boxes of shape (n, m), where each row contains
bounding box coordinates and confidence scores. The confidence scores
are expected to be in the column specified by `scores_column`.
scores_column (int): The index of the column that contains the confidence scores. Default is -1 (last column).
score_threshold (float): The minimum confidence score threshold to select bounding boxes.
Default is 0.7.
Returns:
np.ndarray: An array of bounding boxes that have confidence scores greater than `score_threshold`.
"""
return bboxes_batch[np.where(bboxes_batch[:, scores_column] > score_threshold)[0]]
def to_json(bboxes_batch, images_count, input_as_width_height=False, output_as_width_height=True):
"""
Converts a batch of bounding boxes and facial keypoints into a JSON-friendly format.
This function processes the bounding boxes grouped by unique image IDs, and formats each bounding box
and its associated keypoints (facial landmarks) into a dictionary structure suitable for JSON serialization.
Args:
bboxes_batch (np.ndarray): An array of shape (n, 16) where each row represents a bounding box
and associated keypoints in the following format:
[image_id, x1, y1, x2, y2, confidence, left_eye_x, left_eye_y, right_eye_x,
right_eye_y, nose_x, nose_y, mouth_left_x, mouth_left_y, mouth_right_x, mouth_right_y].
images_count (int): Number of different images composed by the batch.
input_as_width_height (bool, optional): True if format of input bounding boxes is [x1, x2, width, height].
False if format is [x1, y1, x2, y2].
output_as_width_height (bool, optional): True to format bounding boxes as [x1, x2, width, height].
False to format as [x1, y1, x2, y2].
Returns:
list: A list of lists, where each inner list contains dictionaries for bounding boxes and keypoints
for a specific image. Each dictionary has the following structure:
{
"box": [x, y, width, height],
"keypoints": {
"nose": [nose_x, nose_y],
"mouth_right": [mouth_right_x, mouth_right_y],
"right_eye": [right_eye_x, right_eye_y],
"left_eye": [left_eye_x, left_eye_y],
"mouth_left": [mouth_left_x, mouth_left_y]
},
"confidence": confidence_score
}
"""
single_element = len(bboxes_batch.shape) == 1
if single_element:
bboxes_batch = np.expand_dims(bboxes_batch, axis=0)
#unique_ids = np.unique(bboxes_batch[:, 0])
result_batch = []
# Loop over each unique image ID
for unique_id in range(images_count):
result = []
bboxes_subset = bboxes_batch[bboxes_batch[:, 0] == unique_id]
# Loop over each bounding box in the subset
for bbox in bboxes_subset:
row = {
"box": parse_bbox(bbox,
output_as_width_height=output_as_width_height,
input_as_width_height=input_as_width_height).tolist(),
"confidence": bbox[5]
}
result.append(row)
# If the stages combination allows landmarks, then we append them. Otherwise we don't
try:
row["keypoints"] = parse_landmarks(bbox)
except IndexError:
pass
result_batch.append(result)
return result_batch
def limit_bboxes(bboxes_batch, images_shapes, limit_landmarks=True):
"""
Adjusts bounding boxes so that they fit within the boundaries of their corresponding images.
If any bounding box exceeds the image dimensions, it will be corrected to stay within the limits.
Args:
bboxes_batch (np.ndarray): An array of bounding boxes of shape (n, 5), where each row
represents [batch_index, x1, y1, x2, y2].
images_shapes (np.ndarray): A tensor of image shapes of shape (batch, 3), where each row
represents [width, height, channels] of each image in the batch.
limit_landmarks (bool): A flag to specify whether the limit should also apply to landmarks or not.
Returns:
np.ndarray: The adjusted bounding boxes where no coordinate exceeds the image dimensions.
"""
bboxes_batch_fitted = bboxes_batch.copy()
# Get the original shapes (height, width) for each image in the batch
expected_shapes = images_shapes[bboxes_batch_fitted[:, 0].astype(int)]
# Adjust x1 and x2 to be within [0, width-1]
bboxes_batch_fitted[:, 1] = np.minimum(np.maximum(bboxes_batch_fitted[:, 1], 0), expected_shapes[:, 1] - 1)
bboxes_batch_fitted[:, 3] = np.minimum(np.maximum(bboxes_batch_fitted[:, 3], 0), expected_shapes[:, 1] - 1)
# Adjust y1 and y2 to be within [0, height-1]
bboxes_batch_fitted[:, 2] = np.minimum(np.maximum(bboxes_batch_fitted[:, 2], 0), expected_shapes[:, 0] - 1)
bboxes_batch_fitted[:, 4] = np.minimum(np.maximum(bboxes_batch_fitted[:, 4], 0), expected_shapes[:, 0] - 1)
if limit_landmarks:
# Adjust x1..x5 of the landmarks to not surpass boundaries
bboxes_batch_fitted[:, 6:11] = np.minimum(np.maximum(bboxes_batch_fitted[:, 6:11], 0), expected_shapes[:, 1:2] - 1)
# Adjust y1..y5 of the landmarks to not surpass boundaries
bboxes_batch_fitted[:, 11:16] = np.minimum(np.maximum(bboxes_batch_fitted[:, 11:16], 0), expected_shapes[:, 0:1] - 1)
return bboxes_batch_fitted
def parse_bbox(bbox, output_as_width_height=True, input_as_width_height=True):
"""
Parses a bounding box from different formats (dict, list, or ndarray) into a standardized format.
Args:
bbox (dict, list, np.ndarray): Bounding box in one of the following formats:
- dict with key 'box': [x1, y1, x2, y2]
- list: [x1, y1, x2, y2] or [x1, y1, width, height]
- np.ndarray: Shape (4,) or (5,) where the first value might be an index.
output_as_width_height (bool): Whether to return the bounding box as [x1, y1, width, height] (default True) or [x1, y1, x2, y2] if False.
input_as_width_height (bool): Whether the input format of the bounding box is [x1, y1, width, height] (default True) or
[x1, y1, x2, y2] if False.
Returns:
np.ndarray: Parsed bounding box in format [x1, y1, width, height] or [x1, y1, x2, y2].
"""
# Extract box if input is a dict
if isinstance(bbox, dict):
bbox = bbox['box']
# Parse list format
if isinstance(bbox, list):
x1, y1, width, height = bbox
if not input_as_width_height:
width = width - x1
height = height - y1
x2_or_w = width if output_as_width_height else x1 + width
y2_or_h = height if output_as_width_height else y1 + height
return np.asarray([x1, y1, x2_or_w, y2_or_h]).round().astype(int)
# Parse ndarray format
if isinstance(bbox, np.ndarray):
offset = 1 if bbox.shape[0] > 4 else 0 # Handle optional first element
x1, y1, width, height = bbox[offset:offset+4]
if not input_as_width_height:
width = width - x1
height = height - y1
x2_or_w = width if output_as_width_height else x1 + width
y2_or_h = height if output_as_width_height else y1 + height
return np.asarray([x1, y1, x2_or_w, y2_or_h]).round().astype(int)
raise ValueError("Invalid bbox format. Expected dict, list, or ndarray.")
def fix_bboxes_offsets(bboxes_batch, pad_param):
"""
Adjusts the bounding boxes and landmarks by subtracting the corresponding padding offsets applied during image padding.
This function corrects the bounding boxes' coordinates and facial landmarks after padding the images, ensuring that
the boxes and landmarks are aligned with the padded images by subtracting the appropriate offsets.
Args:
bboxes_batch (np.ndarray): An array of bounding boxes and landmarks of shape (n, m), where each row represents
[image_id, x1, y1, x2, y2, confidence, landmark_x1, landmark_y1, ..., landmark_x5, landmark_y5].
The first column (index 0) corresponds to the image ID.
pad_param (np.ndarray): An array of padding parameters of shape (n, 2, 2), where each entry represents the amount
of padding applied to each image along the width and height dimensions.
Returns:
np.ndarray: A modified copy of `bboxes_batch` with updated bounding box and landmark coordinates adjusted for padding.
"""
bboxes_batch = bboxes_batch.copy()
images_ids = np.unique(bboxes_batch[:, 0]) # Get unique image IDs
indexes_bbox_x = [1,3]
indexes_bbox_y = [2,4]
indexes_landmarks_x = [6, 7, 8, 9, 10]
indexes_landmarks_y = [11, 12, 13, 14, 15]
# Adjust bounding boxes and landmarks for each image based on its padding parameters
for image_id, pad in zip(images_ids, pad_param):
selector = bboxes_batch[:, 0] == image_id
# Adjust the x-coordinates of bounding boxes by subtracting width padding
bboxes_batch[np.ix_(selector, indexes_bbox_x)] -= pad[1, 0]
# Adjust the y-coordinates of bounding boxes by subtracting height padding
bboxes_batch[np.ix_(selector, indexes_bbox_y)] -= pad[0, 0]
# If stages combinations contain landmarks, we adjust them too
try:
# Adjust the x-coordinates of landmarks by subtracting width padding
bboxes_batch[np.ix_(selector, indexes_landmarks_x)] -= pad[1, 0]
# Adjust the y-coordinates of landmarks by subtracting height padding
bboxes_batch[np.ix_(selector, indexes_landmarks_y)] -= pad[0, 0]
except IndexError:
pass
return bboxes_batch