SSD-TensorFlow 訓練自己的數據並可視化

---

1. 源碼下載

• GitHub下載地址：https://github.com/balancap/SSD-Tensorflow
• 解壓 根目錄中 checkpoint 文件夾下的 ssd_300_vgg.ckpt.zip 文件

解壓方法：

unzip ssd_300_vgg.ckpt.zip

2. 數據準備

2.1. 標註數據

• 將自己的數據按照Pascal VOC格式進行標註
• 標註方法：可以 百度 / Google 搜索，網上有很多教程
  標註完如下圖：
  
  爲了後續處理的方便，數據集命名爲 VOC2007

2.2. 劃分數據集

• 按照一定的比例劃分爲訓練集、測試集和驗證集
  通過Python代碼腳本自動生成txt文件，train.txt, trainval.txt, test.txt, val.txt
• 劃分代碼如下：

# -*- coding: UTF-8 -*-
#程序功能：按照比例隨機劃分訓練集、測試集和驗證集
#讀取xml文件存入xml_list
#隨機讀取xml_list的內容，按照比例存儲在對應的txt中
#保存txt文件
import os, random
import copy#use copy.deepcopy to copy list
#xml的路徑以及分配後的txt存儲路徑
xml_path = 'E:\Research\ShellDataset_Plus\ShellDataset\Annotations'
txt_path = 'E:\Research\ShellDataset_Plus\ShellDataset\JPEGImages'

test_percent = 0.1#測試集所佔比例
val_percent = 0.1#驗證集所佔比例
xml_list = os.listdir(xml_path)

#remove the extend .xml
xml_list = [xml.split('.')[0] for xml in xml_list]

xml_len = len(xml_list)
#allot randomly
test = random.sample(xml_list, int(xml_len * test_percent))
trainval = copy.deepcopy(xml_list)
for xml in test:
    trainval.remove(xml)
val = random.sample(trainval, int(len(trainval) * val_percent))
train = copy.deepcopy(trainval)
for xml in val:
    train.remove(xml)

#renew txt and write
trainval_path = open(os.path.join(txt_path,'trainval.txt'),'w+')
train_path = open(os.path.join(txt_path,'train.txt'),'w+')
val_path = open(os.path.join(txt_path,'val.txt'),'w+')
test_path = open(os.path.join(txt_path,'test.txt'),'w+')
# add enter to each line
for trainval_xml in trainval:
    trainval_path.write(trainval_xml+'\n')
for train_xml in train:
    train_path.write(train_xml+'\n')
for val_xml in val:
    val_path.write(val_xml+'\n')
for test_xml in test:
    test_path.write(test_xml+'\n')

trainval_path.close()
train_path.close()
val_path.close()
test_path.close()

執行完會自動生成 train.txt, trainval.txt, test.txt, val.txt 文件~

2.3. 轉換數據集格式

• 將圖像數據轉換爲 .tfrecods 格式

2.3.1. 修改 pascalvoc_common.py

• 將訓練類別更改爲自己數據集中的類別
• 文件位置：SSD-Tensorflow-master —> datasets —> pascalvoc_common.py

# class表示自己數據集的標籤，有幾個寫幾個，此處列舉3個
# CLASS表示大類，所有的CLASS可以相同，可以隨便寫
VOC_LABELS = {
    'none': (0, 'Background'),
    'class1': (1, 'CLASS1'),
    'class2': (2, 'CLASS2'),
    'class3': (3, 'CLASS3'),
}

2.3.2. 修改 pascalvoc_to_tfrecords.py

• 更改文件的83行讀取方式爲 ’rb’
• 文件位置： SSD-Tensorflow-master —> datasets —> pascalvoc_to_tfrecords.py

# 修改67行，可以修改幾張圖片轉爲一個tfrecords，此處100表示每100張圖片生成1個.tfrecords文件
SAMPLES_PER_FILES = 100   

# 然後更改文件的83行讀取方式爲’rb’，修改爲如下形式
image_data = tf.gfile.FastGFile(filename, 'rb').read()

2.3.3. 創建 tf_conver_data.sh

• 創建位置：在 SSD-Tensorflow-master (根目錄) 文件夾下創建該文件
• 此外，在根下創建 tfrecords 文件夾，用於存放轉換完的數據

DATASET_DIR=./VOC2007/
OUTPUT_DIR=./tfrecords
python tf_convert_data.py \
    --dataset_name=pascalvoc \
    --dataset_dir=${DATASET_DIR} \
    --output_name=voc_2007_train \
    --output_dir=${OUTPUT_DIR}

2.3.4. 運行 tf_conver_data.sh
首先，賦予 tf_conver_data.sh 文件執行權限，再將格式標準化.
使用 vim 編輯器打開文件，並執行以下命令：

set ff=unix

保存並退出（命令：ESC ——> :wq）
在控制檯執行：

chmod +x tf_conver_data.sh
sed -i 's/\r$//g' tf_conver_data.sh
./tf_conver_data.sh

執行效果：

ligy@lab729:~/SSD-Tensorflow-master$ ./tf_convert_data.sh 
Dataset directory: ./VOC2007/
Output directory: ./tfrecords
>> Converting image 10080/10080
Finished converting the Pascal VOC dataset!

進入 tfrecords 文件夾，可以發現訓練數據已經轉換爲 .tfrecords 格式.

3. 模型訓練

3.1. 修改參數

3.1.1. 修改 train_ssd_network.py

• 修改說明：如果不將None改爲一個具體的數值，訓練會一直進行下去。
• 其他網絡參數配置，可以根據需要在此文件中調整。（初學者建議先不要調整）

# 第63行，可以將600修改，表示多長時間保存一次模型
tf.app.flags.DEFINE_integer(
    'save_interval_secs', 600,
    'The frequency with which the model is saved, in seconds.')

# 修改第154行的最大訓練步數，將None改爲比如40000
（tf.contrib.slim.learning.training函數中max-step爲None時訓練會無限進行。）

3.2.2. 修改 nets/ssd_vgg_300.py

• 因爲使用此網絡結構，修改 87 和 88 行的類別

 default_params = SSDParams(
        img_shape=(300, 300),
        num_classes=21,   #根據自己的數據修改爲類別+1，此處的 1 表示背景標籤
        no_annotation_label=21, #根據自己的數據修改爲類別+1
        feat_layers=['block4', 'block7', 'block8', 'block9', 'block10', 'block11'],
        feat_shapes=[(38, 38), (19, 19), (10, 10), (5, 5), (3, 3), (1, 1)],
        anchor_size_bounds=[0.15, 0.90],
        # anchor_size_bounds=[0.20, 0.90],

3.2.3. 修改 train_ssd_network.py

• 修改類別120行，GPU佔用量，學習率，batch_size等
• 修改第135 行，num_classes 改爲自己的種類數 + 1

# 根據自己的數據修改爲類別+1
tf.app.flags.DEFINE_integer(
    'num_classes', 21, 'Number of classes to use in the dataset.')

3.2.4. 修改 eval_ssd_network.py

• 第66行，修改爲自己的類別數 + 1

# =========================================================================== #
# Main evaluation flags.
# =========================================================================== #
tf.app.flags.DEFINE_integer(
    'num_classes', 9, 'Number of classes to use in the dataset.')

3.2.5. 修改 datasets/pascalvoc_2007.py

TRAIN_STATISTICS = {
    'none': (0, 0),
    'car': (192, 234),   192表示圖片的數量，234表示邊界框的數量
    'person': (1554, 2862),
    'dog': (558, 1230),
    ...
    'total': (10080, 18798),   10080表示訓練圖片的綜述，18798表示邊界框的總數
}
TEST_STATISTICS = {
    'none': (0, 0),
    'car': (1, 1),
    'person': (1, 1),
    'dog': (1, 1),
    ...
    'total': (20, 20),
}
SPLITS_TO_SIZES = {
    'train': 8165,   #訓練的圖片數量
    'test': 1008,    #測試的圖片數量
}
SPLITS_TO_STATISTICS = {
    'train': TRAIN_STATISTICS,
    'test': TEST_STATISTICS,
}
NUM_CLASSES = 20    #數據集的種類，此處不用加上背景

3.2.6. 執行訓練

在SSD根目錄下編寫訓練腳本 train.sh

DATASET_DIR=./tfrecords
TRAIN_DIR=./log/
CHECKPOINT_PATH=./checkpoints/ssd_300_vgg.ckpt
python train_ssd_network.py \
    --train_dir=${TRAIN_DIR} \
    --dataset_dir=${DATASET_DIR} \
    --dataset_name=pascalvoc_2007 \
    --dataset_split_name=train \
    --model_name=ssd_300_vgg \
    --checkpoint_path=${CHECKPOINT_PATH} \
    --checkpoint_exclude_scopes =ssd_300_vgg/conv6,ssd_300_vgg/conv7,ssd_300_vgg/block8,ssd_300_vgg/block9,ssd_300_vgg/block10,ssd_300_vgg/block11,ssd_300_vgg/block4_box,ssd_300_vgg/block7_box,ssd_300_vgg/block8_box,ssd_300_vgg/block9_box,ssd_300_vgg/block10_box,ssd_300_vgg/block11_box \
    --save_summaries_secs=60 \
    --save_interval_secs=600 \
    --weight_decay=0.0005 \
    --optimizer=adam \
    --learning_rate=0.001 \
    --batch_size=8 \
    --ignore_missing_vars = True \

賦予 train.sh 文件執行權限，再將格式標準化.
使用 vim 編輯器打開文件，並執行以下命令：

set ff=unix

保存並退出（命令：ESC ——> :wq）
在控制檯執行：

chmod +x train.sh
sed -i 's/\r$//g' train.sh
./train.sh

此處訓練時間由訓練迭代次數、GPU等決定。

---

4. 測試評價

4.1. 轉換數據格式

將需要測試的數據集轉換爲 .tfrecords 格式
在 VOC2007 中創建一個test文件夾，用於存放需要驗證的數據

• 在SSD根目錄創建一個 tf_convert_testdata.sh 文件

DATASET_DIR=./VOC2007/
OUTPUT_DIR=./tfrecords/test
python tf_convert_data.py \
    --dataset_name=pascalvoc \
    --dataset_dir=${DATASET_DIR} \
    --output_name=voc_2007_test \
    --output_dir=${OUTPUT_DIR}

4.2. 驗證精度

• 在SSD-Tensorflow-master文件夾下建立一個test.sh文件

# -*- coding: UTF-8 -*-
DATASET_DIR=./tfrecords/test/    #保存的轉換爲tfrcodes格式的數據
EVAL_DIR=./log_eval/   # Directory where the results are saved to    
CHECKPOINT_PATH=./log/model.ckpt-40000   #換爲自己訓練的模型
 python3 ./eval_ssd_network.py \
        --eval_dir=${EVAL_DIR} \
        --dataset_dir=${DATASET_DIR} \
        --dataset_name=pascalvoc_2007 \
        --dataset_split_name=test \
        --model_name=ssd_300_vgg \
        --checkpoint_path=${CHECKPOINT_PATH} \
        --batch_size=1

執行結果：

···
INFO:tensorflow:Evaluation [100/1008]
INFO:tensorflow:Evaluation [200/1008]
INFO:tensorflow:Evaluation [300/1008]
INFO:tensorflow:Evaluation [400/1008]
INFO:tensorflow:Evaluation [500/1008]
INFO:tensorflow:Evaluation [600/1008]
INFO:tensorflow:Evaluation [700/1008]
INFO:tensorflow:Evaluation [800/1008]
INFO:tensorflow:Evaluation [900/1008]
INFO:tensorflow:Evaluation [1000/1008]
INFO:tensorflow:Evaluation [1008/1008]
2019-03-29 10:40:43.982762: W ./tensorflow/core/grappler/optimizers/graph_optimizer_stage.h:230] Failed to run optimizer ArithmeticOptimizer, stage HoistCommonFactor. Error: Node average_precision_voc07/ArithmeticOptimizer/HoistCommonFactor_Add_AddN is missing output properties at position :0 (num_outputs=0)
AP_VOC07/mAP[0.80815487204921244]
AP_VOC12/mAP[0.84965517094243659]
INFO:tensorflow:Finished evaluation at 2019-03-29-02:40:44
Time spent : 53.065 seconds.
Time spent per BATCH: 0.053 seconds.

可以看到執行完畢！

4.2. 單張圖片檢測

• 利用 notebooks 文件夾下 ssd_notebook.ipynb 顯示訓練測試模型的結果

4.2.1. 創建並修改 ssd_notebook.py

• 新建一個 ssd_notebook.py 文件，並將ssd_notebook.ipynb中的代碼寫入並修改爲自己需要的。
  代碼如下：（需要修改的位置已加註釋）

import os
import math
import random

import numpy as np
import tensorflow as tf
import cv2

slim = tf.contrib.slim

#matplotlib inline
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

import sys
sys.path.append('../')

from nets import ssd_vgg_300, ssd_common, np_methods
from preprocessing import ssd_vgg_preprocessing
from notebooks import visualization

# TensorFlow session: grow memory when needed. TF, DO NOT USE ALL MY GPU MEMORY!!!
gpu_options = tf.GPUOptions(allow_growth=True)
config = tf.ConfigProto(log_device_placement=False, gpu_options=gpu_options)
isess = tf.InteractiveSession(config=config)

# Input placeholder.
net_shape = (300, 300)
data_format = 'NHWC'
img_input = tf.placeholder(tf.uint8, shape=(None, None, 3))
# Evaluation pre-processing: resize to SSD net shape.
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)
image_4d = tf.expand_dims(image_pre, 0)

# Define the SSD model.
reuse = True if 'ssd_net' in locals() else None
ssd_net = ssd_vgg_300.SSDNet()
with slim.arg_scope(ssd_net.arg_scope(data_format=data_format)):
    predictions, localisations, _, _ = ssd_net.net(image_4d, is_training=False, reuse=reuse)

# Restore SSD model.
ckpt_filename = '../log/model.ckpt-40000'  #此處改爲自己的模型
isess.run(tf.global_variables_initializer())
saver = tf.train.Saver()
saver.restore(isess, ckpt_filename)

# SSD default anchor boxes.
ssd_anchors = ssd_net.anchors(net_shape)

# Main image processing routine.
def process_image(img, select_threshold=0.5, nms_threshold=.45, net_shape=(300, 300)):
    # Run SSD network.
    rimg, rpredictions, rlocalisations, rbbox_img = isess.run([image_4d, 
    predictions, localisations, bbox_img],feed_dict={img_input: img})
    
    # Get classes and bboxes from the net outputs.
    rclasses, rscores, rbboxes = np_methods.ssd_bboxes_select(
            rpredictions, rlocalisations, ssd_anchors,
            select_threshold=select_threshold, img_shape=net_shape, num_classes=9, 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)
    # Resize bboxes to original image shape. Note: useless for Resize.WARP!
    rbboxes = np_methods.bboxes_resize(rbbox_img, rbboxes)
    return rclasses, rscores, rbboxes

# Test on some demo image and visualize output.
path = '../demo/'    #存放驗證圖片的位置
image_names = sorted(os.listdir(path))

for i in range(19):    #此處的19爲本人的demo下放了19張圖片，可以一次性檢測19張圖片
	img = mpimg.imread(path + image_names[-i+1])
	rclasses, rscores, rbboxes =  process_image(img)
	# visualization.bboxes_draw_on_img(img, rclasses, rscores, rbboxes, visualization.colors_plasma)
	visualization.plt_bboxes(i, img, rclasses, rscores, rbboxes)

4.2.2. 修改 visualization.py

====================================================================================
# 修改45行，將函數定義中的參數添加一個 img
def draw_lines(img, lines, color=[255, 0, 0], thickness=2):

====================================================================================
# 將最後一行的
plt.show()

# 修改爲
plt.savefig("%d.jpg"%(nm+1))
====================================================================================

執行 ssd_notebook.py ，即可在notebooks文件夾下看到檢測的結果。

---

參考資料：
[1]. SSD-tensorflow 訓練自己的數據並顯示訓練結果
[2]. 目標檢測SSD+Tensorflow 訓練自己的數據集