VOC格式转COCO格式

VOC和coco格式数据集结构如下图：

 

 其中最主要的是 将VOC的xml文件转换成json格式，train.json中的信息由VOC-main-train.txt中的所有索引文件组成，val和test类似。本demo生成的json如下图：

附代码xml2json.py：

"""
    本脚本用于将VOC数据集转换成coco格式 ***仅限目标检测***
    代码适用于一个类别的情况
    如有多类需进行微调
"""
import json
import os
import xml.etree.ElementTree as ET

CATEGORIES = [
    {
        "category": "1",
        "id": 0,
        "name": "fire"
    },
]


def voc2coco(xml_path, txt_path, json_path):
    """
    本函数从txt_path中的txt文件读取需要处理的文件，在xml_path中找到对应文件写入json
    :param xml_path: xml文件的目录路径
    :param txt_path: txt的绝对路径，主要存放需要处理的文件名索引
    :param json_path: json保存的路径
    :return:
    """
    xml_files = os.listdir(xml_path)
    files_process = []
    with open(txt_path, 'r') as ft:
        for line in ft:
            line = line.rstrip('\n')
            files_process.append(line)

    coco_output = {
        "images": [],
        "categories": CATEGORIES,
        "annotations": []
    }

    annotation_id = 0
    image_id = 0
    is_crowd = 0
    for xml_file in xml_files:
        xml_t = os.path.join(xml_path, xml_file)
        tree = ET.parse(xml_t)
        root = tree.getroot()

        filename = root.find('filename').text
        if filename.split('.')[0] in files_process:
            print(f"Processing {image_id} pic: {filename}...")
            width = root.find('size').find('width').text
            height = root.find('size').find('height').text

            image_info = {
                "id": image_id,
                "file_name": filename,
                "width": width,
                "height": height,
            }
            coco_output["images"].append(image_info)

            category_id = 0  # 注意这里对多个类别需小作修改
            for obj in root.findall('object'):
                bndbox = obj.find('bndbox')
                xmin = int(bndbox.find('xmin').text)
                xmax = int(bndbox.find('xmax').text)
                ymin = int(bndbox.find('ymin').text)
                ymax = int(bndbox.find('ymax').text)
                x = int(xmin)
                y = int(ymin)
                w = int(xmax-xmin)
                h = int(ymax-ymin)
                bounding_box = [x,
                                y,
                                w,
                                h]
                annotation_info = {
                    "id": annotation_id,
                    "image_id": image_id,
                    "category_id": category_id,
                    "iscrowd": is_crowd,
                    "bbox": bounding_box,
                }
                coco_output['annotations'].append(annotation_info)
                save_path = os.path.join(json_path, "test.json")
                with open(save_path, 'w') as output_json_file:
                    json.dump(coco_output, output_json_file, indent=4)
                annotation_id += 1
            image_id += 1


if __name__ == '__main__':
    xml_path = r"E:\python_workdir\dataset\VOC2007\Annotations"
    txt_path = r"E:\python_workdir\dataset\VOC2007\ImageSets\Main\test.txt"
    json_path = r"E:\python_workdir\dataset\tmp_test\tmp_pth2"

    voc2coco(xml_path=xml_path, txt_path=txt_path, json_path=json_path)

注：本demo适用于单类的情况，如涉及多类，需要修改CATEGORIES以及增加一个记录class信息的输入。