目标检测 的标注数据 .xml 转为 tfrecord 的格式用于 TensorFlow 训练

将目标检测 的标注数据 .xml 转为 tfrecord 的格式用于 TensorFlow 训练。

import xml.etree.ElementTree as ET

import numpy as np

import os

import tensorflow as tf

from PIL import Image

classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable",

           "dog", "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]

def convert(size, box):

    dw = 1./size[0]

    dh = 1./size[1]

    x = (box[0] + box[1])/2.0

    y = (box[2] + box[3])/2.0

    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_id):

    in_file = open('F:/xml/%s.xml'%(image_id))

    tree = ET.parse(in_file)

    root = tree.getroot()

    size = root.find('size')

    w = int(size.find('width').text)

    h = int(size.find('height').text)

    bboxes = []

    for i, obj in enumerate(root.iter('object')):

        if i > 29:

            break

        difficult = obj.find('difficult').text

        cls = obj.find('name').text

        if cls not in classes or int(difficult) == 1:

            continue

        cls_id = classes.index(cls)

        xmlbox = obj.find('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) + [cls_id]

        bboxes.extend(bb)

    if len(bboxes) < 30*5:

        bboxes = bboxes + [0, 0, 0, 0, 0]*(30-int(len(bboxes)/5))

    return np.array(bboxes, dtype=np.float32).flatten().tolist()

def convert_img(image_id):

    image = Image.open('F:/snow leopard/test_im/%s.jpg' % (image_id))

    resized_image = image.resize((416, 416), Image.BICUBIC)

    image_data = np.array(resized_image, dtype='float32')/255

    img_raw = image_data.tobytes()

    return img_raw

filename = os.path.join('test'+'.tfrecords')

writer = tf.python_io.TFRecordWriter(filename)

# image_ids = open('F:/snow leopard/test_im/%s.txt' % (

#     year, year, image_set)).read().strip().split()

image_ids = os.listdir('F:/snow leopard/test_im/')

# print(filename)

for image_id in image_ids:

    print (image_id)

    image_id = image_id.split('.')[0]

    print (image_id)

    xywhc = convert_annotation(image_id)

    img_raw = convert_img(image_id)

    example = tf.train.Example(features=tf.train.Features(feature={

        'xywhc':

                tf.train.Feature(float_list=tf.train.FloatList(value=xywhc)),

        'img':

                tf.train.Feature(bytes_list=tf.train.BytesList(value=[img_raw])),

        }))

    writer.write(example.SerializeToString())

writer.close()

　　

Python读取文件夹下图片的两种方法：

import os

imagelist = os.listdir('./images/')      #读取images文件夹下所有文件的名字

import glob

imagelist= sorted(glob.glob('./images/' + 'frame_*.png'))      #读取带有相同关键字的图片名字，比上一中方法好

参考：

https://blog.csdn.net/CV_YOU/article/details/80778392

https://github.com/raytroop/YOLOv3_tf