将voc2007数据写成TFRecord文件
import tensorflow as tf
import os.path
from lxml import etree
DEBUG=True
DIRECTORY_IMAGES = 'JPEGImages/'
VOC_LABELS = {
'none': (0, 'Background'),
'aeroplane': (1, 'Vehicle'),
'bicycle': (2, 'Vehicle'),
'bird': (3, 'Animal'),
'boat': (4, 'Vehicle'),
'bottle': (5, 'Indoor'),
'bus': (6, 'Vehicle'),
'car': (7, 'Vehicle'),
'cat': (8, 'Animal'),
'chair': (9, 'Indoor'),
'cow': (10, 'Animal'),
'diningtable': (11, 'Indoor'),
'dog': (12, 'Animal'),
'horse': (13, 'Animal'),
'motorbike': (14, 'Vehicle'),
'person': (15, 'Person'),
'pottedplant': (16, 'Indoor'),
'sheep': (17, 'Animal'),
'sofa': (18, 'Indoor'),
'train': (19, 'Vehicle'),
'tvmonitor': (20, 'Indoor'),
}
def _process_xml(annotations_path,xml_file_name):
# Read the XML annotation file.
filename = os.path.join(annotations_path, xml_file_name )
tree = etree.parse(filename)
# Image shape.
width = int(tree.xpath('//size/width/text()')[0])
height = int(tree.xpath('//size/height/text()')[0])
depth = int(tree.xpath('//size/depth/text()')[0])
shape = [height,width,depth]
# Find annotations.
bboxes = []
labels = []
labels_text = []
difficult=[]
truncated=[]
for obj in tree.xpath('//object'):
label = obj.xpath('name/text()')[0]
labels.append(int(VOC_LABELS[label][0]))
labels_text.append(label.encode('ascii'))
difficult_text=obj.xpath('difficult/text()')
if difficult_text:
difficult.append(int(difficult_text[0]))
else:
difficult.append(0)
truncated_text=obj.xpath('truncated/text()')
if truncated_text:
truncated.append(int(truncated_text[0]))
else:
truncated.append(0)
bbox = obj.xpath('bndbox')[0]
bboxes.append((float(bbox.xpath('ymin/text()')[0]) / shape[0],
float(bbox.xpath('xmin/text()')[0]) / shape[1],
float(bbox.xpath('ymax/text()')[0]) / shape[0],
float(bbox.xpath('xmax/text()')[0]) / shape[1]
))
if DEBUG:
print("shape:",shape)
print("labels:",labels)
print("labels_text:",labels_text)
print("difficult:",difficult)
print("truncated:",truncated)
print("bboxes:",bboxes)
return shape, bboxes, labels, labels_text, difficult, truncated
def _build_example(image_data, labels, labels_text, bboxes, shape,
difficult, truncated):
"""
构建example
Args:
image_data: string, JPEG encoding of RGB image;
labels: list of integers, identifier for the ground truth;
labels_text: list of strings, human-readable labels;
bboxes: list of bounding boxes; each box is a list of integers;
specifying [xmin, ymin, xmax, ymax]. All boxes are assumed to belong
to the same label as the image label.
shape: 3 integers, image shapes in pixels.
Returns:
Example proto
"""
xmin = []
ymin = []
xmax = []
ymax = []
for b in bboxes:
assert len(b) == 4
# pylint: disable=expression-not-assigned
[l.append(point) for l, point in zip([ymin, xmin, ymax, xmax], b)]
# pylint: enable=expression-not-assigned
image_format = b'JPEG'
example = tf.train.Example(features=tf.train.Features(feature={
'image/height': int64_feature(shape[0]),
'image/width': int64_feature(shape[1]),
'image/channels': int64_feature(shape[2]),
'image/shape': int64_feature(shape),
'image/object/bbox/xmin': float_feature(xmin),
'image/object/bbox/xmax': float_feature(xmax),
'image/object/bbox/ymin': float_feature(ymin),
'image/object/bbox/ymax': float_feature(ymax),
'image/object/bbox/label': int64_feature(labels),
'image/object/bbox/label_text': bytes_feature(labels_text),
'image/object/bbox/difficult': int64_feature(difficult),
'image/object/bbox/truncated': int64_feature(truncated),
'image/format': bytes_feature(image_format),
'image/encoded': bytes_feature(image_data)}))
return example
def _get_output_filename(output_dir, output_name, index):
"""
构建tfrecord文件的全路径名
output_dir：tfrecord文件保存路径
output_name：生成文件的名字
index：索引
"""
name='%s_%03d.tfrecord' % (output_name, index)
return os.path.join(output_dir,name)
def _process_image(annotations_path,xml_file_name):
"""
根据annotations路径，xml文件名，读取图片并返回图片的bytes型数据
"""
img_name = xml_file_name[:-4]+".jpg"
# print(img_name)
path=os.path.dirname(annotations_path)
print("path：",path)
path=os.path.join(path,DIRECTORY_IMAGES,img_name)
# image_data = tf.io.gfile.GFile(path, 'rb').read()
with tf.io.gfile.GFile(path, 'rb') as f:
image_data = f.read()
return image_data
def int64_feature(value):
"""
Wrapper for inserting int64 features into Example proto.
"""
if not isinstance(value, list):
value = [value]
return tf.train.Feature(int64_list=tf.train.Int64List(value=value))
def float_feature(value):
"""
Wrapper for inserting float features into Example proto.
"""
if not isinstance(value, list):
value = [value]
return tf.train.Feature(float_list=tf.train.FloatList(value=value))
def bytes_feature(value):
"""
Wrapper for inserting bytes features into Example proto.
"""
if not isinstance(value, list):
value = [value]
return tf.train.Feature(bytes_list=tf.train.BytesList(value=value))
def run(annotations_path, output_dir,output_name):
"""
将voc2007数据的一个图片（和对应xml用到信息）写成一个tfrecord文件
annotations_path：Annotations文件夹的路径
output_dir：tfrecord文件保存路径
output_name：生成文件的名字
"""
if annotations_path[-1:]=='/':
annotations_path=annotations_path[0:-1]
if output_dir[-1:]=='/':
output_dir=output_dir[0:-1]
if DEBUG:
print("annotations_path：",annotations_path)
print("output_dir：",output_dir)
xml_file_names = sorted(os.listdir(annotations_path))
name_index=0
for xml_file_name in xml_file_names[:6]:
tf_filename=_get_output_filename(output_dir,output_name,name_index)
image_data=_process_image(annotations_path,xml_file_name)
if DEBUG:
print(xml_file_name)
print(tf_filename)
print("image_data:",type(image_data),len(image_data))
shape, bboxes, labels, labels_text, difficult, truncated=_process_xml(annotations_path,xml_file_name)
example=_build_example(image_data, labels, labels_text, bboxes, shape,difficult, truncated)
with tf.io.TFRecordWriter(tf_filename) as tfrecord_writer:
tfrecord_writer.write(example.SerializeToString())
name_index+=1
def run2(annotations_path, output_dir,output_name):
"""
将voc2007中所有图片（和对应的xml有用信息）写成一个tfrecord文件
annotations_path：Annotations文件夹的路径
output_dir：tfrecord文件保存路径
output_name：生成文件的名字
"""
if annotations_path[-1:]=='/':
annotations_path=annotations_path[0:-1]
if output_dir[-1:]=='/':
output_dir=output_dir[0:-1]
if DEBUG:
print("annotations_path：",annotations_path)
print("output_dir：",output_dir)
xml_file_names = sorted(os.listdir(annotations_path))
name_index=0
with tf.io.TFRecordWriter(_get_output_filename(output_dir,output_name,0)) as tfrecord_writer:
for xml_file_name in xml_file_names[:6]:
tf_filename=_get_output_filename(output_dir,output_name,name_index)
image_data=_process_image(annotations_path,xml_file_name)
if DEBUG:
print(xml_file_name)
print(tf_filename)
print("image_data:",type(image_data),len(image_data))
shape, bboxes, labels, labels_text, difficult, truncated=_process_xml(annotations_path,xml_file_name)
example=_build_example(image_data, labels, labels_text, bboxes, shape,difficult, truncated)
# with tf.io.TFRecordWriter(tf_filename) as tfrecord_writer:
# tfrecord_writer.write(example.SerializeToString())
tfrecord_writer.write(example.SerializeToString())
name_index+=1
annotations_path="/Users/mym/Desktop/数据集/VOCdevkit/VOC2007/Annotations/"
output_dir="/Users/mym/Desktop/数据集/"
output_name="voc_2007_train"
run2(annotations_path,output_dir,output_name)