在自己的数据集上测试coco评价指标——以Mar20为例

参考：
1.在自己的数据集上调用cocoapi计算map
2. COCO Result Format
3.COCO result json
之前的模型都是在COCO数据集上训练，数据集的标注以及结果的生成格式都是按照官方的格式组织的，调用cocoapi和官方下载的instance_val2017.json计算就可以了。
现在需要在其他数据集上测试map等指标，这些图片都是标注好的，但是格式和coco要求不一样，因此需要进行转换。
分为四个步骤：1. 数据集划分和标签转换；2.将标注转为coco的result格式；3. 将模型推理结果保存为result格式；4.调用cocoapi计算；

我要测试的数据集为Mar20，数据的标注格式为未归一化的（xmin, xmax, ymin, ymax），COCO的标注格式为未归一化的（xmin, ymin, width, height）。数据集的组织形式如下：

注意：这里测试的MAR20数据集类别为20种飞机类，测试过程中我将这20类全部映射为了COCO的飞机类别。如果需要测试其他数据集，在标签转换过程中需要注意cls_id这个属性。

[‘A1’,‘A2’,‘A3’,‘A4’,‘A5’,‘A6’,‘A7’,‘A8’,‘A9’,‘A10’,‘A11’,‘A12’,‘A13’,‘A14’,‘A15’,‘A16’,‘A17’,‘A18’,‘A19’,‘A20’]

一、数据集划分和标签转换

1.xml标签转为txt

首先将xml标签转化为txt。注意不同的数据集修改数据集类别，convert函数，convert_annotation函数里的cls_id，以及数据的路径。转换后的标签保存在MAR20/coco_Labels目录下。

import xml.etree.ElementTree as ET
import os
import cv2
import random
random.seed(0)

# 数据集类别
classes = ['A1','A2','A3','A4','A5','A6','A7','A8','A9','A10','A11','A12','A13','A14','A15','A16','A17','A18','A19','A20' ]
 
def convert(box):
    # 修改 box : xmin, xmax, ymin, ymax -- xmin, ymin, w, h
    y= box[2]
    x= box[0]
    w = box[1] - box[0]
    h = box[3] - box[2]
    return (int(x), int(y), int(w), int(h))
 
#  修改 数据集地址
dataset_path = './datasets/MAR20'

def convert_annotation(image_id):
    in_file = open(os.path.join(dataset_path, f'Annotations/Horizontal Bounding Boxes/{image_id}.xml'))  # 修改 xml所在路径
    img_file = cv2.imread(os.path.join(dataset_path, f'JPEGImages/{image_id}.jpg'))  # 修改 图片所在路径
    out_file = open(os.path.join(dataset_path, f'coco_Labels/{image_id}.txt' ),'w+')  # 修改 转换后的txt保存路径
    tree = ET.parse(in_file)
    root = tree.getroot()

    assert img_file is not None
    size = img_file.shape[0:-1]
    h = int(size[0])
    w = int(size[1])
    for obj in root.iter('object'):
        cls = obj.find('name').text
        if cls not in classes :
            continue
        # cls_id = classes.index(cls)
        cls_id = 4  # 修改 Mar20是飞机目标识别，细分为10类，这里将飞机目标统一为COCO的飞机目标类别，即4
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        ZIP_ONE = convert(b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in ZIP_ONE]) + '\n')
 
wd = getcwd()

coco_Labels_out = os.path.join(dataset_path, 'coco_Labels')  # 修改 保存图片绝对路径的txt文件的路径

if not os.path.exists(coco_Labels_out):
    os.makedirs(coco_Labels_out)

images = os.listdir(os.path.join(dataset_path, 'JPEGImages'))  # 修改 图片所在文件夹
files = [file for file in images if file.endswith('.jpg')]
image_ids = [file.split('.')[0] for file in files]

for image_id in image_ids:
    try:
        print(image_id)        
        convert_annotation(image_id)
    except:
        print('error img:', image_id)

运行以上代码后会在coco_Labels文件夹下生成以下文本：

2.划分数据集

然后划分数据集的图片和标签，注意修改划分的比例，输入和输出的地址。划分后的数据保存在MAR20/split目录下。

import os
import random
from shutil import copyfile
random.seed(0)

def split_dataset(input_images_dir, input_labels_dir, output_dir, split_ratio=(0.7, 0.05, 0.25)):
    # 创建输出目录结构
    os.makedirs(output_dir, exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'images', 'train'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'images', 'val'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'images', 'test'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'labels', 'train'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'labels', 'val'), exist_ok=True)
    os.makedirs(os.path.join(output_dir, 'labels', 'test'), exist_ok=True)

    # 获取所有图片文件
    image_files = [f for f in os.listdir(input_images_dir) if f.endswith('.jpg')]
    num_images = len(image_files)

    # 随机打乱图片顺序
    random.shuffle(image_files)

    # 计算划分的数量
    num_train = int(num_images * split_ratio[0])
    num_val = int(num_images * split_ratio[1])
    num_test = num_images - num_train - num_val

    # 分割图片和标签文件
    for i, image_file in enumerate(image_files):
        if i < num_train:
            set_name = 'train'
        elif i < num_train + num_val:
            set_name = 'val'
        else:
            set_name = 'test'

        # 复制图片文件
        copyfile(os.path.join(input_images_dir, image_file), os.path.join(output_dir, 'images', set_name, image_file))

        # 构建对应的标签文件名
        label_file = os.path.splitext(image_file)[0] + '.txt'

        # 复制标签文件
        copyfile(os.path.join(input_labels_dir, label_file), os.path.join(output_dir, 'labels', set_name, label_file))

# 修改 数据集地址
dataset_path = './datasets/MAR20'

# 修改输出地址
output_dir = os.path.join(dataset_path, 'split')
os.makedirs(output_dir, exist_ok=True)

# 修改输入图片和标签地址
input_images_dir = os.path.join(dataset_path, 'JPEGImages')
input_labels_dir = os.path.join(dataset_path,'coco_Labels')

split_ratio=(0.7, 0.05, 0.25)
# 调用划分函数 划分比例为70%训练集，5%验证集，25%测试集
split_dataset(input_images_dir, input_labels_dir, output_dir, split_ratio)

划分好后，在MAR20/split文件夹下生成以下文件：

二、将标注转为coco的result格式

首先将test数据集的图片路径保存到test.txt文件中：

import xml.etree.ElementTree as ET
import os

# test图片路径
test_path = './datasets/MAR20/split/images/test'
# 保存txt路径
saved_txt_path = './datasets/MAR20/test.txt'

for img in os.listdir(test_path):
    img_path = os.path.join(test_path, img)
    with open(saved_txt_path, 'a') as f:
        f.write(img_path + '\n')

MAR20/test.txt文件内容如下：

然后将MAR20/labels/test文件夹下的标注转换为coco格式，输出为annotations.json:

import json
import cv2
import os


if __name__=='__main__':
    cats = list()
    # 输出的json文件路径
    out_path = 'annotations.json'
    # test.txt路径
    test_path = './datasets/MAR20/test.txt'

    with open('obj.names', 'r') as f:
        for line in f.readlines():
            line = line.strip('\n')
            cats.append(line)
    cat_info = []
    for i, cat in enumerate(cats):
        cat_info.append({'name': cat, 'id': i})

    ret = {'images': [], 'annotations': [], "categories": cat_info}
    i = 0
    for line in open(test_path, 'r'):
        line = line.strip('\n')
        i += 1
        image_id = eval(os.path.basename(line).split('.')[0])
        image_info = {'file_name': '{}'.format(line), 'id': image_id}
        ret['images'].append(image_info)
        anno_path = line.replace('.jpg', '.txt')
        anno_path = anno_path.replace('images', 'labels')
        anns = open(anno_path, 'r')
        img = cv2.imread(line)
        height, width = img.shape[0], img.shape[1]
        for ann_id, txt in enumerate(anns):
            tmp = txt[:-1].split(' ')
            cat_id = tmp[0]
            bbox = [float(x) for x in tmp[1:]]  # 注意box格式，已经提前转换成coco格式了
            area = round(bbox[2] * bbox[3], 2)
            # coco annotation format
            ann = {'image_id': image_id,
                'id': int(len(ret['annotations']) + 1),
                'category_id': int(cat_id),
                'bbox': bbox,
                'iscrowd': 0,
                'area': area}
            ret['annotations'].append(ann)
    
    json.dump(ret, open(out_path, 'w'))

以上转换需要用到的coco标签和id对应关系如下，文件名为obj.names，复制以下内容保存到obj.names中:

0: person
1: bicycle
2: car
3: motorcycle
4: airplane
5: bus
6: train
7: truck
8: boat
9: traffic light
10: fire hydrant
11: stop sign
12: parking meter
13: bench
14: bird
15: cat
16: dog
17: horse
18: sheep
19: cow
20: elephant
21: bear
22: zebra
23: giraffe
24: backpack
25: umbrella
26: handbag
27: tie
28: suitcase
29: frisbee
30: skis
31: snowboard
32: sports ball
33: kite
34: baseball bat
35: baseball glove
36: skateboard
37: surfboard
38: tennis racket
39: bottle
40: wine glass
41: cup
42: fork
43: knife
44: spoon
45: bowl
46: banana
47: apple
48: sandwich
49: orange
50: broccoli
51: carrot
52: hot dog
53: pizza
54: donut
55: cake
56: chair
57: couch
58: potted plant
59: bed
60: dining table
61: toilet
62: tv
63: laptop
64: mouse
65: remote
66: keyboard
67: cell phone
68: microwave
69: oven
70: toaster
71: sink
72: refrigerator
73: book
74: clock
75: vase
76: scissors
77: teddy bear
78: hair drier
79: toothbrush

三、将推理结果转换为coco格式

推理的时候将单帧结果保存在items，所有的推理结果保存在result，然后将result保存到results.txt文件中。
保存的格式可以参考https://cocodataset.org/#format-results 和 https://github.com/cocodataset/cocoapi/tree/master/results

然后手动将results.txt后缀改为.json即可（保存为json总是报错，麻了）。

#  items为每一帧的检测结果
    for i in range(len(classes)):
        items.append({"image_id": eval(image_name),"category_id":classes[i],"bbox":boxes[i].tolist(), "score":1.0})
    # 检测结果为空也要保存，否则会导致后续的评估出错
    if len(items)==0:
        items.append({"image_id": eval(image_name),"category_id":0,"bbox":[0,0,0,0], "score":0})

# 以上代码保存了单帧检测结果，result保存了所有的结果
result = []
# ...
result.extend(items)

json_file_path = 'results.txt'
# 字典键值会自动变为单引号，json格式必须为双引号，所以需要用json.dumps()函数转换字符
json_str = json.dumps(result, ensure_ascii=False, default=default_dump) 
with open(json_file_path, 'w') as file:
     file.write(str(json_str))

四、调用cocoapi计算coco指标

直接调用接口即可计算coco指标：

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

def main():
    results_file ='result.json'
    annotations = 'annotations.json'
   
    cocoGt = COCO(annotations)
    cocoDt = cocoGt.loadRes(results_file)

    cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')

    cocoEval.params.catIds = [4] # 你可以根据需要增减类别
    cocoEval.evaluate()
    cocoEval.accumulate()
    cocoEval.summarize()

if __name__ == '__main__':
    main()

五、YOLO系列调用cocoapi

根据前面一、二步骤划分好数据集，转换好annotations.json，可以直接运行以下.py文件获得coco指标：

import os
import json

from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
from ultralytics import YOLO

def generate_results(yolo, imgs_dir, jpgs, results_file):
    """Run detection on each jpg and write results to file."""
    results = []
    for jpg in jpgs:
        img_path = os.path.join(imgs_dir, jpg)
        image_id = int(jpg.split('.')[0])
        det = yolo.predict(img_path, conf=0.25,save=True)
        boxes = det[0].boxes
        for i in range(len(boxes)):
            box = boxes[i]
            # 注意ultralytics中的xywh坐标中xy是中心点坐标，coco中的xy是左上角坐标
            x_c, y_c, w, h = box.xywh.tolist()[0]   
            
            x_min = x_c - w / 2
            y_min = y_c - h / 2
            conf = box.conf.tolist()[0]
            cls = int(box.cls.tolist()[0])
            results.append({'image_id': image_id,
                            'category_id': cls,
                            'bbox': [x_min, y_min, w, h],
                            'score': float(conf)})
    with open(results_file, 'w') as f:
        f.write(json.dumps(results, indent=4))
        
def main():
    results_file ='result.json'  # yolo推理结果保存文件
    imgs_dir = './datasets/MAR20/split/images/test'  # 测试集图片路径
    annotations = 'annotations.json'  # gt标注文件
    model=YOLO('yolov8l.yaml').load("/home/jingjia/sdb/liaocheng/ultralytics-main/yolov8l.pt")

    jpgs = [j for j in os.listdir(imgs_dir) if j.endswith('.jpg')]

    generate_results(model, imgs_dir, jpgs, results_file)
    

    # Run COCO mAP evaluation
    # Reference: https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb
    cocoGt = COCO(annotations)
    cocoDt = cocoGt.loadRes(results_file)

    cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')

    cocoEval.params.catIds = [4] # 你可以根据需要增减类别
    cocoEval.evaluate()
    cocoEval.accumulate()
    cocoEval.summarize()

if __name__ == '__main__':
    main()

运行结果：