【AI實戰】手把手教你訓練自己的目標檢測模型(SSD篇)

目標檢測是AI的一項重要應用,通過目標檢測模型能在圖像中把人、動物、汽車、飛機等目標物體檢測出來,甚至還能將物體的輪廓描繪出來,就像下面這張圖,是不是很酷炫呢,嘿嘿

在動手訓練自己的目標檢測模型之前,建議先了解一下目標檢測模型的原理(見文章:大話目標檢測經典模型RCNN、Fast RCNN、Faster RCNN以及Mark R-CNN),這樣纔會更加清楚模型的訓練過程。

本文將在我們前面搭建好的AI實戰基礎環境上(見文章:AI基礎環境搭建),基於SSD算法,介紹如何使用自己的數據訓練目標檢測模型。SSD,全稱Single Shot MultiBox Detector(單鏡頭多盒檢測器),是Wei Liu在ECCV 2016上提出的一種目標檢測算法,是目前流行的主要檢測框架之一。

本案例要做的識別便是在圖像中識別出熊貓,可愛吧,呵呵
 
下面按照以下過程介紹如何使用自己的數據訓練目標檢測模型:
 
1、安裝標註工具
要使用自己的數據來訓練模型,首先得先作數據標註,也就是先要告訴機器圖像裏面有什麼物體、物體在位置在哪裏,有了這些信息後才能來訓練模型。
(1)標註數據文件
目前流行的數據標註文件格式主要有VOC_2007、VOC_2012,該文本格式來源於Pascal VOC標準數據集,這是衡量圖像分類識別能力的重要基準之一。本文采用VOC_2007數據格式文件,以xml格式存儲,如下:

其中重要的信息有:
filename:圖片的文件名
name:標註的物體名稱
xmin、ymin、xmax、ymax:物體位置的左上角、右下角座標

(2)安裝標註工具
如果要標註的圖像有很多,那就需要一張一張手動去計算位置信息,製作xml文件,這樣的效率就太低了。
所幸,有一位大神開源了一個數據標註工具labelImg,可以通過可視化的操作界面進行畫框標註,就能自動生成VOC格式的xml文件了。該工具是基於Python語言編寫的,這樣就支持在Windows、Linux的跨平臺運行,實在是良心之作啊。安裝方式如下:
a. 下載源代碼
通過訪問labelImg的github頁面(https://github.com/tzutalin/labelImg),下載源代碼。可通過git進行clone,也可以直接下載成zip壓縮格式的文件。
 
在本案例中直接下載成zip文件。
b.安裝編譯
解壓labelImg的zip文件,得到LabelImg-master文件夾。
labelImg的界面是使用PyQt編寫的,由於我們搭建的基礎環境使用了最新版本的anaconda已經自帶了PyQt5,在python3的環境下,只需再安裝lxml即可,進入LabelImg-master目錄進行編譯,代碼如下:

#激活虛擬環境
source activate tensorflow
#在python3環境中安裝PyQt5(anaconda已自帶),如果是在python2環境下,則要安裝PyQt4,PyQt4的安裝方式如下
#conda install -c anaconda pyqt=4.11.4
#安裝xml
conda install xml
#編譯
make qt5py3
#打開標註工具
python3 labelImg.py

成功打開labelImg標註工具的界面如下:

2、標註數據
成功安裝了標註工具後,現在就來開始標註數據了。
(1)創建文件夾
按照VOC數據集的要求,創建以下文件夾
Annotations:用於存放標註後的xml文件
ImageSets/Main:用於存放訓練集、測試集、驗收集的文件列表
JPEGImages:用於存放原始圖像

(2)標註數據
將熊貓圖片集放在JPEGImages文件夾裏面(熊貓的美照請找度娘要哦~),注意圖片的格式必須是jpg格式的。
打開labelImg標註工具,然後點擊左側的工具欄“Open Dir”按鈕,選擇剛纔放熊貓的JPEGImages文件夾。這時,主界面將會自動加載第一張熊貓照片。

點擊左側工具欄的“Create RectBox”按鈕,然後在主界面上點擊拉個矩形框,將熊貓圈出來。圈定後,將會彈出一個對話框,用於輸入標註物體的名稱,輸入panda作爲熊貓的名稱。

然後點擊左側工具欄的“Save”按鈕,選擇剛纔創建的Annotations作爲保存目錄,系統將自動生成voc_2007格式的xml文件保存起來。這樣就完成了一張熊貓照片的物體標註了。

接下來點擊左側工具欄的“Next Image”進入下一張圖像,按照以上步驟,畫框、輸入名稱、保存,如此反覆,直到把所有照片都標註好,保存起來。

(3)劃分訓練集、測試集、驗證集
完成所有熊貓照片的標註後,還要將數據集劃分下訓練集、測試集和驗證集。
在github上下載一個自動劃分的腳本(https://github.com/EddyGao/make_VOC2007/blob/master/make_main_txt.py)
然後執行以下代碼

python make_main_txt.py

將會按照腳本里面設置的比例,自動拆分訓練集、測試集和驗證集,將相應的文件名列表保存在裏面。

3、配置SSD
(1)下載SSD代碼
由於本案例是基於tensorflow的,因此,在github上下載一個基於tensorflow的SSD,地址是 https://github.com/balancap/SSD-Tensorflow
 
以zip文件的方式下載下來,然後解壓,得到SSD-Tensorflow-master文件夾
(2)轉換文件格式
將voc_2007格式的文件轉換爲tfrecord格式,tfrecord數據文件tensorflow中的一種將圖像數據和標籤統一存儲的二進制文件,能更加快速地在tensorflow中複製、移動、讀取和存儲等。
SSD-Tensorflow-master提供了轉換格式的腳本,轉換代碼如下:

DATASET_DIR=./panda_voc2007/
OUTPUT_DIR=./panda_tfrecord/
python SSD-Tensorflow-master/tf_convert_data.py --dataset_name=pascalvoc --dataset_dir=${DATASET_DIR} --output_name=voc_2007_train --output_dir=${OUTPUT_DIR}

(3)修改物體類別
由於是我們自定義的物體,因此,要修改SSD-Tensorflow-master中關於物體類別的定義,打開SSD-Tensorflow-master/datasets/pascalvoc_common.py文件,進行修改,將VOC_LABELS中的其它無關類別全部刪掉,增加panda的名稱、ID、類別,如下:

VOC_LABELS = {
    'none': (0, 'Background'),
'panda': (1, 'Animal'),
}

4、下載預訓練模型
SSD-Tensorflow提供了預訓練好的模型,基於VGG模型(要了解VGG模型詳情,請閱讀文章:大話經典CNN經典模型VGG),如下表:
 
但這些預訓練的模型文件都是存儲在drive.google.com上,因此,無法直接下載。只能通過“你懂的”方式進行下載,在這裏下載SSD-300 VGG-based預訓練模型,得到文件:VGG_VOC0712_SSD_300x300_ft_iter_120000.ckpt.zip,然後進行解壓

5、訓練模型
終於把標註文件、SSD模型都準備好了,現在準備開始來訓練了。
在訓練模型之前,有個參數要修改下,打開SSD-Tensorflow-master/train_ssd_network.py找到裏面的DATA_FORMAT參數項,如果是使用cpu訓練則值爲NHWC,如果是使用gpu訓練則值爲NCHW,如下:

DATA_FORMAT = 'NCHW'  # gpu
# DATA_FORMAT = 'NHWC'    # cpu

現在終於可以開始來訓練了,打開終端,切換conda虛擬環境

source activate tensorflow

然後執行以下命令,開始訓練

# 使用預訓練好的 vgg_ssd_300 模型 
DATASET_DIR=./ panda_tfrecord
TRAIN_DIR=./panda_model
CHECKPOINT_PATH=./model_pre_train/VGG_VOC0712_SSD_300x300_ft_iter_120000.ckpt/VGG_VOC0712_SSD_300x300_ft_iter_120000.ckpt
python3 SSD-Tensorflow-master/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} \
    --save_summaries_secs=60 \
    --save_interval_secs=600 \
    --weight_decay=0.0005 \
    --optimizer=adam \
    --learning_rate=0.0001 \
    --batch_size=16

其中,根據自己電腦的性能情況,設置batch_size的值,值越大表示批量處理的數量越大,對機器性能的要求越高。如果電腦性能普通的,則可以設置爲8,甚至4,土豪請忽略。
學習率learning_rate也可以根據實際情況調整,學習率越小則越精確,訓練的時間也越長,學習率越大則可縮短訓練時間,但就會降低精準度。

在這裏使用預訓練好的模型,SSD將會鎖定VGG模型的一些參數進行訓練,這樣能在較短的時間內完成訓練。

6、使用模型
SSD模型訓練好了,現在要來使用了,使用的方式也很簡單。
SSD-Tensorflow-master自帶了一個notebooks腳本,可通過jupyter直接使用模型。
先安裝jupyter,安裝方式如下:

conda install jupyter

然後啓動jupyter-notebook,代碼如下:

jupyter-notebook SSD-Tensorflow-master/notebooks/ssd_notebook.ipynb

啓動後在SSD 300 Model的代碼塊設置模型的路徑和名稱

然後在最後的代碼塊中,設置要測試的圖像路徑path

然後點擊菜單“Cell”,點擊子菜單“Run All”,便能按順序全部執行代碼,並顯示出結果出來

執行後,可愛的熊貓就被圈出來了

經過以上步驟,我們便使用了自己的數據完成了目標檢測模型的訓練。只要以後還有物體檢測的需求,然後找相關的圖片集進行標註,標註後進行模型訓練,就能完成一個定製化的目標檢測模型了,非常方便,希望本案例對大家能有所幫助。

後面將陸續推出更多【AI實戰】內容,敬請留意。

 

推薦相關閱讀

 

關注本人公衆號“大數據與人工智能Lab”(BigdataAILab),獲取更多信息

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章