訓練YOLO v4模型時,xml格式轉txt格式

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)

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章