【Pytorch】我的 SSD 專案筆記(一)：讀取 XML 標註檔

因為某些需求，我在一個專案裡需要訓練 Single Shot MultiBox Detector (SSD) 來做物件偵測。考量過各種選項之後，我最後選擇使用 torchvision 官方的 SSD。

不過，由於專案中得用到自訂資料集，而且每張圖片都對應一個 XML 標註檔，所以第一步就得先搞懂該怎麼用 Python 去讀取這些 XML 檔！

首先我需要準備一個 XML 檔，內容如下 :

<annotation>
	<folder>images</folder>
	<filename>g0001.jpg</filename>
	<path>D:\Glaucoma Detection\code\archive\REFUGE2\train\images\g0001.jpg</path>
	<source>
		<database>Unknown</database>
	</source>
	<size>
		<width>2124</width>
		<height>2056</height>
		<depth>3</depth>
	</size>
	<segmented>0</segmented>
	<object>
		<name>optic_disc</name>
		<pose>Unspecified</pose>
		<truncated>0</truncated>
		<difficult>0</difficult>
		<bndbox>
			<xmin>62</xmin>
			<ymin>679</ymin>
			<xmax>518</xmax>
			<ymax>1136</ymax>
		</bndbox>
	</object>
</annotation>

引入所需套件

為了讓程式能夠順利讀取 XML 檔的內容，最方便的做法就是使用 xml.etree.ElementTree 這個套件，它專門處理 XML 格式的資料。這個套件是 Python 內建的函式庫，這代表我們不用另外安裝它 !

import xml.etree.ElementTree as ET

載入 XML 檔

我們可以把指定路徑的 XML 檔讀取進來，並建立一個 ElementTree 物件。這個 tree 就好比整顆 XML 的「樹狀結構」載體。

tree = ET.parse("./Annotations/g0001.xml")

取得這棵樹的「根節點」

什麼是「根節點」? 以我的檔案為例，根節點就是 <annotation>，之後我就可以用 root.find(“XXX”)，來抓底下的所有子元素，例如<object>、<name>、<bndbox>等。

root = tree.getroot()

取得 XML 檔案中的重點資訊

根據 Pytorch 官方的資訊，在訓練一個 SSD 模型的時候，除了要輸入影像之外，還要輸入「每個 Ground Truth 框框的左上角、右下角座標」，以及分類名稱。這兩個資訊可以從 XML 檔的 object/name 以及 object/bndbox/xmin 等子元素獲得。

對我來說這兩個資訊當然就是最重要的，其餘的例如 filename、folder 就還好，只要每張圖(JPG)對應一個標註檔(XML)就好(檔案名稱相同)。

obj = root.find("object")
name = obj.find("name").text
xmin = int(obj.find("bndbox/xmin").text)
ymin = int(obj.find("bndbox/ymin").text)
xmax = int(obj.find("bndbox/xmax").text)
ymax = int(obj.find("bndbox/ymax").text)

print("Class:", name)
print("BndBox:", xmin, ymin, xmax, ymax)

題外話，由於我要訓練的資料只要辨識眼底圖的視盤(Optic Disc)即可，所以可以確定 object 元素只有一個。

最後的輸出結果

Root tag: annotation
Class: optic_disc
BndBox: 62 679 518 1136