SSD: Signle Shot Detector 用於自然場景文字檢測

SSD: Signle Shot Detector 用於自然場景文字檢測

前言

之前我在 論文閱讀:SSD: Single Shot MultiBox Detector 中,講了這個最新的 Object Detection 演算法。

既然 SSD 是用來檢測物體的,那麼可不可以將 SSD 用來檢測自然場景影象中的文字呢?答案肯定是可以的~

同時,受到浙大 solace_hyh 同學的 ssd-plate_detection 工作,這篇文章記錄我自己將 SSD 用於文字檢測的過程。

全部的程式碼上傳到 Github 了:https://github.com/chenxinpeng/SSD_scene-text-detection程式碼質量不太高,還請高手指點 。^_^

準備與轉換資料集

ICDAR 2011 資料集訓練集共有 229 張影象,我將其分為 159 張、70張影象兩部分。前者用作訓練,後者用於訓練時進行測試。

下面就是要將這些影象,轉換成 lmdb 格式,用於 caffe 訓練;將文字區域的標籤,轉換為 Pascal VOC 的 XML 格式。

將 ground truth 轉換為 Pascal VOC XML 檔案

先將 ICDAR 2011 給定的 gt_**.txt 標籤檔案轉換為 Pascal VOC XML 格式。

先看下原來的 gt_**.txt 格式,如下圖,有一張原始影象:

下面是其 ground truth 檔案:

158,128,412,182,"Footpath"
442,128,501,170,"To"
393,198,488,240,"and"
63,200,363,242,"Colchester"
71,271,383,313,"Greenstead"

ground truth 檔案格式為:xmin, ymin, xmax, ymax, labelx_{min},\ y_{min},\ x_{max},\ y_{max},\ label。同時,要注意,這裡的座標系是如下襬放:

將 ground truth 的 txt 檔案轉換為 Pascal VOC 的 XML 格式的程式碼如下:

#! /usr/bin/python
import os, sys
import glob
from PIL import Image
# ICDAR 影象儲存位置
src_img_dir = "/media/chenxp/Datadisk/ocr_dataset/ICDAR2011/train-textloc"
# ICDAR 影象的 ground truth 的 txt 檔案存放位置
src_txt_dir = "/media/chenxp/Datadisk/ocr_dataset/ICDAR2011/train-textloc"
img_Lists = glob.glob(src_img_dir   '/*.jpg')
img_basenames = [] # e.g. 100.jpg
for item in img_Lists:
img_basenames.append(os.path.basename(item))
img_names = [] # e.g. 100
for item in img_basenames:
temp1, temp2 = os.path.splitext(item)
img_names.append(temp1)
for img in img_names:
im = Image.open((src_img_dir   '/'   img   '.jpg'))
width, height = im.size
# open the crospronding txt file
gt = open(src_txt_dir   '/gt_'   img   '.txt').read().splitlines()
# write in xml file
os.mknod(src_txt_dir   '/'   img   '.xml')
xml_file = open((src_txt_dir   '/'   img   '.xml'), 'w')
xml_file.write('<annotation>\n')
xml_file.write('    <folder>VOC2007</folder>\n')
xml_file.write('    <filename>'   str(img)   '.jpg'   '</filename>\n')
xml_file.write('    <size>\n')
xml_file.write('        <width>'   str(width)   '</width>\n')
xml_file.write('        <height>'   str(height)   '</height>\n')
xml_file.write('        <depth>3</depth>\n')
xml_file.write('    </size>\n')
# write the region of text on xml file
for img_each_label in gt:
spt = img_each_label.split(',')
xml_file.write('    <object>\n')
xml_file.write('        <name>text</name>\n')
xml_file.write('        <pose>Unspecified</pose>\n')
xml_file.write('        <truncated>0</truncated>\n')
xml_file.write('        <difficult>0</difficult>\n')
xml_file.write('        <bndbox>\n')
xml_file.write('            <xmin>'   str(spt[0])   '</xmin>\n')
xml_file.write('            <ymin>'   str(spt[1])   '</ymin>\n')
xml_file.write('            <xmax>'   str(spt[2])   '</xmax>\n')
xml_file.write('            <ymax>'   str(spt[3])   '</ymax>\n')
xml_file.write('        </bndbox>\n')
xml_file.write('    </object>\n')
xml_file.write('</annotation>')

x上面程式碼執行結果是得到如下的 XML 檔案,同樣用上面的 100.jpg 影象示例,其轉換結果如下:

<annotation>
<folder>VOC2007</folder>
<filename>100.jpg</filename>
<size>
<width>640</width>
<height>480</height>
<depth>3</depth>
</size>
<object>
......
</annotation>

上面程式碼生成的 XML 檔案,與影象檔案儲存在一個地方。

生成訓練影象與 XML 標籤的位置檔案

這一步,按照 SSD 訓練的需求,將影象位置,及其對應的 XML 檔案位置寫入一個 txt 檔案,供訓練時讀取,一個檔名稱叫做:trainval.txt 檔案,另一個叫做:test.txt 檔案。形式如下:

scenetext/JPEGImages/106.jpg scenetext/Annotations/106.xml
scenetext/JPEGImages/203.jpg scenetext/Annotations/203.xml
scenetext/JPEGImages/258.jpg scenetext/Annotations/258.xml
scenetext/JPEGImages/122.jpg scenetext/Annotations/122.xml
scenetext/JPEGImages/103.jpg scenetext/Annotations/103.xml
scenetext/JPEGImages/213.jpg scenetext/Annotations/213.xml
scenetext/JPEGImages/149.jpg scenetext/Annotations/149.xml
......

生成的程式碼如下:

#! /usr/bin/python
import os, sys
import glob
trainval_dir = "/home/chenxp/data/VOCdevkit/scenetext/trainval"
test_dir = "/home/chenxp/data/VOCdevkit/scenetext/test"
trainval_img_lists = glob.glob(trainval_dir   '/*.jpg')
trainval_img_names = []
for item in trainval_img_lists:
temp1, temp2 = os.path.splitext(os.path.basename(item))
trainval_img_names.append(temp1)
test_img_lists = glob.glob(test_dir   '/*.jpg')
test_img_names = []
for item in test_img_lists:
temp1, temp2 = os.path.splitext(os.path.basename(item))
test_img_names.append(temp1)
dist_img_dir = "scenetext/JPEGImages"
dist_anno_dir = "scenetext/Annotations"
trainval_fd = open("/home/chenxp/caffe/data/scenetext/trainval.txt", 'w')
test_fd = open("/home/chenxp/caffe/data/scenetext/test.txt", 'w')
for item in trainval_img_names:
trainval_fd.write(dist_img_dir   '/'   str(item)   '.jpg'   ' '   dist_anno_dir   '/'   str(item)   '.xml\n')
for item in test_img_names:
test_fd.write(dist_img_dir   '/'   str(item)   '.jpg'   ' '   dist_anno_dir   '/'   str(item)   '.xml\n')

生成 test name size 文字檔案

這一步,SSD 還需要一個名叫:test_name_size.txt 的檔案,裡面記錄訓練影象、測試影象的影象名稱、height、width。內容形式如下:

106 480 640
203 480 640
258 480 640
318 480 640
122 480 640
103 480 640
320 640 480
......

生成這個文字檔案的程式碼如下:

#! /usr/bin/python
import os, sys
import glob
from PIL import Image
img_dir = "/home/chenxp/data/VOCdevkit/scenetext/JPEGImages"
img_lists = glob.glob(img_dir   '/*.jpg')
test_name_size = open('/home/chenxp/caffe/data/scenetext/test_name_size.txt', 'w')
for item in img_lists:
img = Image.open(item)
width, height = img.size
temp1, temp2 = os.path.splitext(os.path.basename(item))
test_name_size.write(temp1   ' '   str(height)   ' '   str(width)   '\n')

準備標籤對映檔案 labelmap

這個 prototxt 檔案是記錄 label 與 name 之間的對應關係的,內容如下:

item {
name: "none_of_the_above"
label: 0
display_name: "background"
}
item {
name: "object"
label: 1
display_name: "text"
}

我的 prototxt 檔名稱,被我重新命名為:labelmap_voc.prototxt

生成 lmdb 資料庫

準備好上述的幾個文字檔案,將其放置在如下位置:

/home/chenxp/caffe/data/scenetext

這時候,需要修改呼叫 SSD 原始碼中提供的 create_data.sh 指令碼檔案(我將檔案重新命名為:create_data_scenetext.sh):

cur_dir=$(cd $( dirname ${BASH_SOURCE[0]} ) && pwd )
root_dir=$cur_dir/../..
cd $root_dir
redo=1
data_root_dir="$HOME/data/VOCdevkit"
dataset_name="scenetext"
mapfile="$root_dir/data/$dataset_name/labelmap_voc_scenetext.prototxt"
anno_type="detection"
db="lmdb"
min_dim=0
max_dim=0
width=0
height=0
extra_cmd="--encode-type=jpg --encoded"
if [ $redo ]
then
extra_cmd="$extra_cmd --redo"
fi
for subset in test trainval
do
python $root_dir/scripts/create_annoset.py --anno-type=$anno_type --label-map-file=$mapfile \
--min-dim=$min_dim --max-dim=$max_dim --resize-width=$width --resize-height=$height \
--check-label $extra_cmd $data_root_dir $root_dir/data/$dataset_name/$subset.txt \
$data_root_dir/$dataset_name/$db/$dataset_name"_"$subset"_"$db examples/$dataset_name
done

上面的 bash 指令碼會自動將訓練的 ICDAR 2011 的影象檔案與對應 label 轉換為 lmdb 檔案。轉換後的檔案位置可參見上面指令碼的內容,我的位置為:

/home/chenxp/caffe/examples/scenetext_trainval_lmdb
/home/chenxp/caffe/examples/scenetext_test_lmdb

訓練模型

將 SSD 用於自己的檢測任務,是需要 Fine-tuning a pretrained network 的。

具體的,需要載入 SSD 作者提供的 VGG_ILSVRC_16_layers_fc_reduced.caffemodel,在這個預訓練的模型上,繼續用我們的資料訓練。

下載下來後,放在如下位置下面:

/home/chenxp/caffe/models/VGGNet

之後,修改作者提供的訓練 Python 程式碼:ssd_pascal.py,這份程式碼會自動建立訓練所需要的如下幾個檔案:

  • deploy.prototxt
  • solver.prototxt
  • trainval.prototxt
  • test.prototxt

我們需要按照自己的情況,修改如下幾處地方:

# Modify the job name if you want.
job_name = "SSD_{}".format(resize)
# The name of the model. Modify it if you want.
model_name = "VGG_VOC0712_{}".format(job_name)
# Directory which stores the model .prototxt file.
save_dir = "models/VGGNet/VOC0712/{}".format(job_name)
# Directory which stores the snapshot of models.
snapshot_dir = "models/VGGNet/VOC0712/{}".format(job_name)
# Directory which stores the job script and log file.
job_dir = "jobs/VGGNet/VOC0712/{}".format(job_name)
# Directory which stores the detection results.
output_result_dir = "{}/data/VOCdevkit/results/VOC2007/{}/Main".format(os.environ['HOME'], job_name)
# model definition files.
train_net_file = "{}/train.prototxt".format(save_dir)
test_net_file = "{}/test.prototxt".format(save_dir)
deploy_net_file = "{}/deploy.prototxt".format(save_dir)
solver_file = "{}/solver.prototxt".format(save_dir)
# snapshot prefix.
snapshot_prefix = "{}/{}".format(snapshot_dir, model_name)
# job script path.
job_file = "{}/{}.sh".format(job_dir, model_name)
# Stores the test image names and sizes. Created by data/VOC0712/create_list.sh
name_size_file = "data/VOC0712/test_name_size.txt"
# The pretrained model. We use the Fully convolutional reduced (atrous) VGGNet.
pretrain_model = "models/VGGNet/VGG_ILSVRC_16_layers_fc_reduced.caffemodel"
# Stores LabelMapItem.
label_map_file = "data/VOC0712/labelmap_voc.prototxt"
num_classes = 21
num_test_image = 4952

我的訓練引數

其實還需要修改一些,如訓練時的引數。因為一開始若直接用作者 ssd_pascal.py 檔案中的預設的 solver.prototxt 引數,會出現如下情況:

這裡寫圖片描述

跑著跑著,loss 就變成 nan 了,發散了,不收斂。

我除錯了一段時間,我的 solver.prototxt 引數設定如下,可保證收斂:

base_lr: 0.0001

其餘引數可看自己設定。學習率一定要小,原先的 0.001 就會發散。

訓練結束:

這裡寫圖片描述

可以看見,最後的測試精度為 0.776573,感覺 SSD 效果還可以。

我自己訓練好的模型,上傳到雲端了:連結:http://share.weiyun.com/1c544de66be06ea04774fd11e820a780 (密碼:ERid5Y)

這個需要在下一階段的測試中用到。

用訓練好的 model 進行 predict

SSD 的作者也給我們寫好了 predict 的程式碼,我們只需要該引數就可以了。

用 jupyter notebook 開啟 ~/caffe/examples/ssd_detect.ipynb 檔案,這是作者為我們寫好的將訓練好的 caffemodel 用於檢測的檔案。

指定好 caffemodeldeploy.txt,詳細的看我上傳的程式碼吧。

測試幾張影象,結果如下:

這裡寫圖片描述
這裡寫圖片描述
這裡寫圖片描述

參考

  1. ECCV2016 Paper: 《SSD: Single Shot MultiBox Detector》
  2. SSD 原始碼
  3. SSD-plate_detection from solace_hyh
  4. SSD框架訓練自己的資料集