这几天学习了tensorflow中的Object Detection API,虽然网上有很多资料,但还是采了很多坑,没有找到完整详细的一整套实战记录,所以自己整理这几天的成果。
之前一篇博客OpenCV实现人脸检测详解实现了孙红雷和黄渤的人脸目标检测,这里我们依旧使用该数据集进行深度学习的实战。数据集分为训练集和测试集(只是记录实战过程所以数据集很少只有50张左右)。
数据集准备完毕!!!
标注工具LabelImg的下载地址:https://github.com/tzutalin/labelImg下载,解压,得到LabelImg-master文件。
然后打开终端,进入到LabelImg-master文件下,依次执行以下命令:
conda install pyqt = 5
pyrcc5 -o resources.py resources.qrc
python labelImg.py
既可以打开标注工具进行标注(以后只需要打开终端,进入到LabelImg-master文件夹执行python labelImg.py即可打开标注工具)
点击Open Dir打开图像数据所在的位置,点击Change Save Dir选择标注XML文件的存放位置,即可开始标注。(对训练集和测试集数据均进行标注)
为了标注方便,可以更改data文件夹下的predefined_classes.txt文件,将其更改为自己的类别名称,就可以在方框中选择相应类别,比较方便。
最后的得到相应的XML文件。
数据标注完成!!!
标注数据之后我们得到训练集和数据集的XML文件,并存放在相应的文件夹。下面要把文件转换为Object Detection API可以识别的数据格式。一共分为两步:xml转cvs;cvs转tfrecord。
首先附上xml转cvs的代码xml_to_cvs.py
import os
import glob
import pandas as pd
import xml.etree.ElementTree as ET
def xml_to_csv(path):
xml_list = []
for xml_file in glob.glob(path + '/*.xml'):
tree = ET.parse(xml_file)
root = tree.getroot()
for member in root.findall('object'):
value = (root.find('filename').text,
int(root.find('size')[0].text),
int(root.find('size')[1].text),
member[0].text,
int(member[4][0].text),
int(member[4][1].text),
int(member[4][2].text),
int(member[4][3].text)
)
xml_list.append(value)
column_name = ['filename', 'width', 'height', 'class', 'xmin', 'ymin', 'xmax', 'ymax']
xml_df = pd.DataFrame(xml_list, columns=column_name)
return xml_df
def main():
# XML文件的位置
image_path = 'C:/Users/pang/Desktop/Python_Code/tensorflow_object/xqImg_xml/test'
xml_df = xml_to_csv(image_path)
# 输出位置及输出的文件名
xml_df.to_csv('C:/Users/pang/Desktop/Python_Code/tensorflow_object/xqImg_xml/test.cvs', index=None)
print('Successfully converted xml to csv.')
main()
需要更改的地方有两处:
(1)第28行:XML文件的位置改称自己的XML文件位置
(2)第31行:输出位置及输出的文件名
对train和test都做这样的操作,最后得到test.cvs和train.cvs两个文件。
直接上代码cvs_to_tfrecord.py:
"""
Usage:
# From tensorflow/models/
# Create train data:
python generate_tfrecord.py --csv_input=data/train_labels.csv --output_path=train.record
# Create test data:
python generate_tfrecord.py --csv_input=data/test_labels.csv --output_path=test.record
"""
from __future__ import division
from __future__ import print_function
from __future__ import absolute_import
import os
import io
import pandas as pd
import tensorflow as tf
from PIL import Image
from object_detection.utils import dataset_util
from collections import namedtuple, OrderedDict
flags = tf.app.flags
flags.DEFINE_string('csv_input', '', 'Path to the CSV input')
flags.DEFINE_string('output_path', '', 'Path to output TFRecord')
FLAGS = flags.FLAGS
# TO-DO replace this with label map
def class_text_to_int(row_label):
if row_label == 'sun':
return 1
elif row_label == 'huang':
return 2
else:
None
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.gfile.GFile(os.path.join(path, '{}'.format(group.filename)), 'rb') as fid:
encoded_jpg = fid.read()
encoded_jpg_io = io.BytesIO(encoded_jpg)
image = Image.open(encoded_jpg_io)
width, height = image.size
filename = group.filename.encode('utf8')
image_format = b'jpg'
xmins = []
xmaxs = []
ymins = []
ymaxs = []
classes_text = []
classes = []
for index, row in group.object.iterrows():
xmins.append(row['xmin'] / width)
xmaxs.append(row['xmax'] / width)
ymins.append(row['ymin'] / height)
ymaxs.append(row['ymax'] / height)
classes_text.append(row['class'].encode('utf8'))
classes.append(class_text_to_int(row['class']))
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_jpg),
'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
def main(csv_input, output_path, image_path):
writer = tf.python_io.TFRecordWriter(output_path)
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())
writer.close()
print('Successfully created the TFRecords: {}'.format(output_path))
if __name__ == '__main__':
# CSV文件的位置
csv_input = 'C:/Users/pang/Desktop/Python_Code/tensorflow_object/xqImg_xml/test.cvs'
# TFRecords的输出位置及文件名
output_path = 'C:/Users/pang/Desktop/Python_Code/tensorflow_object/xqImg_xml/test.record'
# 图像数据的位置
image_path = 'C:/Users/pang/Desktop/Python_Code/tensorflow_object/xqImg/test'
main(csv_input, output_path, image_path)
需要更改的地方有四处:
(1)第30行:class_text_to_int函数,根据自己的具体任务增加或删除相应的标签与整数
(2)第100行:CSV文件的位置,即上一步生成的文件
(3)第102行:TFRecords的输出保存位置及文件名
(4)第104行:图像数据的位置,即原始图片的位置
同样对train和test都做这样的操作,最后得到test.record和train.record两个文件,也是我们用于模型训练的数据。
数据集格式转换完成!!!
下一步就可以用自己的数据集训练自己的深度学习目标检测模型。