【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,代码如下：

  1 from __future__ import division  
  2 from __future__ import print_function  
  3 from __future__ import absolute_import  
  4   
  5 import os  
  6 import io  
  7 import pandas as pd  
  8 import tensorflow as tf  
  9   
 10 from PIL import Image  
 11 from object_detection.utils import dataset_util  
 12 from collections import namedtuple, OrderedDict  
 13   
 14 flags = tf.app.flags  
 15 flags.DEFINE_string('csv_input', '', 'Path to the CSV input')  
 16 flags.DEFINE_string('output_path', '', 'Path to output TFRecord')  
 17 FLAGS = flags.FLAGS  
 18 # TO-DO replace this with label map  
 19 def class_text_to_int(row_label,filename):
 20     if row_label == 'person':
 21         return 1  
 22     elif row_label == 'investigator':
 23         return 2 
 24     elif row_label == 'collector':
 25         return 3
 26     elif row_label == 'wolf':
 27         return 4
 28     elif row_label == 'skull':
 29         return 5
 30     elif row_label == 'inferno':
 31         return 6
 32     elif row_label == 'stone_blame':
 33         return 7
 34     elif row_label == 'green_jelly':
 35         return 8
 36     elif row_label == 'blue_jelly':
 37         return 9
 38     elif row_label == 'box':
 39         return 10
 40     elif row_label == 'golden_box':
 41         return 11
 42     elif row_label == 'silver_box':
 43         return 12
 44     elif row_label == 'jar':
 45         return 13
 46     elif row_label == 'purple_jar':
 47         return 14
 48     elif row_label == 'purple_weapon':
 49         return 15
 50     elif row_label == 'blue_weapon':
 51         return 16
 52     elif row_label == 'blue_shoe':
 53         return 17
 54     elif row_label == 'blue_barde':
 55         return 18
 56     elif row_label == 'blue_ring':
 57         return 19
 58     elif row_label == 'badge':
 59         return 20
 60     elif row_label == 'dragon_stone':
 61         return 21
 62     elif row_label == 'lawn':
 63         return 22
 64     elif row_label == 'mine':
 65         return 23
 66     elif row_label == 'portal':
 67         return 24
 68     elif row_label == 'tower':
 69         return 25
 70     elif row_label == 'hero_stone':
 71         return 26
 72     elif row_label == 'oracle_stone':
 73         return 27
 74     elif row_label == 'arena':
 75         return 28
 76     elif row_label == 'gold_ore':
 77         return 29
 78     elif row_label == 'relic':
 79         return 30
 80     elif row_label == 'ancient':
 81         return 31
 82     elif row_label == 'house':
 83         return 32
 84     else:
 85         print("------------------nonetype:", filename)
 86         None
 87   
 88 def split(df, group):  
 89     data = namedtuple('data', ['filename', 'object'])  
 90     gb = df.groupby(group)  
 91     return [data(filename, gb.get_group(x)) for filename, x in zip(gb.groups.keys(), gb.groups)]  
 92   
 93   
 94 def create_tf_example(group, path):  
 95     with tf.gfile.GFile(os.path.join(path, '{}'.format(group.filename)), 'rb') as fid:  
 96         encoded_jpg = fid.read()  
 97     encoded_jpg_io = io.BytesIO(encoded_jpg)  
 98     image = Image.open(encoded_jpg_io)  
 99     width, height = image.size  
100   
101     filename = group.filename.encode('utf8')  
102     image_format = b'png'  
103     xmins = []  
104     xmaxs = []  
105     ymins = []  
106     ymaxs = []  
107     classes_text = []  
108     classes = []  
109   
110     for index, row in group.object.iterrows():  
111         xmins.append(row['xmin'] / width)  
112         xmaxs.append(row['xmax'] / width)  
113         ymins.append(row['ymin'] / height)  
114         ymaxs.append(row['ymax'] / height)  
115         classes_text.append(row['class'].encode('utf8'))  
116         classes.append(class_text_to_int(row['class'], group.filename))
117   
118     tf_example = tf.train.Example(features=tf.train.Features(feature={  
119         'image/height': dataset_util.int64_feature(height),  
120         'image/width': dataset_util.int64_feature(width),  
121         'image/filename': dataset_util.bytes_feature(filename),  
122         'image/source_id': dataset_util.bytes_feature(filename),  
123         'image/encoded': dataset_util.bytes_feature(encoded_jpg),  
124         'image/format': dataset_util.bytes_feature(image_format),  
125         'image/object/bbox/xmin': dataset_util.float_list_feature(xmins),  
126         'image/object/bbox/xmax': dataset_util.float_list_feature(xmaxs),  
127         'image/object/bbox/ymin': dataset_util.float_list_feature(ymins),  
128         'image/object/bbox/ymax': dataset_util.float_list_feature(ymaxs),  
129         'image/object/class/text': dataset_util.bytes_list_feature(classes_text),  
130         'image/object/class/label': dataset_util.int64_list_feature(classes),  
131     }))  
132     return tf_example  
133   
134   
135 def main(_):  
136     writer = tf.python_io.TFRecordWriter(FLAGS.output_path)  
137     path = os.path.join(os.getcwd(), 'images')  
138     examples = pd.read_csv(FLAGS.csv_input)  
139     grouped = split(examples, 'filename')  
140     num=0  
141     for group in grouped:  
142         num+=1  
143         tf_example = create_tf_example(group, path)  
144         writer.write(tf_example.SerializeToString())  
145         if(num%100==0):  #每完成100个转换，打印一次  
146             print(num)  
147   
148     writer.close()  
149     output_path = os.path.join(os.getcwd(), FLAGS.output_path)  
150     print('Successfully created the TFRecords: {}'.format(output_path))  
151   
152   
153 if __name__ == '__main__':  
154     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文件夹下，修改之后代码如下：

  1 # SSD with Mobilenet v1, configured for Oxford-IIIT Pets Dataset.
  2 # Users should configure the fine_tune_checkpoint field in the train config as
  3 # well as the label_map_path and input_path fields in the train_input_reader and
  4 # eval_input_reader. Search for "PATH_TO_BE_CONFIGURED" to find the fields that
  5 # should be configured.
  6 
  7 model {
  8   ssd {
  9     num_classes: 32
 10     box_coder {
 11       faster_rcnn_box_coder {
 12         y_scale: 10.0
 13         x_scale: 10.0
 14         height_scale: 5.0
 15         width_scale: 5.0
 16       }
 17     }
 18     matcher {
 19       argmax_matcher {
 20         matched_threshold: 0.45
 21         unmatched_threshold: 0.35
 22         ignore_thresholds: false
 23         negatives_lower_than_unmatched: true
 24         force_match_for_each_row: true
 25       }
 26     }
 27     similarity_calculator {
 28       iou_similarity {
 29       }
 30     }
 31     anchor_generator {
 32       ssd_anchor_generator {
 33         num_layers: 6
 34         min_scale: 0.2
 35         max_scale: 0.95
 36         aspect_ratios: 1.0
 37         aspect_ratios: 2.0
 38         aspect_ratios: 0.5
 39         aspect_ratios: 3.0
 40         aspect_ratios: 0.3333
 41       }
 42     }
 43     image_resizer {
 44       fixed_shape_resizer {
 45         height: 300
 46         width: 300
 47       }
 48     }
 49     box_predictor {
 50       convolutional_box_predictor {
 51         min_depth: 0
 52         max_depth: 0
 53         num_layers_before_predictor: 0
 54         use_dropout: false
 55         dropout_keep_probability: 0.8
 56         kernel_size: 1
 57         box_code_size: 4
 58         apply_sigmoid_to_scores: false
 59         conv_hyperparams {
 60           activation: RELU_6,
 61           regularizer {
 62             l2_regularizer {
 63               weight: 0.00004
 64             }
 65           }
 66           initializer {
 67             truncated_normal_initializer {
 68               stddev: 0.03
 69               mean: 0.0
 70             }
 71           }
 72           batch_norm {
 73             train: true,
 74             scale: true,
 75             center: true,
 76             decay: 0.9997,
 77             epsilon: 0.001,
 78           }
 79         }
 80       }
 81     }
 82     feature_extractor {
 83       type: 'ssd_mobilenet_v1'
 84       min_depth: 16
 85       depth_multiplier: 1.0
 86       conv_hyperparams {
 87         activation: RELU_6,
 88         regularizer {
 89           l2_regularizer {
 90             weight: 0.00004
 91           }
 92         }
 93         initializer {
 94           truncated_normal_initializer {
 95             stddev: 0.03
 96             mean: 0.0
 97           }
 98         }
 99         batch_norm {
100           train: true,
101           scale: true,
102           center: true,
103           decay: 0.9997,
104           epsilon: 0.001,
105         }
106       }
107     }
108     loss {
109       classification_loss {
110         weighted_sigmoid {
111         }
112       }
113       localization_loss {
114         weighted_smooth_l1 {
115         }
116       }
117       hard_example_miner {
118         num_hard_examples: 3000
119         iou_threshold: 0.99
120         loss_type: CLASSIFICATION
121         max_negatives_per_positive: 3
122         min_negatives_per_image: 0
123       }
124       classification_weight: 1.0
125       localization_weight: 1.0
126     }
127     normalize_loss_by_num_matches: true
128     post_processing {
129       batch_non_max_suppression {
130         score_threshold: 1e-8
131         iou_threshold: 0.6
132         max_detections_per_class: 100
133         max_total_detections: 100
134       }
135       score_converter: SIGMOID
136     }
137   }
138 }
139 
140 train_config: {
141   batch_size: 24
142   optimizer {
143     rms_prop_optimizer: {
144       learning_rate: {
145         exponential_decay_learning_rate {
146           initial_learning_rate: 0.004
147           decay_steps: 1000
148           decay_factor: 0.95
149         }
150       }
151       momentum_optimizer_value: 0.9
152       decay: 0.9
153       epsilon: 1.0
154     }
155   }
156   #fine_tune_checkpoint: "PATH_TO_BE_CONFIGURED/model.ckpt"
157   from_detection_checkpoint: false
158   # Note: The below line limits the training process to 200K steps, which we
159   # empirically found to be sufficient enough to train the pets dataset. This
160   # effectively bypasses the learning rate schedule (the learning rate will
161   # never decay). Remove the below line to train indefinitely.
162   num_steps: 40000
163   data_augmentation_options {
164     random_horizontal_flip {
165     }
166   }
167   data_augmentation_options {
168     ssd_random_crop {
169     }
170   }
171 }
172 
173 train_input_reader: {
174   tf_record_input_reader {
175     input_path: "E:/Project/object-detection-Game-yellow/data/whsyxt_train.tfrecord"
176   }
177   label_map_path: "E:/Project/object-detection-Game-yellow/data/label_map.pbtxt"
178 }
179 
180 eval_config: {
181   num_examples: 2000
182   # Note: The below line limits the evaluation process to 10 evaluations.
183   # Remove the below line to evaluate indefinitely.
184   max_evals: 10
185 }
186 
187 eval_input_reader: {
188   tf_record_input_reader {
189     input_path: "E:/Project/object-detection-Game-yellow/data/whsyxt_validation.tfrecord"
190   }
191   label_map_path: "E:/Project/object-detection-Game-yellow/data/label_map.pbtxt"
192   shuffle: false
193   num_readers: 1
194 }

这里需要修改的几处有：第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文件，并加入以下代码，如下：

 1 from PIL import Image
 2 import os.path
 3 import glob
 4 
 5 annotations_test_dir = "E:\\Project\\object-detection-Game-yellow\\annotations\\test\\"
 6 Images_dir = "E:\\Project\\object-detection-Game-yellow\\Images"
 7 test_images_dir = "E:\\Project\\object-detection-Game-yellow\\test_images"
 8 i = 0
 9 for xmlfile in os.listdir(annotations_test_dir):
10     (filepath, tempfilename) = os.path.split(xmlfile)
11     (shotname, extension) = os.path.splitext(tempfilename)
12     xmlname = shotname
13     for pngfile in os.listdir(Images_dir):
14         (filepath, tempfilename) = os.path.split(pngfile)
15         (pngname, extension) = os.path.splitext(tempfilename)
16         if pngname == xmlname:
17              img = Image.open(Images_dir+"\\" + pngname + ".png")
18              img.save(os.path.join(test_images_dir, os.path.basename(pngfile)))
19              print(pngname)
20              i += 1
21 print(i)

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

批量保存测试结果

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

  1 # -*- coding: utf-8 -*-
  2 import os
  3 from PIL import Image
  4 import time
  5 import tensorflow as tf
  6 from PIL import Image
  7 import numpy as np
  8 import os
  9 import six.moves.urllib as urllib
 10 import sys
 11 import tarfile
 12 import zipfile
 13 import time
 14 
 15 from collections import defaultdict
 16 from io import StringIO
 17 from matplotlib import pyplot as plt
 18 # plt.switch_backend('Agg')
 19 from utils import label_map_util
 20 
 21 from utils import visualization_utils as vis_util
 22 
 23 PATH_TO_TEST_IMAGES = "E:\\Project\\object-detection-Game-yellow\\test_images\\"
 24 MODEL_NAME = 'E:/Project/object-detection-Game-yellow/data'
 25 PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb'
 26 PATH_TO_LABELS = MODEL_NAME+'/label_map.pbtxt'
 27 NUM_CLASSES = 32
 28 PATH_TO_RESULTS = "E:\\Project\\object-detection-Game-yellow\\results2\\"
 29 
 30 
 31 def load_image_into_numpy_array(image):
 32     (im_width, im_height) = image.size
 33     return np.array(image.getdata()).reshape((im_height, im_width, 3)).astype(np.uint8)
 34 
 35 
 36 def save_object_detection_result():
 37     IMAGE_SIZE = (12, 8)
 38     # Load a (frozen) Tensorflow model into memory.
 39     detection_graph = tf.Graph()
 40     with detection_graph.as_default():
 41         od_graph_def = tf.GraphDef()
 42         # loading ckpt file to graph
 43         with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
 44             serialized_graph = fid.read()
 45             od_graph_def.ParseFromString(serialized_graph)
 46             tf.import_graph_def(od_graph_def, name='')
 47     # Loading label map
 48     label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
 49     categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES,
 50                                                                 use_display_name=True)
 51     category_index = label_map_util.create_category_index(categories)
 52     # Helper code
 53     with detection_graph.as_default():
 54         with tf.Session(graph=detection_graph) as sess:
 55             start = time.time()
 56             for test_image in os.listdir(PATH_TO_TEST_IMAGES):
 57                 image = Image.open(PATH_TO_TEST_IMAGES + test_image)
 58                 # the array based representation of the image will be used later in order to prepare the
 59                 # result image with boxes and labels on it.
 60                 image_np = load_image_into_numpy_array(image)
 61                 # Expand dimensions since the model expects images to have shape: [1, None, None, 3]
 62                 image_np_expanded = np.expand_dims(image_np, axis=0)
 63                 image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
 64                 # Each box represents a part of the image where a particular object was detected.
 65                 boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
 66                 # Each score represent how level of confidence for each of the objects.
 67                 # Score is shown on the result image, together with the class label.
 68                 scores = detection_graph.get_tensor_by_name('detection_scores:0')
 69                 classes = detection_graph.get_tensor_by_name('detection_classes:0')
 70                 num_detections = detection_graph.get_tensor_by_name('num_detections:0')
 71                 # Actual detection.
 72                 (boxes, scores, classes, num_detections) = sess.run(
 73                     [boxes, scores, classes, num_detections],
 74                     feed_dict={image_tensor: image_np_expanded})
 75                 # Visualization of the results of a detection.
 76                 vis_util.visualize_boxes_and_labels_on_image_array(
 77                     image_np,
 78                     np.squeeze(boxes),
 79                     np.squeeze(classes).astype(np.int32),
 80                     np.squeeze(scores),
 81                     category_index,
 82                     use_normalized_coordinates=True,
 83                     line_thickness=8)
 84 
 85                 final_score = np.squeeze(scores)
 86                 count = 0
 87                 for i in range(100):
 88                     if scores is None or final_score[i] > 0.5:
 89                         count = count + 1
 90                 print()
 91                 print("the count of objects is: ", count)
 92                 (im_width, im_height) = image.size
 93                 for i in range(count):
 94                     # print(boxes[0][i])
 95                     y_min = boxes[0][i][0] * im_height
 96                     x_min = boxes[0][i][1] * im_width
 97                     y_max = boxes[0][i][2] * im_height
 98                     x_max = boxes[0][i][3] * im_width
 99                     x = int((x_min + x_max) / 2)
100                     y = int((y_min + y_max) / 2)
101                     if category_index[classes[0][i]]['name'] == "tower":
102                         print("this image has a tower!")
103                         y = int((y_max - y_min) / 4 * 3 + y_min)
104                     print("object{0}: {1}".format(i, category_index[classes[0][i]]['name']),
105                           ',Center_X:', x, ',Center_Y:', y)
106                     # print(x_min,y_min,x_max,y_max)
107                 plt.figure(figsize=IMAGE_SIZE)
108                 plt.imshow(image_np)
109                 picName = test_image.split('/')[-1]
110                 # print(picName)
111                 plt.savefig(PATH_TO_RESULTS + picName)
112                 print(test_image + ' succeed')
113 
114             end = time.time()
115             seconds = end - start
116             print("Time taken : {0} seconds".format(seconds))
117 
118 
119 save_object_detection_result()