tensorflow2.x版本的xml转csv转record

xml转csv

import os
import xml.dom.minidom

path_img = 'C:/Users/PaulY/Desktop/testdata/testimgs'
path_xml = 'C:/Users/PaulY/Desktop/testdata/testxml'

xml_list = []
for xml1 in os.listdir(path_xml):
    if xml1.endswith(".xml") and xml1 not in "desktop.xml":
        xml_list.append(xml1)
print(len(xml_list))
VOC_LABELS = {'car', 'bus', 'person', 'bike', 'truck', 'motor'}
# csv_labels = open("C:/Users/PaulY/Desktop/cocotrain/train.csv", "w")
csv_labels = open("C:/Users/PaulY/Desktop/testdata/test.csv", "w")
li = 'filename,width,height,class,xmin,ymin,xmax,ymax' + "\n"
csv_labels.write(li)
for xml_file in xml_list:
    print(xml_file)
    image_id, _ = os.path.splitext(xml_file)

    DomTree = xml.dom.minidom.parse(os.path.join(path_xml, xml_file))
    annotation = DomTree.documentElement
    objectlist = annotation.getElementsByTagName('object')
    sizelist = annotation.getElementsByTagName('size')
    for o in sizelist:
        si = o.getElementsByTagName('width')
        width = si[0].childNodes[0].data
        hi = o.getElementsByTagName('height')
        height = hi[0].childNodes[0].data
    for objects in objectlist:
        namelist = objects.getElementsByTagName('name')
        objectname = namelist[0].childNodes[0].data
        bndbox = objects.getElementsByTagName("bndbox")
        for box in bndbox:
            x1_list = box.getElementsByTagName('xmin')
            x1 = int(eval(x1_list[0].childNodes[0].data))
            y1_list = box.getElementsByTagName('ymin')
            y1 = int(eval(y1_list[0].childNodes[0].data))
            x2_list = box.getElementsByTagName('xmax')
            x2 = int(eval(x2_list[0].childNodes[0].data))
            y2_list = box.getElementsByTagName('ymax')
            y2 = int(eval(y2_list[0].childNodes[0].data))
        if objectname in str(VOC_LABELS):

            line = str(image_id) + "," + width + "," + height + "," + objectname + "," + str(x1) + "," + str(y1) \
                   + "," + str(x2) + "," + str(y2) + "\n"
            # print(line)
            csv_labels.write(line)
        else:
            print(objectname,image_id)
csv_labels.close()

csv转record

import tensorflow as tf
import pandas as pd
import os
from PIL import Image
from collections import namedtuple
import io

def class_text_to_int(row_label):
    if row_label == 'car':
        return 1
    elif row_label == 'bus':
        return 2
    elif row_label == 'person':
        return 3
    elif row_label == 'bike':
        return 4
    elif row_label == 'truck':
        return 5
    elif row_label == 'motor':
        return 6
    else:
        print('NONE: ' + row_label)
        return 0

def split(df, group):
    data = namedtuple('data', ['filename', 'object'])
    gb = df.groupby(group)
    return [data(filename, gb.get_group(x)) for filename, x in zip(gb.groups.keys(), gb.groups)]

def create_tf_example(group, path):
    with tf.io.gfile.GFile(os.path.join(path, '{}.jpg'.format(group.filename)), 'rb') as f:
      encoded_jpg = f.read()
    encoded_jpg_io = io.BytesIO(encoded_jpg)
    image = Image.open(encoded_jpg_io)
    width, height = image.size
    filename = (group.filename + '.jpg').encode('utf8')
    image_format = b'jpg'
    xmins = []
    xmaxs = []
    ymins = []
    ymaxs = []
    classes_text = []
    classes = []

    for index, row in group.object.iterrows():
        xmins.append(int(row['xmin']) / int(width))
        xmaxs.append(int(row['xmax']) / int(width))
        ymins.append(int(row['ymin']) / int(height))
        ymaxs.append(int(row['ymax']) / int(height))
        classes_text.append(row['class'].encode('utf8'))
        classes.append(class_text_to_int(row['class']))

    feature_dict = {
        'image/height': tf.train.Feature(int64_list=tf.train.Int64List(value=[height])),
        'image/width': tf.train.Feature(int64_list=tf.train.Int64List(value=[width])),
        'image/filename': tf.train.Feature(bytes_list=tf.train.BytesList(value=[filename])),
        'image/source_id': tf.train.Feature(bytes_list=tf.train.BytesList(value=[filename])),
        'image/encoded': tf.train.Feature(bytes_list=tf.train.BytesList(value=[encoded_jpg])),
        'image/format': tf.train.Feature(bytes_list=tf.train.BytesList(value=[image_format])),
        'image/object/bbox/xmin': tf.train.Feature(float_list=tf.train.FloatList(value=xmins)),
        'image/object/bbox/xmax': tf.train.Feature(float_list=tf.train.FloatList(value=xmaxs)),
        'image/object/bbox/ymin': tf.train.Feature(float_list=tf.train.FloatList(value=ymins)),
        'image/object/bbox/ymax': tf.train.Feature(float_list=tf.train.FloatList(value=ymaxs)),
        'image/object/class/text': tf.train.Feature(bytes_list=tf.train.BytesList(value=classes_text)),
        'image/object/class/label': tf.train.Feature(int64_list=tf.train.Int64List(value=classes)),
    }

    tf_example = tf.train.Example(features=tf.train.Features(feature=feature_dict))
    return tf_example

def main(csv_input, output_path, image_path):
    with tf.io.TFRecordWriter(output_path) as writer:
        path = image_path
        examples = pd.read_csv(csv_input)
        grouped = split(examples, 'filename')
        for group in grouped:
            tf_example = create_tf_example(group, path)
            writer.write(tf_example.SerializeToString())

    print('Successfully created the TFRecords: {}'.format(output_path))


if __name__ == '__main__':
    # imgPath = 'C:/Users/PaulY/Desktop/cocotrain/mydata'
    # output_path = 'C:/Users/PaulY/Desktop/cocotrain/train.record'
    # csv_input = 'C:/Users/PaulY/Desktop/cocotrain/train.csv'
    imgPath = 'C:/Users/PaulY/Desktop/testdata/testimgs'
    output_path = 'C:/Users/PaulY/Desktop/testdata/test.record'
    csv_input = 'C:/Users/PaulY/Desktop/testdata/test.csv'
    main(csv_input, output_path, imgPath)