【TensorFlow】基于ssd_mobilenet模型实现目标检测

时间:2021-01-20 04:05:32

  最近工作的项目使用了TensorFlow中的目标检测技术,通过训练自己的样本集得到模型来识别游戏中的物体,在这里总结下。

  本文介绍在Windows系统下,使用TensorFlow的object detection API来训练自己的数据集,所用的模型为ssd_mobilenet,当然也可以使用其他模型,包括ssd_inception、faster_rcnn、rfcnn_resnet等,其中,ssd模型在各种模型中性能最好,所以便采用它来进行训练。

配置环境

  1. 在GitHub上下载所需的models文件,地址:https://github.com/tensorflow/models

  2. 安装pillow、Jupyter、matplotlib、lxml,打开anaconda prompt输入以下命令,并安装成功

pip install pillow
pip install jupyter
pip install matplotlib
pip install lxml

  3. 编译protobuf,object detection API是使用protobuf来训练模型和配置参数的,所以得先编译protobuf,下载地址:https://github.com/google/protobuf/releases,具体配置过程可参考:https://blog.csdn.net/dy_guox/article/details/79081499 。

制作自己的样本集

  1. 下载labelImg,并标注自己收集的图片样本,标注的标签自动保存为xml格式,

【TensorFlow】基于ssd_mobilenet模型实现目标检测

<annotation>
<folder>images1</folder>
<filename>0.png</filename>
<path>C:\Users\White\Desktop\images1\0.png</path>
<source>
<database>Unknown</database>
</source>
<size>
<width>1080</width>
<height>1920</height>
<depth>3</depth>
</size>
<segmented>0</segmented>
<object>
<name>box</name>
<pose>Unspecified</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<bndbox>
<xmin>345</xmin>
<ymin>673</ymin>
<xmax>475</xmax>
<ymax>825</ymax>
</bndbox>
</object>
<object>
<name>box</name>
<pose>Unspecified</pose>
<truncated>0</truncated>
<difficult>0</difficult>
<bndbox>
<xmin>609</xmin>
<ymin>1095</ymin>
<xmax>759</xmax>
<ymax>1253</ymax>
</bndbox>
</object>
</annotation>

  2. 在工程文件夹下新建以下目录,并将所有的样本图片放入images文件夹,将标注保存的xml文件保存到merged_xml文件夹,

  【TensorFlow】基于ssd_mobilenet模型实现目标检测【TensorFlow】基于ssd_mobilenet模型实现目标检测【TensorFlow】基于ssd_mobilenet模型实现目标检测

将样本数据转换为TFRecord格式

  1. 新建train_test_split.py把xml数据集分为了train 、test、 validation三部分,并存储在annotations文件夹中,train为训练集占76.5%,test为测试集10%,validation为验证集13.5%,train_test_split.py代码如下:

import os
import random
import time
import shutil xmlfilepath=r'merged_xml'
saveBasePath=r"./annotations" trainval_percent=0.9
train_percent=0.85
total_xml = os.listdir(xmlfilepath)
num=len(total_xml)
list=range(num)
tv=int(num*trainval_percent)
tr=int(tv*train_percent)
trainval= random.sample(list,tv)
train=random.sample(trainval,tr)
print("train and val size",tv)
print("train size",tr)
# print(total_xml[1])
start = time.time()
# print(trainval)
# print(train)
test_num=0
val_num=0
train_num=0
# for directory in ['train','test',"val"]:
# xml_path = os.path.join(os.getcwd(), 'annotations/{}'.format(directory))
# if(not os.path.exists(xml_path)):
# os.mkdir(xml_path)
# # shutil.copyfile(filePath, newfile)
# print(xml_path)
for i in list:
name=total_xml[i]
# print(i)
if i in trainval: #train and val set
# ftrainval.write(name)
if i in train:
# ftrain.write(name)
# print("train")
# print(name)
# print("train: "+name+" "+str(train_num))
directory="train"
train_num+=1
xml_path = os.path.join(os.getcwd(), 'annotations/{}'.format(directory))
if(not os.path.exists(xml_path)):
os.mkdir(xml_path)
filePath=os.path.join(xmlfilepath,name)
newfile=os.path.join(saveBasePath,os.path.join(directory,name))
shutil.copyfile(filePath, newfile) else:
# fval.write(name)
# print("val")
# print("val: "+name+" "+str(val_num))
directory="validation"
xml_path = os.path.join(os.getcwd(), 'annotations/{}'.format(directory))
if(not os.path.exists(xml_path)):
os.mkdir(xml_path)
val_num+=1
filePath=os.path.join(xmlfilepath,name)
newfile=os.path.join(saveBasePath,os.path.join(directory,name))
shutil.copyfile(filePath, newfile)
# print(name)
else: #test set
# ftest.write(name)
# print("test")
# print("test: "+name+" "+str(test_num))
directory="test"
xml_path = os.path.join(os.getcwd(), 'annotations/{}'.format(directory))
if(not os.path.exists(xml_path)):
os.mkdir(xml_path)
test_num+=1
filePath=os.path.join(xmlfilepath,name)
newfile=os.path.join(saveBasePath,os.path.join(directory,name))
shutil.copyfile(filePath, newfile)
# print(name) # End time
end = time.time()
seconds=end-start
print("train total : "+str(train_num))
print("validation total : "+str(val_num))
print("test total : "+str(test_num))
total_num=train_num+val_num+test_num
print("total number : "+str(total_num))
print( "Time taken : {0} seconds".format(seconds))

【TensorFlow】基于ssd_mobilenet模型实现目标检测

【TensorFlow】基于ssd_mobilenet模型实现目标检测

  2. 把xml转换成csv文件,新建xml_to_csv.py,,运行代码前,需要建一个data目录,用来放生成的csv文件,结果和代码如下:

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()
# print(root)
print(root.find('filename').text)
for member in root.findall('object'):
value = (root.find('filename').text,
int(root.find('size')[0].text), #width
int(root.find('size')[1].text), #height
member[0].text,
int(member[4][0].text),
int(float(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():
for directory in ['train','test','validation']:
xml_path = os.path.join(os.getcwd(), 'annotations/{}'.format(directory))
# image_path = os.path.join(os.getcwd(), 'merged_xml')
xml_df = xml_to_csv(xml_path)
# xml_df.to_csv('whsyxt.csv', index=None)
xml_df.to_csv('data/whsyxt_{}_labels.csv'.format(directory), index=None)
print('Successfully converted xml to csv.') main()

运行结果如下:

【TensorFlow】基于ssd_mobilenet模型实现目标检测

在data文件夹下生成的csv文件:

【TensorFlow】基于ssd_mobilenet模型实现目标检测

  3. 生成tfrecords文件,python文件名为generate_tfrecord.py,代码如下:

 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,filename):
if row_label == 'person':
return 1
elif row_label == 'investigator':
return 2
elif row_label == 'collector':
return 3
elif row_label == 'wolf':
return 4
elif row_label == 'skull':
return 5
elif row_label == 'inferno':
return 6
elif row_label == 'stone_blame':
return 7
elif row_label == 'green_jelly':
return 8
elif row_label == 'blue_jelly':
return 9
elif row_label == 'box':
return 10
elif row_label == 'golden_box':
return 11
elif row_label == 'silver_box':
return 12
elif row_label == 'jar':
return 13
elif row_label == 'purple_jar':
return 14
elif row_label == 'purple_weapon':
return 15
elif row_label == 'blue_weapon':
return 16
elif row_label == 'blue_shoe':
return 17
elif row_label == 'blue_barde':
return 18
elif row_label == 'blue_ring':
return 19
elif row_label == 'badge':
return 20
elif row_label == 'dragon_stone':
return 21
elif row_label == 'lawn':
return 22
elif row_label == 'mine':
return 23
elif row_label == 'portal':
return 24
elif row_label == 'tower':
return 25
elif row_label == 'hero_stone':
return 26
elif row_label == 'oracle_stone':
return 27
elif row_label == 'arena':
return 28
elif row_label == 'gold_ore':
return 29
elif row_label == 'relic':
return 30
elif row_label == 'ancient':
return 31
elif row_label == 'house':
return 32
else:
print("------------------nonetype:", filename)
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'png'
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'], group.filename)) 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(os.getcwd(), 'images')
examples = pd.read_csv(FLAGS.csv_input)
grouped = split(examples, 'filename')
num=0
for group in grouped:
num+=1
tf_example = create_tf_example(group, path)
writer.write(tf_example.SerializeToString())
if(num%100==0): #每完成100个转换,打印一次
print(num) 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()

其中,20~83行应改成在样本集中标注的类别,我这里总共有32个类别,字符串row_label应与labelImg中标注的名称相同。

现将训练集转换为tfrecord格式,输入如下命令:

python generate_tfrecord.py --csv_input=data/whsyxt_train_labels.csv --output_path=data/whsyxt_train.tfrecord  

【TensorFlow】基于ssd_mobilenet模型实现目标检测

【TensorFlow】基于ssd_mobilenet模型实现目标检测

类似的,我们可以输入如下命令,将验证集和测试集也转换为tfrecord格式,

python generate_tfrecord.py --csv_input=data/whsyxt_validation_labels.csv --output_path=data/whsyxt_validation.tfrecord 
python generate_tfrecord.py --csv_input=data/whsyxt_test_labels.csv --output_path=data/whsyxt_test.tfrecord

都执行成功后,获得如下文件,

【TensorFlow】基于ssd_mobilenet模型实现目标检测

训练

  1. 在工程文件夹data目录下创建标签分类的配置文件(label_map.pbtxt),需要检测几种目标,将创建几个id,代码如下:

item {
id: 1 # id从1开始编号
name: 'person'
}
item {
id: 2
name: 'investigator'
}
item {
id: 3
name: 'collector'
}
item {
id: 4
name: 'wolf'
}
item {
id: 5
name: 'skull'
}
item {
id: 6
name: 'inferno'
}
item {
id: 7
name: 'stone_blame'
}
item {
id: 8
name: 'green_jelly'
}
item {
id: 9
name: 'blue_jelly'
}
item {
id: 10
name: 'box'
}
item {
id: 11
name: 'golden_box'
}
item {
id: 12
name: 'silver_box'
}
item {
id: 13
name: 'jar'
}
item {
id: 14
name: 'purple_jar'
}
item {
id: 15
name: 'purple_weapon'
}
item {
id: 16
name: 'blue_weapon'
}
item {
id: 17
name: 'blue_shoe'
}
item {
id: 18
name: 'blue_barde'
}
item {
id: 19
name: 'blue_ring'
}
item {
id: 20
name: 'badge'
}
item {
id: 21
name: 'dragon_stone'
}
item {
id: 22
name: 'lawn'
}
item {
id: 23
name: 'mine'
}
item {
id: 24
name: 'portal'
}
item {
id: 25
name: 'tower'
}
item {
id: 26
name: 'hero_stone'
}
item {
id: 27
name: 'oracle_stone'
}
item {
id: 28
name: 'arena'
}
item {
id: 29
name: 'gold_ore'
}
item {
id: 30
name: 'relic'
}
item {
id: 31
name: 'ancient'
}
item {
id: 32
name: 'house'
}

  2. 配置管道配置文件,找到 models\research\object_detection\samples\configs\ssd_inception_v2_pets.config文件,复制到data文件夹下,修改之后代码如下:

 # SSD with Mobilenet v1, configured for Oxford-IIIT Pets Dataset.
# Users should configure the fine_tune_checkpoint field in the train config as
# well as the label_map_path and input_path fields in the train_input_reader and
# eval_input_reader. Search for "PATH_TO_BE_CONFIGURED" to find the fields that
# should be configured. model {
ssd {
num_classes: 32
box_coder {
faster_rcnn_box_coder {
y_scale: 10.0
x_scale: 10.0
height_scale: 5.0
width_scale: 5.0
}
}
matcher {
argmax_matcher {
matched_threshold: 0.45
unmatched_threshold: 0.35
ignore_thresholds: false
negatives_lower_than_unmatched: true
force_match_for_each_row: true
}
}
similarity_calculator {
iou_similarity {
}
}
anchor_generator {
ssd_anchor_generator {
num_layers: 6
min_scale: 0.2
max_scale: 0.95
aspect_ratios: 1.0
aspect_ratios: 2.0
aspect_ratios: 0.5
aspect_ratios: 3.0
aspect_ratios: 0.3333
}
}
image_resizer {
fixed_shape_resizer {
height: 300
width: 300
}
}
box_predictor {
convolutional_box_predictor {
min_depth: 0
max_depth: 0
num_layers_before_predictor: 0
use_dropout: false
dropout_keep_probability: 0.8
kernel_size: 1
box_code_size: 4
apply_sigmoid_to_scores: false
conv_hyperparams {
activation: RELU_6,
regularizer {
l2_regularizer {
weight: 0.00004
}
}
initializer {
truncated_normal_initializer {
stddev: 0.03
mean: 0.0
}
}
batch_norm {
train: true,
scale: true,
center: true,
decay: 0.9997,
epsilon: 0.001,
}
}
}
}
feature_extractor {
type: 'ssd_mobilenet_v1'
min_depth: 16
depth_multiplier: 1.0
conv_hyperparams {
activation: RELU_6,
regularizer {
l2_regularizer {
weight: 0.00004
}
}
initializer {
truncated_normal_initializer {
stddev: 0.03
mean: 0.0
}
}
batch_norm {
train: true,
scale: true,
center: true,
decay: 0.9997,
epsilon: 0.001,
}
}
}
loss {
classification_loss {
weighted_sigmoid {
}
}
localization_loss {
weighted_smooth_l1 {
}
}
hard_example_miner {
num_hard_examples: 3000
iou_threshold: 0.99
loss_type: CLASSIFICATION
max_negatives_per_positive: 3
min_negatives_per_image: 0
}
classification_weight: 1.0
localization_weight: 1.0
}
normalize_loss_by_num_matches: true
post_processing {
batch_non_max_suppression {
score_threshold: 1e-8
iou_threshold: 0.6
max_detections_per_class: 100
max_total_detections: 100
}
score_converter: SIGMOID
}
}
} train_config: {
batch_size: 24
optimizer {
rms_prop_optimizer: {
learning_rate: {
exponential_decay_learning_rate {
initial_learning_rate: 0.004
decay_steps: 1000
decay_factor: 0.95
}
}
momentum_optimizer_value: 0.9
decay: 0.9
epsilon: 1.0
}
}
#fine_tune_checkpoint: "PATH_TO_BE_CONFIGURED/model.ckpt"
from_detection_checkpoint: false
# Note: The below line limits the training process to 200K steps, which we
# empirically found to be sufficient enough to train the pets dataset. This
# effectively bypasses the learning rate schedule (the learning rate will
# never decay). Remove the below line to train indefinitely.
num_steps: 40000
data_augmentation_options {
random_horizontal_flip {
}
}
data_augmentation_options {
ssd_random_crop {
}
}
} train_input_reader: {
tf_record_input_reader {
input_path: "E:/Project/object-detection-Game-yellow/data/whsyxt_train.tfrecord"
}
label_map_path: "E:/Project/object-detection-Game-yellow/data/label_map.pbtxt"
} eval_config: {
num_examples: 2000
# 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: "E:/Project/object-detection-Game-yellow/data/whsyxt_validation.tfrecord"
}
label_map_path: "E:/Project/object-detection-Game-yellow/data/label_map.pbtxt"
shuffle: false
num_readers: 1
}

这里需要修改的几处有:第9行,改为自己标注的总类别数;第175行,改为训练集tfrecord文件的路径;第177行和191行,改为自己label_map的路径;第189行,改为验证集tfrecord文件的路径。

我们可以在这个管道文件中设置网络的各种学习参数,如:第141行设置批次大小,第145~148行设置学习率和退化率,第162行设置训练的总步数等等。

  3. 开始训练,将object_detection\train.py文件复制到工程目录下进行训练即可,命令如下:

python train.py --logtostderr --pipeline_config_path=E:/Project/object-detection-Game-yellow/data/ssd_mobilenet_v1_pets.config --train_dir=E:/Project/object-detection-Game-yellow/data

  无错误则开始训练,等待训练结束,如下:

【TensorFlow】基于ssd_mobilenet模型实现目标检测

使用TensorBoard进行监测

  1.在输入训练的命令后,data文件夹下会生成如下文件,该文件存放训练过程中的中间数据,并可以用图形化的方式展现出来。

  【TensorFlow】基于ssd_mobilenet模型实现目标检测

  2. 新打开一个命令提示符窗口,首先激活TensorFlow,然后输入如下命令:

tensorboard --logdir==training:your_log_dir --host=127.0.0.1

其中,your_log_dir为工程目录中存放训练结果的文件夹目录,把目录地址拷贝出来将其替代。

  3.打开浏览器,在地址栏输入:localhost:6006,即可显示tensorboard:

【TensorFlow】基于ssd_mobilenet模型实现目标检测

导出训练结果

  1.训练过程中将在training目录下生成一堆model.ckpt-*的文件,如下:

【TensorFlow】基于ssd_mobilenet模型实现目标检测

选择相应步数的模型,使用export_inference_graph.py(其在object detection目录下)导出pb文件,命令如下:

python export_inference_graph.py --pipeline_config_path=E:\Project\object-detection-Game-yellow\data\ssd_mobilenet_v1_pets.config --trained_checkpoint_prefix ./data/model.ckpt-30000 --output_directory ./data/exported_model_directory

【TensorFlow】基于ssd_mobilenet模型实现目标检测

运行命令后,会在工程的data目录下生成名为exported_model_directory文件夹,如下:

【TensorFlow】基于ssd_mobilenet模型实现目标检测

文件夹内容如下:

【TensorFlow】基于ssd_mobilenet模型实现目标检测

其中,frozen_inference_graph.pb就是我们以后将要使用的模型结果。

获取测试图片

  1. 新建test_images文件夹和get_testImages.py文件,并加入以下代码,如下:

 from PIL import Image
import os.path
import glob annotations_test_dir = "E:\\Project\\object-detection-Game-yellow\\annotations\\test\\"
Images_dir = "E:\\Project\\object-detection-Game-yellow\\Images"
test_images_dir = "E:\\Project\\object-detection-Game-yellow\\test_images"
i = 0
for xmlfile in os.listdir(annotations_test_dir):
(filepath, tempfilename) = os.path.split(xmlfile)
(shotname, extension) = os.path.splitext(tempfilename)
xmlname = shotname
for pngfile in os.listdir(Images_dir):
(filepath, tempfilename) = os.path.split(pngfile)
(pngname, extension) = os.path.splitext(tempfilename)
if pngname == xmlname:
img = Image.open(Images_dir+"\\" + pngname + ".png")
img.save(os.path.join(test_images_dir, os.path.basename(pngfile)))
print(pngname)
i += 1
print(i)

第5、6、7行,分别是annotations\test文件夹路径、Images文件夹路径和test_images文件夹路径,运行python文件,获取测试图片并存储到test_images文件夹目录下。

【TensorFlow】基于ssd_mobilenet模型实现目标检测

【TensorFlow】基于ssd_mobilenet模型实现目标检测

批量保存测试结果

  1. 在工程目录下新建results文件夹和get_allTestResults.py文件并加入如下代码,我们将使用前面训练出的模型批量测试test_images文件夹中的图片并保存到results文件夹中,

 # -*- coding: utf-8 -*-
import os
from PIL import Image
import time
import tensorflow as tf
from PIL import Image
import numpy as np
import os
import six.moves.urllib as urllib
import sys
import tarfile
import zipfile
import time from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
# plt.switch_backend('Agg')
from utils import label_map_util from utils import visualization_utils as vis_util PATH_TO_TEST_IMAGES = "E:\\Project\\object-detection-Game-yellow\\test_images\\"
MODEL_NAME = 'E:/Project/object-detection-Game-yellow/data'
PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb'
PATH_TO_LABELS = MODEL_NAME+'/label_map.pbtxt'
NUM_CLASSES = 32
PATH_TO_RESULTS = "E:\\Project\\object-detection-Game-yellow\\results2\\" 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) def save_object_detection_result():
IMAGE_SIZE = (12, 8)
# Load a (frozen) Tensorflow model into memory.
detection_graph = tf.Graph()
with detection_graph.as_default():
od_graph_def = tf.GraphDef()
# loading ckpt file to graph
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
with detection_graph.as_default():
with tf.Session(graph=detection_graph) as sess:
start = time.time()
for test_image in os.listdir(PATH_TO_TEST_IMAGES):
image = Image.open(PATH_TO_TEST_IMAGES + test_image)
# the array based representation of the image will be used later in order to prepare the
# result image with boxes and labels on it.
image_np = load_image_into_numpy_array(image)
# Expand dimensions since the model expects images to have shape: [1, None, None, 3]
image_np_expanded = np.expand_dims(image_np, axis=0)
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.
boxes = 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.
scores = detection_graph.get_tensor_by_name('detection_scores:0')
classes = detection_graph.get_tensor_by_name('detection_classes:0')
num_detections = detection_graph.get_tensor_by_name('num_detections:0')
# Actual detection.
(boxes, scores, classes, num_detections) = sess.run(
[boxes, scores, classes, num_detections],
feed_dict={image_tensor: image_np_expanded})
# Visualization of the results of a detection.
vis_util.visualize_boxes_and_labels_on_image_array(
image_np,
np.squeeze(boxes),
np.squeeze(classes).astype(np.int32),
np.squeeze(scores),
category_index,
use_normalized_coordinates=True,
line_thickness=8) final_score = np.squeeze(scores)
count = 0
for i in range(100):
if scores is None or final_score[i] > 0.5:
count = count + 1
print()
print("the count of objects is: ", count)
(im_width, im_height) = image.size
for i in range(count):
# print(boxes[0][i])
y_min = boxes[0][i][0] * im_height
x_min = boxes[0][i][1] * im_width
y_max = boxes[0][i][2] * im_height
x_max = boxes[0][i][3] * im_width
x = int((x_min + x_max) / 2)
y = int((y_min + y_max) / 2)
if category_index[classes[0][i]]['name'] == "tower":
print("this image has a tower!")
y = int((y_max - y_min) / 4 * 3 + y_min)
print("object{0}: {1}".format(i, category_index[classes[0][i]]['name']),
',Center_X:', x, ',Center_Y:', y)
# print(x_min,y_min,x_max,y_max)
plt.figure(figsize=IMAGE_SIZE)
plt.imshow(image_np)
picName = test_image.split('/')[-1]
# print(picName)
plt.savefig(PATH_TO_RESULTS + picName)
print(test_image + ' succeed') end = time.time()
seconds = end - start
print("Time taken : {0} seconds".format(seconds)) save_object_detection_result()

【TensorFlow】基于ssd_mobilenet模型实现目标检测

【TensorFlow】基于ssd_mobilenet模型实现目标检测

最后,我们就可以使用results中的测试结果进行准确率的计算,查看训练效果及后期优化。

总结

【TensorFlow】基于ssd_mobilenet模型实现目标检测

转载请注明出处:https://www.cnblogs.com/White-xzx/p/9503203.html