天天看點

SSD: Single Shot MultiBox Detector 訓練KITTI資料集(1)

前言

之前介紹了SSD的基本用法和檢測單張圖檔的方法,那麼本篇部落格将詳細記錄如何使用SSD檢測架構訓練KITTI資料集。SSD項目中自帶了用于訓練PASCAL VOC資料集的腳本,基本不用做修改就可以輕松完成訓練;但是想要訓練其他資料集比如KITTI,則需做很大的調整。本文所有工具源碼都已公開,請根據實際情況自行修改。

下載下傳資料集

部落客打算将SSD算法用于檢測車載視訊,用到的是 KITTI資料集 。簡單介紹一下,KITTI資料集由德國卡爾斯魯厄理工學院和豐田美國技術研究院聯合創辦,是目前國際上最大的自動駕駛場景下的計算機視覺算法評測資料集。用于評測目标(機動車、非機動車、行人等)檢測、目标跟蹤、路面分割等計算機視覺技術在車載環境下的性能。KITTI包含市區、鄉村和高速公路等場景采集的真實圖像資料,每張圖像中最多達15輛車和30個行人,還有各種程度的遮擋(ps:歐洲道路狀況和中國還是很不相同,期待國内早日能有同類資料集)。

進入官網,找到object一欄,準備下載下傳資料集:

SSD: Single Shot MultiBox Detector 訓練KITTI資料集(1)

根據下載下傳情況(部落客把前四個都下載下傳了,點開看過),進行SSD訓練隻需要下載下傳第1個圖檔集 Download left color images of object data set (12 GB)和标注檔案 Download training labels of object data set (5 MB) 就夠了。然後将其解壓,發現其中7481張訓練圖檔有标注資訊,而測試圖檔沒有,這就是本次訓練所使用的圖檔數量。由于SSD中訓練腳本是基于VOC資料集格式的,是以我們需要把KITTI資料集做成PASCAL VOC的格式,其基本架構可以參看這篇部落格:PASCAL VOC資料集分析 。根據SSD訓練要求,部落客在/home/mx/data/中目錄中建立一系列檔案夾存放所需資料集和工具檔案,具體如下:

PS.參看截圖,資料要放在home目錄下的data檔案夾,不是caffe中的data檔案夾,這個要注意,否則後續腳本出錯。

# 在data/檔案夾下建立KITTIdevkit/KITTI兩層子目錄,所需檔案放在KITTI/中
Annotations/
└── xml 
ImageSets/
└── main/
      └── trainval.txt
      └── test.txt # 等等
JPEGImages/
└── png
Labels/
└── txt # 自建檔案夾,存放原始标注資訊,待轉化為xml,不屬于VOC格式
create_train_test_txt.py # 3個python工具,後面有詳細介紹
modify_annotations_txt.py
txt_to_xml.py
           
SSD: Single Shot MultiBox Detector 訓練KITTI資料集(1)

(截圖來源于小規模試驗,圖檔隻有400張)

轉換資料集

為了友善SSD進行訓練,我們需要将KITTI資料集轉換成PASCAL VOC的格式,細心的朋友可能已經發現,KITTI官網提供了一個工具: code to convert from KITTI to PASCAL VOC file format ,為啥不用呢?因為我覺得很難用,缺乏靈活性,還不如自己的Python轉換工具好使。

KITTI标注資訊說明

KITTI資料集中标注資訊是存放在txt文本中的,我們随便複制一些标注語句,看看都包含了那些資訊:

Car 0.00 0 -1.67 642.24 178.50 680.14 208.68 1.38 1.49 3.32 2.41 1.66 34.98 -1.60

Car 0.00 0 -1.75 685.77 178.12 767.02 235.21 1.50 1.62 3.89 3.27 1.67 21.18 -1.60

具體的含義在官網沒找到,但是部落客偶然在DIGITS項目中看到了KITTI标注資訊的明确含義:

SSD: Single Shot MultiBox Detector 訓練KITTI資料集(1)

上圖連結:Object Detection Data Extension ,可以看到,KITTI的标注資訊中,SSD訓練需要使用的隻有類别’Car‘和物體外框的坐标‘387.63 181.54 423.81 203.12’,其餘的字段都可以忽略。

轉換KITTI類别

PASCAL VOC資料集總共20個類别,如果用于特定場景,20個類别确實多了。此次部落客為資料集設定3個類别, ‘Car’,’Cyclist’,’Pedestrian’,隻不過标注資訊中還有其他類型的車和人,直接略過有點浪費,部落客希望将 ‘Van’, ‘Truck’, ‘Tram’ 合并到 ‘Car’ 類别中去,将 ‘Person_sitting’ 合并到 ‘Pedestrian’ 類别中去(‘Misc’ 和 ‘Dontcare’ 這兩類直接忽略)。這裡使用的是modify_annotations_txt.py工具,源碼如下:

# modify_annotations_txt.py
import glob
import string

txt_list = glob.glob('./Labels/*.txt') # 存儲Labels檔案夾所有txt檔案路徑
def show_category(txt_list):
    category_list= []
    for item in txt_list:
        try:
            with open(item) as tdf:
                for each_line in tdf:
                    labeldata = each_line.strip().split(' ') # 去掉前後多餘的字元并把其分開
                    category_list.append(labeldata[]) # 隻要第一個字段,即類别
        except IOError as ioerr:
            print('File error:'+str(ioerr))
    print(set(category_list)) # 輸出集合

def merge(line):
    each_line=''
    for i in range(len(line)):
        if i!= (len(line)-):
            each_line=each_line+line[i]+' '
        else:
            each_line=each_line+line[i] # 最後一條字段後面不加空格
    each_line=each_line+'\n'
    return (each_line)

print('before modify categories are:\n')
show_category(txt_list)

for item in txt_list:
    new_txt=[]
    try:
        with open(item, 'r') as r_tdf:
            for each_line in r_tdf:
                labeldata = each_line.strip().split(' ')
                if labeldata[] in ['Truck','Van','Tram']: # 合并汽車類
                    labeldata[] = labeldata[].replace(labeldata[],'Car')
                if labeldata[] == 'Person_sitting': # 合并行人類
                    labeldata[] = labeldata[].replace(labeldata[],'Pedestrian')
                if labeldata[] == 'DontCare': # 忽略Dontcare類
                    continue
                if labeldata[] == 'Misc': # 忽略Misc類
                    continue
                new_txt.append(merge(labeldata)) # 重新寫入新的txt檔案
        with open(item,'w+') as w_tdf: # w+是打開原檔案将内容删除,另寫新内容進去
            for temp in new_txt:
                w_tdf.write(temp)
    except IOError as ioerr:
        print('File error:'+str(ioerr))

print('\nafter modify categories are:\n')
show_category(txt_list)
           

執行指令:

python modify_annotations_txt.py

來運作py程式,這裡以000400.txt為例,顯示轉換前後的對比效果:

# 轉換前
Car   -           -
Car   -           -
Car   -           -
Truck   -           -
Car   -           -
DontCare - - -     - - - - - - -
DontCare - - -     - - - - - - -

# 轉換後
Car   -           -
Car   -           -
Car   -           -
Car   -           -
Car   -           -
           

轉換txt标注資訊為xml格式

對原始txt檔案進行上述處理後,接下來需要将标注檔案從txt轉化為xml,并去掉标注資訊中用不上的部分,隻留下3類,還有把坐标值從float型轉化為int型,最後所有生成的xml檔案要存放在Annotations檔案夾中。這裡使用的是txt_to_xml.py工具,此處是由 KITTI_SSD 的代碼修改而來,感謝作者的貢獻。

# txt_to_xml.py
# encoding:utf-8
# 根據一個給定的XML Schema,使用DOM樹的形式從空白檔案生成一個XML
from xml.dom.minidom import Document
import cv2
import os

def generate_xml(name,split_lines,img_size,class_ind):
    doc = Document()  # 建立DOM文檔對象

    annotation = doc.createElement('annotation')
    doc.appendChild(annotation)

    title = doc.createElement('folder')
    title_text = doc.createTextNode('KITTI')
    title.appendChild(title_text)
    annotation.appendChild(title)

    img_name=name+'.png'

    title = doc.createElement('filename')
    title_text = doc.createTextNode(img_name)
    title.appendChild(title_text)
    annotation.appendChild(title)

    source = doc.createElement('source')
    annotation.appendChild(source)

    title = doc.createElement('database')
    title_text = doc.createTextNode('The KITTI Database')
    title.appendChild(title_text)
    source.appendChild(title)

    title = doc.createElement('annotation')
    title_text = doc.createTextNode('KITTI')
    title.appendChild(title_text)
    source.appendChild(title)

    size = doc.createElement('size')
    annotation.appendChild(size)

    title = doc.createElement('width')
    title_text = doc.createTextNode(str(img_size[]))
    title.appendChild(title_text)
    size.appendChild(title)

    title = doc.createElement('height')
    title_text = doc.createTextNode(str(img_size[]))
    title.appendChild(title_text)
    size.appendChild(title)

    title = doc.createElement('depth')
    title_text = doc.createTextNode(str(img_size[]))
    title.appendChild(title_text)
    size.appendChild(title)

    for split_line in split_lines:
        line=split_line.strip().split()
        if line[] in class_ind:
            object = doc.createElement('object')
            annotation.appendChild(object)

            title = doc.createElement('name')
            title_text = doc.createTextNode(line[])
            title.appendChild(title_text)
            object.appendChild(title)

            bndbox = doc.createElement('bndbox')
            object.appendChild(bndbox)
            title = doc.createElement('xmin')
            title_text = doc.createTextNode(str(int(float(line[]))))
            title.appendChild(title_text)
            bndbox.appendChild(title)
            title = doc.createElement('ymin')
            title_text = doc.createTextNode(str(int(float(line[]))))
            title.appendChild(title_text)
            bndbox.appendChild(title)
            title = doc.createElement('xmax')
            title_text = doc.createTextNode(str(int(float(line[]))))
            title.appendChild(title_text)
            bndbox.appendChild(title)
            title = doc.createElement('ymax')
            title_text = doc.createTextNode(str(int(float(line[]))))
            title.appendChild(title_text)
            bndbox.appendChild(title)

    # 将DOM對象doc寫入檔案
    f = open('Annotations/'+name+'.xml','w')
    f.write(doc.toprettyxml(indent = ''))
    f.close()

if __name__ == '__main__':
    class_ind=('Pedestrian', 'Car', 'Cyclist')
    cur_dir=os.getcwd()
    labels_dir=os.path.join(cur_dir,'Labels')
    for parent, dirnames, filenames in os.walk(labels_dir): # 分别得到根目錄,子目錄和根目錄下檔案   
        for file_name in filenames:
            full_path=os.path.join(parent, file_name) # 擷取檔案全路徑
            f=open(full_path)
            split_lines = f.readlines()
            name= file_name[:-] # 後四位是擴充名.txt,隻取前面的檔案名
            img_name=name+'.png' 
            img_path=os.path.join('/home/mx/KITTI/train_image',img_name) # 路徑需要自行修改            
            img_size=cv2.imread(img_path).shape
            generate_xml(name,split_lines,img_size,class_ind)
print('all txts has converted into xmls')
           

執行指令:

python txt_to_xml.py

來運作py程式,轉換效果如下:

# 原始的000400.txt
Car 0.00 0 -1.67 642.24 178.50 680.14 208.68 1.38 1.49 3.32 2.41 1.66 34.98 -1.60
Car 0.00 0 -1.75 685.77 178.12 767.02 235.21 1.50 1.62 3.89 3.27 1.67 21.18 -1.60
Car 0.67 0 -2.15 885.80 160.44 1241.00 374.00 1.69 1.58 3.95 3.64 1.65 5.47 -1.59
Car 0.00 0 -1.89 755.82 101.65 918.16 230.75 3.55 2.56 7.97 7.06 1.63 23.91 -1.61
Car 0.00 1 -2.73 928.61 177.14 1016.83 209.77 1.48 1.36 3.51 17.33 1.71 34.63 -2.27
# 生成的000400.xml(部分)
This XML file does not appear to have any style information associated with it. The document tree is shown below.
<annotation>
  <folder>KITTI</folder>
  <filename>000400.png</filename>
  <source>
    <database>The KITTI Database</database>
    <annotation>KITTI</annotation>
  </source>
  <size>
    <width>1242</width>
    <height>375</height>
    <depth>3</depth>
  </size>
  <object>
    <name>Car</name>
    <bndbox>
      <xmin>642</xmin>
      <ymin>178</ymin>
      <xmax>680</xmax>
      <ymax>208</ymax>
    </bndbox>
  </object>
  <object>
    <name>Car</name>
    <bndbox>
      <xmin>685</xmin>
      <ymin>178</ymin>
      <xmax>767</xmax>
      <ymax>235</ymax>
    </bndbox>
  </object>
......
</annotation>
           

生成訓練驗證集和測試集清單

用于SSD訓練的Pascal VOC格式的資料集總共就是三大塊:首先是JPEGImages檔案夾,放入了所有png圖檔;然後是Annotations檔案夾,上述步驟已經生成了相應的xml檔案;最後就是imagesSets檔案夾,裡面有一個Main子檔案夾,這個檔案夾存放的是訓練驗證集,測試集的相關清單檔案,如下圖所示:

SSD: Single Shot MultiBox Detector 訓練KITTI資料集(1)

這裡使用create_train_test_txt.py工具,自動生成上述16個txt檔案,其中訓練測試部分的比例可以自行修改,由于這個工具是用Python3寫的,是以執行的時候應該是:

python3 create_train_test_txt.py

# create_train_test_txt.py
# encoding:utf-8
import pdb
import glob
import os
import random
import math

def get_sample_value(txt_name, category_name):
    label_path = './Labels/'
    txt_path = label_path + txt_name+'.txt'
    try:
        with open(txt_path) as r_tdf:
            if category_name in r_tdf.read():
                return ' 1'
            else:
                return '-1'
    except IOError as ioerr:
        print('File error:'+str(ioerr))

txt_list_path = glob.glob('./Labels/*.txt')
txt_list = []

for item in txt_list_path:
    temp1,temp2 = os.path.splitext(os.path.basename(item))
    txt_list.append(temp1)
txt_list.sort()
print(txt_list, end = '\n\n')

# 有部落格建議train:val:test=8:1:1,先嘗試用一下
num_trainval = random.sample(txt_list, math.floor(len(txt_list)*/)) # 可修改百分比
num_trainval.sort()
print(num_trainval, end = '\n\n')

num_train = random.sample(num_trainval,math.floor(len(num_trainval)*/)) # 可修改百分比
num_train.sort()
print(num_train, end = '\n\n')

num_val = list(set(num_trainval).difference(set(num_train)))
num_val.sort()
print(num_val, end = '\n\n')

num_test = list(set(txt_list).difference(set(num_trainval)))
num_test.sort()
print(num_test, end = '\n\n')

pdb.set_trace()

Main_path = './ImageSets/Main/'
train_test_name = ['trainval','train','val','test']
category_name = ['Car','Pedestrian','Cyclist']

# 循環寫trainvl train val test
for item_train_test_name in train_test_name:
    list_name = 'num_'
    list_name += item_train_test_name
    train_test_txt_name = Main_path + item_train_test_name + '.txt' 
    try:
        # 寫單個檔案
        with open(train_test_txt_name, 'w') as w_tdf:
            # 一行一行寫
            for item in eval(list_name):
                w_tdf.write(item+'\n')
        # 循環寫Car Pedestrian Cyclist
        for item_category_name in category_name:
            category_txt_name = Main_path + item_category_name + '_' + item_train_test_name + '.txt'
            with open(category_txt_name, 'w') as w_tdf:
                # 一行一行寫
                for item in eval(list_name):
                    w_tdf.write(item+' '+ get_sample_value(item, item_category_name)+'\n')
    except IOError as ioerr:
        print('File error:'+str(ioerr))
           

執行程式過程中,如遇到pdb提示,可按c鍵,再按enter鍵。

如果想把标注資料全部作為trainval,而把未标注的資料(大約有7000多圖檔)作為test,需要重新修改腳本,待續。

資料集的後續處理

下面進行資料集的後續處理,在/home.mx/caffe/data之下建立KITTI檔案夾,用于存放本次訓練所需的腳本工具,如下圖所示。

SSD: Single Shot MultiBox Detector 訓練KITTI資料集(1)

生成訓練所需清單檔案

SSD訓練的時候除了需要LMDB格式的資料以外,還需要讀取三個清單檔案,分别是:trainval.txt,test.txt和test_name_size.txt。前兩個txt檔案存放訓練、測試圖檔的png路徑和xml路徑,第三個txt檔案存放測試圖檔的名稱和尺寸。所需工具可以由/home/mx/caffe/data/VOC0712/create_list.sh腳本修改而來。

複制一份上述腳本,并重命名為create_list_kitti.sh,存放在KITTI檔案夾中。經過修改後的腳本檔案如下(雙#号注釋處為部落客修改過的地方):

# create_list_kitti.sh
#!/bin/bash
root_dir=$HOME/data/KITTIdevkit/ ## 自行修改
sub_dir=ImageSets/Main
bash_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
for dataset in trainval test
do
  dst_file=$bash_dir/$dataset.txt
  if [ -f $dst_file ]
  then
    rm -f $dst_file
  fi
  for name in KITTI ## 自行修改
  do
    #if [[ $dataset == "test" && $name == "VOC2012" ]] ## 這段可以注釋掉
    #then
        #continue
    #fi
    echo "Create list for $name $dataset..."
    dataset_file=$root_dir/$name/$sub_dir/$dataset.txt

    img_file=$bash_dir/$dataset"_img.txt"
    cp $dataset_file $img_file
    sed -i "s/^/$name\/JPEGImages\//g" $img_file
    sed -i "s/$/.png/g" $img_file ## 從jpg改為png

    label_file=$bash_dir/$dataset"_label.txt"
    cp $dataset_file $label_file
    sed -i "s/^/$name\/Annotations\//g" $label_file
    sed -i "s/$/.xml/g" $label_file

    paste -d' ' $img_file $label_file >> $dst_file

    rm -f $label_file
    rm -f $img_file
  done

  # Generate image name and size infomation.
  if [ $dataset == "test" ]
  then
    $bash_dir/../../build/tools/get_image_size $root_dir $dst_file $bash_dir/$dataset"_name_size.txt"
  fi

  # Shuffle trainval file.
  if [ $dataset == "trainval" ]
  then
    rand_file=$dst_file.random
    cat $dst_file | perl -MList::Util=shuffle -e 'print shuffle(<STDIN>);' > $rand_file
    mv $rand_file $dst_file
  fi
done
           

執行下面指令,可在/home/mx/caffe/data/KITTI檔案夾下生成3個訓練所需txt檔案。

$ cd ~/caffe
$ ./data/KITTI/create_list_kitti.sh
           

而生成的txt清單格式如下:

# trainval.txt和test.txt檔案格式
KITTI/JPEGImages/png KITTI/Annotations/xml
KITTI/JPEGImages/png KITTI/Annotations/xml
KITTI/JPEGImages/png KITTI/Annotations/xml
KITTI/JPEGImages/png KITTI/Annotations/xml
......
           
# test_name_size.txt檔案格式
  
  
  
  
......
           

準備标簽映射檔案

由于隻有3類,是以可以仿照例子,寫一個labelmap_kitti.prototxt檔案,用于記錄label和name的對應關系,存放在/home/mx/caffe/data/KITTI檔案夾中,具體内容如下:

item {
  name: "none_of_the_above"
  label: 
  display_name: "background"
}
item {
  name: "Car"
  label: 
  display_name: "Car"
}
item {
  name: "Pedestrian"
  label: 
  display_name: "Pedestrian"
}
item {
  name: "Cyclist"
  label: 
  display_name: "Cyclist"
}
           

生成LMDB資料庫

如果前面一切順利,現在就可以生成LMDB檔案了,所需工具可以由/home/mx/caffe/data/VOC0712/create_data.sh腳本修改而來。仍然複制一份上述腳本,并重命名為create_data_kitti.sh,存放在KITTI檔案夾中。經過修改後的腳本檔案如下:

# create_data_kitti.sh
cur_dir=$(cd $( dirname ${BASH_SOURCE[0]} ) && pwd )
root_dir=$cur_dir/../..

cd $root_dir

redo=
data_root_dir="$HOME/data/KITTIdevkit" ## 自行修改
dataset_name="KITTI" ## 自行修改
mapfile="$root_dir/data/$dataset_name/labelmap_kitti.prototxt" ## 自行修改
anno_type="detection"
db="lmdb"
min_dim=
max_dim=
width=
height=

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
           

執行指令

./data/KITTI/create_data_kitti.sh

來運作腳本,将會生成兩份LMDB檔案,路徑分别如下:

$ /home/mx/caffe/examples/KITTI/KITTI_test_lmdb
$ /home/mx/caffe/examples/KITTI/KITTI_trainval_lmdb
           

至此,訓練資料可以說已經準備好了,下一篇部落格将記錄訓練SSD模型的過程,敬請期待。

繼續閱讀