MSCOCO数据集转VOC数据集格式

背景

VOC数据集类别为20类, 在很多比赛中经常应用VOC数据集，一方面数据集质量比较好，一方面数据集标注比较好。缺点是数据量上不如MSCOCO数据集，这里介绍一种方法，利用python把MSCOCO数据集转换成为VOC数据集格式，分为3步，每一步为一个py文件。

step1

import json

'''
    挑选出想要的类别
'''

className = {
    1: 'person', 
    63: 'tv',
    64: 'laptop',
    78: 'human face',
    79: 'hair drier',
    80: 'license plate',
}

# coco数据集标签的路径
coco_json_path = 'D:/FLIR/1/FLIR_ADAS_1_3.tar/FLIR_ADAS_1_3/train/thermal_annotations.json'
# 保存为新的标签（只含有想要的类别）
save_path = "COCO_train.json"

classNum = [1, 63, 64, 78, 79, 80]

def writeNum(Num):
    with open(save_path, "a+") as f:
        f.write(str(Num))

inputfile = []
inner = {}
# 向test.json文件写入内容
with open(coco_json_path, "r+") as f:
    allData = json.load(f)
    data = allData["annotations"]
    print(data[1])
    print("read ready")

for i in data:
    if i['category_id'] in classNum:
        inner = {
            "filename": str(int(i["image_id"]) + 1).zfill(6),  # 文件名长度为6 比如 001265
            "name": className[i["category_id"]],
            "bndbox": i["bbox"]
        }
        inputfile.append(inner)
inputfile = json.dumps(inputfile)
writeNum(inputfile)

step2

import json
import os

nameStr = []

"""
    推荐备份一下COCO数据集，方便以后使用，我这里备份了COCO数据集，
    这个py文件的作用就是把COCO数据中，不想要的图片删掉，只留想要的图片。
    （这里想要的图片是指step1中选的类别图片）
"""
# COCO_train.json 刚才保存的想要的类别
with open("COCO_train.json", "r+") as f:
    data = json.load(f)
    print("read ready")

for i in data:
    imgName = "FLIR_" + str(i["filename"]) + ".jpg"
    nameStr.append(imgName)

nameStr = set(nameStr)

# 你们保存的MSCOCO数据集的路径（图片文件路径）
path = "D:/FLIR/1/FLIR_ADAS_1_3.tar/FLIR_ADAS_1_3/train/train_jpg_file/"

for file in os.listdir(path):
    if (file not in nameStr): os.remove(path + file)

step3

import xml.dom
import xml.dom.minidom
import os
# from PIL import Image
import cv2
import json

'''
    生成xml 文件
'''

_ANNOTATION_SAVE_PATH = 'Annotations'  # 生成的xml保存的路径
_INDENT = ' ' * 4
_NEW_LINE = '\n'
_FOLDER_NODE = 'COCO2014'
_ROOT_NODE = 'annotation'
_DATABASE_NAME = 'LOGODection'
_ANNOTATION = 'COCO2014'
_AUTHOR = 'MEISI'
_SEGMENTED = '0'
_DIFFICULT = '0'
_TRUNCATED = '0'
_POSE = 'Unspecified'

def createElementNode(doc, tag, attr): 
    element_node = doc.createElement(tag)
    text_node = doc.createTextNode(attr)
    element_node.appendChild(text_node)
    return element_node


def createChildNode(doc, tag, attr, parent_node):
    child_node = createElementNode(doc, tag, attr)
    parent_node.appendChild(child_node)


def createObjectNode(doc, attrs):
    object_node = doc.createElement('object')
    createChildNode(doc, 'name', attrs['name'], object_node)
    createChildNode(doc, 'pose', _POSE, object_node)
    createChildNode(doc, 'truncated',  _TRUNCATED, object_node)
    createChildNode(doc, 'difficult', _DIFFICULT, object_node)
    bndbox_node = doc.createElement('bndbox')
    createChildNode(doc, 'xmin', str(int(attrs['bndbox'][0])), bndbox_node)
    createChildNode(doc, 'ymin', str(int(attrs['bndbox'][1])), bndbox_node)
    createChildNode(doc, 'xmax', str(int(attrs['bndbox'][0] + attrs['bndbox'][2])), bndbox_node)
    createChildNode(doc, 'ymax', str(int(attrs['bndbox'][1] + attrs['bndbox'][3])), bndbox_node)
    object_node.appendChild(bndbox_node)
    return object_node


# 将documentElement写入XML文件
def writeXMLFile(doc, filename):
    tmpfile = open('tmp.xml', 'w')
    doc.writexml(tmpfile, addindent=' ' * 4, newl='\n', encoding='utf-8')
    tmpfile.close()
    fin = open('tmp.xml')
    fout = open(filename, 'w')
    lines = fin.readlines()
    for line in lines[1:]:
        if line.split():
            fout.writelines(line)
    fin.close()
    fout.close()


if __name__ == "__main__":
    # 图片的路径
    img_path = "D:/FLIR/1/FLIR_ADAS_1_3.tar/FLIR_ADAS_1_3/train/train_jpg_file/"
    fileList = os.listdir(img_path)
    if fileList == 0: os._exit(-1)
    with open("COCO_train.json", "r") as f:
        ann_data = json.load(f)
    current_dirpath = os.path.dirname(os.path.abspath('__file__'))
    if not os.path.exists(_ANNOTATION_SAVE_PATH):
        os.mkdir(_ANNOTATION_SAVE_PATH)
    for imageName in fileList:
        saveName = imageName.strip(".jpg")
        xml_file_name = os.path.join(_ANNOTATION_SAVE_PATH, (saveName + '.xml'))
        img = cv2.imread(os.path.join(img_path, imageName))
        height, width, channel = img.shape
        my_dom = xml.dom.getDOMImplementation()
        doc = my_dom.createDocument(None, _ROOT_NODE, None)
        root_node = doc.documentElement
        createChildNode(doc, 'folder', _FOLDER_NODE, root_node)
        createChildNode(doc, 'filename', saveName + '.jpg', root_node)
        source_node = doc.createElement('source')
        createChildNode(doc, 'database', _DATABASE_NAME, source_node)
        createChildNode(doc, 'annotation', _ANNOTATION, source_node)
        createChildNode(doc, 'image', 'flickr', source_node)
        createChildNode(doc, 'flickrid', 'NULL', source_node)
        root_node.appendChild(source_node)
        owner_node = doc.createElement('owner')
        createChildNode(doc, 'flickrid', 'NULL', owner_node)
        createChildNode(doc, 'name', _AUTHOR, owner_node)
        root_node.appendChild(owner_node)
        size_node = doc.createElement('size')
        createChildNode(doc, 'width', str(width), size_node)
        createChildNode(doc, 'height', str(height), size_node)
        createChildNode(doc, 'depth', str(channel), size_node)
        root_node.appendChild(size_node)
        createChildNode(doc, 'segmented', _SEGMENTED, root_node)
        for ann in ann_data:
            """
   				[5:]
            """
            if (saveName == ann["filename"]):
                # object节点
                object_node = createObjectNode(doc, ann)
                root_node.appendChild(object_node)
            else:
                continue
        print(xml_file_name)
        writeXMLFile(doc, xml_file_name)

通过上面3步的到的xml