YOLO訓練時需要的標籤數據是txt格式的
而以前做Faster-RCNN的時候,標籤文件是xml格式的
特別當下載別人標記好的數據集時,很有可能提供的是xml格式的文件,那麼如何將xml格式轉爲txt格式?
代碼如下:
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
sets = []
classes = ["face","face_mask"] ##修改爲自己的類別
#原樣保留。size爲圖片大小
# 將ROI的座標轉換爲yolo需要的座標
# size是圖片的w和h
# box裏保存的是ROI的座標(x,y的最大值和最小值)
# 返回值爲ROI中心點相對於圖片大小的比例座標,和ROI的w、h相對於圖片大小的比例
def convert(size, box):
dw = 1./(size[0])
dh = 1./(size[1])
x = (box[0] + box[1])/2.0 - 1
y = (box[2] + box[3])/2.0 - 1
w = box[1] - box[0]
h = box[3] - box[2]
x = x*dw
w = w*dw
y = y*dh
h = h*dh
return (x,y,w,h)
def convert_annotation(image_add):
#image_add進來的是帶地址的.jpg
image_add = os.path.split(image_add)[1] #截取文件名帶後綴
image_add = image_add[0:image_add.find('.',1)] #刪除後綴,現在只有文件名沒有後綴
#現在傳進來的只有圖片名沒有後綴
in_file = open('G:/1/xml/' + image_add + '.xml') #修改爲你自己的輸入目錄
out_file = open('G:/1/labels/%s.txt'%(image_add), 'w') #修改爲你自己的輸出目錄
tree=ET.parse(in_file)
root = tree.getroot()
if root.find('size'):
size = root.find('size')
w = int(size.find('width').text) #偶爾xml標記出錯,width或height設置爲0了
h = int(size.find('height').text) #需要標記出來,便於單獨處理
if w==0:
print("出錯! width或height爲0: "+image_add)
os.remove("G:/set/"+image_add+".jpg")
os.remove("G:/set/"+image_add+".xml")
return
#在一個XML中每個Object的迭代
for obj in root.iter('object'):
#iter()方法可以遞歸遍歷元素/樹的所有子元素
difficult = obj.find('difficult').text
cls = obj.find('name').text
#如果訓練標籤中的品種不在程序預定品種,或者difficult = 1,跳過此object
if cls not in classes or int(difficult)==1:
continue
#cls_id 只等於1
cls_id = classes.index(cls)
xmlbox = obj.find('bndbox')
#b是每個Object中,一個bndbox上下左右像素的元組
b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text))
bb = convert((w,h), b)
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
else:
print("出錯!xml缺少size: "+image_add) #偶爾xml缺少size,需要標記出來,便於單獨處理
os.remove("G:/set/"+image_add+".jpg")
os.remove("G:/set/"+image_add+".xml")
image_adds = open("G:/1/train.txt") #修改爲你自己的訓練數據集目錄
for image_add in image_adds:
#print(image_add)
#image_add = image_add.strip()
#print (image_add)
convert_annotation(image_add)