【Tensorflow】SSD_Mobilenet_v2实现目标检测（二）：测试

继续上篇博客介绍的
【Tensorflow】SSD_Mobilenet_v2实现目标检测（一）：环境配置+训练
接下来SSD_Mobilenet_v2实现目标检测之训练后实现测试。
训练后会在指定的文件夹内生成如下文件

1. 可视化训练过程

tensorboard --logdir=C:\Users\znjt\Desktop\loss  # 储存.tfevents的路径

将获得的网址复制到火狐或谷歌浏览器进行查看：

根据以上的信息，可以选择最佳的模型，笔者简单测试了6000epoch，因此训练结果并不好，可以增加epoch的数量以及适当调整训练参数。
具体的可视化过程可以参考：【tensorboard】可视化events.out.tfevents文件

2. 导出模型文件

在object detection目录下的export_inference_graph.py，使用如下指令导出模型结果：

python export_inference_graph.py --pipeline_config_path=/home/lianlirong/models-master/research/object_detection/data/hengfeng/ssd_mobilenet_v2_coco.config  --trained_checkpoint_prefix /home/lianlirong/models-master/logs/hengfeng/model.ckpt-60000 --output_directory /home/lianlirong/models-master/logs/hengfeng/model

执行完命令后，在model文件夹下生成如下文件

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

3. 获取测试图片

在Annotations的同级目录下新建一个文件夹test（用于储存测试文件）以及python文件get_test.py（用于获取测试图片），代码如下:

"""
2020.09.23:alian
tensorflow-ssd-test-img

"""
from PIL import Image
import os.path
import glob

annotations_test_dir = "/home/lianlirong/models-master/research/object_detection/images/VOC2007-hengfeng/Annotations/test/" # 训练时生成的用于测试划分出来的xml文件
Images_dir = "/home/lianlirong/models-master/research/object_detection/images/VOC2007-hengfeng/JPEGImages"  # 原图路径
test_images_dir = "/home/lianlirong/models-master/research/object_detection/images/VOC2007-hengfeng/test"   # 空文件夹，保存获取的测试图片
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 + ".jpg")
             img.save(os.path.join(test_images_dir, os.path.basename(pngfile)))
             print(pngname)
             i += 1
print(i)

运行代码，在test文件夹下获得测试图片

4. 测试图片并保存测试结果

在./logs/hengfeng/(即model的同级目录)下新建results_imgs文件夹；在/models-master/research/object_detection/文件夹下get_results.py文件并加入如下代码，我们将使用前面训练出的模型批量测试test文件夹中的图片并保存到results_imgs文件夹中。

"""
2020.09.23:alian
获得测试结果
"""
# -*- 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 = "/home/lianlirong/models-master/research/object_detection/images/VOC2007-hengfeng/JPEGImages/"  # 测试图片的路径
MODEL_NAME = '/home/lianlirong/models-master/logs/hengfeng/model'  #  模型的储存路径
PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb'  # 模型文件
PATH_TO_LABELS = '/home/lianlirong/models-master/research/object_detection/data/hengfeng/hengfeng_label_map.pbtxt'  #  标签文件的路径
NUM_CLASSES = 2  # 检测的目标数
PATH_TO_RESULTS = "/home/lianlirong/models-master/logs/hengfeng/results_imgs/"  # 测试结果图的保存路径


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.io.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:  ## tf.gfile.GFile更新为tf.io.gfile.GFile，否则会出错‘utf8’
            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()

运行代码，获得测试结果图。

文件夹树形结构如下：

├──models-master （tensorflow项目文件）
　　　　├── logs # 笔者存放训练模型的目录
　　　　　　　　├──hengfeng
　　　　　　　　　　　　├──model # 转换生成的模型文件
　　　　　　　　　　　　├──results_imgs # 保存结果图片
　　　　├── research
　　　　　　　　├──generate_tfrecord.py # 生成tfrecord文件的代码
　　　　　　　　├──object_detection
　　　　　　　　　　　　├──ssd_mobilenet_v2_coco_2018_03_29
　　　　　　　　　　　　　　　　├──model.ckpt # 预训练模型文件
　　　　　　　　　　　　├──model_main.py # 训练代码文件
　　　　　　　　　　　　├──get_redults.py # 测试运行代码
　　　　　　　　　　　　├──images
　　　　　　　　　　　　　　　　├──hengfeng # 目标数据集
　　　　　　　　　　　　　　　　　　　　├──Annotations # xml文件
　　　　　　　　　　　　　　　　　　　　├──JPEGImages # 原图片
　　　　　　　　　　　　　　　　　　　　├──train_test_split.py # 划分训练测试集的代码
　　　　　　　　　　　　　　　　　　　　├──xml_to_csv.py # 生成csv的代码文件
　　　　　　　　　　　　　　　　　　　　├──others_object …# 其他目标的数据集
　　　　　　　　　　　　　　　　　　　　├──get_test.py # 获取测试图片
　　　　　　　　　　　　　　　　　　　　├──test # 储存测试图片的文件夹
　　　　　　　　　　　　├──data
　　　　　　　　　　　　　　　　├──hengfeng # 目标训练的必要文件
　　　　　　　　　　　　　　　　　　　　├──hengfeng_label_map.pbtxt # 目标标签文件
　　　　　　　　　　　　　　　　　　　　├──hengfeng_train.tfrecord #生成的tfrecord文件
　　　　　　　　　　　　　　　　　　　　├──hengfeng_val.tfrecord
　　　　　　　　　　　　　　　　　　　　├──hengfeng_test.tfrecord
　　　　　　　　　　　　　　　　　　　　├──hengfeng_train_labels.csv # 生成的csv文件
　　　　　　　　　　　　　　　　　　　　├──hengfeng_val_labels.csv
　　　　　　　　　　　　　　　　　　　　├──hengfeng_test_labels.csv
　　　　　　　　　　　　　　　　　　　　├──ssd_mobilenet_v2_coco.config # 训练的配置文件
　　　　　　　　　　　　　　　　├──others_object #其他目标
以上是基于tensorflow +SSD_Mobilenet_v2实现目标检测的全部内容