iou2d_calculator.py
yolov6\assigners\iou2d_calculator.py
目录
2.def cast_tensor_type(x, scale=1., dtype=None):
3.def fp16_clamp(x, min=None, max=None):
4.def iou2d_calculator(bboxes1, bboxes2, mode='iou', is_aligned=False, scale=1., dtype=None):
5.def bbox_overlaps(bboxes1, bboxes2, mode='iou', is_aligned=False, eps=1e-6):
1.所需的库和模块
#This code is based on
#https://github.com/fcjian/TOOD/blob/master/mmdet/core/bbox/iou_calculators/iou2d_calculator.py
import torch
2.def cast_tensor_type(x, scale=1., dtype=None):
def cast_tensor_type(x, scale=1., dtype=None):
if dtype == 'fp16':
# scale 是为了防止溢出。
# scale is for preventing overflows
x = (x / scale).half()
return x3.def fp16_clamp(x, min=None, max=None):
def fp16_clamp(x, min=None, max=None):
if not x.is_cuda and x.dtype == torch.float16:
# cpu float16 的 clamp，tensor fp16 没有 clamp 实现
# clamp for cpu float16, tensor fp16 has no clamp implementation
# torch.clamp(input, min, max, out=None) → Tensor
# clamp() 函数的功能将输入 input 张量每个元素的值压缩到区间 [min,max]，并返回结果到一个新张量。
# input ：输入张量。
# min ：限制范围下限。
# max ：限制范围上限。
# out ：输出张量。
# pytorch中，一般来说如果对tensor的一个函数后加上了下划线，则表明这是一个 in-place 类型。
# in-place 类型是指，当在一个tensor上操作了之后，是直接修改了这个tensor，而不是返回一个新的tensor并不修改旧的tensor。
return x.float().clamp(min, max).half()
return x.clamp(min, max)4.def iou2d_calculator(bboxes1, bboxes2, mode='iou', is_aligned=False, scale=1., dtype=None):
def iou2d_calculator(bboxes1, bboxes2, mode='iou', is_aligned=False, scale=1., dtype=None):
# 2D 重叠（例如 IoU、GIoU）计算器。
# 计算 2D bbox 之间的 IoU。
"""2D Overlaps (e.g. IoUs, GIoUs) Calculator."""
"""Calculate IoU between 2D bboxes.
Args:
bboxes1 (Tensor): bboxes have shape (m, 4) in <x1, y1, x2, y2>
format, or shape (m, 5) in <x1, y1, x2, y2, score> format. bboxes 在 <x1, y1, x2, y2> 格式中具有形状 (m, 4)，或在 <x1, y1, x2, y2, score> 格式中具有形状 (m, 5)。
bboxes2 (Tensor): bboxes have shape (m, 4) in <x1, y1, x2, y2>
format, shape (m, 5) in <x1, y1, x2, y2, score> format, or be
empty. If ``is_aligned `` is ``True``, then m and n must be
equal. bboxes 具有 <x1, y1, x2, y2> 格式的形状 (m, 4)，具有 <x1, y1, x2, y2, score> 格式的形状 (m, 5)，或者为空。如果 ``is_aligned`` 为 ``True``，则 m 和 n 必须相等。
mode (str): "iou" (intersection over union), "iof" (intersection
over foreground), or "giou" (generalized intersection over
union). “iou”（并集交集）、“iof”（前景交集）或“giou”（广义并集交集）。
is_aligned (bool, optional): If True, then m and n must be equal.
Default False. 如果为 True，则 m 和 n 必须相等。默认为 False。
Returns:
Tensor: shape (m, n) if ``is_aligned `` is False else shape (m,) 形状（m，n）如果``is_aligned``为False，则形状（m，）
"""
assert bboxes1.size(-1) in [0, 4, 5]
assert bboxes2.size(-1) in [0, 4, 5]
if bboxes2.size(-1) == 5:
bboxes2 = bboxes2[..., :4]
if bboxes1.size(-1) == 5:
bboxes1 = bboxes1[..., :4]
if dtype == 'fp16':
# 更改张量类型以节省 CPU 和 Cuda 内存并保持速度。
# change tensor type to save cpu and cuda memory and keep speed
# def cast_tensor_type(x, scale=1., dtype=None): -> x = (x / scale).half() -> return x
bboxes1 = cast_tensor_type(bboxes1, scale, dtype)
bboxes2 = cast_tensor_type(bboxes2, scale, dtype)
overlaps = bbox_overlaps(bboxes1, bboxes2, mode, is_aligned) # aligned对齐
if not overlaps.is_cuda and overlaps.dtype == torch.float16:
# resume cpu float32 恢复 CPU float32。
overlaps = overlaps.float()
return overlaps
# def bbox_overlaps(bboxes1, bboxes2, mode='iou', is_aligned=False, eps=1e-6):
# -> 计算两组 bbox 之间的重叠。
# if rows * cols == 0: if is_aligned: return bboxes1.new(batch_shape + (rows, )) else: return bboxes1.new(batch_shape + (rows, cols))
# -> if mode in ['iou', 'iof']: return ious / return gious
return bbox_overlaps(bboxes1, bboxes2, mode, is_aligned)
5.def bbox_overlaps(bboxes1, bboxes2, mode='iou', is_aligned=False, eps=1e-6):
def bbox_overlaps(bboxes1, bboxes2, mode='iou', is_aligned=False, eps=1e-6):
# 计算两组 bbox 之间的重叠。
# FP16 由 https://github.com/open-mmlab/mmdetection/pull/4889 贡献
# 如果“ is_aligned ”为“ False ”，则计算 bboxes1 和 bboxes2 的每个 bbox 之间的重叠，否则计算 每对对齐的 bboxes1 和 bboxes2 之间的重叠。
# 如果“ is_aligned ”为“ True ”，则 m 和 n 必须相等。
"""Calculate overlap between two set of bboxes.
FP16 Contributed by https://github.com/open-mmlab/mmdetection/pull/4889
Note:
Assume bboxes1 is M x 4, bboxes2 is N x 4, when mode is 'iou',
there are some new generated variable when calculating IOU
using bbox_overlaps function:
1) is_aligned is False
area1: M x 1
area2: N x 1
lt: M x N x 2
rb: M x N x 2
wh: M x N x 2
overlap: M x N x 1
union: M x N x 1
ious: M x N x 1
Total memory:
S = (9 x N x M + N + M) * 4 Byte,
When using FP16, we can reduce:
R = (9 x N x M + N + M) * 4 / 2 Byte
R large than (N + M) * 4 * 2 is always true when N and M >= 1.
Obviously, N + M <= N * M < 3 * N * M, when N >=2 and M >=2,
N + 1 < 3 * N, when N or M is 1.
Given M = 40 (ground truth), N = 400000 (three anchor boxes
in per grid, FPN, R-CNNs),
R = 275 MB (one times)
A special case (dense detection), M = 512 (ground truth),
R = 3516 MB = 3.43 GB
When the batch size is B, reduce:
B x R
Therefore, CUDA memory runs out frequently.
Experiments on GeForce RTX 2080Ti (11019 MiB):
| dtype | M | N | Use | Real | Ideal |
|:----:|:----:|:----:|:----:|:----:|:----:|
| FP32 | 512 | 400000 | 8020 MiB | -- | -- |
| FP16 | 512 | 400000 | 4504 MiB | 3516 MiB | 3516 MiB |
| FP32 | 40 | 400000 | 1540 MiB | -- | -- |
| FP16 | 40 | 400000 | 1264 MiB | 276MiB | 275 MiB |
2) is_aligned is True
area1: N x 1
area2: N x 1
lt: N x 2
rb: N x 2
wh: N x 2
overlap: N x 1
union: N x 1
ious: N x 1
Total memory:
S = 11 x N * 4 Byte
When using FP16, we can reduce:
R = 11 x N * 4 / 2 Byte
So do the 'giou' (large than 'iou').
Time-wise, FP16 is generally faster than FP32.
When gpu_assign_thr is not -1, it takes more time on cpu
but not reduce memory.
There, we can reduce half the memory and keep the speed.
If ``is_aligned`` is ``False``, then calculate the overlaps between each
bbox of bboxes1 and bboxes2, otherwise the overlaps between each aligned
pair of bboxes1 and bboxes2.
Args:
bboxes1 (Tensor): shape (B, m, 4) in <x1, y1, x2, y2> format or empty.
bboxes2 (Tensor): shape (B, n, 4) in <x1, y1, x2, y2> format or empty.
B indicates the batch dim, in shape (B1, B2, ..., Bn).
If ``is_aligned`` is ``True``, then m and n must be equal.
mode (str): "iou" (intersection over union), "iof" (intersection over
foreground) or "giou" (generalized intersection over union).
Default "iou".
is_aligned (bool, optional): If True, then m and n must be equal.
Default False.
eps (float, optional): A value added to the denominator for numerical
stability. Default 1e-6.
Returns:
Tensor: shape (m, n) if ``is_aligned`` is False else shape (m,)
Example:
>>> bboxes1 = torch.FloatTensor([
>>> [0, 0, 10, 10],
>>> [10, 10, 20, 20],
>>> [32, 32, 38, 42],
>>> ])
>>> bboxes2 = torch.FloatTensor([
>>> [0, 0, 10, 20],
>>> [0, 10, 10, 19],
>>> [10, 10, 20, 20],
>>> ])
>>> overlaps = bbox_overlaps(bboxes1, bboxes2)
>>> assert overlaps.shape == (3, 3)
>>> overlaps = bbox_overlaps(bboxes1, bboxes2, is_aligned=True)
>>> assert overlaps.shape == (3, )
Example:
>>> empty = torch.empty(0, 4)
>>> nonempty = torch.FloatTensor([[0, 0, 10, 9]])
>>> assert tuple(bbox_overlaps(empty, nonempty).shape) == (0, 1)
>>> assert tuple(bbox_overlaps(nonempty, empty).shape) == (1, 0)
>>> assert tuple(bbox_overlaps(empty, empty).shape) == (0, 0)
"""
# 不支持的模式 {mode}。
assert mode in ['iou', 'iof', 'giou'], f'Unsupported mode {mode}'
# 要么边界框是空的，要么边界框的最后一个维度的长度为 4。
# Either the boxes are empty or the length of boxes' last dimension is 4
# bboxes1.size(-1) == 4 ： bboxes1.size(-1) 获取张量 bboxes1 的最后一个维度的大小。在边界框的上下文中，最后一个维度通常表示每个边界框的坐标数量，这应该是4（即 x_min, y_min, x_max, y_max ）。 这个条件检查 bboxes1 的最后一个维度是否为4，确保每个边界框都有四个坐标值。
# bboxes1.size(0) == 0 ： bboxes1.size(0) 获取张量 bboxes1 的第一个维度的大小。如果这个维度为0，意味着没有边界框。 这个条件检查 bboxes1 是否为空，即没有边界框需要处理。
# 总的来说，这个断言确保了 bboxes1 要么包含具有4个坐标值的边界框，要么为空。这是为了在处理边界框数据之前确保数据的一致性和正确性。如果数据不符合这些条件，程序将抛出 AssertionError ，提示可能存在数据错误或不一致的问题。
assert (bboxes1.size(-1) == 4 or bboxes1.size(0) == 0)
assert (bboxes2.size(-1) == 4 or bboxes2.size(0) == 0)
# 批量维度必须相同。
# Batch dim must be the same
# 批量维度：（B1，B2，...Bn）
# Batch dim: (B1, B2, ... Bn)
assert bboxes1.shape[:-2] == bboxes2.shape[:-2]
# 这是一个元组，表示张量 bboxes1 的完整形状（shape）。例如，如果 bboxes1 是一个包含多个边界框的张量，其形状可能是 [batch_size, num_boxes, 4] ，其中 batch_size 是批次中的图像数量， num_boxes 是每个图像中的边界框数量， 4 代表每个边界框的坐标数（x_min, y_min, x_max, y_max）。
# [:-2] 是一个切片操作，它获取从开始到倒数第二个维度之前的所有维度。这意味着它会忽略最后两个维度，通常这两个维度分别代表边界框的数量和每个边界框的坐标数。
# 对于上面的例子， bboxes1.shape 可能是 [batch_size, num_boxes, 4] ，那么 bboxes1.shape[:-2] 将是 [batch_size] 。
# batch_shape ： 这个变量存储了 bboxes1 的批量形状，即不包括边界框数量和坐标数的所有维度。这通常用于确定批次中的图像数量。
batch_shape = bboxes1.shape[:-2]
# bboxes1.size(-2) 获取张量 bboxes1 的倒数第二个维度的大小。在边界框的上下文中，这个维度通常表示边界框的数量或批次中的图像数量。 rows 存储了这个大小，代表 bboxes1 张量中的边界框数量。
rows = bboxes1.size(-2)
# 同上， cols 代表 bboxes2 张量中的边界框数量。
cols = bboxes2.size(-2)
if is_aligned: # aligned对齐
assert rows == cols
if rows * cols == 0:
if is_aligned:
# batch_shape + (rows, ) ： 这是一个元组拼接操作，它将 batch_shape （一个元组）与另一个元组 ( rows, ) 连接起来。 batch_shape 包含了批次的形状，但不包括边界框的数量，而 ( rows, ) 是一个包含单个元素 rows 的元组。rows 是一个整数，表示 bboxes1 中边界框的数量。
# bboxes1.new(...) ：bboxes1.new 使用与 bboxes1 相同的数据类型和设备，根据提供的元组形状创建一个新的空张量。 这个新张量的形状将是 batch_shape + (rows, ) ，这意味着它的维度将与 batch_shape 相同，并且会新增一个维度，其大小为 rows 。
# 例如，如果 batch_shape 是 [32] （表示批次中有32张图像），而 rows 是80（表示每个图像中有80个边界框），那么返回的张量形状将是 [32, 80] 。这个张量可以用于存储与每个边界框相关的信息，例如每个边界框是否包含某个目标（即前景掩码）。
return bboxes1.new(batch_shape + (rows, ))
else:
return bboxes1.new(batch_shape + (rows, cols))
# area1 是一个张量，包含了 bboxes1 中所有边界框的面积。这个面积信息在计算交并比（IoU）和进行非极大值抑制（NMS）等操作时非常有用。
area1 = (bboxes1[..., 2] - bboxes1[..., 0]) * (
bboxes1[..., 3] - bboxes1[..., 1])
# 同理，area1 是一个张量，包含了 bboxes1 中所有边界框的面积。
area2 = (bboxes2[..., 2] - bboxes2[..., 0]) * (
bboxes2[..., 3] - bboxes2[..., 1])
if is_aligned:
# torch.max(input, dim, *, keepdim=False, out=None) → max, max_indices
# input ：输入的张量。
# dim ：整数，指定要计算最大值的维度。
# keepdim ：布尔值，如果为 True ，则输出张量会保持输入张量的维度，即不会移除计算最大值的维度。
# out ：一个可选的张量，用于存储输出结果。
# 返回两个张量： max 包含指定维度上的最大值， max_indices 包含最大值的索引。
# 用于找到两组边界框的交集的左上角和右下角坐标。
# lt 张量包含了两组边界框在每个维度上的交集的左上角坐标， rb 张量包含了交集的右下角坐标。
# 如果 lt 的任一坐标大于 rb 的对应坐标，这意味着没有交集，交集的面积将为零。这些坐标通常用于计算交集的面积，或者进一步计算交并比（IoU）。
lt = torch.max(bboxes1[..., :2], bboxes2[..., :2]) # [B, rows, 2]
# torch.min(input, dim, *, keepdim=False, out=None) → min, min_indices
# input ：输入的张量。
# dim ：整数，指定要计算最小值的维度。
# keepdim ：布尔值，如果为 True ，则输出张量会保持输入张量的维度，即不会移除计算最小值的维度。
# out ：一个可选的张量，用于存储输出结果。
# 返回两个张量： min 包含指定维度上的最小值， min_indices 包含最小值的索引。
rb = torch.min(bboxes1[..., 2:], bboxes2[..., 2:]) # [B, rows, 2]
wh = fp16_clamp(rb - lt, min=0)
overlap = wh[..., 0] * wh[..., 1]
if mode in ['iou', 'giou']:
union = area1 + area2 - overlap
else:
union = area1
if mode == 'giou':
enclosed_lt = torch.min(bboxes1[..., :2], bboxes2[..., :2])
enclosed_rb = torch.max(bboxes1[..., 2:], bboxes2[..., 2:])
else:
lt = torch.max(bboxes1[..., :, None, :2],
bboxes2[..., None, :, :2]) # [B, rows, cols, 2]
rb = torch.min(bboxes1[..., :, None, 2:],
bboxes2[..., None, :, 2:]) # [B, rows, cols, 2]
wh = fp16_clamp(rb - lt, min=0)
overlap = wh[..., 0] * wh[..., 1]
if mode in ['iou', 'giou']:
union = area1[..., None] + area2[..., None, :] - overlap
else:
union = area1[..., None]
if mode == 'giou':
enclosed_lt = torch.min(bboxes1[..., :, None, :2],
bboxes2[..., None, :, :2])
enclosed_rb = torch.max(bboxes1[..., :, None, 2:],
bboxes2[..., None, :, 2:])
# 这段代码的意图是将这个数值转换为一个新的张量。
# 这将创建一个新的张量，包含 eps 的值，并且数据类型和设备将默认为CPU上的浮点数。
eps = union.new_tensor([eps])
union = torch.max(union, eps)
ious = overlap / union
if mode in ['iou', 'iof']:
return ious
# calculate gious 计算giou
enclose_wh = fp16_clamp(enclosed_rb - enclosed_lt, min=0)
enclose_area = enclose_wh[..., 0] * enclose_wh[..., 1]
enclose_area = torch.max(enclose_area, eps)
gious = ious - (enclose_area - union) / enclose_area
return gious