Google發(fā)布了新的TensorFlow物體檢測(cè)API，包含了預(yù)訓(xùn)練模型，一個(gè)發(fā)布模型的jupyter notebook，一些可用于使用自己數(shù)據(jù)集對(duì)模型進(jìn)行重新訓(xùn)練的有用腳本。

使用該API可以快速的構(gòu)建一些圖片中物體檢測(cè)的應(yīng)用。這里我們一步一步來(lái)看如何使用預(yù)訓(xùn)練模型來(lái)檢測(cè)圖像中的物體。

首先我們載入一些會(huì)使用的庫(kù)

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 

接下來(lái)進(jìn)行環(huán)境設(shè)置

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

物體檢測(cè)載入

from utils import label_map_util  from utils import visualization_utils as vis_util 

準(zhǔn)備模型

變量  任何使用export_inference_graph.py工具輸出的模型可以在這里載入，只需簡(jiǎn)單改變PATH_TO_CKPT指向一個(gè)新的.pb文件。這里我們使用“移動(dòng)網(wǎng)SSD”模型。

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()) 

將（frozen）TensorFlow模型載入內(nèi)存

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='') 

載入標(biāo)簽圖

標(biāo)簽圖將索引映射到類名稱，當(dāng)我們的卷積預(yù)測(cè)5時(shí)，我們知道它對(duì)應(yīng)飛機(jī)。這里我們使用內(nèi)置函數(shù)，但是任何返回將整數(shù)映射到恰當(dāng)字符標(biāo)簽的字典都適用。

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) 

檢測(cè)

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 copywith detection_graph.as_default():   with tf.Session(graph=detection_graph) as sess:   for image_path in TEST_IMAGE_PATHS:    image = Image.open(image_path)    # 這個(gè)array在之后會(huì)被用來(lái)準(zhǔn)備為圖片加上框和標(biāo)簽    image_np = load_image_into_numpy_array(image)    # 擴(kuò)展維度，應(yīng)為模型期待: [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')    # 每個(gè)框代表一個(gè)物體被偵測(cè)到.    boxes = detection_graph.get_tensor_by_name('detection_boxes:0')    # 每個(gè)分值代表偵測(cè)到物體的可信度.    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')    # 執(zhí)行偵測(cè)任務(wù).    (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)