30秒輕松實現(xiàn)TensorFlow物體檢測
Google發(fā)布了新的TensorFlow物體檢測API,包含了預(yù)訓(xùn)練模型,一個發(fā)布模型的jupyter notebook,一些可用于使用自己數(shù)據(jù)集對模型進(jìn)行重新訓(xùn)練的有用腳本。
使用該API可以快速的構(gòu)建一些圖片中物體檢測的應(yīng)用。這里我們一步一步來看如何使用預(yù)訓(xùn)練模型來檢測圖像中的物體。
首先我們載入一些會使用的庫
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
接下來進(jìn)行環(huán)境設(shè)置
%matplotlib inline
sys.path.append("..")
物體檢測載入
from utils import label_map_util from utils import visualization_utils as vis_util
準(zhǔn)備模型
變量 任何使用export_inference_graph.py工具輸出的模型可以在這里載入,只需簡單改變PATH_TO_CKPT指向一個新的.pb文件。這里我們使用“移動網(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ù)測5時,我們知道它對應(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)
檢測
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在之后會被用來準(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')
# 每個框代表一個物體被偵測到.
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')
# 執(zhí)行偵測任務(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)
在載入模型部分可以嘗試不同的偵測模型以比較速度和準(zhǔn)確度,將你想偵測的圖片放入TEST_IMAGE_PATHS中運行即可。
以上就是本文的全部內(nèi)容,希望對大家的學(xué)習(xí)有所幫助,也希望大家多多支持腳本之家。
相關(guān)文章
python使用pandas處理excel文件轉(zhuǎn)為csv文件的方法示例
這篇文章主要介紹了python使用pandas處理excel文件轉(zhuǎn)為csv文件的方法示例,文中通過示例代碼介紹的非常詳細(xì),對大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價值,需要的朋友們下面隨著小編來一起學(xué)習(xí)學(xué)習(xí)吧2019-07-07
VScode連接遠(yuǎn)程服務(wù)器上的jupyter notebook的實現(xiàn)
這篇文章主要介紹了VScode連接遠(yuǎn)程服務(wù)器上的jupyter notebook的實現(xiàn),文中通過示例代碼介紹的非常詳細(xì),對大家的學(xué)習(xí)或者工作具有一定的參考學(xué)習(xí)價值,需要的朋友們下面隨著小編來一起學(xué)習(xí)學(xué)習(xí)吧2020-04-04
python列表使用實現(xiàn)名字管理系統(tǒng)
這篇文章主要為大家詳細(xì)介紹了python列表使用實現(xiàn)名字管理系統(tǒng),具有一定的參考價值,感興趣的小伙伴們可以參考一下2019-01-01
Python中time模塊與datetime模塊在使用中的不同之處
這篇文章主要介紹了Python中time模塊與datetime模塊在使用中的不同之處,是Python入門學(xué)習(xí)中的基礎(chǔ)知識,需要的朋友可以參考下2015-11-11

