本博客的内容是:在 BlackWalnut Labs 完成交通标志标志识别实验的学习笔记
其中主要的部分是 BlackWalnut Labs 的实验过程介绍,少部分自己实验过程中的心得,所以就作为转载自 BlackWalnut 的一篇博文吧
说明:以下内容均在 BlackWalnut Labs 提供的环境下进行,如果想在自己的环境下运行,可参考TensorFlow Object Detection
项目
交通标志识别系统是 ADAS(高级驾驶辅助系统)场景的组成部分之一,能使小车能自动识别前方的交通标志并提醒驾驶人。
Google 在 github 上提供的了一个对象识别的项目,可以利用该项目训练各种对象识别模型。这些模型不需要进行预处理,可以直接识别图像中的交通标志,并返回交通标志类型和在图中的位置。
对象识别模型是一种混合技术,先用图像分割技术把每一个对象分割出来,然后用图像分类技术识别分割出的对象是否是训练集中的对象。
在次实验中,将要测试 TensorFlow Object Detection
项目,使用自己拍摄的交通标志训练集重训练对象识别模型,并将建立以下技能:
• 收集交通标志图片
• 使用 LabelImg 为图像添加标签
• 将图像和标签转换为 TFRecord
• 选择并部署 ssd_mobilenet_v1 模型
• 训练模型
• 使用测试图片执行预测
• BlackWalnut Labs. AI Cloud 访问账号
• Windows 10(可选)
Object Detection
的数据集和分类器的数据集不同,图像中需要包含识别物体在内的整个场景。例如需要识别的交通标志类型是 “前方行人”,那么训练集图像只要包含 “前方行人” 这一部分即可,并且同一张图片中可以出现多个交通标志。
在真正场景中使用对象识别技术需要上百甚至上千张图片作为训练集,但在 Codelab中因为是个人手动收集,而且场景固定,大约每组交通标志 100 张图像就可以有效果了。
新建名为Codelab 的目录作为项目目录,在该目录下新建名为 data/TrafficSign 的目录用于存放数据集。由 Codelab 提供收集完毕的训练集,可以在 Tools/datasets/image 目录下找到。该数据集包含 200 张图片,每张图片中都包含 2 个较为明显的交通标志。切换到 Terminal 界面,进入 Codelab 目录,复制图片文件到 Codelab/data/TrafficSign 目录下。
cp -r ../Tools/datasets/image/* data/TrafficSign
当然你也可以使用自己拍摄的数据集,只要拍摄的图片清晰且符合下面的要求即可。
① 拍摄图片的分辨率为 480*480
② 每张图片中至少包含 1 个清晰的交通标志
③ 每种交通标志至少出现 100 次
训练对象识别模型时,不仅需要图像作为数据集,还需要告诉模型图像中对象的位置以及其类型。也就是需要生成图片对应的XML文件。为了准确和方便的标出对象在图中的位置,可以使用Windows10平台下的 LabelImg 工具。点击 这里 下载工具的对应版本
这里使用的是 Windows_v1.6 版本。解压下载的文件后,直接运行 labelImg.exe
文件即可。软件的界面是英文的,不过简单易懂。基本功能解释如下:
选项 | 功能 |
---|---|
Open | 表示打开单张图片 |
Open Dir | 表示打开一个目录,该软件支持读取一个目录中所有的图片。 |
Change Save Dir | 表示更改保存的目录,该软件支持自动保存。 |
Next Image | 表示下一张图片。 |
Prev Image | 表示上一张图片 |
Save | 表示保存图片 |
Create RectBox | 表示创建标签,对象标签是用矩形框表示的 |
Duplicate RectBox | 表示复制标签 |
Delete RectBox | 表示删除标签 |
使用方法:
用鼠标拖动的方式创建标签框后,程序会弹出定义标签名称的对话框。当一个标签设置完毕时,右上角会出现添加的标签,此时保存标签文件时会自动生成 XML 类型的标签文件。对于标签文件,这里有个建议 —— 图片文件与标签文件前缀名相同。
标定数据集是一项很大的工作,为了节省时间,Codelab 提供了标定好的标签文件,前缀名与图片相同,可以在 Tools/datasets/label 目录下找到。该标定好的数据集可以搭配自己拍摄的交通标志数据集使用。切换到 Terminal 界面,进入 Codelab 目录,复制标签文件到 Codelab/data/TrafficSign 目录下。
cp -r ../Tools/datasets/label/* data/TrafficSign
在 Codelab/data 目录下新建名为 label.pbtxt
的文件,内容根据实际使用的数据集定,其文件内容如下:
item {
id: 1
name: 'Car'
}
item {
id: 2
name: '10T'
}
item {
id: 3
name: 'BanRight'
}
item {
id: 4
name: 'Pedestrian'
}
item {
id: 5
name: 'TurnLeft'
}
item {
id: 6
name: 'Speaker'
}
item {
id: 7
name: 'Crosswalk'
}
item {
id: 8
name: 'TurnAround'
}
item {
id: 9
name: 'GoStraight'
}
item {
id: 10
name: 'GoStraightOrRight'
}
item {
id: 11
name: 'GoStraightOrLeft'
}
item {
id: 12
name: 'BanLeft'
}
item {
id: 13
name: 'TurnRight'
}
item {
id: 14
name: 'BanSpeaker'
}
item {
id: 15
name: 'NoParking'
}
item {
id: 16
name: 'BanTurnAround'
}
item {
id: 17
name: 'BanCar'
}
item {
id: 18
name: 'BanStraightAndLeft'
}
item {
id: 19
name: 'SlowDown'
}
item {
id: 20
name: 'BanLeftAndRight'
}
item {
id: 21
name: 'Limit40'
}
item {
id: 22
name: 'BanStraightAndRight'
}
图像和标签都可以通过流式传输。Google 规范了训练集的格式 TFREcord
,在 Codelab 目录下新建名为 Image2TFRecord.py
的 Python3 文件。在开始编写生成程序前,先做本地依赖目录结构的导入和配置。
import sys
sys.path.append('/home/jovyan/Appendix/tensorflow_models/research')
sys.path.append('/home/jovyan/Appendix/tensorflow_models/research/slim')
在导入本地依赖目录之后,再进行依赖包的导入。
import numpy as np
import tensorflow as tf
import os
import xml.dom.minidom
from object_detection.utils import dataset_util
在 Object Detection
项目中, Google 提供了一个生成 TFRecord
的函数,不过这个函数一次执行只能处理一张图片,但可以基于此批量处理数据集。
def create_cat_tf_example(encoded_cat_image_data):
"""Creates a tf.Example proto from sample cat image.
Args:
encoded_cat_image_data: The jpg encoded data of the cat image.
Returns:
example: The created tf.Example.
"""
height = 1032.0
width = 1200.0
filename = 'example_cat.jpg'
image_format = b'jpg'
xmins = [322.0 / 1200.0]
xmaxs = [1062.0 / 1200.0]
ymins = [174.0 / 1032.0]
ymaxs = [761.0 / 1032.0]
classes_text = ['Cat']
classes = [1]
tf_example = tf.train.Example(features=tf.train.Features(feature={
'image/height': dataset_util.int64_feature(height),
'image/width': dataset_util.int64_feature(width),
'image/filename': dataset_util.bytes_feature(filename),
'image/source_id': dataset_util.bytes_feature(filename),
'image/encoded': dataset_util.bytes_feature(encoded_image_data),
'image/format': dataset_util.bytes_feature(image_format),
'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),
'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),
'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),
'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),
'image/object/class/text': dataset_util.bytes_list_feature(classes_text),
'image/object/class/label': dataset_util.int64_list_feature(classes),
}))
return tf_example
在这个函数中,只需要把 height
、width
、filename
、image_format
、xmins
、xmaxs
、ymins
、ymaxs
、classes_text
、classes
修改为依据传入参数动态获取即可将其用于批量处理。
参数 | 含义 |
---|---|
height | 表示待处理图片的高度 |
width | 表示待处理图片的宽度 |
filename | 表示待处理图片的名称,用于区分不同的图片 |
image_format | 表示图片的格式,通常为 jpeg 或者 png |
xmins | 表示图片中对象标注框的左边框的横坐标值,该变量为数组,数组中的一个值表示一个对象标签。 |
xmax | 表示图片中对象标注框的右边框的横坐标值,该变量为数组,数组中的一个值表示一个对象标签。 |
ymin | 表示图片中对象标注框的上边框的纵坐标值,该变量为数组,数组中的一个值表示一个对象标签。 |
ymax | 表示图片中对象标注框的下边框的纵坐标值,该变量为数组,数组中的一个值表示一个对象标签。 |
classes_text | 表示图片中对象的名称,该变量为数组,数组中的一个值表示一个对象标签。 |
classes | 表示图片中对象名称所对应的 ID 号,这个编号针对所有用于训练的数据集来说,从 1 开始编号,该变量为数组,数组中的一个值表示一个对象标签。 |
修改转化函数,新增函数的传入参数,并将其作为上述参数的字典。
def create_tf_example(image,labelInfo,imageInfo):
height = imageInfo['height'] # Image height
width = imageInfo['width'] # Image width
filename = labelInfo['name'] # Filename of the image. Empty if image is not from file
encoded_image_data = image # Encoded image bytes
image_format = b'jpeg' # b'jpeg' or b'png'
xmins = labelInfo['xmins'] # List of normalized left x coordinates in bounding box (1 per box)
xmaxs = labelInfo['xmaxs'] # List of normalized right x coordinates in bounding box
# (1 per box)
ymins = labelInfo['ymins'] # List of normalized top y coordinates in bounding box (1 per box)
ymaxs = labelInfo['ymaxs'] # List of normalized bottom y coordinates in bounding box
# (1 per box)
classes_text = labelInfo['classes_text'] # List of string class name of bounding box (1 per box)
classes = labelInfo['classes'] # List of integer class id of bounding box (1 per box)
tf_example = tf.train.Example(features=tf.train.Features(feature={
'image/height': dataset_util.int64_feature(height),
'image/width': dataset_util.int64_feature(width),
'image/filename': dataset_util.bytes_feature(filename),
'image/source_id': dataset_util.bytes_feature(filename),
'image/encoded': dataset_util.bytes_feature(encoded_image_data),
'image/format': dataset_util.bytes_feature(image_format),
'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),
'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),
'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),
'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),
'image/object/class/text': dataset_util.bytes_list_feature(classes_text),
'image/object/class/label': dataset_util.int64_list_feature(classes),
}))
return tf_example
在调用该函数前,设置全局图像参数,即图像的宽高像素。
imageInfo = {
'height':480,
'width':480
}
每调用一次转换函数都会返回一个存储着二进制文件的变量,针对每一张图片都做一次转换,把所有转换出来的变量累积形成一个新的变量后输出成文件,这个文件即为 TFRecord
。TensorFlow 提供了能够将其保存到本地的函数,如下:
trainFileWriter = tf.python_io.TFRecordWriter('data/train.record')
trainDirString = 'data/TrafficSign'
trainDir = os.listdir(trainDirString)
for fileName in trainDir:
if(fileName.endswith('.jpg')):
# ... 稍后将在这里插入代码
with tf.gfile.GFile(img_path, 'rb') as fid:
encoded_jpg = fid.read()
tf_example = create_tf_example(encoded_jpg,labelInfo,imageInfo)
trainFileWriter.write(tf_example.SerializeToString())
trainFileWriter.close()
使用 Python 的 XML 库读取标签文件并填充进字典文件中(即前面 … 的地方)。
DOMTree = xml.dom.minidom.parse(trainDirString + '/' + fileName[0:len(fileName)-4] + '.xml')
collection = DOMTree.documentElement
objects = collection.getElementsByTagName("object")
labelInfo = {
'name':bytes(fileName, encoding = "utf8"),
'xmins':[],
'xmaxs':[],
'ymins':[],
'ymaxs':[],
'classes_text':[],
'classes':[]
}
textDict = {
'Car':1,
'10T':2,
'BanRight':3,
'Pedestrian':4,
'TurnLeft':5,
'Speaker':6,
'Crosswalk':7,
'TurnAround':8,
'GoStraight':9,
'GoStraightOrRight':10,
'GoStraightOrLeft':11,
'BanLeft':12,
'TurnRight':13,
'BanSpeaker':14,
'NoParking':15,
'BanTurnAround':16,
'BanCar':17,
'BanStraightAndLeft':18,
'SlowDown':19,
'BanLeftAndRight':20,
'Limit40':21,
'BanStraightAndRight':22
}
for object in objects:
labelInfo['xmins'].append(float(object.getElementsByTagName('xmin')[0].childNodes[0].data)/imageInfo['width'])
labelInfo['xmaxs'].append(float(object.getElementsByTagName('xmax')[0].childNodes[0].data)/imageInfo['width'])
labelInfo['ymins'].append(float(object.getElementsByTagName('ymin')[0].childNodes[0].data)/imageInfo['height'])
labelInfo['ymaxs'].append(float(object.getElementsByTagName('ymax')[0].childNodes[0].data)/imageInfo['height'])
labelInfo['classes_text'].append(bytes(object.getElementsByTagName('name')[0].childNodes[0].data, encoding = "utf8"))
labelInfo['classes'].append(textDict[object.getElementsByTagName('name')[0].childNodes[0].data])
print(labelInfo)
img_path = trainDirString + '/' + fileName
如果使用的是自己的数据集,需要根据自己的数据集中的对象编号进行编号,修改 textDict
的部分,该对象结构如下。
{
"对象名1":1,//对象ID,数字类型,从 1 开始编号
"对象名2":2,//对象ID,数字类型,从 1 开始编号,
//...
}
完整示例代码见下面的附录。执行程序,即可在 data 目录中找到 train.record
。
现在可以查看一下 LabelImg
生成的 XML 文件,其中包含了图片尺寸的描述和标签的描述。
参数 | 含义 |
---|---|
width | 表示图片的宽度 |
height | 表示图片的高度 |
filename | 表示图片的名称,每一个对象都用 object 标签表示。 |
name | 对象名 |
xmin | 表示图片中对象标注框的左边框的横坐标值。 |
xmax | 表示图片中对象标注框的又边框的横坐标值. |
ymin | 表示图片中对象标注框的上边框的纵坐标值. |
ymax | 表示图片中对象标注框的下边框的纵坐标值. |
<annotation>
<folder>TrafficSignfolder>
<filename>1521536642525.jpgfilename>
<path>D:\Project\Python_Projects\ObjectDetectionTraining\data\TrafficSign\1521536642525.jpgpath>
<source>
<database>Unknowndatabase>
source>
<size>
<width>480width>
<height>480height>
<depth>3depth>
size>
<segmented>0segmented>
<object>
<name>Carname>
<pose>Unspecifiedpose>
<truncated>0truncated>
<difficult>0difficult>
<bndbox>
<xmin>112xmin>
<ymin>188ymin>
<xmax>185xmax>
<ymax>262ymax>
bndbox>
object>
<object>
<name>10Tname>
<pose>Unspecifiedpose>
<truncated>0truncated>
<difficult>0difficult>
<bndbox>
<xmin>380xmin>
<ymin>189ymin>
<xmax>457xmax>
<ymax>265ymax>
bndbox>
object>
annotation>
实现对象识别的模型有很多,Google 根据这些模型的论文用 TensorFlow 进行了重现,并向广大开发者开放了预训练的模型。开发者使用自己的训练集将这些模型进行重训练后,即可使用对象识别技术。 在本次实验中,我们选择 ssd_mobilenet_v1_coco
模型。
Codelab 在 Tools/model 目录下提供了 ssd_mobilenet_v1_coco
的镜像,在 Codelab 目录下新建名为 model的目录,切换到 Terminal 界面,进入 Codelab 目录,把模型复制到该目录下。
cp -r ../Tools/model/* model
model 目录下自带训练配置文件 pipeline.config
,用 jupyter
进入编辑。
修改
num_classes
后的参数为本次训练中对象的总种类数。(如果使用的是提供的数据集,这里修改为 22)
修改fine_tune_checkpoint
后的路径为/home/jovyan/Codelab/model/model.ckpt
修改train_input_reader/input_path
后的参数为/home/jovyan/Codelab/data/train.record
修改eval_input_reader/input_path
后的参数为/home/jovyan/Codelab/data/train.record
为了方便,这里是直接使用训练集作为测试集。
修改label_map_path
后的参数为/home/jovyan/Codelab/data/label.pbtxt
保存并退出编辑后,用 Terminal 进入项目目录,进入 /home/jovyan/Appendix/tensorflow_models/research 目录。
cd ~/Appendix/tensorflow_models/research
在该目录下执行下面的训练程序。默认情况下,程序会训练 50000 步后结束,可以通过修改NUM_TRAIN_STEPS
参数,改变执行次数。
PIPELINE_CONFIG_PATH='/home/jovyan/Codelab/model/pipeline.config'
MODEL_DIR='/home/jovyan/Codelab/data'
NUM_TRAIN_STEPS=50000
SAMPLE_1_OF_N_EVAL_EXAMPLES=1
python3 object_detection/model_main.py \
--pipeline_config_path=${PIPELINE_CONFIG_PATH} \
--model_dir=${MODEL_DIR} \
--num_train_steps=${NUM_TRAIN_STEPS} \
--sample_1_of_n_eval_examples=$SAMPLE_1_OF_N_EVAL_EXAMPLES \
--alsologtostderr
最新版的 TensorFlow Object Detection
项目删去了每步的 loss 输出,但在程序开始执行之后,可以通过 tensorboard 查看训练的情况。默认情况下,程序会训练 50000 步后结束,可以通过修改NUM_TRAIN_STEPS
参数,改变执行次数。
默认情况下,程序会训练 50000 步后结束,但在训练的过程中会输出中间训练的 checkpoint
,可以不等待训练完全结束就提前查看训练的中间结果。
在开始编写预测程序之前,先将中间训练的 checkpoint
转化为可调用的 tensorflow graph
文件。其中的TRAINED_CKPT_PREFIX
参数,需要填写在上面训练过程中,生成的ckpt
文件路径,如TRAINED_CKPT_PREFIX='/home/jovyan/Codelab/data/model.ckpt-12345'
。文件代码如下:
INPUT_TYPE=image_tensor
PIPELINE_CONFIG_PATH='/home/jovyan/Codelab/model/pipeline.config'
TRAINED_CKPT_PREFIX='/home/jovyan/Codelab/data/model.ckpt-选择的步数'
EXPORT_DIR='/home/jovyan/Codelab/output'
python3 object_detection/export_inference_graph.py \
--input_type=${INPUT_TYPE} \
--pipeline_config_path=${PIPELINE_CONFIG_PATH} \
--trained_checkpoint_prefix=${TRAINED_CKPT_PREFIX} \
--output_directory=${EXPORT_DIR}
在 Codelab 目录下新建名为 Inference.py
的 Python3 文件,按照惯例,先添加系统依赖库。
import sys
sys.path.append('/home/jovyan/Appendix/tensorflow_models/research')
sys.path.append('/home/jovyan/Appendix/tensorflow_models/research/slim')
sys.path.append('/home/jovyan/Appendix/tensorflow_models/research/object_detection')
import numpy as np
import os
import tensorflow as tf
from distutils.version import StrictVersion
from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
from PIL import Image
from object_detection.utils import ops as utils_ops
if StrictVersion(tf.__version__) < StrictVersion('1.9.0'):
raise ImportError('Please upgrade your TensorFlow installation to v1.9.* or later!')
%matplotlib inline
from utils import label_map_util
from utils import visualization_utils as vis_util
定义路径变量,指向导出的图文件和标签所对应的文件。
PATH_TO_FROZEN_GRAPH = 'output/frozen_inference_graph.pb'
PATH_TO_LABELS = os.path.join('data', 'label.pbtxt')
读取图文件,保存在 detection_graph
变量中。
detection_graph = tf.Graph()
with detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(PATH_TO_FROZEN_GRAPH, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
读取标签文件,保存在 category_index
文件中。
category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)
编写工具函数,用于读取图片并保存到 numpy
数组中。
def load_image_into_numpy_array(image):
(im_width, im_height) = image.size
return np.array(image.getdata()).reshape(
(im_height, im_width, 3)).astype(np.uint8)
设置预测图像的路径变量。
PATH_TO_TEST_IMAGES_DIR = 'test_images'
TEST_IMAGE_PATHS = [ os.path.join(PATH_TO_TEST_IMAGES_DIR, 'image{}.jpg'.format(i)) for i in range(1, 2) ]
IMAGE_SIZE = (12, 8)
编写预测函数,执行计算图并返回原始结果。
def run_inference_for_single_image(image, graph):
with graph.as_default():
with tf.Session() as sess:
ops = tf.get_default_graph().get_operations()
all_tensor_names = {output.name for op in ops for output in op.outputs}
tensor_dict = {}
for key in [
'num_detections', 'detection_boxes', 'detection_scores',
'detection_classes', 'detection_masks'
]:
tensor_name = key + ':0'
if tensor_name in all_tensor_names:
tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(
tensor_name)
if 'detection_masks' in tensor_dict:
detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])
detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])
real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)
detection_boxes = tf.slice(detection_boxes, [0, 0], [real_num_detection, -1])
detection_masks = tf.slice(detection_masks, [0, 0, 0], [real_num_detection, -1, -1])
detection_masks_reframed = utils_ops.reframe_box_masks_to_image_masks(
detection_masks, detection_boxes, image.shape[0], image.shape[1])
detection_masks_reframed = tf.cast(
tf.greater(detection_masks_reframed, 0.5), tf.uint8)
tensor_dict['detection_masks'] = tf.expand_dims(
detection_masks_reframed, 0)
image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')
output_dict = sess.run(tensor_dict,
feed_dict={image_tensor: np.expand_dims(image, 0)})
output_dict['num_detections'] = int(output_dict['num_detections'][0])
output_dict['detection_classes'] = output_dict[
'detection_classes'][0].astype(np.uint8)
output_dict['detection_boxes'] = output_dict['detection_boxes'][0]
output_dict['detection_scores'] = output_dict['detection_scores'][0]
if 'detection_masks' in output_dict:
output_dict['detection_masks'] = output_dict['detection_masks'][0]
return output_dict
针对每一张测试图片都执行一次预测,并可视化显示在jupyter
for image_path in TEST_IMAGE_PATHS:
image = Image.open(image_path)
image_np = load_image_into_numpy_array(image)
image_np_expanded = np.expand_dims(image_np, axis=0)
output_dict = run_inference_for_single_image(image_np, detection_graph)
vis_util.visualize_boxes_and_labels_on_image_array(
image_np,
output_dict['detection_boxes'],
output_dict['detection_classes'],
output_dict['detection_scores'],
category_index,
instance_masks=output_dict.get('detection_masks'),
use_normalized_coordinates=True,
line_thickness=8)
plt.figure(figsize=IMAGE_SIZE)
plt.imshow(image_np)
在 Codelab 目录下新建名为 test_images 的目录,把要测试的图片放在该目录下并以 image1.jpg
、image2.jpg
、image3.jpg
(以此类推)的方式命名。
切换到 Inference.py
的 Python3 文件中,参照下面的方式修改程序。
#找到这段代码
TEST_IMAGE_PATHS = [ os.path.join(PATH_TO_TEST_IMAGES_DIR, 'image{}.jpg'.format(i)) for i in range(1, 2) ]
#修改为这段代码
TEST_IMAGE_PATHS = [ os.path.join(PATH_TO_TEST_IMAGES_DIR,'image{}.jpg'.format(i)) for i in range(1, 替换为测试图片的数量+1) ]
执行程序,即可在 jupyter
界面看到当前训练程度下预测的结果。
#file_name: Image2TFRecord.py
import sys
sys.path.append('/home/jovyan/Appendix/tensorflow_models/research')
sys.path.append('/home/jovyan/Appendix/tensorflow_models/research/slim')
import numpy as np
import tensorflow as tf
import os
import xml.dom.minidom
from object_detection.utils import dataset_util
imageInfo = {
'height':480,
'width':480
}
def create_tf_example(image,labelInfo,imageInfo):
height = imageInfo['height'] # Image height
width = imageInfo['width'] # Image width
filename = labelInfo['name'] # Filename of the image. Empty if image is not from file
encoded_image_data = image # Encoded image bytes
image_format = b'jpeg' # b'jpeg' or b'png'
xmins = labelInfo['xmins'] # List of normalized left x coordinates in bounding box (1 per box)
xmaxs = labelInfo['xmaxs'] # List of normalized right x coordinates in bounding box
# (1 per box)
ymins = labelInfo['ymins'] # List of normalized top y coordinates in bounding box (1 per box)
ymaxs = labelInfo['ymaxs'] # List of normalized bottom y coordinates in bounding box
# (1 per box)
classes_text = labelInfo['classes_text'] # List of string class name of bounding box (1 per box)
classes = labelInfo['classes'] # List of integer class id of bounding box (1 per box)
tf_example = tf.train.Example(features=tf.train.Features(feature={
'image/height': dataset_util.int64_feature(height),
'image/width': dataset_util.int64_feature(width),
'image/filename': dataset_util.bytes_feature(filename),
'image/source_id': dataset_util.bytes_feature(filename),
'image/encoded': dataset_util.bytes_feature(encoded_image_data),
'image/format': dataset_util.bytes_feature(image_format),
'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),
'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),
'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),
'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),
'image/object/class/text': dataset_util.bytes_list_feature(classes_text),
'image/object/class/label': dataset_util.int64_list_feature(classes),
}))
return tf_example
trainFileWriter = tf.python_io.TFRecordWriter('data/train.record')
trainDirString = 'data/TrafficSign'
trainDir = os.listdir(trainDirString)
for fileName in trainDir:
if(fileName.endswith('.jpg')):
DOMTree = xml.dom.minidom.parse(trainDirString + '/' + fileName[0:len(fileName)-4] + '.xml')
collection = DOMTree.documentElement
objects = collection.getElementsByTagName("object")
labelInfo = {
'name':bytes(fileName, encoding = "utf8"),
'xmins':[],
'xmaxs':[],
'ymins':[],
'ymaxs':[],
'classes_text':[],
'classes':[]
}
textDict = {
'Car':1,
'10T':2,
'BanRight':3,
'Pedestrian':4,
'TurnLeft':5,
'Speaker':6,
'Crosswalk':7,
'TurnAround':8,
'GoStraight':9,
'GoStraightOrRight':10,
'GoStraightOrLeft':11,
'BanLeft':12,
'TurnRight':13,
'BanSpeaker':14,
'NoParking':15,
'BanTurnAround':16,
'BanCar':17,
'BanStraightAndLeft':18,
'SlowDown':19,
'BanLeftAndRight':20,
'Limit40':21,
'BanStraightAndRight':22
}
for object in objects:
labelInfo['xmins'].append(float(object.getElementsByTagName('xmin')[0].childNodes[0].data)/imageInfo['width'])
labelInfo['xmaxs'].append(float(object.getElementsByTagName('xmax')[0].childNodes[0].data)/imageInfo['width'])
labelInfo['ymins'].append(float(object.getElementsByTagName('ymin')[0].childNodes[0].data)/imageInfo['height'])
labelInfo['ymaxs'].append(float(object.getElementsByTagName('ymax')[0].childNodes[0].data)/imageInfo['height'])
labelInfo['classes_text'].append(bytes(object.getElementsByTagName('name')[0].childNodes[0].data, encoding = "utf8"))
labelInfo['classes'].append(textDict[object.getElementsByTagName('name')[0].childNodes[0].data])
print(labelInfo)
img_path = trainDirString + '/' + fileName
with tf.gfile.GFile(img_path, 'rb') as fid:
encoded_jpg = fid.read()
tf_example = create_tf_example(encoded_jpg,labelInfo,imageInfo)
trainFileWriter.write(tf_example.SerializeToString())
trainFileWriter.close()