Tensorflow-SSD之原模型測試
今天下午嘗試了一下tensorflow-ssd的模型，具體代碼還沒細讀，調試了一番之後，把官方提供的ssd_300_vgg的模型測試了一下，感覺比yolo準确率高一些，下一步準備訓練一下自己的模型，嘗試對比一下。
以下代碼來自網絡，如侵權，聯系我。
圖檔測試代碼：
# -*- coding：utf-8 -*-
import os
import cv2
import math
import random
import tensorflow as tf
import matplotlib.pyplot as plt
import matplotlib.cm as mpcm
import matplotlib.image as mpimg
from nets import ssd_vgg_300, ssd_common, np_methods
from preprocessing import ssd_vgg_preprocessing
import sys
slim = tf.contrib.slim
# TensorFlow session
gpu_options = tf.GPUOptions(allow_growth=True)
config = tf.ConfigProto(log_device_placement=False, gpu_options=gpu_options)
isess = tf.InteractiveSession(config=config)
l_VOC_CLASS =( ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle',
'bus', 'car', 'cat', 'chair', 'cow',
'diningTable', 'dog', 'horse', 'motorbike', 'person',
'pottedPlant', 'sheep', 'sofa', 'train', 'TV'])
# 定義資料格式，設定占位符
net_shape = (300, 300)
# 預處理，以Tensorflow backend, 将輸入圖檔大小改成 300x300，作為下一步輸入
img_input = tf.placeholder(tf.uint8, shape=(None, None, 3))
# 輸入圖像的通道排列形式，'NHWC'表示 [batch_size,height,width,channel]
data_format = 'NHWC'
# 資料預處理，将img_input輸入的圖像resize為300大小，labels_pre,bboxes_pre,bbox_img待解析
image_pre, labels_pre, bboxes_pre, bbox_img = ssd_vgg_preprocessing.preprocess_for_eval(
img_input, None, None, net_shape, data_format,
resize=ssd_vgg_preprocessing.Resize.WARP_RESIZE)
# 拓展為4維變量用于輸入
image_4d = tf.expand_dims(image_pre, 0)
# 定義SSD模型
# 是否複用，目前我們沒有在訓練是以為None
reuse = True if 'ssd_net' in locals() else None
# 調出基于VGG神經網絡的SSD模型對象，注意這是一個自定義類對象
ssd_net = ssd_vgg_300.SSDNet()
# 得到預測類和預測坐标的Tensor對象，這兩個就是神經網絡模型的計算流程
with slim.arg_scope(ssd_net.arg_scope(data_format=data_format)):
predictions, localisations, _, _ = ssd_net.net(image_4d, is_training=False, reuse=reuse)
# 導入官方給出的 SSD 模型參數
ckpt_filename = 'checkpoints\\ssd_300_vgg.ckpt'
# ckpt_filename = '../checkpoints/VGG_VOC0712_SSD_300x300_ft_iter_120000.ckpt'
isess.run(tf.global_variables_initializer())
saver = tf.train.Saver()
saver.restore(isess, ckpt_filename)
# 在網絡模型結構中，提取搜尋網格的位置
# 根據模型超參數，得到每個特征層（這裡用了6個特征層，分别是4，7，8，9，10，11）的anchors_boxes
ssd_anchors = ssd_net.anchors(net_shape)
# 加載輔助作圖函數
def colors_subselect(colors, num_classes=21):
dt = len(colors) // num_classes
sub_colors = []
for i in range(num_classes):
color = colors[i * dt]
if isinstance(color[0], float):
sub_colors.append([int(c * 255) for c in color])
else:
sub_colors.append([c for c in color])
return sub_colors
def bboxes_draw_on_img(img, classes, scores, bboxes, colors, thickness=2):
shape = img.shape
for i in range(bboxes.shape[0]):
bbox = bboxes[i]
color = colors[classes[i]]
# Draw bounding box...
p1 = (int(bbox[0] * shape[0]), int(bbox[1] * shape[1]))
p2 = (int(bbox[2] * shape[0]), int(bbox[3] * shape[1]))
cv2.rectangle(img, p1[::-1], p2[::-1], color, 2)
# Draw text...
s = '%s/%.3f' % (l_VOC_CLASS[int(classes[i]) - 1], scores[i])
p1 = (p1[0] - 5, p1[1])
cv2.putText(img, s, p1[::-1], cv2.FONT_HERSHEY_DUPLEX, 1, color,2)
colors_plasma = colors_subselect(mpcm.plasma.colors, num_classes=21)
# 主流程函數
def process_image(img, case, select_threshold=0.15, nms_threshold=.1, net_shape=(300, 300)):
# select_threshold：box門檻值——每個像素的box分類預測資料的得分會與box門檻值比較，高于一個box門檻值則認為這個box成功框到了一個對象
# nms_threshold：重合度門檻值——同一對象的兩個框的重合度高于該門檻值，則運作下面去重函數
# 執行SSD模型，得到4維輸入變量，分類預測，坐标預測，rbbox_img參數為最大檢測範圍，本文固定為[0,0,1,1]即全圖
rimg, rpredictions, rlocalisations, rbbox_img = isess.run([image_4d, predictions,
localisations, bbox_img], feed_dict={img_input: img})
# ssd_bboxes_select()函數根據每個特征層的分類預測分數，歸一化後的映射坐标，
# ancohor_box的大小，通過設定一個門檻值計算得到每個特征層檢測到的對象以及其分類和坐标
rclasses, rscores, rbboxes = np_methods.ssd_bboxes_select(rpredictions, rlocalisations, ssd_anchors,
select_threshold=select_threshold,
img_shape=net_shape,
num_classes=21, decode=True)
# 檢測有沒有超出檢測邊緣
rbboxes = np_methods.bboxes_clip(rbbox_img, rbboxes)
rclasses, rscores, rbboxes = np_methods.bboxes_sort(rclasses, rscores, rbboxes, top_k=400)
# 去重，将重複檢測到的目标去掉
rclasses, rscores, rbboxes = np_methods.bboxes_nms(rclasses, rscores, rbboxes, nms_threshold=nms_threshold)
# 将box的坐标重新映射到原圖上（上文所有的坐标都進行了歸一化，是以要逆操作一次）
rbboxes = np_methods.bboxes_resize(rbbox_img, rbboxes)
if case == 1:
bboxes_draw_on_img(img, rclasses, rscores, rbboxes, colors_plasma, thickness=8)
return img
else:
return rclasses, rscores, rbboxes
case = 1
img = cv2.imread("demo\\000010.jpg")
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
plt.title('SSD_test')
plt.imshow(process_image(img, case))
plt.show()
以下為視訊測試代碼：
import os
import cv2
import math
import random
import tensorflow as tf
import matplotlib.pyplot as plt
import matplotlib.cm as mpcm
import matplotlib.image as mpimg
from notebooks import visualization
from nets import ssd_vgg_300, ssd_common, np_methods
from preprocessing import ssd_vgg_preprocessing
import sys
# 當引用子產品和運作的腳本不在同一個目錄下，需在腳本開頭添加如下代碼：
#sys.path.append('./SSD-Tensorflow/')
slim = tf.contrib.slim
# TensorFlow session
gpu_options = tf.GPUOptions(allow_growth=True)
config = tf.ConfigProto(log_device_placement=False, gpu_options=gpu_options)
isess = tf.InteractiveSession(config=config)
l_VOC_CLASS = ['aeroplane', 'bicycle', 'bird', 'boat', 'bottle',
'bus', 'car', 'cat', 'chair', 'cow',
'diningTable', 'dog', 'horse', 'motorbike', 'person',
'pottedPlant', 'sheep', 'sofa', 'train', 'TV']
# 定義資料格式，設定占位符
net_shape = (300, 300)
# 預處理，以Tensorflow backend, 将輸入圖檔大小改成 300x300，作為下一步輸入
img_input = tf.placeholder(tf.uint8, shape=(None, None, 3))
# 輸入圖像的通道排列形式，'NHWC'表示 [batch_size,height,width,channel]
data_format = 'NHWC'
# 資料預處理，将img_input輸入的圖像resize為300大小，labels_pre,bboxes_pre,bbox_img待解析
image_pre, labels_pre, bboxes_pre, bbox_img = ssd_vgg_preprocessing.preprocess_for_eval(
img_input, None, None, net_shape, data_format,
resize=ssd_vgg_preprocessing.Resize.WARP_RESIZE)
# 拓展為4維變量用于輸入
image_4d = tf.expand_dims(image_pre, 0)
# 定義SSD模型
# 是否複用，目前我們沒有在訓練是以為None
reuse = True if 'ssd_net' in locals() else None
# 調出基于VGG神經網絡的SSD模型對象，注意這是一個自定義類對象
ssd_net = ssd_vgg_300.SSDNet()
# 得到預測類和預測坐标的Tensor對象，這兩個就是神經網絡模型的計算流程
with slim.arg_scope(ssd_net.arg_scope(data_format=data_format)):
predictions, localisations, _, _ = ssd_net.net(image_4d, is_training=False, reuse=reuse)
# 導入官方給出的 SSD 模型參數
ckpt_filename = 'checkpoints/ssd_300_vgg.ckpt'
# ckpt_filename = '../checkpoints/VGG_VOC0712_SSD_300x300_ft_iter_120000.ckpt'
isess.run(tf.global_variables_initializer())
saver = tf.train.Saver()
saver.restore(isess, ckpt_filename)
# 在網絡模型結構中，提取搜尋網格的位置
# 根據模型超參數，得到每個特征層（這裡用了6個特征層，分别是4，7，8，9，10，11）的anchors_boxes
ssd_anchors = ssd_net.anchors(net_shape)
"""
每層的anchors_boxes包含4個arrayList，前兩個List分别是該特征層下x，y坐标軸對于原圖（300x300）大小的映射
第三,四個List為anchor_box的長度和寬度，同樣是經過歸一化映射的，根據每個特征層box數量的不同，這兩個List元素
個數會變化。其中，長寬的值根據超參數anchor_sizes和anchor_ratios制定。
"""
# 加載輔助作圖函數
def colors_subselect(colors, num_classes=21):
dt = len(colors) // num_classes
sub_colors = []
for i in range(num_classes):
color = colors[i * dt]
if isinstance(color[0], float):
sub_colors.append([int(c * 255) for c in color])
else:
sub_colors.append([c for c in color])
return sub_colors
def bboxes_draw_on_img(img, classes, scores, bboxes, colors, thickness=1):
shape = img.shape
for i in range(bboxes.shape[0]):
bbox = bboxes[i]
color = colors[classes[i]]
# Draw bounding box...
p1 = (int(bbox[0] * shape[0]), int(bbox[1] * shape[1]))
p2 = (int(bbox[2] * shape[0]), int(bbox[3] * shape[1]))
cv2.rectangle(img, p1[::-1], p2[::-1], color, thickness)
# Draw text...
s = '%s/%.3f' % (l_VOC_CLASS[int(classes[i]) - 1], scores[i])
p1 = (p1[0] - 5, p1[1])
cv2.putText(img, s, p1[::-1], cv2.FONT_HERSHEY_DUPLEX, 0.5, color, 1)
colors_plasma = colors_subselect(mpcm.plasma.colors, num_classes=21)
# 主流程函數
def process_image(img, select_threshold=0.2, nms_threshold=.1, net_shape=(300, 300)):
# select_threshold：box門檻值——每個像素的box分類預測資料的得分會與box門檻值比較，高于一個box門檻值則認為這個box成功框到了一個對象
# nms_threshold：重合度門檻值——同一對象的兩個框的重合度高于該門檻值，則運作下面去重函數
# 執行SSD模型，得到4維輸入變量，分類預測，坐标預測，rbbox_img參數為最大檢測範圍，本文固定為[0,0,1,1]即全圖
rimg, rpredictions, rlocalisations, rbbox_img = isess.run([image_4d, predictions, localisations, bbox_img],
feed_dict={img_input: img})
# ssd_bboxes_select函數根據每個特征層的分類預測分數，歸一化後的映射坐标，
# ancohor_box的大小，通過設定一個門檻值計算得到每個特征層檢測到的對象以及其分類和坐标
rclasses, rscores, rbboxes = np_methods.ssd_bboxes_select(rpredictions, rlocalisations, ssd_anchors,
select_threshold=select_threshold,
img_shape=net_shape,
num_classes=21, decode=True)
# 檢測有沒有超出檢測邊緣
rbboxes = np_methods.bboxes_clip(rbbox_img, rbboxes)
rclasses, rscores, rbboxes = np_methods.bboxes_sort(rclasses, rscores, rbboxes, top_k=400)
# 去重，将重複檢測到的目标去掉
rclasses, rscores, rbboxes = np_methods.bboxes_nms(rclasses, rscores, rbboxes, nms_threshold=nms_threshold)
# 将box的坐标重新映射到原圖上（上文所有的坐标都進行了歸一化，是以要逆操作一次）
rbboxes = np_methods.bboxes_resize(rbbox_img, rbboxes)
bboxes_draw_on_img(img, rclasses, rscores, rbboxes, colors_plasma, thickness=2)
return img
# 視訊物體定位
import imageio # pip install imageio==2.4.1
imageio.plugins.ffmpeg.download()
from moviepy.editor import VideoFileClip
def get_process_video(input_path, output_path):
video = VideoFileClip(input_path)
result = video.fl_image(process_image)
result.write_videofile(output_path, fps=25)
try:
os.makedirs("Video/input/")
os.makedirs("Video/output/")
except:
input_folder = "Video/input/"
input_video_name = "66.mp4"
input_video_path = input_folder + input_video_name
output_video_path = "Video/output/output_" + input_video_name
get_process_video(input_video_path, output_video_path)