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

　　最近工作的项目使用了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格式，

<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文件夹，

　　

将样本数据转换为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))  

　　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()  

运行结果如下：

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

　　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  

类似的，我们可以输入如下命令，将验证集和测试集也转换为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

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

训练

　　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

　　无错误则开始训练，等待训练结束，如下：

使用TensorBoard进行监测

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

　　

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

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

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

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

导出训练结果

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

选择相应步数的模型，使用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

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

文件夹内容如下：

其中，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文件夹目录下。

批量保存测试结果

　　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()

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

总结

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