目标检测中的mAP-代码实现

mAP

mAP，全称为mean Average Precision，在目标检测任务中被用于衡量检测器的好坏。本文第一部分讲解mAP的概念以及计算过程，第二部分专注于用代码实现mAP的计算。

在做目标检测时，每个类别对应有一个AP，全部类别的AP求平均就是mAP。

AP是P-R曲线下方的面积。

P-R曲线的横轴是Recall，纵轴是Precision。

因此，欲计算mAP，得先计算每个类别对应的AP，进一步，得先计算每个类别对应的Recall和Precision。

下面正文开始~

一、计算某个类别对应的AP

1. 获取测试集上所有预测的bbox的结果(TP or FP)

这里，假设预测得到的bbox已经做过NMS处理，如果不清楚什么是NMS，请查看这篇文章。

GT总数(绿色框)为4，而红色框是预测得到的，红色数字表示置信度(confidence)。

接下来需要计算每张图片中的每一个红色预测框与其附近的绿色真实框的IoU，当IoU大于预先设定的IoU阈值时，设置该预测框为TP，否则设置为FP。

之所以使用TP和FP，是因为我们感兴趣的是 预测框的位置是否接近真实框的位置 (接近程度用两者的IoU来衡量)，接近的(IoU大于某个阈值)就视为正例(Positive，简写为P)。

对第一张图片：

对第二张图片：

对第三张图片：

2. 按照置信度从大到小排序

3. 计算Precision和Recall

即：
$$Precision=\frac{TP}{TP+FP}$$
$$Recall=\frac{TP}{TP+FN}$$

而P-R曲线由一系列的点(Recall,Precision)连接而成，因此需要若干这样的点，那么如何获取这些点呢？

Recall和Precision的计算公式已经在上面写出了。

其中，$TP+FN$就是全部绿色真实框的个数(即：GT总数)，这是恒定不变的。

而$TP$和$FP$是动态变化的，这一点可以通过类比来理解。在普通的分类任务中，通过设置不同的阈值，可以获取多个不同的分类结果，这些分类结果对应不同的混淆矩阵，据此可以获得多个不同的(Recall, Precision)，它们构成了P-R曲线上的一个个点。而在目标检测任务中，置信度充当了上述阈值的角色，不同的置信度下，可获取不同的(Recall,Precision)，于是就可以绘制P-R曲线，进一步计算AP(P-R曲线下方面积)了。

这里，将置信度从大到小排序，可以方便的通过累加方式统计不同置信度下的TP和FP。

不同置信度下的Precision和Recall统计如下：

其中，TP+FN=4，即总GT数。

第一行，累加TP=1，累加FP=0，$Precision=\frac{TP}{TP+FP}=\frac1{1+0}$，$Recall=\frac{TP}{TP+FN}=\frac14$

第二行，累加TP=1，累加FP=1，$Precision=\frac{TP}{TP+FP}=\frac1{1+1}$，$Recall=\frac{TP}{TP+FN}=\frac14$

第三行，累加TP=2，累加FP=1，$Precision=\frac{TP}{TP+FP}=\frac2{2+1}$，$Recall=\frac{TP}{TP+FN}=\frac24$

同理可计算后续所有行。

得到这些点之后，就可以绘制P-R曲线了。

4. 绘制P-R曲线

5. 计算P-R曲线下方面积，即AP

二、计算所有类别的AP

对每一个类别，都执行(一)中的步骤，得到每个类别的AP。

三、计算mAP

将(二)中计算得到的AP求平均，就得到了当前IoU阈值下的mAP。

以上仅仅是对于某一个特定的IoU阈值，我们还需要尝试不同的IoU阈值，然后再取平均【sum(每一个IoU阈值下得到的mAP)/不同IoU阈值的个数】，得到最终的mAP。

PyTorch实现mAP

定义计算mAP的函数mean_average_precision，该函数共4个参数。其中，pred_bboxes 代表所有预测框，true_boxes代表所有真实框，iou_threshold代表设定的IoU阈值，num_classes是总类别数。

pred_bboxes 和true_boxes都包含多条数据。

pred_bboxes中的每条数据格式如下：

[train_idx,class_pred,prob_score,x1,y1,x2,y2]

train_idx：指示图片编号，用于区分不同的图片

class_pred：预测的类别

prob_score：置信度

(x1,y1)：bbox左上角坐标（可能已经归一化）

(x2,y2)：bbox右下角坐标（可能已经归一化）

true_boxes中每条数据的格式与之类似，只不过它是确定的（GT）。

注意：为了描述方便，接下来会将上面所说的”一条数据”直接说成是”一个(真实/预测)框”，这是不准确的，因为这一条数据除了包含框的位置信息外，还包含上面讲到的train_idx等信息。

开始写代码：

def mean_average_precision(pred_bboxes,true_boxes,iou_threshold,num_classes=20):
    
    #pred_bboxes(list): [[train_idx,class_pred,prob_score,x1,y1,x2,y2], ...]
    
    average_precisions=[]#存储每一个类别的AP
    epsilon=1e-6#防止分母为0

我们需要分别求解每个类别的AP，然后对所有类别的AP求平均就能得到mAP，因此接下来将专注于求解一个类别的AP。继续写代码：

#对于每一个类别
for c in range(num_classes):
    detections=[]#存储预测为该类别的bbox
    ground_truths=[]#存储本身就是该类别的bbox(GT)
    
    for detection in pred_bboxes:
        if detection[1]==c:
            detections.append(detection)
    
    for true_box in true_boxes:
        if true_box[1]==c:
            ground_truths.append(true_box)

上面的代码将预测为该类别的框存储在detections列表中，将本身就是该类别的真实框存储在ground_truths列表中。

继续写代码：

amount_bboxes=Counter(gt[0] for gt in ground_truths)

for key,val in amount_bboxes.items():
    amount_bboxes[key]=torch.zeros(val)#置0，表示这些真实框初始时都没有与任何预测框匹配

第一行代码用于统计每一张图片中真实框的个数，gt[0]是train_idx，它指示了用于区分每张图片的一个编号。

举个栗子：

假设：
img 0 has 3 bboxes
img 1 has 5 bboxes

那么这句代码的作用是得到如下结果：
amount_bboxes={0:3,1:5}

后面的代码改变字典的value，继续上面的栗子，改变后结果如下：

amount_bboxes={0:torch.tensor([0,0,0]),1:torch.tensor([0,0,0,0,0])}

刚开始，amout_bboxes包含了每张图片中真实框的个数，它是一个字典，其中key为图片编号，value为该图片包含的真实框的个数；之后，将真实框的个数用全0向量来替代，有几个真实框，全0向量就包含几个0。

比如这里，编号为0的图片中共包含3个真实框，编号为1的图片中共包含5个真实框。

继续写代码：

#将预测框按照置信度从大到小排序
detections.sort(key=lambda x:x[2],reverse=True)

#初始化TP,FP
TP=torch.zeros(len(detections))
FP=torch.zeros(len(detections))

#TP+FN就是当前类别GT框的总数，是固定的
total_true_bboxes=len(ground_truths)

#如果当前类别一个GT框都没有，那么直接跳过即可
if total_true_bboxes == 0:
    continue

按照之前所讲，把预测框按照置信度从大到小进行排序，并初始化TP和FP。由于每一个预测框对应一个TP或者FP，而当前类别下总共有 len(detections)个预测框，因此初始化为长度为len(detections)的全0向量。

继续写代码：

#对于每个预测框，先找到它所在图片中的所有真实框，然后计算预测框与每一个真实框之间的IoU，大于IoU阈值且该真实框没有与其他预测框匹配，则置该预测框的预测结果为TP，否则为FP
       for detection_idx,detection in enumerate(detections):
           #在计算IoU时，只能是同一张图片内的框做，不同图片之间不能做
           #图片的编号存在第0个维度
           #于是下面这句代码的作用是：找到当前预测框detection所在图片中的所有真实框，用于计算IoU
           ground_truth_img=[bbox for bbox in ground_truths if bbox[0]==detections[0]]
           
           num_gts=len(ground_truth_img)
           
           best_iou=0
           for idx,gt in emnumerate(ground_truth_img):
               #计算当前预测框detection与它所在图片内的每一个真实框的IoU
               iou=insert_over_union(torch.tensor(detection[3:]),torch.tensor(gt[3:]))
               if iou >best_iou:
                   best_iou=iou
                   best_gt_idx=idx
           if best_iou>iou_threshold:
               #这里的detection[0]是amount_bboxes的一个key，best_gt_idx是该key对应的真实框中的train_idx
               if amount_bboxes[detection[0]][best_gt_idx]==0:#只有没被占用的真实框才能用，0表示未被占用（占用：该真实框与某预测框匹配【两者IoU大于设定的IoU阈值】）
                   TP[detection_idx]=1#该预测框为TP
                   amount_bboxes[detection[0]][best_gt_idx]=1#将该真实框标记为已经用过了，不能再用于其他预测框。因为一个预测框最多只能对应一个真实框（最多：IoU小于IoU阈值时，预测框没有对应的真实框)
               else:
                   FP[detection_idx]=1#虽然该预测框与真实框中的一个框之间的IoU大于IoU阈值，但是这个真实框已经与其他预测框匹配，因此该预测框为FP
           else:
               FP[detection_idx]=1#该预测框与真实框中的每一个框之间的IoU都小于IoU阈值，因此该预测框直接为FP

这段代码完成了IoU的计算与比较，并由此确定每一个预测框是TP还是FP。

首先，找到当前预测框detection所在图片中的所有真实框，存储在ground_truth_img中，因为某一个预测框只能与该预测框所在图片内的真实框之间计算IoU。

amount_bboxes[detection[0]][best_gt_idx]用于表示编号（train_idx）为detection[0]的图片中，第best_gt_idx个真实框是否已经与某一个预测框匹配（两者IoU大于设定的IoU阈值），如果是，则置为1，否则置0。初始化时置为全零向量，表示都还未匹配。

为什么需要这样做？举个例子：下面的绿色框为真实框，红色框为预测框。假设1号预测框已经与左上角的绿色框匹配，那么对于2号预测框来说，它会分别计算自己与左上角以及右下角的两个真实框之间的IoU，虽然结果是它与左上角的真实框之间的IoU更大，但是左上角的真实框已经与1号预测框匹配，因此2号预测框就无法再与左上角的真实框进行匹配了。上述代码保证了这一点。

对了，在计算IoU时调用了insert_over_union函数，它是我们之前实现过的，这里直接搬过来了，忘记的可以查看这篇文章。

继续写代码，接下来的代码就没那么难理解了：

TP_cumsum=torch.cumsum(TP,dim=0)
FP_cumsum=torch.cumsum(FP,dim=0)

#套公式
recalls=TP_cumsum/(total_true_bboxes+epsilon)
precisions=torch.divide(TP_cumsum,(TP_cumsu+FP_cumsum+epsilon))

#把[0,1]这个点加入其中
precisions=torch.cat((torch.tensor([1]),precision))
recalls=torch.cat((torch.tensor([0]),recalls))
#使用trapz计算AP
average_precisions.append(torch.trapz(precisions,recalls)) 

计算每个置信度下的TP和FP，然后套公式计算每个置信度下的Recall和Precision，最后使用数值积分(torch.trapz)算P-R曲线下方的面积，即当前类别对应的AP。

最后，退出for循环，就得到了每个类别对应的AP，返回均值即可：

return sum(average_precisions)/len(average_precisions) 

这样就得到了当前IoU_threshold下的mAP。

正如前面所说，你还可以更换不同的IoU_threshold，重复调用上述函数进行计算，每次都将结果保存下来，最后取它们的均值，就得到了最终的mAP。

完整代码如下：

import torch
from collections import Counter

def mean_average_precision(pred_bboxes,true_boxes,iou_threshold,num_classes=20):
    
    #pred_bboxes(list): [[train_idx,class_pred,prob_score,x1,y1,x2,y2], ...]
    
    average_precisions=[]#存储每一个类别的AP
    epsilon=1e-6#防止分母为0
    
    #对于每一个类别
    for c in range(num_classes):
        detections=[]#存储预测为该类别的bbox
        ground_truths=[]#存储本身就是该类别的bbox(GT)
        
        for detection in pred_bboxes:
            if detection[1]==c:
                detections.append(detection)
        
        for true_box in true_boxes:
            if true_box[1]==c:
                ground_truths.append(true_box)
                
        #img 0 has 3 bboxes
        #img 1 has 5 bboxes
        #就像这样：amount_bboxes={0:3,1:5}
        #统计每一张图片中真实框的个数,train_idx指示了图片的编号以区分每张图片
        amount_bboxes=Counter(gt[0] for gt in ground_truths)
        
        for key,val in amount_bboxes.items():
            amount_bboxes[key]=torch.zeros(val)#置0，表示这些真实框初始时都没有与任何预测框匹配
        #此时，amount_bboxes={0:torch.tensor([0,0,0]),1:torch.tensor([0,0,0,0,0])}
        
        #将预测框按照置信度从大到小排序
        detections.sort(key=lambda x:x[2],reverse=True)
        
        #初始化TP,FP
        TP=torch.zeros(len(detections))
        FP=torch.zeros(len(detections))
        
        #TP+FN就是当前类别GT框的总数，是固定的
        total_true_bboxes=len(ground_truths)
        
        #如果当前类别一个GT框都没有，那么直接跳过即可
        if total_true_bboxes == 0:
            continue
        
        #对于每个预测框，先找到它所在图片中的所有真实框，然后计算预测框与每一个真实框之间的IoU，大于IoU阈值且该真实框没有与其他预测框匹配，则置该预测框的预测结果为TP，否则为FP
        for detection_idx,detection in enumerate(detections):
            #在计算IoU时，只能是同一张图片内的框做，不同图片之间不能做
            #图片的编号存在第0个维度
            #于是下面这句代码的作用是：找到当前预测框detection所在图片中的所有真实框，用于计算IoU
            ground_truth_img=[bbox for bbox in ground_truths if bbox[0]==detections[0]]
            
            num_gts=len(ground_truth_img)
            
            best_iou=0
            for idx,gt in emnumerate(ground_truth_img):
                #计算当前预测框detection与它所在图片内的每一个真实框的IoU
                iou=insert_over_union(torch.tensor(detection[3:]),torch.tensor(gt[3:]))
                if iou >best_iou:
                    best_iou=iou
                    best_gt_idx=idx
            if best_iou>iou_threshold:
                #这里的detection[0]是amount_bboxes的一个key，表示图片的编号，best_gt_idx是该key对应的value中真实框的下标
                if amount_bboxes[detection[0]][best_gt_idx]==0:#只有没被占用的真实框才能用，0表示未被占用（占用：该真实框与某预测框匹配【两者IoU大于设定的IoU阈值】）
                    TP[detection_idx]=1#该预测框为TP
                    amount_bboxes[detection[0]][best_gt_idx]=1#将该真实框标记为已经用过了，不能再用于其他预测框。因为一个预测框最多只能对应一个真实框（最多：IoU小于IoU阈值时，预测框没有对应的真实框)
                else:
                    FP[detection_idx]=1#虽然该预测框与真实框中的一个框之间的IoU大于IoU阈值，但是这个真实框已经与其他预测框匹配，因此该预测框为FP
            else:
                FP[detection_idx]=1#该预测框与真实框中的每一个框之间的IoU都小于IoU阈值，因此该预测框直接为FP
                
        TP_cumsum=torch.cumsum(TP,dim=0)
        FP_cumsum=torch.cumsum(FP,dim=0)
        
        #套公式
        recalls=TP_cumsum/(total_true_bboxes+epsilon)
        precisions=torch.divide(TP_cumsum,(TP_cumsu+FP_cumsum+epsilon))
        
        #把[0,1]这个点加入其中
        precisions=torch.cat((torch.tensor([1]),precision))
        recalls=torch.cat((torch.tensor([0]),recalls))
        #使用trapz计算AP
        average_precisions.append(torch.trapz(precisions,recalls))
        
    return sum(average_precisions)/len(average_precisions) 


def insert_over_union(boxes_preds,boxes_labels):
    
    box1_x1=boxes_preds[...,0:1]
    box1_y1=boxes_preds[...,1:2]
    box1_x2=boxes_preds[...,2:3]
    box1_y2=boxes_preds[...,3:4]#shape:[N,1]
    
    box2_x1=boxes_labels[...,0:1]
    box2_y1=boxes_labels[...,1:2]
    box2_x2=boxes_labels[...,2:3]
    box2_y2=boxes_labels[...,3:4]
    
    x1=torch.max(box1_x1,box2_x1)
    y1=torch.max(box1_y1,box2_y1)
    x2=torch.min(box1_x2,box2_x2)
    y2=torch.min(box1_y2,box2_y2)
    
    
    #计算交集区域面积
    intersection=(x2-x1).clamp(0)*(y2-y1).clamp(0)
    
    box1_area=abs((box1_x2-box1_x1)*(box1_y1-box1_y2))
    box2_area=abs((box2_x2-box2_x1)*(box2_y1-box2_y2))
    
    return intersection/(box1_area+box2_area-intersection+1e-6)

参考：

• [1] https://www.youtube.com/watch?v=FppOzcDvaDI&list=PLhhyoLH6Ijfw0TpCTVTNk42NN08H6UvNq&index=4

• [2] https://jonathan-hui.medium.com/map-mean-average-precision-for-object-detection-45c121a31173