效果图

以下内容为tensorflow官方测试代码，仅用于测试自己模型训练的效果，不是训练代码

导包

import warnings
warnings.filterwarnings("ignore")
import numpy as np
import os
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile

from distutils.version import StrictVersion
from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
from PIL import Image

# 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

# tf版本需要大于1.9
if StrictVersion(tf.__version__) < StrictVersion('1.9.0'):
    raise ImportError('Please upgrade your TensorFlow installation to v1.9.* or later!')

导下到本地的tensorflow工具包

from utils import label_map_util

from utils import visualization_utils as vis_util

环境

# pb模型存放位置
PATH_TO_FROZEN_GRAPH = 'models/frozen_inference_graph.pb'

# coco数据集的label映射文件
PATH_TO_LABELS = 'label_map.pbtxt'

# 载入训练好的pd模型
detection_graph = tf.Graph()
with detection_graph.as_default():
    od_graph_def = tf.GraphDef()
    with tf.gfile.GFile(PATH_TO_FROZEN_GRAPH, 'rb') as fid:
        serialized_graph = fid.read()
        od_graph_def.ParseFromString(serialized_graph)
        tf.import_graph_def(od_graph_def, name='')

# 得到一个保存编号和类别描述映射关系的列表
category_index = label_map_util.create_category_index_from_labelmap(PATH_TO_LABELS, use_display_name=True)

# 把图片数据变成3维的数据，定义数据类型为uint8
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 run_inference_for_single_image(image, graph):
    with graph.as_default():
        with tf.Session() as sess:
            # 获得图中所有op
            ops = tf.get_default_graph().get_operations()
            # 获得输出op的名字
            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'
                # 如果tensor_name在all_tensor_names中
                if tensor_name in all_tensor_names:
                    # 则获取到该tensor
                    tensor_dict[key] = tf.get_default_graph().get_tensor_by_name(
                      tensor_name)

            if 'detection_masks' in tensor_dict:
                # 获得检测框数据
                detection_boxes = tf.squeeze(tensor_dict['detection_boxes'], [0])
                # 获得mask数据
                detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])
                # 检测框数量
                real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)
                # detection_boxes数据维度为(100,4)，部分数据有值，部分数据值为0
                # detection_masks数据维度为(num_mask,image_h,image_w)
                # 比如detection_boxes和detection_masks中都有10个有效数据，
                # 但是real_num_detection的数量可能会少于10，所以需要根据real_num_detection的数量取数据
                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[1], image.shape[2])
                # 阈值设置为0.5，大于0.5则为True->1.0，小于0.5则为Fasle->0
                detection_masks_reframed = tf.cast(tf.greater(detection_masks_reframed, 0.5), tf.uint8)
                # 增加一个维度
                tensor_dict['detection_masks'] = tf.expand_dims(
                    detection_masks_reframed, 0)       


            # 图片输入的tensor
            image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')

            # 传入图片运行模型获得结果
            output_dict = sess.run(tensor_dict,
                             feed_dict={image_tensor: image})

            # 所有的结果都是float32类型的，有些数据需要做数据格式转换
            # 检测到目标的数量
            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

配置jupyter显示图片

# 在jupyter里面显示图片
%matplotlib inline

执行

for root,dirs,files in os.walk('test_images/'):
    for image_path in files:
        # 读取图片
        image = Image.open(os.path.join(root,image_path))
        # 把图片数据变成3维的数据，定义数据类型为uint8
        image_np = load_image_into_numpy_array(image)
        # 增加一个维度，数据变成: [1, None, None, 3]
        image_np_expanded = np.expand_dims(image_np, axis=0)
        # 目标检测
        output_dict = run_inference_for_single_image(image_np_expanded, detection_graph)
        print(output_dict['detection_boxes'].shape)
        print(output_dict['detection_masks'].shape)
        # 给原图加上预测框，置信度和类别信息
        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)
        # 画图
        plt.figure(figsize=(20,10))
        plt.imshow(image_np)
        plt.axis('off')
        plt.show()

测试训练结果

训练2000次的loss情况