tensorflow objectct_detection_tutorial 对象检测演示

 本笔记本将逐步介绍使用预先训练好的模型来检测图像中的物体的过程

在开始之前，请务必遵循[安装说明]（https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/installation.md）。

import numpy as np
import os
import six.moves.urllib as urllib
import sys
import tarfile
import tensorflow as tf
import zipfile


from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
from PIL import Image

sys.path.append("..")
from object_detection.utils import ops as utils_ops


if tf.__version__ < '1.4.0':
    raise ImportError('Please upgrade your tensorflow installation to v1.4.* or later!')

Env setup

# 这是显示图像所需的。
%matplotlib inline

对象检测导入

以下是从对象检测模块导入的内容

from utils import label_map_util

from utils import visualization_utils as vis_util

模型准备：

变量：使用`export_inference_graph.py`工具导出的任何模型都可以在这里加载，只需将`PATH_TO_CKPT`改为指向一个新的.pb文件即可。默认情况下，我们在这里使用“带Mobilenet的SSD”模型。 请参阅[detection model zoo]]（https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/detection_model_zoo.md）以获取可以用尽的其他模型的列表， 具有不同的速度和精度的盒子。

# 下载什么样的模型。
MODEL_NAME = 'ssd_mobilenet_v1_coco_2017_11_17'
MODEL_FILE = MODEL_NAME + '.tar.gz'
DOWNLOAD_BASE = 'http://download.tensorflow.org/models/object_detection/'

#冻结检测图的路径。 这是用于对象检测的实际模型。
PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb'

#用于为每个框添加正确标签的字符串列表。
PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt')

NUM_CLASSES = 90            #90中类别

下载模式：

opener = urllib.request.URLopener()                     #网络请求操作
opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE)
tar_file = tarfile.open(MODEL_FILE)
for file in tar_file.getmembers():
    file_name = os.path.basename(file.name)
    if 'frozen_inference_graph.pb' in file_name:
        tar_file.extract(file, os.getcwd())

Tensorflow模型加载到内存中:

detection_graph = tf.Graph()
with detection_graph.as_default():
    od_graph_def = tf.GraphDef()
    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='')

加载标签：标签映射将索引映射到类别名称，以便当我们的卷积网络预测'5'时，我们知道这对应于'飞机'。 这里我们使用内部实用函数，但任何返回字典映射整数到适当的字符串标签的东西都可以。

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：

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)

Detection：

#为了简单起见，我们将只使用2张图片：
# image1.jpg
# image2.jpg
# 如果您想用图像测试代码，只需将图像的路径添加到TEST_IMAGE_PATHS即可。
PATH_TO_TEST_IMAGES_DIR = 'test_images'
TEST_IMAGE_PATHS = [ os.path.join(PATH_TO_TEST_IMAGES_DIR, 'image{}.jpg'.format(i)) for i in range(1, 3) ]

# Size, in inches, of the output images.输出图像的大小（以英寸为单位）。
IMAGE_SIZE = (12, 8)

def run_inference_for_single_image(image, graph):
    with graph.as_default():
        with tf.Session() as sess:
            
      # 获取输入和输出张量的句柄
            ops = tf.get_default_graph().get_operations()
            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'
        if tensor_name in all_tensor_names:
            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])
               detection_masks = tf.squeeze(tensor_dict['detection_masks'], [0])
        # 需要重构来将蒙版从框坐标转换为图像坐标并适合图像大小。
               real_num_detection = tf.cast(tensor_dict['num_detections'][0], tf.int32)
               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[0], image.shape[1])
                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)
                image_tensor = tf.get_default_graph().get_tensor_by_name('image_tensor:0')

      # Run inference运行推理
                output_dict = sess.run(tensor_dict,
                             feed_dict={image_tensor: np.expand_dims(image, 0)})

      # all outputs are float32 numpy arrays, so convert types as appropriate
    #所有的输出都是float32 numpy数组，所以根据需要转换类型
                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

for image_path in TEST_IMAGE_PATHS:
    
    image = Image.open(image_path)
  # 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]
#由于模型预期图像具有形状，因此扩展尺寸：[1，None，None，3]

    image_np_expanded = np.expand_dims(image_np, axis=0)
  # Actual detection.实际检测。
    output_dict = run_inference_for_single_image(image_np, detection_graph)
  # Visualization of the results of a detection.可视化检测结果。
    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=IMAGE_SIZE)
    plt.imshow(image_np)