TensorFlow Object Detection API 只检测特定的目标以及摄像头实时检测

前言

本人小白一枚最近在做一个检测行人的项目，参考了各位前辈的文章跑通了官方的demo。大佬在B站上的视频很好大家可以学习下 https://blog.csdn.net/dy_guox/article/details/79111949 官方的demo会对图片中所有可以检测的目标进行检测，但是我的目标是只检测行人。查了很久没有合适的材料，折腾了很久才把这个问题解决了。在这里给大家分享一下。

一、 以只进行行人的检测为例

首先进入object_detection文件夹找到utils文件夹，找到名为visualization_utils.py的python文件。我的路径如下：

A:\anaconda\Lib\site-packages\research\object_detection\utils

大家找到自己的visualization_utils.py文件后，在第735行左右会看到如下代码段：

 draw_bounding_box_on_image_array(
    image,
    ymin,
    xmin,
    ymax,
    xmax,
    color= color,
    thickness= line_thickness,
    display_str_list=box_to_display_str_map[box],
    use_normalized_coordinates=use_normalized_coordinates)
  if keypoints is not None:
    draw_keypoints_on_image_array(
      image,
      box_to_keypoints_map[box],
      color= color,
      radius= line_thickness /.2,
      use_normalized_coordinates= use_normalized_coordinates)

对这段代码进行修改如下：

display_str_list = box_to_display_str_map[box]
if (("bottle" in display_str_list[0]) or ("person") in display_str_list[0]):
  draw_bounding_box_on_image_array(
    image,
    ymin,
    xmin,
    ymax,
    xmax,
    color= color,
    thickness= line_thickness,
    display_str_list=box_to_display_str_map[box],
    use_normalized_coordinates=use_normalized_coordinates)
  if keypoints is not None:
    draw_keypoints_on_image_array(
      image,
      box_to_keypoints_map[box],
      color= color,
      radius= line_thickness /.2,
      use_normalized_coordinates= use_normalized_coordinates)

在这里一定要注意缩进的问题，官方demo里面的缩进格式我不是特别熟悉。我把代码放在下面一些大家按照我的格式来应该就没有问题了。

在运行官方demo之后会发现只会对图片中的行人进行了检测：

修改之后的检测结果

修改之前的检测结果

大家在自己实现自己的项目时只需要进行修改即可，改成你需要检测的类型即可。

if (("bottle" in display_str_list[0]) or ("person") in display_str_list[0]):

参考了：https://stackoverflow.com/questions/46515358/how-to-get-only-boxes-for-specific-categories-in-tensorflow-object-detection?rq=1

二、 利用摄像头进行实时行人监测

copy一下官方demo：

对官方demo进行修改，利用OpenCV打开摄像头，运行默认的model即ssd_mobilenet_v1_coco_2017_11_17进行实时监测（这里是修改后的因此只能检测人）。

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

from distutils.version import StrictVersion
from collections import defaultdict
from io import StringIO
from matplotlib import pyplot as plt
from PIL import Image
from object_detection.utils import label_map_util
from object_detection.utils import visualization_utils as vis_util

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

if StrictVersion(tf.__version__) < StrictVersion('1.12.0'):
    raise ImportError('Please upgrade your TensorFlow installation to v1.12.*.')
    
cap = cv2.VideoCapture(0)

MODEL_NAME = 'ssd_mobilenet_v1_coco_2017_11_17'
PATH_TO_FROZEN_GRAPH = MODEL_NAME + '/frozen_inference_graph.pb'
PATH_TO_LABELS = os.path.join('data', 'mscoco_label_map.pbtxt')

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)

with detection_graph.as_default():
    with tf.Session(graph=detection_graph) as sess:
        while True:
            ret, image_np = cap.read()
            # 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)

            cv2.imshow('object detection', image_np)
            if cv2.waitKey(25) & 0xFF == ord('q'):
                cv2.destroyAllWindows()
                break
cap.release()
cv2.destroyAllWindows()

运行上述代码就可以实时的进行检测啦。参考了：https://blog.csdn.net/weixin_42499236/article/details/83826472

总结

通过对visualization_utils.py的修改，实现了利用官方提供的ssd_mobilenet_v1_coco_2017_11_17实现了实时的行人检测。这里的模型可以根据自己的需要去官方下载不一定非要和文章中用的一样。这样就避免自己去训练模型（当然自己训练就不需要这么样修改了，要检测几个目标在训练之前就已经确定了）。只需要简单的几步就可以直接用官方demo进行实时监测啦。（前提是官方的model中包含你需要检测的类型）。


本人第一次写博客，自己也是接触深度学习不久。如有疏漏之处，还请各位包涵。大家有什么问题欢迎互相讨论，共同学习!

参考：
1.https://blog.csdn.net/dy_guox/article/details/79111949
2.https://blog.csdn.net/weixin_42499236/article/details/83826472
3.https://stackoverflow.com/questions/46515358/how-to-get-only-boxes-for-specific-categories-in-tensorflow-object-detection?rq=1