如何在COCOEval中计算指定关键点的mAP指标

pycocotools默认基于全部17个COCO人体关键点计算OKS指标下的mAP，要计算自定义指定关键点集合的mAP，只需要在评估前过滤不需要参与计算的关键点即可，具体修改逻辑如下：

• 先建立关键点名称到COCO默认索引的映射，提取需要计算的关键点对应的索引
• 评估前将非目标关键点标记为「未标注」状态，让这部分点不参与OKS匹配和指标统计
• 按原有流程执行评估、累加、统计即可，最终输出的就是指定关键点集合的mAP

修改后可直接运行的代码 #

from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
import copy

# 初始化COCO 基础API
cocoGt = COCO('annotations/person_keypoints_val2017.json')
cocoDt = cocoGt.loadRes('detections/results.json')

cat_ids = cocoGt.getCatIds(catNms=['person'])
imgIds = cocoGt.getImgIds(catIds=cat_ids)

# ========== 自定义配置：指定要计算的关键点 ==========
target_kpt_names = ['nose', 'left_hip', 'right_hip']
# 从类别元信息读取COCO官方关键点顺序，建立名称到索引的映射
person_kpt_list = cocoGt.cats[cat_ids[0]]['keypoints']
target_kpt_idx = [person_kpt_list.index(name) for name in target_kpt_names]
# ==============================================

cocoEval = COCOeval(cocoGt, cocoDt, 'keypoints')
cocoEval.params.imgIds = imgIds

# ========== 过滤不需要参与评估的关键点 ==========
# 遍历所有GT标注，将非目标关键点的visibility设为0（标记为未标注，评估时自动跳过）
for gt_id, gt_ann in cocoEval.cocoGt.anns.items():
    # 深拷贝关键点数组，避免修改原始加载的标注内容
    gt_ann['keypoints'] = copy.deepcopy(gt_ann['keypoints'])
    kpt_arr = gt_ann['keypoints']
    for idx in range(len(person_kpt_list)):
        if idx not in target_kpt_idx:
            # COCO关键点格式为(x,y,v)三元组，v=0代表点未标注，不参与评估
            kpt_arr[idx*3 + 2] = 0
# ======================================================

cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()

# 输出的stats[0]就是指定关键点集合在OKS=0.5:0.95阈值下的mAP
print(cocoEval.stats[0])

注意事项 #

• COCO关键点存储为长度3*N的一维数组，每3位对应一个关键点的(x坐标, y坐标, 可见性标记)，可见性为0时评估器会直接跳过该点，不参与任何匹配和指标计算
• 代码中用copy.deepcopy拷贝关键点数组，不会修改原始标注加载的内容，不会影响后续其他全量指标计算
• 如果需要计算单个关键点的AP，只需要把target_kpt_names设为对应单个关键点的名称即可

内容的提问来源于stack exchange，提问作者Armand Assault