Caffe-SSD 训练自己的数据集教程

SSD Train Own Data Tutorial

This tutorial written by Jin Tian, Master in Tsinghua University, if you have any question about this post, contact me via wechat: jintianiloveu. Repost is also welcomed, but please remain this copyright info, enjoy :)

Prepare

从我的github中获取kitti-ssd源代码，注意，我将ssd训练和生成数据的代码和原始的ssd代码分离了，你只需要把原始ssdclone下来编译即可，我们训练数据集只需要调用它，以后的操作都和它没有关系了。

Part A Data Orginization

Before we get started, I have say that SSD original source code data orginization is really a mess. If you want train your own data you don't know where to go. But now, I am going change it, reshape it to a simple and clear way. you just clone source code and make it, the rest thing is all about my code, using my code you can sperate caffe-ssd source code from your dataset folder in a more clear way

Part B Get Your Images And Labels

First of all, get your images and labels, I assume that you have 7000 images and same count labels in txt format, orginize them in 2 folder, called Images which contains all images, and Labels which contains all labels. And, most important at all is that, every single image must have same name mapped label txt file, means if you have a image 0001.jpg you must labeled it in 0001.txt. And all txt files must in this format:

class_index xmin ymin xmax ymax

It's simple enough! In my way, I place all images inside ~/data/MyDataset/Images and all my labels in ~/data/MyDataset/Labels, hopefully please do not change Images and Labels folder name, we gonna use it.

Part C Create Your Work Folder

Seperate from caffe-ssd source code directory, you can create a invidual folder named MyDataset, our all work will compelet in this folder. OK, clone my kitti-ssd into your's anywhere you like. you can change this folder name as you like(example. face-ssd). Here inside we got this things:

data
models
train_ssd_kitti.py

ok, next we are going work with data first, because we have to generate lmdb file first.

Part D Create lmdb Database

OK, in this step, we are going put all data into lmdb database., this will generate a lmdb folder inside ~/data/MyDataset folder which contains KITTI_trainval and KITTI_test data.

cd data
bash create_list.sh
bash create_data.sh

Done! now you get trainval.txt test.txt and test_name_size.txt
But you have to get your labelmap_dataset.prototxt, here is suggestion:

if you have 5 classed named the 6th class in name background

item {
  name: "none_of_the_above"
  label: 6
  display_name: "background"
}

And later in train_ssd.py you gonna change two value : num_classes and background_index_id.

Part E Get VGG Pretrain Model and Start Train SSD

Download VGG pretrain model and place into models/VGGNet , the everything was done! Just a little change in train_ssd_kitti.py you can train ssd ready! Here is something you have to change:

sys.path.insert(0, "/home/chenqi-didi/Documents/work/caffe/python")
data_root_dir = "/home/chenqi-didi/data/"
caffe_root = "/home/chenqi-didi/Documents/work/caffe"
train_data = data_root_dir + "KITTI/lmdb/KITTI_trainval_lmdb"
# The database file for testing data. Created by data/KITTI/create_data.sh
test_data = data_root_dir + "KITTI/lmdb/KITTI_test_lmdb"
model_name = "VGG_KITTI_{}".format(job_name)

# Directory which stores the model .prototxt file.
save_dir = "models/VGGNet/KITTI/{}".format(job_name)
# Directory which stores the snapshot of models.
snapshot_dir = "models/VGGNet/KITTI/{}".format(job_name)
# Directory which stores the job script and log file.
job_dir = "jobs/VGGNet/KITTI/{}".format(job_name)
# Directory which stores the detection results.
output_result_dir = "{}/data/KITTI/results/{}/Main".format(os.environ['HOME'], job_name)
label_map_file = "{}/data/labelmap_kitti.prototxt".format(current_dir)

# Defining which GPUs to use.
gpus = "0,1"

Find above code change all KITTI into your dataset name, save it and you are ready to go!

python train_ssd_kitti.py

Finally, For Test Result

got your image path in data, for example data/test2.jpg and then change ssd_detection.py file path, then run:

python ssd_detection.py

Here is the result:

image