前言
模型训练过程中每一轮都会计算P,R,mAP,mAP@0.5等数值,本篇分析这些数值的计算过程,分析最核心部分。我的感受是计算的过程比想象的复杂。
主要的流程在yolov5/val.py文件的process_batch处理函数中。
if nl:
tbox = xywh2xyxy(labels[:, 1:5]) # target boxes
scale_boxes(im[si].shape[1:], tbox, shape, shapes[si][1]) # native-space labels
labelsn = torch.cat((labels[:, 0:1], tbox), 1) # native-space labels
correct = process_batch(predn, labelsn, iouv, im[si])
if plots:
confusion_matrix.process_batch(predn, labelsn)
计算map@0.5的计算基本原理是:
- 给定一个阈值数组,从0.5-0.95每间隔0.05,生成一共10个数据
- 获取预测结果和标注,计算预测框和gt框的IOU,判断类别是否一致。
- 用阈值数组中的每一个阈值去筛选IOU,大于阈值的预测为True, 小于IOU的预测为False
- 统计结果,输出结果
入参分析
detections:bbox,预测框信息,格式为:x1, y1, x2, y2, conf, class
labels:gt框信息,格式为:class, x1, y1, x2, y2
iouv:阈值数组,内容是:[0.50000, 0.55000, 0.60000, 0.65000, 0.70000, 0.75000, 0.80000, 0.85000, 0.90000, 0.95000]
def process_batch(detections, labels, iouv, img=None):
"""
Return correct prediction matrix
Arguments:
detections (array[N, 6]), x1, y1, x2, y2, conf, class
labels (array[M, 5]), class, x1, y1, x2, y2
Returns:
correct (array[N, 10]), for 10 IoU levels
"""
预测框的可视化:
构造返回结果数组
对于每一个预测框,在不同的阈值下会判断成True 或 False。构造一个所有预测框在所有阈值下的结果二维数组。
# 预测正确的记录,所有iouv [0.5, 0.55, 0.6 ... 0.95]阈值下iou大于阈值且类型匹配的结果。
# shape:(300, 10) 300个预测,10个阈值
correct = np.zeros((detections.shape[0], iouv.shape[0])).astype(bool)
- detections.shape[0]: 300
- iouv.shape[0]:10
最终的结果可以理解成300个预测框在10个不同阈值下一共会有3000个结果,结果为布尔数组。
计算所有预测框和标注框的IOU结果
将gt框和预测框传入box_iou函数,得到结果。
"""
计算gt框和预测的iou,这个过程会出现广播机制,5个gt框分别和300个预测框计算iou,得到的结果为5,300
(Pdb) pp detections[:, :4].shape
torch.Size([300, 4])
(Pdb) pp labels[:, 1:].shape
torch.Size([5, 4])
(Pdb) pp iou.shape
torch.Size([5, 300])
"""
iou = box_iou(labels[:, 1:], detections[:, :4])
- labels[:, 1:]:(5, 4)
- detections[:, :4]:(300, 4)
5个标注框和300个预测框匹配,也就是拿每一个标注框都和300个预测框匹配,得到5*300个结果。相当于每一行代表这一个标注框和预测框的iou匹配结果。
tensor([[0.00000, 0.00000, 0.78682, 0.02270, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.02236, 0.00000, 0.00000, 0.00000, 0.00000, 0.43876, 0.70469, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000,
0.00000, 0.00000, 0.00000, 0.04540, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.46606, 0.00000, 0.00000, 0.00000, 0.00000, 0.13432, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.00000, 0.53255,
匹配标签
上一步计算了预测框和标注框的iou,这一步计算预测类别和标注类别的匹配。
"""
计算标签匹配。labels[:, 0:1]代表所有标注的标签。用每一个gt的标签和所有的预测的标签匹配。
过程存在广播机制,结果为:torch.Size([5, 300])
结果中 0维也就是5是gt的label 1维也就是300是预测label
(Pdb) labels[:, 0:1].shape
torch.Size([5, 1])
(Pdb) detections[:, 5].shape
torch.Size([300])
correct_class.shape
torch.Size([5, 300])
"""
correct_class = labels[:, 0:1] == detections[:, 5]
labels[:, 0:1]:[5, 1]
detections[:, 5]: [300]
和上一步类似,每一个标注类表和预测类别匹配,5个标注类别和300个预测类别匹配,得到一个5*300的结果。得到的是一个布尔数组
(Pdb) correct_class
tensor([[ True, True, True, True, True, True, True, False, True, True, False, False, True, False, True, True, False, False, True, False, True, False, True, True, True, True, True, True, False, False, True, True, True, False, True, True, False, True, True, False, True, False, False, True,
False, True, True, True, True, True, True, True, True, False, False, False, True, True, False, True, True, False, False, True, True, True, False, True, True, True, True, False, True, True, True, True, False, False, True, True, True, True, True, True, True, True, True, False
遍历所有阈值计算在不同的阈值下预测正确的个数
阈值数组是[0.50000, 0.55000, 0.60000, 0.65000, 0.70000, 0.75000, 0.80000, 0.85000, 0.90000, 0.95000],计算在不同的阈值下预测框和标注框的iou匹配结果。以0.5为例,大于0.5为True,小于0.5为False
# 循环遍历所有的iou阈值,计算在不同的阈值下预测正确的个数
for i in range(len(iouv)):
# (iou >= iouv[i]) & correct_class 得到一个布尔数组
# torch.where((iou >= iouv[i]) & correct_class) 筛选出布尔数组中为True的元素,返回和布尔数组结构一致,返回为True元素在当前张量中的坐标。当前为二维数组。
# 返回两个元素,元素1是数组,代表为True元素的x轴下标, 元素2也是数组,代表为True元素的y轴下标。
# 类似:(tensor([0, 0, 1, 1]), tensor([0, 1, 1, 3]))
x = torch.where((iou >= iouv[i]) & correct_class) # IoU > threshold and classes match
(iou >= iouv[i]) & correct_class
iou 是 [5, 300]的iou结果
correct_class 是[5, 300]的类别结果
将两者的判断与操作,得到的就是iou大于阈值,且类别正确的结果。结果是5*300的布尔数组
torch.where((iou >= iouv[i]) & correct_class)
筛选出布尔数组中为True的元素,返回Ture在张量中的坐标。
x的结果如下,结果代表的含义是第一个张量代表标注标签的ID,第二个张量代表匹配的对应的预测标签的ID。
(tensor([0, 0, 0, 1, 1, 2, 2, 2, 3, 3, 4, 4, 4, 5], device='cuda:0'),
tensor([ 2, 67, 68, 3, 79, 5, 31, 90, 4, 86, 1, 85, 92, 0], device='cuda:0'))
判断是否有正确匹配的结果
# 判断是否有结果
if x[0].shape[0]:
"""
将结果汇总成[label, detect, iou]的格式。
label:gt的标签的下标
detect: 预测结果标签的下标
iou: 预测和gt的iou值
torch.stack(x, 1) 将x轴和y轴表示的坐标转换成[x, y]的格式
iou[x[0], x[1]][:, None]) 获取iou中匹配上的元素值
"""
matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy() # [label, detect, iou]
torch.stack(x, 1)
将标注框的ID和预测框的ID组合在一次
iou[x[0], x[1]][:, None]
筛选出对应的IOU值
将两者组合在一起,得到的就是[label_id, detect_id, iou]
去除重复匹配
一个gt框可能和多个预测框的iou都大于阈值,这是完全可能的。这里只保留最好的匹配,所以把那些多余的匹配都去掉。一个gt框只保留一个预测框,同时一个预测框只对应一个gt框。
# 匹配存在多个元素,可能是因为广播机制导致的重复,去除重复。
if x[0].shape[0] > 1:
# 对置信度排序,倒序排列
matches = matches[matches[:, 2].argsort()[::-1]]
# 对预测detect进行去重,一个预测框只保留一个iou匹配。返回去重之后的matches
matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
# 对label进行去重,一个gt框只保留一个最高iou匹配,返回去重之后的matches
# matches = matches[matches[:, 2].argsort()[::-1]]
matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
保存结果
correct是300*10的二维数组,将对应行和列的元素设置为True。
# 将correct中相应的元素置为True
correct[matches[:, 1].astype(int), i] = True
matches[:, 1].astype(int)
预测框的下标i
某一个阈值
将某一个阈值下对应的预测框设置成True,也就是代表着在该阈值下,预测框预测的类别正确,同时IOU也超过阈值。
完整代码
def process_batch(detections, labels, iouv, img=None):
"""
Return correct prediction matrix
Arguments:
detections (array[N, 6]), x1, y1, x2, y2, conf, class
labels (array[M, 5]), class, x1, y1, x2, y2
Returns:
correct (array[N, 10]), for 10 IoU levels
"""
# 预测正确的记录,所有iouv [0.5, 0.55, 0.6 ... 0.95]阈值下iou大于阈值且类型匹配的结果。
# shape:(300, 10) 300个预测,10个阈值
correct = np.zeros((detections.shape[0], iouv.shape[0])).astype(bool)
"""
计算gt框和预测的iou,这个过程会出现广播机制,5个gt框分别和300个预测框计算iou,得到的结果为5,300
(Pdb) pp detections[:, :4].shape
torch.Size([300, 4])
(Pdb) pp labels[:, 1:].shape
torch.Size([5, 4])
(Pdb) pp iou.shape
torch.Size([5, 300])
"""
iou = box_iou(labels[:, 1:], detections[:, :4])
breakpoint()
# 计算标签匹配。labels[:, 0:1]代表所有标注的标签。用每一个gt的标签和所有的预测的标签匹配。
# 过程存在广播机制,结果为:torch.Size([5, 300])
# 结果中 0维也就是5是gt的label 1维也就是300是预测label
"""
(Pdb) labels[:, 0:1].shape
torch.Size([5, 1])
(Pdb) detections[:, 5].shape
torch.Size([300])
correct_class.shape
torch.Size([5, 300])
"""
correct_class = labels[:, 0:1] == detections[:, 5]
# 循环遍历所有的iou阈值,计算在不同的阈值下预测正确的个数
for i in range(len(iouv)):
# (iou >= iouv[i]) & correct_class 得到一个布尔数组
# torch.where((iou >= iouv[i]) & correct_class) 筛选出布尔数组中为True的元素,返回和布尔数组结构一致,返回为True元素在当前张量中的坐标。当前为二维数组。
# 返回两个元素,元素1是数组,代表为True元素的x轴下标, 元素2也是数组,代表为True元素的y轴下标。
# 类似:(tensor([0, 0, 1, 1]), tensor([0, 1, 1, 3]))
x = torch.where((iou >= iouv[i]) & correct_class) # IoU > threshold and classes match
# 判断是否有结果
if x[0].shape[0]:
"""
将结果汇总成[label, detect, iou]的格式。
label:gt的标签的下标
detect: 预测结果标签的下标
iou: 预测和gt的iou值
torch.stack(x, 1) 将x轴和y轴表示的坐标转换成[x, y]的格式
iou[x[0], x[1]][:, None]) 获取iou中匹配上的元素值
"""
matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy() # [label, detect, iou]
# 匹配存在多个元素,可能是因为广播机制导致的重复,去除重复。
if x[0].shape[0] > 1:
# 对置信度排序,倒序排列
matches = matches[matches[:, 2].argsort()[::-1]]
# 对预测detect进行去重,一个预测框只保留一个iou匹配。返回去重之后的matches
matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
# 对label进行去重,一个gt框只保留一个最高iou匹配,返回去重之后的matches
matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
# 将correct中相应的元素置为True
correct[matches[:, 1].astype(int), i] = True
"""
(Pdb) pp correct
array([[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False],
...,
[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False],
[False, False, False, ..., False, False, False]])
(Pdb) pp correct.shape
(300, 10)
"""
return torch.tensor(correct, dtype=torch.bool, device=iouv.device)