COCOAPI计算mAP

Posted 2022-10-19 洪流之源

cal_mAP.py 

import re
import os
import json
import cv2
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
from collections import OrderedDict
import argparse
from pathlib import Path


class COCOResults(object):
    METRICS = 
        "bbox": ["AP", "AP50", "AP75", "APs", "APm", "APl"],
        "segm": ["AP", "AP50", "AP75", "APs", "APm", "APl"],
        "box_proposal": [
            "AR@100",
            "ARs@100",
            "ARm@100",
            "ARl@100",
            "AR@1000",
            "ARs@1000",
            "ARm@1000",
            "ARl@1000",
        ],
        "keypoints": ["AP", "AP50", "AP75", "APm", "APl"],
    

    def __init__(self, *iou_types):
        allowed_types = ("box_proposal", "bbox", "segm", "keypoints")
        assert all(iou_type in allowed_types for iou_type in iou_types)
        results = OrderedDict()
        for iou_type in iou_types:
            results[iou_type] = OrderedDict(
                [(metric, -1) for metric in COCOResults.METRICS[iou_type]]
            )
        self.results = results

    def update(self, coco_eval):
        if coco_eval is None:
            return
        from pycocotools.cocoeval import COCOeval

        assert isinstance(coco_eval, COCOeval)
        s = coco_eval.stats
        iou_type = coco_eval.params.iouType
        res = self.results[iou_type]
        metrics = COCOResults.METRICS[iou_type]
        for idx, metric in enumerate(metrics):
            res[metric] = s[idx]

    def __repr__(self):
        results = '\\n'
        for task, metrics in self.results.items():
            results += 'Task: \\n'.format(task)
            metric_names = metrics.keys()
            metric_vals = [':.4f'.format(v) for v in metrics.values()]
            results += (', '.join(metric_names) + '\\n')
            results += (', '.join(metric_vals) + '\\n')
        return results


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--log')
    parser.add_argument('--anno')
    parser.add_argument('--image-dir')
    args = parser.parse_args()
    coco = COCO(args.anno)
    results = COCOResults('bbox')
    coco_dt = coco.loadRes(args.log)
    coco_eval = COCOeval(coco, coco_dt, 'bbox')
    imagelist = os.listdir(args.image_dir)
    coco_eval.params.imgIds = [int(Path(x).stem) for x in imagelist if x.endswith('jpg')]
    coco_eval.evaluate()
    coco_eval.accumulate()
    coco_eval.summarize()
    results.update(coco_eval)
    print(results)
    

 执行：

python cal_mAP.py --anno ./instances_val2017.json --log ./result.json --image-dir ../../coco2017val/images/val2017

上述命令中--anno是coco格式的标注文件， --image-dir是测试集图片路径，--log是预测的结果文件为json格式，可参考如下伪代码生成：

preds = []
for img in js['images']:
    img_p = img['file_name']
    if not os.path.isfile('/'.join((imgdir, img_p))):
       continue
    
    pred = inference(img_p) # 推理接口, 获取预测结果
    coco91class = coco80_to_coco91_class() # 获取coco类别字典
    for pp in pred:
      if pp is None:
        continue
      for p in pp:
        bbox = p[:4]
        prob = float(p[4])
        clse = int(p[5])
        preds.append(dict(
          image_id=img['id'], # 图片id
          category_id=coco91class[clse], # 类别id
          bbox=[float(bbox[0]), float(bbox[1]), float(bbox[2]-bbox[0]), float(bbox[3]-bbox[1])], # 检测框
          score=prob)) # 预测得分

  with open('result.json', 'w') as f: # 保存为json文件
      json.dump(preds, f)

result.json示例格式如下：

[
    "image_id": 397133, "category_id": 1, "bbox": [382.21466064453125, 61.73094177246094, 123.6776123046875, 290.18048095703125], "score": 0.8199472427368164, 
    "image_id": 397133, "category_id": 1, "bbox": [0.8529815673828125, 262.3185729980469, 62.58330535888672, 47.36285400390625], "score": 0.5744290351867676, 
    "image_id": 397134, "category_id": 51, "bbox": [60.40350341796875, 286.3122863769531, 81.48675537109375, 41.669189453125], "score": 0.28584015369415283, 
    "image_id": 397134, "category_id": 64, "bbox": [-3.7143173217773438, -2.5591201782226562, 70.6357421875, 155.88150024414062], "score": 0.20149467885494232, 
    "image_id": 397135, "category_id": 51, "bbox": [28.550830841064453, 344.74737548828125, 70.10545349121094, 40.75140380859375], "score": 0.11809217184782028
]

mAP计算结果如下：