[MMDetection]VOC数据格式转为COCO数据格式

本文链接：https://blog.youkuaiyun.com/qq_41964545/article/details/131612205

以下脚本可以根据创建VOC格式数据集转换为COCO数据集

其中文件组织格式如下

VOC2007

------Annotations

------***********.xml

-------ImageSets

------train.txt

------test.txt

-------JPEGImages

------***********.jpg

COCO2017

------train

------annotations

------train.json

------images

------***********.jpg

-------test

------annotations

------test.json

------images

------***********.jpg

#!/usr/bin/python
# xml是voc的格式
# json是coco的格式
import sys, os, json, glob
import xml.etree.ElementTree as ET
import shutil
INITIAL_BBOXIds = 1
# PREDEF_CLASSE = {}
# 类别名
PREDEF_CLASSE = {'DT_SUBSIDENCE':0, 'PIP':1, 'NET':2, 'GROUND_LAMP':3, 'DT_CRACK':4, 'PIP_S':5, 'DT_GAP':6}



# function
def get(root, name):
    return root.findall(name)


def get_and_check(root, name, length):
    vars = root.findall(name)
    if len(vars) == 0:
        raise NotImplementedError('Can not find %s in %s.' % (name, root.tag))
    if length > 0 and len(vars) != length:
        raise NotImplementedError('The size of %s is supposed to be %d, but is %d.' % (name, length, len(vars)))
    if length == 1:
        vars = vars[0]
    return vars


def convert(xml_paths, out_json, mode = "train"):
    json_dict = {'images': [], 'type': 'instances',
                 'categories': [], 'annotations': []}
    categories = PREDEF_CLASSE
    bbox_id = INITIAL_BBOXIds
    for image_id, xml_f in enumerate(xml_paths):

        # 进度输出
        # sys.stdout.write('\r>> Converting image %d/%d' % (
        #     image_id + 1, len(xml_paths)))
        # sys.stdout.flush()

        tree = ET.parse(xml_f)
        root = tree.getroot()

        # filename = get_and_check(root, 'frame', 1).text
        # print(type(xml_f))
        filename = xml_f.split(r"E:\VOC2007\Annotations")[1]   # VOC的annotations的路径
        filename = filename.split("\\")[1]
        filename = filename.split(".xml")[0]+".jpg"
        print("########################",filename,xml_f)
        if mode == "test":
            src = os.path.join(r"E:\VOC2007\JPEGImages",filename) # VOC的JPEGImages的路径
            dst = os.path.join(r"E:\COCO2017\test\images",filename) # 需要存放的COCOimage路径
            shutil.copyfile(src, dst)
        if mode == "train":
            src = os.path.join(r"E:\VOC2007\JPEGImages",filename)
            dst = os.path.join(r"E:COCO2017\train\images",filename)
            shutil.copyfile(src, dst)

        size = get_and_check(root, 'size', 1)
        width = int(get_and_check(size, 'width', 1).text)
        height = int(get_and_check(size, 'height', 1).text)
        image = {'file_name': filename, 'height': height,
                 'width': width, 'id': image_id + 1}
        json_dict['images'].append(image)
        ## Cruuently we do not support segmentation
        # segmented = get_and_check(root, 'segmented', 1).text
        # assert segmented == '0'

        for obj in get(root, 'object'):
            category = get_and_check(obj, 'name', 1).text
            if category not in categories:
                new_id = max(categories.values()) + 1
                categories[category] = new_id
            category_id = categories[category]
            bbox = get_and_check(obj, 'bndbox', 1)
            xmin = int(get_and_check(bbox, 'xmin', 1).text) - 1
            ymin = int(get_and_check(bbox, 'ymin', 1).text) - 1
            xmax = int(get_and_check(bbox, 'xmax', 1).text)
            ymax = int(get_and_check(bbox, 'ymax', 1).text)
            if xmax <= xmin or ymax <= ymin:
                continue
            o_width = abs(xmax - xmin)
            o_height = abs(ymax - ymin)
            ann = {'area': o_width * o_height, 'iscrowd': 0, 'image_id': image_id + 1,
                   'bbox': [xmin, ymin, o_width, o_height], 'category_id': category_id,
                   'id': bbox_id, 'ignore': 0, 'segmentation': []}
            json_dict['annotations'].append(ann)
            bbox_id = bbox_id + 1

    for cate, cid in categories.items():
        cat = {'supercategory': 'none', 'id': cid, 'name': cate}
        json_dict['categories'].append(cat)

    # json_file = open(out_json, 'w')
    # json_str = json.dumps(json_dict)
    # json_file.write(json_str)
    # json_file.close() # 快
    json.dump(json_dict, open(out_json, 'w'), indent=4)  # indent=4 更加美观显示 慢


if __name__ == '__main__':
    # 1、读取VOC2007\ImagesSets\Main中的文件 得到测试集和训练集的图片名
    test_filenames = []
    train_filenames = []
    with open(r"E:\VOC2007\ImageSets\Main\test.txt") as f:
        for filename in f.readlines():
            test_filenames.append(filename.split('\n')[0])  # 去掉换行符 \n 
    with open(r"E:\VOC2007\ImageSets\Main\train.txt") as f:
        for filename in f.readlines():
            train_filenames.append(filename.split('\n')[0])  
    print(test_filenames)
    print(train_filenames)

    # 2、分别拼凑出 训练 和 测试 所需要的xml 路径
    xml_path = r'E:\VOC2007\Annotations'  # 改一下读取xml文件位置
    train_xml_files = []
    test_xml_files = []
    for train_filename in train_filenames:
        train_xml_file = os.path.join(xml_path,train_filename) + '.xml'   # 拼接路径 加上后缀
        # print(train_xml_file)
        train_xml_files.append(train_xml_file)

    for test_filename in test_filenames:
        test_xml_file = os.path.join(xml_path,test_filename) + '.xml'
        # print(test_xml_file)
        test_xml_files.append(test_xml_file)

    print(test_xml_files)
    print(train_xml_files)

    # 3、转换数据
    convert(test_xml_files, r'E:\COCO2017\test\annotations\test.json',mode="test")  # 这里是生成的json保存位置，改一下
    convert(train_xml_files, r'E:\COCO2017\train\annotations\train.json',mode="train")  # 这里是生成的json保存位置，改一下