xml形式のデータをtxt形式のデータに変換する(yolov3)

#前置き
物体検出をtensorflowを使い、yoloで行おうとしたときにデータがxml形式が多く、yoloに適用できない。それならば自分で作ろう。
#実際に作成し、使用したコード

import xml.etree.ElementTree as ET
import sys , os
import glob

cate_list = ["Car","Pedestrian","Truck","Signal","Signs","Bicycle","Motorbike","Bus","SVehicle","Train"]

with open("voc_classes.txt","w") as f:
    f.write('\n'.join(cate_list))
    
def class_encord(class_name):
    cate_id = {"Car":0,"Pedestrian":1,"Truck":2,"Signal":3,"Signs":4,"Bicycle":5,"Motorbike":6,"Bus":7,"SVehicle":8,"Train":9}
    return cate_id[class_name]

def convert(data_file,list_file):
    in_file = open(data_file)
    tree = ET.parse(in_file)
    root = tree.getroot()
    for obj in root.iter("item"):
        cate = obj.find("category").text
        cate = cate.lstrip("\n").lstrip("   ")
        cate = cate.rstrip("    ").rstrip("\n")
        category_id = class_encord(cate)
        xmlbox = obj.find("box2d")
        data = [int(float(xmlbox.find("x1").text)),int(float(xmlbox.find("x2").text)),int(float(xmlbox.find("y1").text)),int(float(xmlbox.find("y2").text))]
        list_file.write(" " + ",".join([str(a) for a in data]) + "," + str(category_id))

data_file_list = glob.glob("Annotations/*.xml")

list_file = open("2007_train.txt","w")
for data_file in data_file_list:
    jpg_file = "train_" + data_file.rstrip(".xml") + ".jpg"
    list_file.write(jpg_file)
    convert(data_file,list_file)
    list_file.write("\n")
list_file.close()

一般化してないのですごく見ずらいですね(;^ω^)
一応変更することで使用できるパラメータを載せておきます

• cate_list : 検出したい物体のlistになります。
• open("voc_classes.txt") : txtファイルの名前を変えていただければ物体リストをtxtに出力したものになります。
• cate_id : 物体とそれに対応している値を入れてください。
• list_file : 出来上がるトレーニング用のテキストデータです。

#コードの流れ
yoloは物体が記載されたtxtを使用しますのでそれを出力。

cate_list = ["Car","Pedestrian","Truck","Signal","Signs","Bicycle","Motorbike","Bus","SVehicle","Train"]

with open("voc_classes.txt","w") as f:
    f.write('\n'.join(cate_list))

物体名とidをエンコードするための関数。

def class_encord(class_name):
    cate_id = {"Car":0,"Pedestrian":1,"Truck":2,"Signal":3,"Signs":4,"Bicycle":5,"Motorbike":6,"Bus":7,"SVehicle":8,"Train":9}
    return cate_id[class_name]

xmlのデータをtxtファイルにコンバートする関数。
やっていることはtml.etree.ElementTreeで読み込み、各データを取り出し、書き込む、という簡単なものです。

def convert(data_file,list_file):
    in_file = open(data_file)
    tree = ET.parse(in_file)
    root = tree.getroot()
    for obj in root.iter("item"):
        cate = obj.find("category").text
        cate = cate.lstrip("\n").lstrip("   ")
        cate = cate.rstrip("    ").rstrip("\n")
        category_id = class_encord(cate)
        xmlbox = obj.find("box2d")
        data = [int(float(xmlbox.find("x1").text)),int(float(xmlbox.find("x2").text)),int(float(xmlbox.find("y1").text)),int(float(xmlbox.find("y2").text))]
        list_file.write(" " + ",".join([str(a) for a in data]) + "," + str(category_id))

後はそれを実行しているだけ。

data_file_list = glob.glob("Annotations/*.xml")

list_file = open("2007_train.txt","w")
for data_file in data_file_list:
    jpg_file = "train_" + data_file.rstrip(".xml") + ".jpg"
    list_file.write(jpg_file)
    convert(data_file,list_file)
    list_file.write("\n")
list_file.close()

物体検出は面白いですね(^▽^)
データ整理が大変ですが。