mxnet框架自制lst文件

# -*- coding: utf-8 -*-
# Author  : Ming
# File    : gen-mxnet-data.py
# Time    : 2018/9/12 0012 下午 5:14

import glob
import xml.etree.ElementTree as ET

def class_text_to_int(row_label):
    if row_label == 'gun':
        return 1
    else:
        return 2


def xml_to_csv(path):
    idx = 0
    x = []
    for xml_file in glob.glob(path + '/*.xml'):

        tree = ET.parse(xml_file)
        root = tree.getroot()
        f_name = [root.find('filename').text]
        v1 = [idx, 4, 5,
              int(root.find('size')[0].text),
              int(root.find('size')[1].text)]
        a = []
        for member in root.findall('object'):

            value = [
                class_text_to_int(member[0].text),
                     int(member[4][0].text),
                     int(member[4][1].text),
                     int(member[4][2].text),
                     int(member[4][3].text)
                     ]
            a.append(value)
        c = []
        for i in a:
            for j in i:
             c.append(j)
        x.append(v1+c+f_name)
        idx = idx + 1
    return x


def main():
    image_path = path
    x = xml_to_csv(image_path)
    with open ('t.lst','w')as f:
        for i in x:
            line = '\t'.join(str(j) for j in i)
            print(line)
            f.write(line)
            f.write('\n')
    print('Successfully converted xml to csv.')

if __name__ == '__main__':
    path = 'E:\gunbiaozhu\gun_test'
    main()