PascalVOC形式のxmlファイルをCOCO形式のjsonファイルに変換する

コードの大部分は、Convert-Pascal-VOC-to-COCO(github)を流用させていただいています。

個人的な理由になりますが、物体検出の学習にCOCO形式のデータセットを用いる必要がありまして。
だけど、今までアノテーションツールで作成してきた学習データがPascalVOC形式のxmlファイルですぐに使えない。
PascalVOC形式のデータセットをCOCO形式のデータセットに変換する際の情報って少ないんですよね。

自分がみつけられていないだけかもしれませんが。。

Pythonコード

./annotation/ディレクトリの配下にあるxmlファイルすべてを変換し、1つのjsonファイル(train.json)に出力します。

カテゴリーの名前はracket,player,,,としてますが、データに合わせて書き換えてください。

XML2JSON.py

import os
import glob
import xml.etree.ElementTree as ET
import xmltodict
import json
from xml.dom import minidom
from collections import OrderedDict

def XML2JSON(xmlFiles):
    attrDict = dict()
    attrDict["categories"]=[{"supercategory":"none","id":1,"name":"racket"},
                    {"supercategory":"none","id":2,"name":"player"},
                    {"supercategory":"none","id":3,"name":"tennisball"},
                    {"supercategory":"none","id":4,"name":"umpire"},
                {"supercategory":"none","id":5,"name":"ballperson"},
                {"supercategory":"none","id":6,"name":"camera"},
                {"supercategory":"none","id":7,"name":"player"},
                {"supercategory":"none","id":8,"name":"tv"},
                {"supercategory":"none","id":9,"name":"smartphone"}
                  ]
    images = list()
    annotations = list()
    image_id = 0
    for file in xmlFiles:    
        image_id = image_id + 1      
        annotation_path=file
        image = dict()
        doc = xmltodict.parse(open(annotation_path).read(), force_list=('object'))
        image['file_name'] = str(doc['annotation']['filename'])
        image['height'] = int(doc['annotation']['size']['height'])
        image['width'] = int(doc['annotation']['size']['width'])
        image['id'] = image_id
        print ("File Name: {} and image_id {}".format(file, image_id))
        images.append(image)
        id1 = 1
        if 'object' in doc['annotation']:
            for obj in doc['annotation']['object']:
                for value in attrDict["categories"]:
                    annotation = dict()          
                    if str(obj['name']) == value["name"]:
                        annotation["iscrowd"] = 0
                        annotation["image_id"] = image_id
                        x1 = int(obj["bndbox"]["xmin"])  - 1
                        y1 = int(obj["bndbox"]["ymin"]) - 1
                        x2 = int(obj["bndbox"]["xmax"]) - x1
                        y2 = int(obj["bndbox"]["ymax"]) - y1                         
                        annotation["bbox"] = [x1, y1, x2, y2]
                        annotation["area"] = float(x2 * y2)
                        annotation["category_id"] = value["id"]
                        annotation["ignore"] = 0
                        annotation["id"] = id1
                        annotation["segmentation"] = [[x1,y1,x1,(y1 + y2), (x1 + x2), (y1 + y2), (x1 + x2), y1]]
                        id1 +=1
                        annotations.append(annotation)

            else:
                print("File: {} doesn't have any object".format(file))

        else:
            print("File: {} not found".format(file))


    attrDict["images"] = images    
    attrDict["annotations"] = annotations
    attrDict["type"] = "instances"

    jsonString = json.dumps(attrDict)
    with open("train.json", "w") as f:
        f.write(jsonString)


path="./annotations/"
trainXMLFiles=glob.glob(os.path.join(path, '*.xml'))
XML2JSON(trainXMLFiles)