Tensorflow训练自己的SSD目标检测器
2019-05-21 本文已影响0人
技术大渣渣
本文主要讲解如何利用Tensorflow object detection api从0到1训练自己的目标检测器。环境是win10+vscode，ubuntu16.04配置参照官方教程。除第一部分配置在windows和ubuntu下不同外，第二部分的的步骤和代码在两个平台上通用。。。。。。
一.安装Tensorflow Object Detection API
1.配置protobuf，下载protoc-3.6.0-win32.zip（https://github.com/google/protobuf/releases），解压后将bin文件夹加入到系统环境变量，cmd中输入protoc然后回车不报错说明安装ok
2.下载tensorflow models模块，地址https://github.com/tensorflow/models
3.安装tensorflow model 以及slim
a.将protoc编译成py,在models-master/research目录下cmd运行：
protoc object_detection/protos/*.proto --python_out=.
b.进一步安装slim,在models-master/research/slim目录下重命名BUILD文件为BUILD1（运行setup.py时windows下新建的文件夹名build和文件BUILD冲突），然后cmd运行：
python setup.py install
c.配置path环境变量：包括models-master、research和slim三个文件夹的路径。
d.在models-master/research目录下cmd运行：
python setup.py install
安装model模块（一定要首先执行a，因为执行此步时需要将a生成的py文件复制到'C:\Users\user\AppData\Local\Programs\Python\Python36\Lib\site-packages\object_detection-0.1-py3.6.egg\object_detection'，若不执行，则会提示cannot import name 'anchor_generator_pb2' ）
另外，我还遇到了ModuleNotFoundError: No module named 'tensorflow.python.saved_model.model_utils'，出现原因：升级了tensorflow2.0,又降级到1.13会出现该问题，解决方法：删除tensorflow_estimator重新安装
pip uninstall tensorflow_estimator
pip install tensorflow_estimator
4.测试环境，在models-master/research目录下cmd运行：python object_detection/builders/model_builder_test.py，返回ok则表明环境ok！
二.训练自己的目标检测器
1.数据集准备
2.给数据集打标签，采用LabelImg,软件使用可参考《手把手教你图片打标》
原始图像存储在“images”文件夹，打标签后生成的xml文件存储在“annotations”文件夹
1.PNG
3.读取所有的xml列表，运行xml_to_csv.py：
2.PNG
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():
image_path = os.path.join(os.getcwd(), 'annotations')
xml_df = xml_to_csv(image_path)
xml_df.to_csv('data/raccoon_labels.csv', index=None)
print('Successfully converted xml to csv.')
main()
4.将列表数据分为train、validation、test三部分，供训练、验证和测试（也可只分成训练和验证两部分），运行split data.py
import numpy as np
import pandas as pd
np.random.seed(1)
csv_file_url = 'data/raccoon_labels.csv'
full_data = pd.read_csv(csv_file_url)
total_file_number = len(full_data)
print("There are total {} examples in this dataset.".format(total_file_number))
full_data.head() #Viewing the first 5 lines
num_train = int(total_file_number*0.85)
num_validation = int(total_file_number*0.1)
num_test = total_file_number-num_train-num_validation
assert num_train + num_validation + num_test <= total_file_number, "Not enough examples for your choice."
print("Looks good! {} for train, {} for validation and {} for test.".format(num_train, num_validation, num_test))
index_train = np.random.choice(total_file_number, size=num_train, replace=False)
index_validation_test = np.setdiff1d(list(range(total_file_number)), index_train)
index_validation = np.random.choice(index_validation_test, size=num_validation, replace=False)
index_test = np.setdiff1d(index_validation_test, index_validation)
train = full_data.iloc[index_train]
validation = full_data.iloc[index_validation]
test = full_data.iloc[index_test]
train.to_csv('data/data_train.csv', index=None)
validation.to_csv("data/data_validation.csv", index=None)
test.to_csv('data/data_test.csv', index=None)
print("All done!")
5.生成tfrecord，运行generate_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', 'data/data_test.csv', 'Path to the CSV input')
flags.DEFINE_string('output_path', 'data/data_test.record', 'Path to output TFRecord')
flags.DEFINE_string('image_dir', 'images/', 'Path to images')
FLAGS = flags.FLAGS
# TO-DO replace this with label map
def class_text_to_int(row_label):
if row_label == 'raccoon':
return 1
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(_):
writer = tf.python_io.TFRecordWriter(FLAGS.output_path)
path = os.path.join(FLAGS.image_dir)
examples = pd.read_csv(FLAGS.csv_input)
grouped = split(examples, 'filename')
for group in grouped:
tf_example = create_tf_example(group, path)
writer.write(tf_example.SerializeToString())
writer.close()
output_path = os.path.join(os.getcwd(), FLAGS.output_path)
print('Successfully created the TFRecords: {}'.format(output_path))
if __name__ == '__main__':
tf.app.run()
6.框架需要我们定义好我们的类别ID与类别名称的关系，通常用pbtxt格式文件保存，我们在~/master/data/config目录下新建一个名为raccoon_label_map.pbtxt的文本文件（仿照 TensorFlow models/research/object_detection/data 文件夹下的 .pbtxt 文件编写自己的 .pbtxt 文件）,内容如下:
item {
id: 1
name: 'raccoon'
}
7.使用ssd_mobilenet_v1 网络，并且我们需要使用迁移学习的加速我们的训练过程，我们将使用ssd_mobilenet_v1_coco作为预训练模型来进行finetune训练，下载地址 ssd_mobilenet_v1_coco，下载解压后同样放在~/master/data/config/目录下
8.复制TensorFlow
models/research/object_detection/samples/configs 下的ssd_mobilenet_v1_coco.config 到 ~/master/data/config/下，重命名为ssd_mobilenet_v1_raccoon.config,并做如下修改：
num_classes: 90 改为 num_classes: 1
...
type: 'ssd_mobilenet_v1' 改为 type: 'ssd_mobilenet_v1' #若选择其他预训练模型则对应更改
...
fine_tune_checkpoint: "PATH_TO_BE_CONFIGURED/model.ckpt" 更改为步骤7中下载的预训练模型地址
...
train_input_reader: {
tf_record_input_reader {
input_path: "PATH_TO_BE_CONFIGURED/mscoco_train.record-?????-of-00100"
}
label_map_path: "PATH_TO_BE_CONFIGURED/mscoco_label_map.pbtxt"
} #更改为自己的地址
...
eval_config: {
num_examples: 21
# Note: The below line limits the evaluation process to 10 evaluations.
# Remove the below line to evaluate indefinitely.
max_evals: 10
}
eval_input_reader: {
tf_record_input_reader {
input_path: "PATH_TO_BE_CONFIGURED/mscoco_val.record-?????-of-00010"
}
label_map_path: "PATH_TO_BE_CONFIGURED/mscoco_label_map.pbtxt"
shuffle: false
num_readers: 1
} #更改为自己地址
# 很多超参数也可以调
9.在Tensorflow models-master/research/object_detection(tensorflow object detection api)路径下,执行下面的相关命令进行训练：
python model_main.py --model_dir=./master/model/train/
--pipeline_config_path=./master/data/config/ssd_mobilenet_v1_raccoon.config
提示 No module named 'pycocotools',windows 下安装即可：
pip install git+https://github.com/philferriere/cocoapi.git#subdirectory=PythonAPI
若windows下权限不够就用管理员权限运行cmd训练
10.在Tensorflow models-master/research/object_detection/model_main.py
中加入如下代码，进行log打印:
tf.logging.set_verbosity(tf.logging.INFO)
11.查看实时训练曲线
tensorboard --logdir=/home/.../model/train/
12.单张图片测试模型效果，运行test_model_image.py:
import numpy as np
import os
import sys
import tensorflow as tf
import cv2
#add tensorflow object detection api(must!!!!!)
sys.path
sys.path.insert(0, './python/tensorflow-ssd/models-master/research/object_detection/')
# This is needed since the notebook is stored in the object_detection folder.
sys.path.append("..")
from object_detection.utils import ops as utils_ops
print(tf.__version__)
# if tf.__version__ < '1.4.0':
# raise ImportError('Please upgrade your tensorflow installation to v1.4.* or later!')
# This is needed to display the images.
from utils import label_map_util
from utils import visualization_utils as vis_util
# What model to download.
MODEL_NAME = 'ssd_mobilenet_v1_coco_2017_11_17'
# Path to frozen detection graph. This is the actual model that is used for the object detection.
PATH_TO_CKPT = "./python/tensorflow-ssd/raccoon_dataset-master/model/freezed_pb/frozen_inference_graph.pb"
# List of the strings that is used to add correct label for each box.
PATH_TO_LABELS = os.path.join('data/config', 'raccoon_label_map.pbtxt')
NUM_CLASSES = 1
#Load a (frozen) Tensorflow model into memory
detection_graph = tf.Graph()
with detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
#Loading label map
label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
category_index = label_map_util.create_category_index(categories)
#Helper code
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)
#Detection
def run_inference_for_single_image(image, graph):
with graph.as_default():
with tf.Session() as sess:
# Get handles to input and output tensors
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:
# The following processing is only for single image
detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])
detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])
# Reframe is required to translate mask from box coordinates to image coordinates and fit the image size.
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)
# Follow the convention by adding back the batch dimension
tensor_dict['detection_masks'] = tf.expand_dims(
detection_masks_reframed, 0)
image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')
# Run inference
output_dict = sess.run(tensor_dict,
feed_dict={image_tensor: np.expand_dims(image, 0)})
# all outputs are float32 numpy arrays, so convert types as appropriate
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
image_np = cv2.imread("./images/raccoon-96.jpg")
# Expand dimensions since the model expects images to have shape: [1, None, None, 3]
image_np_expanded = np.expand_dims(image_np, axis=0)
# Actual detection.
output_dict = run_inference_for_single_image(image_np, detection_graph)
# Visualization of the results of a detection.
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)
cv2.imshow("test",np.array(image_np,dtype=np.uint8))
cv2.waitKey()
13.视频测试模型效果，运行test_model_video.py(需进一步把类别显示加进去):
import time
import cv2
import numpy as np
import tensorflow as tf
#--------------Model preparation----------------
# Path to frozen detection graph. This is the actual model that is used for
# the object detection.
PATH_TO_CKPT = 'path_to_your_frozen_inference_graph.pb'
# Load a (frozen) Tensorflow model into memory
detection_graph = tf.Graph()
with detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
# Each box represents a part of the image where a particular
# object was detected.
gboxes = detection_graph.get_tensor_by_name('detection_boxes:0')
# Each score represent how level of confidence for each of the objects.
# Score is shown on the result image, together with the class label.
gscores = detection_graph.get_tensor_by_name('detection_scores:0')
gclasses = detection_graph.get_tensor_by_name('detection_classes:0')
gnum_detections = detection_graph.get_tensor_by_name('num_detections:0')
# TODO: Add class names showing in the image
def detect_image_objects(image, sess, detection_graph):
# Expand dimensions since the model expects images to have
# shape: [1, None, None, 3]
image_np_expanded = np.expand_dims(image, axis=0)
# Actual detection.
(boxes, scores, classes, num_detections) = sess.run(
[gboxes, gscores, gclasses, gnum_detections],
feed_dict={image_tensor: image_np_expanded})
# Visualization of the results of a detection.
boxes = np.squeeze(boxes)
scores = np.squeeze(scores)
height, width = image.shape[:2]
for i in range(boxes.shape[0]):
if (scores is None or
scores[i] > 0.5):
ymin, xmin, ymax, xmax = boxes[i]
ymin = int(ymin * height)
ymax = int(ymax * height)
xmin = int(xmin * width)
xmax = int(xmax * width)
score = None if scores is None else scores[i]
font = cv2.FONT_HERSHEY_SIMPLEX
text_x = np.max((0, xmin - 10))
text_y = np.max((0, ymin - 10))
cv2.putText(image, 'Detection score: ' + str(score),
(text_x, text_y), font, 0.4, (0, 255, 0))
cv2.rectangle(image, (xmin, ymin), (xmax, ymax),
(0, 255, 0), 2)
return image
with detection_graph.as_default():
with tf.Session(graph=detection_graph) as sess:
video_path = 'path_to_your_video'
capture = cv2.VideoCapture(video_path)
while capture.isOpened():
if cv2.waitKey(30) & 0xFF == ord('q'):
break
ret, frame = capture.read()
if not ret:
break
t_start = time.clock()
detect_image_objects(frame, sess, detection_graph)
t_end = time.clock()
print('detect time per frame: ', t_end - t_start)
cv2.imshow('detected', frame)
capture.release()
cv2.destroyAllWindows()
后续代码和小样本数据集会放在我的github上，敬请期待。。。。
资料：
1.Tensorflow object detection 官方文档
2.https://www.jianshu.com/p/86894ccaa407
