TensorFlow 製作自己數據集時,xml轉csv

 TensorFlow 製作自己數據集時,xml轉csv千篇一律,把我拐入坑裏了。

如果訓練自己的數據集只有一個類別,用網絡上的xml_to_csv,完全沒有問題,源碼如下:

# -*- coding: utf-8 -*-
import os
import glob
import pandas as pd
import xml.etree.ElementTree as ET

def xml_to_csv(path):
    xml_list = []
    # 讀取註釋文件
    for xml_file in glob.glob(path + '/*.xml'):
        tree = ET.parse(xml_file)
        root = tree.getroot()
        for member in root.findall('object'):
            value = (root.find('filename').text + '.jpg',
                     int(root.find('size')[0].text),
                     int(root.find('size')[1].text),
                     member[0].text,
                     int(member[4][0].text),
                     int(member[4][1].text),
                     int(member[4][2].text),
                     int(member[4][3].text)
                     )
            xml_list.append(value)
    column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']

    # 將所有數據分爲樣本集和驗證集,一般按照3:1的比例
    train_list = xml_list[0: int(len(xml_list) * 0.67)]
    eval_list = xml_list[int(len(xml_list) * 0.67) + 1: ]

    # 保存爲CSV格式
    train_df = pd.DataFrame(train_list, columns=column_name)
    eval_df = pd.DataFrame(eval_list, columns=column_name)
    train_df.to_csv('data/train.csv', index=None)
    eval_df.to_csv('data/eval.csv', index=None)


def main():
    path = './xml'
    xml_to_csv(path)
    print('Successfully converted xml to csv.')

main()

如果你的類別數據集,超過2類以上,再用上述源碼,覺得把所有的數據集3:1的分割,而非一個類別的3:1分割 。

對上述源碼略作調整,完美把每一類數據集按照9:1分割爲訓練數據集和測試數據集,源代碼如下:

# coding: utf-8
import glob
import pandas as pd
import xml.etree.ElementTree as ET

classes = ["20Km_h", "no_passing_35", "no_passing", "keep_left", "keep_right", "mandatory", "straight_or_left", "passing_limits",
           "bicycles", "pedestrians", "stop", "dangerous"]

def xml_to_csv(path):
    train_list = []
    eval_list = []

    for cls in classes:
        xml_list = []
        # 讀取註釋文件
        for xml_file in glob.glob(path + '/*.xml'):
            tree = ET.parse(xml_file)
            root = tree.getroot()
            for member in root.findall('object'):
                if cls == member[0].text:
                    value = (root.find('filename').text,
                             int(root.find('size')[0].text),
                             int(root.find('size')[1].text),
                             member[0].text,
                             int(member[4][0].text),
                             int(member[4][1].text),
                             int(member[4][2].text),
                             int(member[4][3].text)
                             )
                    xml_list.append(value)

        for i in range(0,int(len(xml_list) * 0.9)):
            train_list.append(xml_list[i])
        for j in range(int(len(xml_list) * 0.9) + 1,int(len(xml_list))):
            eval_list.append(xml_list[j])

    column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']


    # 保存爲CSV格式
    train_df = pd.DataFrame(train_list, columns=column_name)
    eval_df = pd.DataFrame(eval_list, columns=column_name)
    train_df.to_csv('data/train.csv', index=None)
    eval_df.to_csv('data/eval.csv', index=None)


def main():
    # path = 'E:\\\data\\\Images'
    path = r'D:\work\PycharmPro\trafficsign\SSD_NET\data\xml_data'  # path參數更具自己xml文件所在的文件夾路徑修改
    xml_to_csv(path)
    print('Successfully converted xml to csv.')


main()

classes = ["20Km_h", "no_passing_35", "no_passing", "keep_left", "keep_right", "mandatory", "straight_or_left", "passing_limits", "bicycles", "pedestrians", "stop", "dangerous"]

該處需要改爲自己數據集類別標籤名。

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