COCO格式数据集转pascal_voc数据集格式

import os
import logging
import argparse
from xml.dom import minidom
from pycocotools.coco import COCO
import json

class Dataset():
    def __init__(self, jsonSrc, annFolder):
        super(Dataset, self).__init__()
        self.setLogger()  # 设置日志
        self.dom = minidom.Document()
        self.jsonSrc = jsonSrc
        self.annFloder = annFolder
        self.coco = COCO(self.jsonSrc)  # 加载数据集到COCO对象

    def setLogger(self):
        self.logger = logging.getLogger(__name__)
        self.logger.setLevel(level=logging.DEBUG)  # 设置日志等级

        # FileHander 将log输出的日志文件
        fileName = os.path.basename(__file__).split('.')[0]  # 设置日志文件名字为执行文件名字
        rootPath = os.path.abspath(".")
        logPath = os.path.join(rootPath, "log")  # 设置日志文件存放路径
        if not os.path.exists(logPath):
            os.makedirs(logPath)
        logFileName = os.path.join(logPath, fileName + ".log")  # 日志全称
        handler = logging.FileHandler(logFileName)  # 设置文件handler以输出到文件
        formatter = logging.Formatter("%(asctime)s - %(filename)s[line:%(lineno)d] - %(levelname)s: %(message)s")
        handler.setFormatter(formatter)
        self.logger.addHandler(handler)

    def createElementNode(self, nodeName, nodeText):
        """
        创建节点
        :param nodeName:节点名称
        :param nodeText:节点内容
        :return:节点
        """
        node = self.dom.createElement(nodeName)
        text = self.dom.createTextNode(str(nodeText))
        node.appendChild(text)
        return node

    def createSourceNode(self, sourceInfo):
        """
        封装source节点
        :param sourceInfo: source节点信息(Dict{database:XXX, annotation:XXX....})
        :return: source节点
        """
        node = self.dom.createElement('source')
        node.appendChild(self.createElementNode('database', sourceInfo['database']))
        node.appendChild(self.createElementNode('annotation', sourceInfo['annotation']))
        node.appendChild(self.createElementNode('image', sourceInfo['image']))
        node.appendChild(self.createElementNode('flickrid', sourceInfo['flickrid']))
        return node

    def createOwnerNode(self, ownerInfo):
        """
        封装owner节点
        :param ownerInfo: owner节点信息(Dict{flickrid:XX, name:XX})
        :return: owner节点
        """
        node = self.dom.createElement('owner')
        node.appendChild(self.createElementNode('flickrid', ownerInfo['flickrid']))
        node.appendChild(self.createElementNode('name', ownerInfo['name']))
        return node

    def createSizeNode(self, sizeInfo):
        """
        封装size节点
        :param sizeInfo: size节点信息(Dict{width:XX, height:XX...})
        :return: size节点
        """
        node = self.dom.createElement('size')
        node.appendChild(self.createElementNode('width', sizeInfo['width']))
        node.appendChild(self.createElementNode('height', sizeInfo['height']))
        node.appendChild(self.createElementNode('depth', sizeInfo['depth']))
        return node

    def createObjectNode(self, objectInfo):
        """
        封装object节点
        :param objectInfo: object节点所需要的信息(Dict{name:XX, pose:XX....})
        :return: boject节点
        """
        node = self.dom.createElement('object')
        node.appendChild(self.createElementNode('name', objectInfo['name']))
        node.appendChild(self.createElementNode('pose', objectInfo['pose']))
        node.appendChild(self.createElementNode('truncated', objectInfo['truncated']))
        node.appendChild(self.createElementNode('difficult', objectInfo['difficult']))
        bndboxNode = self.dom.createElement('bndbox')
        bndboxNode.appendChild(self.createElementNode('xmin', objectInfo['bndbox']['xmin']))
        bndboxNode.appendChild(self.createElementNode('ymin', objectInfo['bndbox']['ymin']))
        bndboxNode.appendChild(self.createElementNode('xmax', objectInfo['bndbox']['xmax']))
        bndboxNode.appendChild(self.createElementNode('ymax', objectInfo['bndbox']['ymax']))
        node.appendChild(bndboxNode)  # 将bndbox加到object节点中
        return node

    def genAnnotation(self, vocInfo):
        """
        创建单张图片的Anntation
        :param vocInfo: 该图片的VOC信息(字典)
        :return: annotation的XML节点
        """
        annotationNode = self.dom.createElement('annotation')
        # folderNode建立并添加
        folderNode = self.createElementNode('folder', vocInfo['folder'])
        annotationNode.appendChild(folderNode)
        # filenameNode建立并添加
        filenameNode = self.createElementNode('filename', vocInfo['filename'])
        annotationNode.appendChild(filenameNode)
        # sourceNode建立并添加
        sourceNode = self.createSourceNode(vocInfo['sourceInfo'])
        annotationNode.appendChild(sourceNode)
        # ownerNode建立并添加
        ownerNode = self.createOwnerNode(vocInfo['ownerInfo'])
        annotationNode.appendChild(ownerNode)
        # sizeNode建立并添加
        sizeNode = self.createSizeNode(vocInfo['sizeInfo'])
        annotationNode.appendChild(sizeNode)
        # segmentedNode建立并添加
        segmentedNode = self.createElementNode('segmented', vocInfo['segmented'])
        annotationNode.appendChild(segmentedNode)
        # 循环建立并添加所有objectNode
        for objectInfo in vocInfo['objectInfo']:
            objectNode = self.createObjectNode(objectInfo)
            annotationNode.appendChild(objectNode)
        return annotationNode

    def genVocInfo(self, imgId):
        """
        根据coco数据集中的imgId生成单张图片的Voc信息
        :param imgId: coco数据集中的图片Id
        :return: imgId对应图片的voc信息
        """
        vocInfo = {
            'folder': None,
            'filename': None,
            'sourceInfo': {
                'database': 'Detection',
                'annotation': 'COCOData',
                'image': 'flickr',
                'flickrid': 'NULL'
            },
            'ownerInfo': {
                'flickrid': 'NULL',
                'name': 'LY'
            },
            'sizeInfo': {
                'width': None,
                'height': None,
                'depth': None
            },
            'segmented': None,
            'objectInfo': []
        }  # 创建vocInfo字典，以储存voc所需要的信息
        imageInfo = self.coco.loadImgs(ids=[imgId])[0]  # image信息

        image_id = self.coco.getAnnIds(imgIds=imgId)
        annotationsInfo = self.coco.loadAnns(ids=image_id)  # 注释信息
        #self.coco.loadAnns(ids=[10243, 10244, 10245, 10246, 10247, 10248, 10249, 10250, 10251, 10252, 10253]))
        #print(self.coco.getAnnIds(imgIds=imgId))
        vocInfo['folder'] = self.annFloder
        vocInfo['filename'] = imageInfo['file_name']  # 填入filename
        vocInfo['sizeInfo']['width'] = imageInfo['width']
        vocInfo['sizeInfo']['height'] = imageInfo['height']
        vocInfo['sizeInfo']['depth'] = 3
        vocInfo['segmented'] = 0
        # 遍历该图片的所有注释
        for anntation in annotationsInfo:
            objectInfo = {
                'name': None,
                'pose': 'Unspecified',
                'truncated': 0,
                'difficult': 0,
                'bndbox': {
                    'xmin': None,
                    'ymin': None,
                    'xmax': None,
                    'ymax': None
                }
            }  # 创建object字典存放信息
            catInfo = self.coco.loadCats(ids=[anntation['category_id']])  # 该注释的分类信息
            print(anntation)
            bbox = anntation['bbox'] # 获取该注释的边界框信息
            objectInfo['name'] = catInfo[0]['name']
            objectInfo['bndbox']['xmin'] = int(bbox[0])
            objectInfo['bndbox']['ymin'] = int(bbox[1])
            objectInfo['bndbox']['xmax'] = int(bbox[2])
            objectInfo['bndbox']['ymax'] = int(bbox[3])
            vocInfo['objectInfo'].append(objectInfo)
        return vocInfo

    def convertXML(self):
        savePath = annFolder + '/' + 'Annotations/'  # 文件生成在当前程序文件夹下
        if not os.path.exists(savePath):  # 若没有文件夹则创建
            os.makedirs(savePath)
        imgIds = self.coco.getImgIds()  # 获取所有图像的ID
        print('There are {} images in {}'.format(len(imgIds), self.jsonSrc))
        for imgId in imgIds:  # 循环遍历每个图像id
            print(imgId)
            vocInfo = self.genVocInfo(imgId)  # 根据coco图像id生成voc数据
            annotationNode = self.genAnnotation(vocInfo)  # 根据voc数据生成xml格式节点
            saveName = vocInfo['filename'].split('.')[0] + '.xml'  # 保存xml名称与图片名称一致
            saveFile = savePath + saveName  # 保存xml的路径
            try:
                with open(saveFile, 'w', encoding='utf-8') as f:  # 保存图片
                    annotationNode.writexml(f, indent='', addindent='\t', newl='\n')
                print(saveFile, ' write ok')
            except Exception as e:
                self.logger.error("Write Xml ERROR: %s failed\n Exception as follow: %s "
                                  % (saveName, e))
                print("\033[31mError: %s wrote failed!\033[0m" % saveName)

if __name__ == '__main__':
    label = 'test'
    annFolder = 'VOC2007'

    jsonPath = 'xxxx_coco.json'
    data = Dataset(jsonPath, annFolder)
    data.convertXML()