史上最全AP/mAP通用代码实现(yolov5 txt版本)-下_yolov5的map

提示：通用map指标框架代码介绍，直接使用yolov5数据格式，实现论文map指标计算代码解读

---

前言

“史上最全AP、mAP详解与代码实现”文章（这里）已经介绍了map相关原理，且给出相应简单代码实现AP方法。然将AP计算融入模型求解AP结果，可能是一个较为复杂的工程量。我也在这里博客分享基于模型构造一个即插即用计算map的相关模块代码，使用者只需复制我的模块，即可使用。然，之前方法是需要使用xml格式，依然对yolo模型不甚友好，我再此修改为txt方式，使用yolo本身数据实现small、medium、large等相关AP或AP0.75等结果预测。本文将直接介绍计算map核心代码简单列子，在此基础上介绍整个即插即用map计算模块使用方法与代码解读。

该版本是直接使用yolo数据格式实现map计算，集成txt转json格式内容。

---

一、map模块整体认识

本文就是一个detection_map即插即用计算map指标模块，可计算任何模型map指标，但有效计算需要稍微修改部分代码，我后面将介绍。基于此，我将整理一份yolo模型的通用map框架代码。那么，本文将介绍2个内容，其一为简单计算map的一个列子，其原理可参考这里博客；其二为基于yolo模型介绍通用map模块计算方法map_yolo。其整体架构如下图：

注：我使用yolov5-6.1模型，仅将detection_map放入该位置，即可使用。

二、map计算应用代码解读

实现mAP计算，我们需要有已知真实标签与模型预测标签，按照pcocotools的格式生成真实标签与预测标签的json格式，即可实现map指标计算。

from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
if __name__ == "__main__":
    cocoGt = COCO('coco_format.json')        #标注文件的路径及文件名，json文件形式
    cocoDt = cocoGt.loadRes('predect_format.json')  #自己的生成的结果的路径及文件名，json文件形式

    cocoEval = COCOeval(cocoGt, cocoDt, "bbox")
    cocoEval.evaluate()
    cocoEval.accumulate()
    cocoEval.summarize()

1
2
3
4
5
6
7
8
9
10
11

介于我在这篇文章这里已有详细介绍，我将不在介绍。

三、通用map计算指标代码解读

介于我在这篇文章这里已有详细介绍，我将不在介绍，文章参考内容如下图：

四、基于yolov5使用通用map计算指标代码解读

这一部分也是本文最重要一部分，实际有关map原理内容或整体模块实现已在我推荐文章中，但推荐文章缺点是没有放置相应代码内容。而该部分就是直接给出基于yolov5模型调用map通用模块实现的相关代码或工程。

1、通用map指标计算模块整体结构说明

构建初始化模型，配置相关参数，直接使用computer_main函数集成，进行推理与map指标计算(整体如下图)。

2、参数构建

我构建模型相关参数，如数据文件夹、权重及推理相关参数，特别是conf阈值与iou阈值需要关注，在yolov5的val.py指标计算设置conf阈值=0.001、iou阈值=0.6，这个根据自己情况而定。

def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--source', type=str, default= r'E:\project\data\voc_data\voc2007_data\images\test', help='dataset.yaml path')
    parser.add_argument('--weights', nargs='+', type=str,
                        default=r'E:\project\project_distilation\experiment\runs\train\yolo_x2s_iou-0.45_conf-0.85/weights/best.pt',
                        help='model.pt path(s)')
    parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--conf_thres', type=float, default=0.001, help='confidence threshold，default=0.001')
    parser.add_argument('--iou_thres', type=float, default=0.6, help='NMS IoU threshold，default=0.6')
    parser.add_argument('--imgsz', '--img', '--img_size', type=int, default=640, help='inference size (pixels)')
    parser.add_argument('--save_dir',  default='runs/val_map/exp', help='图像保存路径')
    parser.add_argument('--save_img', default=False, help='保存框图像查看')
    opt = parser.parse_args()

    return opt
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

3、数据准备

很简单，只要是yolov5格式数据即可，代码会自动将txt内容转为满足pycocotools计算map指标json的json格式。

4、模型初始化

yolov5模型初始化较为简单，直接使用yolov5自带的attempt_load方式初始化模型即可，如下代码：

def init_model(weights):

    model = attempt_load(weights, map_location=device)
    model = model.eval()
    return model
1
2
3
4
5
6

5、map指标计算函数(computer_main)代码解读

我大概描述该函数内容(按步骤说明)：

①、获得图像相关路径及指标计算函数初始化

C = Computer_map()
img_root_lst = C.get_img_root_lst(opt.source)  # 获得图片绝对路径与图片产生image_id映射关系
1
2

②、获得类别

categories = model.names  
C.get_categories(categories)
1
2

③、生成gt的json文件

这个更简单了，和中篇文章不一样就体现在这里，是使用yolo本身txt格式转json，其代码如下：

C.yolov5txt2cocojson(img_root_lst,out_dir=None,save_img=False)
1

④、图像预处理

for img_path in img_root_lst:
   img0 = cv2.imread(img_path)
   img = letterbox(img0, img_size, stride=stride, auto=True)[0]
   img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
   im = np.ascontiguousarray(img)
   im = torch.from_numpy(im).to(device)
   im = im.float()  # uint8 to fp16/32
   im /= 255  # 0 - 255 to 0.0 - 1.0
   if len(im.shape) == 3:
       im = im[None]  # expand for batch dim
1
2
3
4
5
6
7
8
9
10

这里图像预处理是调用yolov5的letterbox函数。

⑤、模型推理与后处理

pred = model(im)[0]  
result = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=None,multi_label=True)
det = result[0]
1
2
3

这里仍然是调用yolov5模型与非极大值后处理函数。

⑥、输出尺寸恢复

if len(det)>0:
    det[:, :4] = scale_coords(im.shape[2:], det[:, :4], img0.shape).round()
1
2

这里也是调用scale_coords函数恢复预测的box到原图尺寸对应box。

⑦、生成预测json格式文件

det = det.cpu().numpy() if det.is_cuda else det.numpy()  # 处理为cuda上的数据或cpu转numpy格式
det = [[d[0],d[1],d[2],d[3],d[4], categories[int(d[5])] ] for d in det] # 给定名称name标签
# det 格式为列表[x1,y1,x2,y2,score,label],若无结果为空
img_name = C.get_strfile(img_path)
C.detect2json(det, img_name)
1
2
3
4
5

这里需要循环推理每个图像预测结果，生成对应满足pycocotools预测json文件格式内容。

当然，我做了是否保存预测图像模块，如果需要使用大致查看预测内容，建议conf与iou阈值试单调整，否则满图都是框。

 if opt.save_img:
     img=draw_bbox(img0,det)
     cv2.imwrite(os.path.join(opt.save_dir,img_name),img)
1
2
3

⑧、map指标计算

循环推理完所有图片，也意味预测json保存完毕，就直接使用gt与pred文件json，调用我集成好的函数，即可实现map指标计算，如下：

C.computer_map()  # 计算map
1

computer_main代码

这边我给出了类似yolov5保存最佳权重判断指标，具体整个代码如下：

def computer_main(opt, model):
    '''
    data_root:任何文件夹，但必须保证每个图片与对应xml必须放在同一个文件夹中
    model:模型，用于预测
    '''

    stride=32
    img_size=[opt.imgsz, opt.imgsz]

    C = Computer_map()
    img_root_lst = C.get_img_root_lst(opt.source)  # 获得图片绝对路径与图片产生image_id映射关系

    # 在self.coco_json中保存categories，便于产生coco_json和predetect_json
    categories = model.names  # 可以给txt路径读取，或直接给列表  #*********************得到classes,需要更改的地方***********##
    C.get_categories(categories)

    C.yolov5txt2cocojson(img_root_lst,out_dir=None,save_img=False)
    # 产生coco_json格式
    # xml_root_lst = [name[:-3] + 'xml' for name in img_root_lst]
    # for xml_root in xml_root_lst: C.xml2cocojson(xml_root)  # 产生coco json 并保存到self.coco_json中





    if opt.save_img:build_dir(opt.save_dir)
    # 产生预测的json
    for img_path in img_root_lst:
        img0 = cv2.imread(img_path)
        img = letterbox(img0, img_size, stride=stride, auto=True)[0]
        img = img.transpose((2, 0, 1))[::-1]  # HWC to CHW, BGR to RGB
        im = np.ascontiguousarray(img)
        print("图片原始尺寸：{}\t模型预测尺寸：{}".format(img0.shape,im.shape))

        im = torch.from_numpy(im).to(device)
        im = im.float()  # uint8 to fp16/32
        im /= 255  # 0 - 255 to 0.0 - 1.0
        if len(im.shape) == 3:
            im = im[None]  # expand for batch dim

        pred = model(im)[0]  ####**********************需要更改的地方***********************####

        result = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=None, multi_label=True)
        det = result[0]
        # result, classes = parse_result['result'], parse_result['classes']
        if len(det)>0:
            det[:, :4] = scale_coords(im.shape[2:], det[:, :4], img0.shape).round()
        det = det.cpu().numpy() if det.is_cuda else det.numpy()  # 处理为cuda上的数据或cpu转numpy格式
        det = [[d[0],d[1],d[2],d[3],d[4], categories[int(d[5])] ] for d in det] # 给定名称name标签
        # det 格式为列表[x1,y1,x2,y2,score,label],若无结果为空
        img_name = C.get_strfile(img_path)
        C.detect2json(det, img_name)

        if opt.save_img:
            img=draw_bbox(img0,det)
            cv2.imwrite(os.path.join(opt.save_dir,img_name),img)
    map_value = C.computer_map()  # 计算map,返回 [mAP@0.5:0.95， mAP@0.5, mAP@0.75, ... ]
    yolo_best = 0.9*map_value[0]+0.1*map_value[1]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

6、基于yolov5的map指标计算代码链接

链接：https://pan.baidu.com/s/1z5tYuYsBtoD3zvipgDKaEA
提取码：map2
其结果如下：

总结

本文核心是介绍自己构建的map通用框架代码，为介绍该框架，我借助yolov5模型作为基准，一步步阐述如何使用map通用框架指标计算。