YOLO形式の座標情報をxmlに変換する

#はじめに
研究中にyoloでの座標情報をSSDなどのPascalVOC形式に直す必要がありました。
しかし、ネットで検索しても出てくるのはその逆である場合が多かったです。
なので、自分用にプログラムを自作してみました。

#コードの紹介
以下に作成したコードを示します。

yolo_to_xml.py

import xml.etree.ElementTree as ET
import xml.dom.minidom as md
import os
import numpy as np

from PIL import  Image

def create_xml(file_name, object_number):

    #保存フォルダへ移動
    os.chdir('./converted_to_xml')

    #xmlファイルを生成する
    annotation = ET.Element('annotation')
    folder = ET.SubElement(annotation, 'folder')
    filename = ET.SubElement(annotation, 'filename')
    path = ET.SubElement(annotation, 'path')
    source = ET.SubElement(annotation, 'source')
    database = ET.SubElement(source, 'database')
    size = ET.SubElement(annotation, 'size')
    width = ET.SubElement(size, 'width')
    height = ET.SubElement(size, 'height')
    depth = ET.SubElement(size, 'depth')
    segment = ET.SubElement(annotation, 'segment')

    for i in range(object_number):
        object = ET.SubElement(annotation, 'object')
        name = ET.SubElement(object, 'name')
        pose = ET.SubElement(object, 'pose')
        truncated = ET.SubElement(object, 'truncated')
        difficult = ET.SubElement(object, 'difficult')
        bndbox = ET.SubElement(object, 'bndbox')
        xmin = ET.SubElement(bndbox, 'xmin')
        ymin = ET.SubElement(bndbox, 'ymin')
        xmax = ET.SubElement(bndbox, 'xmax')
        ymax = ET.SubElement(bndbox, 'ymax')

    tree = ET.ElementTree(annotation)
    fl = file_name
    tree.write(fl)

    os.chdir('../')

    return fl

def add_object(name):

    os.chdir('./converted_to_xml')
    xml_name = name.replace('txt', 'xml')

    # xmlに新たな物体情報を追加する
    tree = ET.parse(xml_name)
    root = tree.getroot()

    for annotation in root.findall('object'):
        object = ET.SubElement(annotation, 'object')
        name = ET.SubElement(object, 'name')
        pose = ET.SubElement(object, 'pose')
        truncated = ET.SubElement(object, 'truncated')
        difficult = ET.SubElement(object, 'difficult')
        bndbox = ET.SubElement(object, 'bndbox')
        xmin = ET.SubElement(bndbox, 'xmin')
        ymin = ET.SubElement(bndbox, 'ymin')
        xmax = ET.SubElement(bndbox, 'xmax')
        ymax = ET.SubElement(bndbox, 'ymax')

    tree = ET.ElementTree(root)
    fl = xml_name
    tree.write(fl)

    os.chdir('../')

    return fl

def read_txt(name):
    txt_path = os.path.join('anottation_data/', name)

    with open(txt_path) as f:
        l_strip = [s.strip() for s in f.readlines()]

        #画像中の物体の個数を出力
        #print(len(l_strip))
        #物体が２つ以上ならxmlファイルを書き換える
        if len(l_strip) != '1':
            xml_name = name.replace('.txt', '.xml')
            create_xml(xml_name, len(l_strip))
        #print(l_strip)
        #txtデータ読み込み
        for i, j in enumerate(l_strip):

            #データを分割する
            # splitは番号、ｘ座標、y座標、Width、Heightの順に格納されているリスト型
            split = j.split()

            #辞書型に代入する
            data = {'number':0,'x_coordinate':0, 'y_coordinate':0, 'width':0, 'height':0}
            data['number'] = split[0]
            data['x_coordinate'] = split[1]
            data['y_coordinate'] = split[2]
            data['width'] = split[3]
            data['height'] = split[4]

            #作成したxmlへ代入
            xml_name = name.replace('.txt', '.xml')
            yolo_to_xml(data, xml_name, name, i)

    return name

def yolo_to_xml(data, xml_name, name, object_counter):

    #保存用の辞書型生成
    coordinated_data = {'x_min': 0, 'x_max': 0, 'y_min': 0, 'y_max': 0}

    #画像データの取得
    os.chdir('./data')
    im_name = name.replace('.txt', '.jpg')
    im = np.array(Image.open(im_name))

    #object_nameの変更(任意の物体名を追加可能)
    if data['number'] == '0':
        data['number'] = 'Can'
    elif data['number'] == '1':
        data['number'] = 'Person'
    elif data['number'] == '2':
        data['number'] = 'Box'
    else:
        data['number'] = 'Bottle'

    #座標の変更(非正規化)
    data['x_coordinate'] = int(float(data['x_coordinate']) * im.shape[1])
    data['y_coordinate'] = int(float(data['y_coordinate']) * im.shape[0])
    data['width'] = int(float(data['width']) * im.shape[1])
    data['height'] = int(float(data['height']) * im.shape[0])
    coordinated_data['x_min'] = int(data['x_coordinate'] - (data['width']/2))
    coordinated_data['x_max'] = int(data['x_coordinate'] + (data['width']/2))
    coordinated_data['y_min'] = int(data['y_coordinate'] - (data['height']/2))
    coordinated_data['y_max'] = int(data['y_coordinate'] + (data['height']/2))

    for k, v in coordinated_data.items():
        print(k, v)

    #xmlへの書き込み開始
    os.chdir('../')
    os.chdir('./converted_to_xml')

    #xmlにデータを書き込んでいる
    tree = ET.parse(xml_name)
    root = tree.getroot()

    root.findall('filename')[0].text = im_name
    root.findall('./*/width')[0].text = str(im.shape[1])
    root.findall('./*/height')[0].text = str(im.shape[0])
    root.findall('./*/depth')[0].text = str(im.shape[2])
    root.findall('./*/name')[object_counter].text = data['number']
    root.findall('./*/*/xmin')[object_counter].text = str(coordinated_data['x_min'])
    root.findall('./*/*/ymin')[object_counter].text = str(coordinated_data['y_min'])
    root.findall('./*/*/xmax')[object_counter].text = str(coordinated_data['x_max'])
    root.findall('./*/*/ymax')[object_counter].text = str(coordinated_data['y_max'])


    tree.write(xml_name)

    os.chdir('../')

    return data


if __name__ == "__main__":
    #txtデータの名前のリストを保存する
    path = os.listdir('./anottation_data')
    #xmlファイルを生成してtxtファイルの内容を書き込む

    # xmlファイル生成
    for i, name in enumerate(path):
        name = name.replace('.txt', '.xml')
        create_xml(name, 1)
        read_txt(path[i])

#補足
フォルダの位置は個人によります。