Ubuntu 16.04下利用tf-faster-rcnn在VOC或其他数据集上训练自己的模型

暑期的时候有朋友问我如何配置tf-faster-rcnn，那会简单写了个steps.pdf。正好最近闲了下来，就把它汇总一下。

Introduction

这是一篇介绍如何使用tf-faster-rcnn在VOC和其他数据集训练自己模型的笔记.
笔记中所有文件的地址: github
我的机器和库版本: GTX 1060, miniconda 4.5.4, CUDA 9.0, CUDNN 7.1.4, tensorflow-gpu 1.8.0.

1 配置tf-faster-rcnn

这里参考https://github.com/endernewton/tf-faster-rcnn，具体细节有一些不同.

1.1 Intall opencv, cython, easydict

我使用的是opencv 3.4.1, cython 0.28.4 and easydict 1.7.

1.2 Clone the repository

git clone https://github.com/endernewton/tf-faster-rcnn.git

1.3 Build the Cython modules

cd tf-faster-rcnn/lib
make clean
make
cd ..

1.4 Install the Python COCO API

git clone https://github.com/pdollar/coco.git
cd coco/PythonAPI
make
cd ../..

1.5 Setup data

cd data
wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar
wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar
wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCdevkit_08-Jun-2007.tar
tar xvf VOCtrainval_06-Nov-2007.tar
tar xvf VOCtest_06-Nov-2007.tar
tar xvf VOCdevkit_08-Jun-2007.tar
ln -s VOCdevkit VOCdevkit2007
cd ..

1.6 Download pre-trained model

Google Driver: link
Baidu Cloud: link
下载模型并放入目录tf-faster-rcnn/, 在当前目录执行命令

tar xvf voc_0712_80k-110k.tgz

1.7 Create a folder and a soft link to use the pre-trained model

NET=res101
TRAIN_IMDB=voc_2007_trainval+voc_2012_trainval
mkdir -p output/$NET/$TRAIN_IMDB
cd output/$NET/$TRAIN_IMDB
ln -s ../../../voc_2007_trainval+voc_2012_trainval ./default
cd ../../..

1.8 Modify python file

修改tf-faster-rcnn/lib/datasets/voc_eval.py, 121 line.

# save
print('Saving cached annotations to {:s}'.format(cachefile))
with open(cachefile, 'w') as f: ---> with open(cachefile, 'wb') as f:
  pickle.dump(recs, f)

1.9 Demo for testing on custom images

修改python文件之后, 执行以下命令.

GPU_ID=0
CUDA_VISIBLE_DEVICES=$GPU_ID ./tools/demo.py

1.10 End

至此为止, 只需要对./tools/demo.py进行轻微修改, 你就可以使用这个模型来预测自己的数据集(数据集中的category必须在VOC数据集里面). 当然mean iou肯定不会很高的, 大约会在0.20~0.60这个区间.

2 Train your own model on the VOC dataset

2.1 Download CNN pre-trained model

mkdir -p data/imagenet_weights
cd data/imagenet_weights
wget -v http://download.tensorflow.org/models/resnet_v1_101_2016_08_28.tar.gz
tar -xzvf resnet_v1_101_2016_08_28.tar.gz
mv resnet_v1_101.ckpt res101.ckpt
cd ../..

2.2 Modify shell file

为了节约时间, 这里将shell文件中的ITES修改为300.

2.2.1 tf-faster-rcnn/experiments/scripts/train_faster_rcnn.sh

pascal_voc)
    TRAIN_IMDB="voc_2007_trainval"
    TEST_IMDB="voc_2007_test"
    STEPSIZE="[50000]"
    ITERS=70000 ---> 300 
    ANCHORS="[8,16,32]"
    RATIOS="[0.5,1,2]"
    ;;

2.2.2 tf-faster-rcnn/experiments/scripts/test_faster_rcnn.sh

pascal_voc)
    TRAIN_IMDB="voc_2007_trainval"
    TEST_IMDB="voc_2007_test"
    ITERS=70000 ---> 300
    ANCHORS="[8,16,32]"
    RATIOS="[0.5,1,2]"
    ;;

2.3 Train your model on pascal_voc

因为train脚本里自动执行test脚本, 而VOC2007的测试集有4952张图片, 所以test过程会花些时间. 为了节约时间, 我们只保留data/VOCdevkit2007/VOC2007/ImageSets/Main/test.txt里前200行. 最后执行下面命令进行模型的训练

./experiments/scripts/train_faster_rcnn.sh 0 pascal_voc res101

2.4 End

到此为止, 你可以在VOC数据集上训练自己的模型.

3 Train your own model on other dataset

这里给大家拿一个车辆检测比赛的数据集举例子, 这个数据集有两类文件(详情见数据集说明). 首先我们把下载的数据集解压后放入一个新目录tf-faster-rcnn/dataset.

3.1 Data Preprocess

3.1.1 清空目录data/VOCdevkit2007/VOC2007/

在目录tf-faster-rcnn/执行命令

rm -rf data/VOCdevkit2007/VOC2007/*

3.1.2 data/VOCdevkit2007/VOC2007/Annotations

这个目录存储训练集所有图像的Annotations文件(image_name.xml)
这里给出一个notebook 来演示如何获取.xml文件

3.1.3 data/VOCdevkit2007/VOC2007/ImageSets/Main

1. train.txt: 训练集中所有图像名(不带后缀).
2. trainval.txt: 训练集和验证集中所有图像名.
3. val.txt: 验证集中所有图像名.
4. test.txt: 测试集中所有图像名 (test脚本以它为基准来评估mean iou的, 所以如果你的测试集没有xml文件, 那你不可以将测试集放入. 相反的, 你应该在此写入验证集中所有图像名).

注意: test.txt和val.txt在这个数据集里是一样的(因为这个数据集的测试集是没有xml文件的), copy一下val.txt然后重命名即可.

3.1.4 获取data/VOCdevkit2007/VOC2007/JPEGImages

所有图像文件(同理如果你的测试集没有xml文件, 你不可以将其放入该目录).
这里给出一个notebook演示如何获取.txt和.jpg文件

3.1.5 End

接下来只需要把dataset/下的Annotations, ImageSets/Main, JPEGImages移动到目录data/VOCdevkit2007/VOC2007/. 至此, 数据集已准备完毕.

3.2 Modify the program

3.2.1 lib/datasets/pascal_voc.py

1. 修改classes, 36 line.

   self._classes = ('__background__', # always index 0
                    'car')
   

2. 如果你的数据集是以0为基准的, 需要移除-1运算, 169 line.

   x1 = float(bbox.find('xmin').text)
   y1 = float(bbox.find('ymin').text)
   x2 = float(bbox.find('xmax').text)
   y2 = float(bbox.find('ymax').text)
   

3.2.2 lib/datasets/imdb.py

修改程序, 105 - 124 line.

def _get_widths(self):
    return [PIL.Image.open(self.image_path_at(i)).size[0]
            for i in range(self.num_images)]

def _get_heights(self):
    return [PIL.Image.open(self.image_path_at(i)).size[1]
            for i in range(self.num_images)]
            
def append_flipped_images(self):
    num_images = self.num_images
    widths = self._get_widths()
    heights = self._get_heights()
    for i in range(num_images):
    	boxes = self.roidb[i]['boxes'].copy()
	oldx1 = boxes[:, 0].copy()
	oldx2 = boxes[:, 2].copy()
	boxes[:, 0] = widths[i] - oldx2 - 1
	boxes[:, 2] = widths[i] - oldx1 - 1
	for ids in range(len(boxes)):
	    if boxes[ids][2] < boxes[ids][0]:
		boxes[ids][0] = 0
	assert (boxes[:, 2] >= boxes[:, 0]).all()
	entry = {'boxes': boxes,
		'gt_overlaps': self.roidb[i]['gt_overlaps'],
		'gt_classes': self.roidb[i]['gt_classes'],
		'flipped': True}
	self.roidb.append(entry)
self._image_index = self._image_index * 2

3.2.3 lib/datasets/voc_eval.py

27 - 30 line, int() -> float().

obj_struct['bbox'] = [float(bbox.find('xmin').text),
                      float(bbox.find('ymin').text),
                      float(bbox.find('xmax').text),
                      float(bbox.find('ymax').text)]

3.3 Train your model

3.3.1 Training

1. Remove old model

   rm -rf output
   

2. Clear cache

   rm data/cache/voc_2007_test_gt_roidb.pkl
   rm data/cache/voc_2007_trainval_gt_roidb.pkl
   rm data/VOCdevkit2007/VOC2007/ImageSets/Main/test.txt_annots.pkl
   

3. Update parameters
   3.1. 在 experiments/scripts/train_faster_rcnn.sh中修改 stepsize 和 iteations.
   3.2. 在 experiments/cfgs中修改模型的参数.
   3.3. 在lib/model/config.py中修改训练参数和测试参数.
   注意: 必须保证train_faster_rcnn.sh和test_faster_rcnn.sh中的ITERS相同.
4. Train model

   ./experiments/scripts/train_faster_rcnn.sh 0 pascal_voc res101
   

3.3.2 Predict

通过./tools/demo.py 可以获取测试图像的边界框. 这里我给出一个例子demo example , 下面是对输入的简单解释:

1. demo_net: 分类网络架构
2. demo_ite: 网络的ITERS, 可以选择不同ITERS时期的网络来做测试
3. demo_dir: 测试集目录
4. demo_vis: 是否可视化测试的图像
5. write_and: 是否将预测的boxes写入csv文件
6. dataset: 选择数据集格式

Note: 为了更好的观察预测效果, 可以先在data/demo里面放几张预测图片, 执行时将demo_vis设置为True. 如果可视化的检测结果没有很大偏差, 这时只需修改demo_dir同时将demo_vis设置为False，即可对全部的数据进行预测,(如果同时可视化很多图片是很可怕的, 切记!!!).

3.4 How to improve the performance of your model

你需要理解透彻这个项目里的所有程序, 只有这样你才可以通过修改参数, 加入数据增强, 选择不同的度量等方法来训练自己的网络.