【DeepLearning】PascalVoc=>keras-yolo3変換

前回の続きです。前回はこちら
buffalokusojima.hatenablog.com

今回は前回作成したファイルのフォーマットをkeras-yolo3に変換するプログラム説明です。
以下、ソース参照

import os, glob
from xml.etree import ElementTree

target_dir = 'image' #対象のフォルダ
output_dir = 'image' #アウトプットフォルダ

classes_path = 'class.txt' #クラスファイル

class_list = []
with open(classes_path, 'r', encoding='UTF-8') as f:
    line = f.readline()
    while line:
        class_list.append(line.replace('\n', ''))
        line = f.readline()
print(class_list)

file_count = 0
error_list = []

for file in glob.glob(os.path.join(target_dir, '*.xml')):
    

    """
     以下、XMLパースしていく
    """
    try:
        tree = ElementTree.parse(file)
        elem = tree.getroot()
        
        objects = elem.findall("object")
        
        annotation_list = []
        
        #アノテーションの数分回す
        for obj in objects:
            
　　     #ラベルをクラスファイルを元に数値化する
            label = class_list.index(obj.find("name").text)

　　　 #bboxを取得し、座標を格納
            for bbox in obj.findall("bndbox"):
                xmin = bbox.find("xmin").text
                xmax = bbox.find("xmax").text
                ymin = bbox.find("ymin").text
                ymax = bbox.find("ymax").text
                annotation_list.append([xmin, ymin, xmax, ymax, str(label)])
                
　　#取得したアノテーションを改行してファイルに記載する
        annotation_file = os.path.join(output_dir, os.path.basename(file).replace("xml", "txt"))
        with open(annotation_file, 'w', encoding='UTF-8') as f:
            for annotation in annotation_list:
                f.write(",".join(annotation))
                f.write('\n')
            print("writing:", annotation_file)
        file_count += 1
    except:
        error_list.append(file)
        
print("file making done:", file_count)
print("error list:", error_list)

以上のコードで指定したフォルダに対して実行すると

<annotation>
	<folder>image</folder>
	<filename>sample1.jpg</filename>
	<path>keras-yolo3/image/sample1.jpg</path>
	<source>
		<database>Unknown</database>
	</source>
	<size>
		<width>510</width>
		<height>340</height>
		<depth>3</depth>
	</size>
	<segmented>0</segmented>
	<object>
		<name>リンゴ</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>57</xmin>
			<ymin>96</ymin>
			<xmax>252</xmax>
			<ymax>311</ymax>
		</bndbox>
	</object>
	<object>
		<name>リンゴ</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>252</xmin>
			<ymin>105</ymin>
			<xmax>452</xmax>
			<ymax>314</ymax>
		</bndbox>
	</object>
	<object>
		<name>リンゴ</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>252</xmin>
			<ymin>25</ymin>
			<xmax>406</xmax>
			<ymax>150</ymax>
		</bndbox>
	</object>
	<object>
		<name>リンゴ</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>101</xmin>
			<ymin>39</ymin>
			<xmax>248</xmax>
			<ymax>141</ymax>
		</bndbox>
	</object>
</annotation>

上記のようなxmlファイルが以下のようなtxtファイルとして新たに作成されます。

57,96,252,311,0
252,105,452,314,0
252,25,406,150,0
101,39,248,141,0

後は、学習リストを作成する際に、学習のフォーマットに合わせてこんファイルから座標を取得すれば学習が出来ます。

【おわりに】
次回からそろそろ学習スクリプトとかをやっていこうと思います。