30秒轻松实现TensorFlow物体检测

Chloe

发布时间：2026-10-13| 2人浏览过

原创

Google发布的TensorFlow物体检测API提供了预训练模型及重新训练脚本，助力快速构建图像检测应用。本文详解如何载入预训练的SSD MobileNet模型，配置环境，下载并加载冻结图，映射标签，并通过辅助代码实现图像检测与可视化，帮助用户在30秒内上手物体检测任务。

Google推出了全新的TensorFlow物体检测API，该API不仅内置了预训练模型，还附带了一个用于发布模型的Jupyter Notebook，以及一系列实用的脚本，支持用户使用自定义数据集对模型进行重新训练。

借助此API，开发者能够迅速搭建出针对图片中特定物体的检测应用。下文将逐步演示如何利用预训练模型来完成图像中物体的检测工作。

首先，我们需要加载运行过程中所需的各类库文件。

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

随后，进行必要的环境配置与设置。

%matplotlib inline 
sys.path.append("..")

接下来进入物体检测模块的载入阶段。

from utils import label_map_util 
 
from utils import visualization_utils as vis_util

在此步骤中，我们需要准备模型相关文件。

任何通过export_inference_graph.py工具生成的模型均可在此处载入，只需将PATH_TO_CKPT变量指向新的.pb文件即可。本例中，我们选用“移动网SSD”（SSD MobileNet）模型进行演示。

MODEL_NAME = 'ssd_mobilenet_v1_coco_11_06_2017' 
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

接下来执行模型的下载操作。

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)

编写必要的辅助代码以支持检测流程。

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)

执行具体的检测逻辑。

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) ] 
IMAGE_SIZE = (12, 8) 
[python] view plain copy
with detection_graph.as_default(): 
 
 with tf.Session(graph=detection_graph) as sess: 
  for image_path in TEST_IMAGE_PATHS: 
   image = Image.open(image_path) 
   # 这个array在之后会被用来准备为图片加上框和标签 
   image_np = load_image_into_numpy_array(image) 
   # 扩展维度，应为模型期待: [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') 
   # 每个框代表一个物体被侦测到. 
   boxes = detection_graph.get_tensor_by_name('detection_boxes:0') 
   # 每个分值代表侦测到物体的可信度. 
   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') 
   # 执行侦测任务. 
   (boxes, scores, classes, num_detections) = sess.run( 
     [boxes, scores, classes, num_detections], 
     feed_dict={image_tensor: image_np_expanded}) 
   # 图形化. 
   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) 
   plt.figure(figsize=IMAGE_SIZE) 
   plt.imshow(image_np)

在载入模型环节，您可以尝试切换不同的检测模型以对比速度与准确度。只需将待检测的图片路径添加至TEST_IMAGE_PATHS列表中并运行代码即可。

本文内容至此结束，希望上述步骤能对您的学习有所帮助，也欢迎大家继续支持本站。