将COCO数据集整合成VOC格式xml

VOC2012的数据集感觉不是很够，因此将部分COCO数据集转为VOC格式。做的是目标检测的，因此主要还是annotation的部分。

因为VOC格式中的需要的数据其实主要是图片id、bbox数据已经分类资料，所以我先从COCO数据集中的json文件中截取出了我需要的内容

其中这些类就是和VOC重合的类别，共20类。

import json

className = {
    1:'person',
    16:'bird',
    17:'cat',
    21:'cow',
    18:'dog',
    19:'horse',
    20:'sheep',
    5:'aeroplane',
    2:'bicycle',
    9:'boat',
    6:'bus',
    3:'car',
    4:'motorbike',
    7:'train',
    44:'bottle',
    62:'chair',
    67:'dining table',
    64:'potted plant',
    63:'sofa',
    72:'tvmonitor'
}

classNum = [1,2,3,4,5,6,7,9,16,17,18,19,20,21,44,62,63,64,67,72]

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

# with open("instances_val2014.json","r+") as f:
#     data = json.load(f)
    # annData = data["annotations"]
    # print(annData[0])
    # for x in annData[0]:
    #     if(x == "image_id"):
    #         print(type(x))
    #         print(x+ ":" + str(annData[0][x]))
    #     if (x == "image_id" or x == "bbox" or x == "category_id"):
    #         print(x + ":" + annData[0][x])
    #     if (x == "image_id" or x == "bbox" or x == "category_id"):
    #         print(x+ ":" + annData[0][x])

# with open("test.json","w") as f:
#     json.dump(annData, f, ensure_ascii=False)

inputfile = []
inner = {}
##向test.json文件写入内容
with open("instances_train2014.json","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(i["image_id"]).zfill(6),
            "name": className[i["category_id"]],
            "bndbox":i["bbox"]
        }
        inputfile.append(inner)
inputfile = json.dumps(inputfile)
writeNum(inputfile)

之后就是根据选取出来的类别中的图片筛选需要的图片，把不需要的删掉。

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2018/1/18 10:46
# @Author  : He Hangjiang
# @Site    : 
# @File    : 根据目标图片id筛选图片.py
# @Software: PyCharm

import json
import os

nameStr = []

with open("COCO_train.json","r+") as f:
    data = json.load(f)
    print("read ready")

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

nameStr = set(nameStr)
print(nameStr)
print(len(nameStr))

path = "D:/迅雷下载/train2014/"

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

第三步就是根据图片名，依次创建xml文件。因为coco数据集并没有将同一个图片的bbox放在一起，所以还要做个判断。我就直接用了遍历，懒得去想了。代码有参考http://blog.csdn.net/yjl9122/article/details/56842098的内容，你们也可以参考他的。

#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2018/1/19 10:05
# @Author  : He Hangjiang
# @Site    : 
# @File    : creatXML.py
# @Software: PyCharm

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

# xml文件规范定义

# _TXT_PATH= '../../SynthText-master/myresults/groundtruth'
_IMAGE_PATH= 'D:/迅雷下载/pytorch/train'

_INDENT= ''*4
_NEW_LINE= '\n'
_FOLDER_NODE= 'COCO2014'
_ROOT_NODE= 'annotation'
_DATABASE_NAME= 'LOGODection'
_ANNOTATION= 'COCO2014'
_AUTHOR= 'HHJ'
_SEGMENTED= '0'
_DIFFICULT= '0'
_TRUNCATED= '0'
_POSE= 'Unspecified'


# _IMAGE_COPY_PATH= 'JPEGImages'
_ANNOTATION_SAVE_PATH= 'Annotations'
# _IMAGE_CHANNEL= 3

# 封装创建节点的过�?
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)


# object节点比较特殊
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')
    # print(filename)
    fout =open(filename, 'w')
    # print(os.path.dirname(fout))

    lines = fin.readlines()

    for line in lines[1:]:

        if line.split():
            fout.writelines(line)

        # new_lines = ''.join(lines[1:])

        # fout.write(new_lines)

    fin.close()

    fout.close()


if __name__ == "__main__":
    ##读取图片列表
    img_path = "D:/train2014/"
    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)

    # if not os.path.exists(_IMAGE_COPY_PATH):
    #     os.mkdir(_IMAGE_COPY_PATH)

    for imageName in fileList:

        saveName= imageName.strip(".jpg")
        print(saveName)
        # pos = fileList[xText].rfind(".")
        # textName = fileList[xText][:pos]

        # ouput_file = open(_TXT_PATH + '/' + fileList[xText])
        # ouput_file =open(_TXT_PATH)

        # lines = ouput_file.readlines()

        xml_file_name = os.path.join(_ANNOTATION_SAVE_PATH, (saveName + '.xml'))
        # with open(xml_file_name,"w") as f:
        #     pass

        img=cv2.imread(os.path.join(img_path,imageName))
        print(os.path.join(img_path,imageName))
        # cv2.imshow(img)
        height,width,channel=img.shape
        print(height,width,channel)



        my_dom = xml.dom.getDOMImplementation()

        doc = my_dom.createDocument(None,_ROOT_NODE,None)

        # 获得根节�?
        root_node = doc.documentElement

        # folder节点

        createChildNode(doc, 'folder',_FOLDER_NODE, root_node)

        # filename节点

        createChildNode(doc, 'filename', saveName+'.jpg',root_node)

        # source节点

        source_node = doc.createElement('source')

        # 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节点

        owner_node = doc.createElement('owner')

        # owner的子节点

        createChildNode(doc, 'flickrid','NULL', owner_node)

        createChildNode(doc, 'name',_AUTHOR, owner_node)

        root_node.appendChild(owner_node)

        # size节点

        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)

        # segmented节点

        createChildNode(doc, 'segmented',_SEGMENTED, root_node)

        for ann in ann_data:
            if(saveName==ann["filename"]):
                # object节点
                object_node = createObjectNode(doc, ann)
                root_node.appendChild(object_node)

            else:
                continue

        # 构建XML文件名称

        print(xml_file_name)

        # 创建XML文件

        # createXMLFile(attrs, width, height, xml_file_name)


        # # 写入文件
        #
        writeXMLFile(doc, xml_file_name)

然后创建出来的就是这样的：