【深度学习】YOLO源码中的mAP计算代码的理解笔记（大部分代码逐行+基础解释）

提示：本篇博客是在阅读了YOLO源码中的mAP计算方法的代码后加上官方解释以及自己的debug调试理解每一步是怎么操作的。由于是大部分代码进行了逐行解释，所以篇幅过长。

---

前言

首先，在理解YOLO源码中的mAP计算过程大部分参考了这篇文章：【目标检测】评价指标：mAP概念及其代码实现(yolo源码/pycocotools)，这篇文章提到了mAP计算的一些基本知识，也提供了代码，这里也是参考的这篇文章里的yolo源码的mAP计算代码（这篇文章是根据YOLO源码中的整理过后的代码）。想要源码可以点链接进去或者本篇最后部分贴了源码。

一、输入格式处理

输入的格式要求如下图：

在进行mAP计算之前需要将YOLO模型预测文件的数据格式转换为绝对位置，并且需要进行相应位置的转换：
YOLO预测文件中：[class, $x_{center}$, $y_{center}$, width, height, conf] $\to$ [$x_{min}$, $y_{min}$, $x_{max}$, $y_{max}$, conf, class]
其中的$x_{center}$, $y_{center}$, width, height为相对位置，都是归一化处理后的百分比位置，（$x_{min}$, $y_{min}$）,（ $x_{max}$, $y_{max}$）则是经过还原后在原图上标注框的左上角和右下角在原图中的坐标
真实结果同理进行转换：我们在标注label的时候保存格式也是经过归一化处理后的数据，在计算mAP的时候也需要还原到原图的坐标上去。

1.1 转换公式

此处的转换原理和公式参考这篇文章所提到的标签格式转换：利用mAP计算yolo精确度，同时这篇文章也有格式转换的代码，大家可以参考，此处就不再贴格式转换代码，只简述原理以及公式。
格式转换的原理如下图：

根据原理可以得到原始公式：
$$\begin{gathered} x_{center}=\frac{x_{min}+x_{max}}{2},y_{center}=\frac{y_{min}+y_{max}}{2}, \\ width=x_{max}-x_{min},height=y_{max}-y_{min} \end{gathered}$$
联立上面四个公式可以得到：
$$\begin{gathered} x_{min}=x_{center}-\frac{width}{2},x_{max}=x_{center}+\frac{width}{2} \\ {y}_{min}=y_{center}-\frac{height}{2},y_{max}=y_{center}+\frac{height}{2} \end{gathered}$$
此时得到的（$x_{min}$, $y_{min}$）,（ $x_{max}$, $y_{max}$）是归一化后的坐标，还需要进行原图的还原：其中W、H为原图的宽度和高度。ps：此处的W、H和width、height不是同一个东西，width和height是经过归一化处理后的值，而W和H是原图的宽高
$$\begin{gathered} x_{min}=x_{min}\ast W,x_{max}=x_{max}\ast W \\ {y}_{min}=y_{min}\ast H,y_{max}=y_{max}\ast H \end{gathered}$$

---

二、init：初始化

2.1 iouv

iouv就是从0.5到0.95分为10个值的数组

2.2 stats

stats是一个列表，包含4个numpy数组
stats[0] $\to$ shape：[$n_{pred}$，10] $\to$ 所有预测框载10个IOU阈值上是TP还是FP，其中$n_{pred}$表述所有预测框的总数量
stats[1] $\to$ shape：[$n_{pred}$] $\to$ 所有预测框的置信度
stats[2] $\to$ shape：[$n_{pred}$] $\to$ 所有预测框的预测类别
stats[3] $\to$ shape：[$n_{label}$] $\to$ 所有真实框（即label标签的）的预测类别，其中$n_{label}$表述所有真实框的总数

三、process_batch：实现预测结果和真实结果的匹配（TP/FP统计）

3.1 输入参数的格式

在一部分已经阐述了格式的转换，此处就不再进行赘述。
其中，N是预测框的总数，M是真实框的总数

3.2 代码注释（逐行）

# 每一个预测结果在不同IoU下的预测结果匹配
correct = np.zeros((detections.shape[0], self.niou)).astype(bool)

初始化correct（bool形式） $\to$ shape：[N, 10] $\to$ 在10个IOU阈值上每个预测框的TP、FP情况，True表示TP，False表示FP。
zeros即初始化全为0 $\to$ False

此处穿插python知识：使用shape可以快速读取矩阵的形状

# 二维矩阵
shape[0] # 读取矩阵第一维度的长度 即数组的行数
shape[1] # 读取矩阵第二维度的长度 即数组的列数

# 图像
image.shape[0] # 图片的高
image.shape[1] # 图片的宽
image.shape[2] # 图片的通道数

# 一般来说，在二维张量里，shape[-1]表示列数，
#注意，即使是一维行向量，shape[-1]表示行向量的元素总数，换言之也是列数：
shape[-1] # 表示最后一个维度

---

if detections is None:
 self.stats.append((correct, *torch.zeros((2, 0), device=self.device), labels[:, 0]))

注意一下：本地调试的时候，这里直接写None会报错
后面本人使用以下解决：当没有产生预测文件的时候：使用tensor生成0行6列的向量。

 else: # 没有预测框生成
      detections = torch.zeros((0, 6))

然后将上面从头开始的两句（correct的初始化以及判空的语句）修改为下面这个即可。

nl, npr = labels.shape[0], detections.shape[0]
correct = torch.zeros(npr, self.niou, dtype=torch.bool, device="cpu")
 if npr == 0:
     if nl:
     self.stats.append((correct, *torch.zeros((2, 0), device=self.device), labels[:, 0]))

---

else:
   # 计算标签与所有预测结果之间的IoU
   iou = box_iou(labels[:, 1:], detections[:, :4])

iou $\to$ shape：[M,N]，以labels真实框为行，detections预测框为列 $\to$ 计算每个真实框与每个预测框之间的交并比IOU

---

   # 计算每一个预测结果可能对应的实际标签
   correct_class = labels[:, 0:1] == detections[:, 5]

correct_class $\to$ shape：[M,N]，以labels真实框为行，detections预测框为列 $\to$ 保存每个真实框与每个预测框之间的类别是否相等，True即为类别一致，False即为类别不一致。
其中，labels[:, 0:1] == detections[:, 5]返回值为bool类型，True or False

例子：比如预测框预测的三个框的类别依次为2，0，1（括号里面表示类别，括号前面表示下标），真实框只有两个框，并且类别依次为0，1。得到的correct_class矩阵如下所示

label\detection	0(2)	1(0)	2(1)
0(1)	F	F	T
1(2)	T	F	F

---

 for i in range(self.niou):  # 在不同IoU置信度下的预测结果匹配结果
      # 根据IoU置信度和类别对应得到预测结果与实际标签的对应关系
      x = torch.where((iou >= self.iouv[i]) & correct_class)

外层的循环是指不同的IOU阈值下的计算
使用torch.where（此处采用的是后面所提的用法2：取Tensor中符合条件的坐标）获得真实框和预测框之间的iou大于此时的IOU阈值并且类别一致的结果 $\to$ shape：[2, $N_{same}$] $\to$ $N_{same}$指的是一致的组总数量，第一行表示行坐标，第二行表示列坐标。

例子：假设前面获得的iou和correct_class如下，此时的iou阈值为0.5，(0，2)(1，0)这两组符合条件，成为返回的结果。故torch.where返回值为[ [0,1] [2,0] ]。（第一行表示真实框索引，第二行表示预测框索引）

iou：

label\detection	0(2)	1(0)	2(1)
0(1)	0.6	0.7	0.8
1(2)	0.9	0.2	0.1

correct_class：

label\detection	0(2)	1(0)	2(1)
0(1)	F	F	T
1(2)	T	F	F

此处穿插python知识：torch.where用法
参考这篇文章：torch.where()的两种用法

用法1：按照指定条件合并两个同维度Tensor：
函数原型：torch.where(condition, x, y) $\to$ Tensor

用法2：取Tensor中符合条件的坐标：
函数原型：torch.where(condition) $\to$ Tensor

---

# 若存在和实际标签相匹配的预测结果
# x[0]:存在为True的索引(实际结果索引), x[1]当前所有True的索引(预测结果索引)
if x[0].shape[0]:  
 # [label, detect, iou]
 matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy()

通俗讲，x[0]指的是x的第一行，即对应真实框的索引号。x[1]指的是x的第二行，即对应的预测框的索引号。如果x[0]存在，就进行一个[label, detect, iou]的拼接。

torch.stack(x, 1) $\to$ x在列上进行拼接，从[ [0,1] [2,0] ] $\to$ [ [0,2] [1,0] ]。相当于就是变回一组坐标的形式。
iou[x[0], x[1]][:, None] $\to$ 将iou矩阵中的对应于x中的索引的iou取出来，以[[0.8],[0.9]]的形式。
torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1) $\to$ 将前面所得到的两个东西在列上进行拼接 $\to$ 得到[ [0,2,0.8], [1,0,0.9] ] （形成[label, detect, iou]）

此处穿插python知识：torch.stack用法
沿一个新维度对输入一系列张量进行连接，序列中所有张量应为相同形状，stack 函数返回的结果会新增一个维度。

# dim = 0 : 在第0维进行连接，相当于在行上进行组合(输入张量为一维，输出张量为两维)
a = torch.tensor([1, 2, 3])
b = torch.tensor([11, 22, 33])
c = torch.stack([a, b],dim=0)
#c => tensor([[ 1,  2,  3], [11, 22, 33]])

# dim=1:在第1维进行连接，相当于在对应行上面对列元素进行组合(输入张量为一维，输出张量为两维)
a = torch.tensor([1, 2, 3])
b = torch.tensor([11, 22, 33])
c = torch.stack([a, b],dim=1)
#c => tensor([[ 1, 11], [ 2, 22], [ 3, 33]])

此处穿插python知识：torch.cat用法
用于在指定的维度上拼接张量。这个函数接收一个张量列表，并在指定的维度上将它们连接起来。它通常用于连接两个或多个张量，以创建一个更大的张量。

tensor1 = torch.tensor([[1, 2], [3, 4]])
tensor2 = torch.tensor([[5, 6], [7, 8]])
# 在第0维（行）上连接这两个张量
result = torch.cat((tensor1, tensor2), dim=0)
# result => tensor([[1, 2], [3, 4], [5, 6], [7, 8]])
# 新的张量，行数=tensor1和tensor2 的行数之和，列数与 tensor1 和 tensor2 的列数相同。

---

   if x[0].shape[0] > 1:  # 存在多个与目标对应的预测结果
        # 根据IoU从高到低排序 [实际结果索引,预测结果索引,结果IoU]
        matches = matches[matches[:, 2].argsort()[::-1]] 

matches[:, 2].argsort()[::-1] $\to$ -1指的是逆序排序，即从IOU高到低排序，2是指matches[2]即iou值，整体返回的是逆序排位后的索引，如[1,0] （因为0.9>0.8，所以第二行应该和第一行交换位置，所以得到的排序后的索引为[1,0]）

再使用matches = matches[排序后的索引] 进行替换重置。

---

 # 每一个预测结果保留一个和实际结果的对应
matches = matches[np.unique(matches[:, 1], return_index=True)[1]] 
# 每一个实际结果和一个预测结果对应
matches = matches[np.unique(matches[:, 0], return_index=True)[1]]

np.unique(matches[:, 0], return_index=True)[0] => 去除重复后的值 [1] => 去除重复后的每个值对应的索引 [2] => dtype
然后matches再根据获得去除重复后的索引进行排列。
ps：因为前面已经按照IOU从高到低进行排序了，故当去除的时候自动保留最高的IOU的那一组

此处穿插python知识：numpy.unique用法
此处参考这篇文章：【Python】np.unique() 介绍与使用
去除其中重复的元素 ，并按元素由小到大返回一个新的无元素重复的元组或者列表。

# 格式：numpy.unique(arr, return_index, return_inverse, return_counts)
# arr：输入数组，如果不是一维数组则会展开
# return_index：如果为 true，返回新列表元素在旧列表中的位置（下标），并以列表形式存储。
# return_inverse：如果为true，返回旧列表元素在新列表中的位置（下标），并以列表形式存储。
# return_counts：如果为 true，返回去重数组中的元素在原数组中的出现次数。
A = [1, 2, 2, 5, 3, 4, 3]
a = np.unique(A) # [1 2 3 4 5]
a, indices = np.unique(A, return_index=True)# 返回新列表元素在旧列表中的位置（下标）
# a => [1 2 3 4 5] indices => [0 1 4 5 3]
a, indices = np.unique(A, return_inverse=True)# 旧列表的元素在新列表的位置
# a => [1 2 3 4 5] indices => [0 1 1 4 2 3 2]
a, indices = np.unique(A, return_counts=True)# 每个元素在旧列表里各自出现了几次
# a => [1 2 3 4 5] indices => [1 2 2 1 1]

---

 # 表明当前预测结果在当前IoU下实现了目标的预测
 correct[matches[:, 1].astype(int), i] = True  

matches[:, 1].astype(int) $\to$ 将matches的1列（第二列，即预测框的索引）作为int类型
matches表明经过筛选后留下的预测框的索引，然后在correct矩阵中将其置为True，即表明在该阈值（i）下，此预测框为正确的预测（TP）。

---

   # 预测结果在不同IoU是否预测正确, 预测置信度, 预测类别, 实际类别
   self.stats.append((correct, detections[:, 4], detections[:, 5], labels[:, 0]))

当所有iou阈值循环完成后，stats添加筛选过后的四个Numpy数组：correct、conf、class、label_class。

---

一个小总结：

数组名	shape	意义
correct	[$N_{pred}$, 10]	$N_{pred}$表示所有预测框的数量，表示所有预测框在该IOU阈值下为TP还是FP
correct_class	[$N_{label}$, $N_{pred}$]	表明每组真实框与预测框之间的类别是否一致
x	[2,X]	经过类别和阈值筛选后剩下的X组数据的索引号，第一行表示行坐标，第二行表示列坐标
matches	[Y,3]	去除重复后剩下的Y组数据，[label, detection, iou]形式，表明预测框和真实框匹配的框的数据
stats	每一维都有四个数组	分别为correct、conf、class、label_class

四、calculate_ap_per_class: 计算每一类别的AP值

4.1 代码注释（逐行）

stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*self.stats)]  # to numpy
# tp:所有预测结果在不同IoU下的预测结果 [n, 10]
# conf: 所有预测结果的置信度
# pred_cls: 所有预测结果得到的类别
# target_cls: 所有图片上的实际类别
tp, conf, pred_cls, target_cls = stats[0], stats[1], stats[2], stats[3]
# 根据类别置信度从大到小排序
i = np.argsort(-conf)  # 根据置信度从大到小排序
tp, conf, pred_cls = tp[i], conf[i], pred_cls[i]

第一行代码是将所有图片的四个数组进行汇总，在YOLO源码中，未拼接前 $\to$ stats[0] 指第一张图的四个数组，拼接后 $\to$ stats[0]指所有图片的所有预测框的TP/FP情况。

然后再根据置信度逆序排序，即按照置信度从高到低进行排序。

---

# 得到所有类别及其对应数量(目标类别数)
unique_classes, nt = np.unique(target_cls, return_counts=True)
nc = unique_classes.shape[0]  # number of classes
# ap: 每一个类别在不同IoU置信度下的AP，shape[nc类别数, 10], 
# p:每一个类别的P曲线(不同类别置信度), r:每一个类别的R(不同类别置信度)
ap, p, r = np.zeros((nc, tp.shape[1])), np.zeros((nc, 1000)), np.zeros((nc, 1000))

np.unique用法见前面，此处不再赘述。

---

for ci, c in enumerate(unique_classes):  # 对每一个类别进行P,R计算，ci为c的index，c为值
            i = pred_cls == c
            n_l = nt[ci]  # number of labels 该类别的实际数量(正样本数量)
            n_p = i.sum()  # number of predictions 预测结果数量
            if n_p == 0 or n_l == 0:
                continue

i = pred_cls == c返回的是bool类型的一维数组。
nt指的是每个类别真实框的总数量，一维数组。

---

# cumsum：轴向的累加和, 计算当前类别在不同的类别置信度下的P,R
fpc = (1 - tp[i]).cumsum(0)  # FP累加和(预测为负样本且实际为负样本)
tpc = tp[i].cumsum(0)  # TP累加和(预测为正样本且实际为正样本)
# 召回率计算(不同的类别置信度下)
recall = tpc / (n_l + eps)
# 精确率计算(不同的类别置信度下)
precision = tpc / (tpc + fpc)

TP和FP的计算方式开头所提到的参考文章有非常清楚的解释，此处不再赘述。
此处说一下tp[i]的用法见下图，前面获得了i这个数组，记录了在tp数组中哪些预测框是当前类别的预测框。
tp[i]会保留对应i为True的数据。

cumsum(）的用法见图

import numpy as np
a = np.asarray([[1, 2, 3],
                [4, 5, 6],
                [7, 8, 9]
                ])
b = a.cumsum(axis=0) # 按行累加
# b=>[[1 2 3] [5 7 9] [12 15 18]]
c = a.cumsum(axis=1) # 按列累加
# c=>[[1 3 6] [4 9 15] [7 15 24]]

---

            # 计算不同类别置信度下的AP(根据P-R曲线计算)
            for j in range(tp.shape[1]):
                ap[ci, j], mpre, mrec = self.compute_ap(recall[:, j], precision[:, j])
        # 所有类别的ap值 @0.5:0.95
       	return ap

最后得到recall和precision值后就是计算ap值。

五、compute_ap：计算PR曲线的面积

此处就是计算每一组PR值形成的曲线的面积，代码本来的注解就很清晰明了，此处就不逐行解释了。
此处贴一个插值积分求面积的博客，里面比较详细介绍了插值积分：numpy.interp()用法

六、源码

class MeanAveragePrecison:
    def __init__(self, device="cpu"):
        '''
        计算mAP: mAP@0.5; mAP @0.5:0.95; mAP @0.75
        '''
        self.iouv = torch.linspace(0.5, 0.95, 10, device=device)  # 不同的IoU置信度 @0.5:0.95
        self.niou = self.iouv.numel()  # IoU置信度数量
        self.stats = []  # 存储预测结果
        self.device = device
    def process_batch(self, detections, labels):
        '''
        预测结果匹配(TP/FP统计)
        :param detections:(array[N,6]) x1,y1,x1,y1,conf,class (原图绝对坐标)
        :param labels:(array[M,5]) class,x1,y1,x2,y2 (原图绝对坐标)
        '''
        # 每一个预测结果在不同IoU下的预测结果匹配
        correct = np.zeros((detections.shape[0], self.niou)).astype(bool)
        if detections is None:
            self.stats.append((correct, *torch.zeros((2, 0), device=self.device), labels[:, 0]))
        else:
        # 计算标签与所有预测结果之间的IoU
            iou = box_iou(labels[:, 1:], detections[:, :4])
            # 计算每一个预测结果可能对应的实际标签
            correct_class = labels[:, 0:1] == detections[:, 5]
            for i in range(self.niou):  # 在不同IoU置信度下的预测结果匹配结果
                # 根据IoU置信度和类别对应得到预测结果与实际标签的对应关系
                x = torch.where((iou >= self.iouv[i]) & correct_class)
                # 若存在和实际标签相匹配的预测结果
                if x[0].shape[0]:  # x[0]:存在为True的索引(实际结果索引), x[1]当前所有True的索引(预测结果索引)
                    # [label, detect, iou]
                    matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy()
                    if x[0].shape[0] > 1:  # 存在多个与目标对应的预测结果
                        matches = matches[matches[:, 2].argsort()[::-1]]  # 根据IoU从高到低排序 [实际结果索引,预测结果索引,结果IoU]
                        matches = matches[np.unique(matches[:, 1], return_index=True)[1]]  # 每一个预测结果保留一个和实际结果的对应
                        matches = matches[np.unique(matches[:, 0], return_index=True)[1]]  # 每一个实际结果和一个预测结果对应
                    correct[matches[:, 1].astype(int), i] = True  # 表面当前预测结果在当前IoU下实现了目标的预测
            # 预测结果在不同IoU是否预测正确, 预测置信度, 预测类别, 实际类别
            self.stats.append((correct, detections[:, 4], detections[:, 5], labels[:, 0]))

    def calculate_ap_per_class(self, save_dir='.', names=(), eps=1e-16):
        stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*self.stats)]  # to numpy
        # tp:所有预测结果在不同IoU下的预测结果 [n, 10]
        # conf: 所有预测结果的置信度
        # pred_cls: 所有预测结果得到的类别
        # target_cls: 所有图片上的实际类别
        tp, conf, pred_cls, target_cls = stats[0], stats[1], stats[2], stats[3]
        # 根据类别置信度从大到小排序
        i = np.argsort(-conf)  # 根据置信度从大到小排序
        tp, conf, pred_cls = tp[i], conf[i], pred_cls[i]

        # 得到所有类别及其对应数量(目标类别数)
        unique_classes, nt = np.unique(target_cls, return_counts=True)
        nc = unique_classes.shape[0]  # number of classes
 
        # ap: 每一个类别在不同IoU置信度下的AP, p:每一个类别的P曲线(不同类别置信度), r:每一个类别的R(不同类别置信度)
        ap, p, r = np.zeros((nc, tp.shape[1])), np.zeros((nc, 1000)), np.zeros((nc, 1000))
        for ci, c in enumerate(unique_classes):  # 对每一个类别进行P,R计算
            i = pred_cls == c
            n_l = nt[ci]  # number of labels 该类别的实际数量(正样本数量)
            n_p = i.sum()  # number of predictions 预测结果数量
            if n_p == 0 or n_l == 0:
                continue

            # cumsum：轴向的累加和, 计算当前类别在不同的类别置信度下的P,R
            fpc = (1 - tp[i]).cumsum(0)  # FP累加和(预测为负样本且实际为负样本)
            tpc = tp[i].cumsum(0)  # TP累加和(预测为正样本且实际为正样本)
            # 召回率计算(不同的类别置信度下)
            recall = tpc / (n_l + eps)
            # 精确率计算(不同的类别置信度下)
            precision = tpc / (tpc + fpc)


            # 计算不同类别置信度下的AP(根据P-R曲线计算)
            for j in range(tp.shape[1]):
                ap[ci, j], mpre, mrec = self.compute_ap(recall[:, j], precision[:, j])
        # 所有类别的ap值 @0.5:0.95
       	return ap


    def compute_ap(self, recall, precision):
        # 增加初始值(P=1.0 R=0.0) 和 末尾值(P=0.0, R=1.0)
        mrec = np.concatenate(([0.0], recall, [1.0]))
        mpre = np.concatenate(([1.0], precision, [0.0]))

        # Compute the precision envelope np.maximun.accumulate
        # (返回一个数组,该数组中每个元素都是该位置及之前的元素的最大值)
        mpre = np.flip(np.maximum.accumulate(np.flip(mpre)))

        # 计算P-R曲线面积
        method = 'interp'  # methods: 'continuous', 'interp'
        if method == 'interp':  # 插值积分求面积
            x = np.linspace(0, 1, 101)  # 101-point interp (COCO))
            # 积分(求曲线面积)
            ap = np.trapz(np.interp(x, mrec, mpre), x)
        elif method == 'continuous':  # 不插值直接求矩阵面积
            i = np.where(mrec[1:] != mrec[:-1])[0]  # points where x axis (recall) changes
            ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])  # area under curve

        return ap, mpre, mrec

---

结束语

浅浅记录这几天阅读YOLO源码中的mAP计算过程，自身python基础不太ok，也算一边查阅python语法知识一边debug看懂每一步中数值的变化。可能以上自身的理解存在一点偏差，如果存在问题欢迎大家在评论区指出~