FasterRCNN训练自己的数据

前言

因为科研需要，折腾了一阵Faster RCNN. 最开始将大佬chenyuntc的simple-faster-rcnn-pytorch鼓捣了一下使之能够在Windows下面跑（有需要的朋友可以联系我），但是该代码的batch size只能设为1，跑一次实验的耗时过长。本渣尝试花了四天时间通读了一遍大佬的代码，奈何功力尚浅，还做不到大范围的修改，因此从github上找到了jwyang的faster-rcnn.pytorch 稍作修改，在Linux下成功实验了自己的数据集。

代码

我所使用的服务器操作系统是Linux，安装了Python3.6和Pytorch1.3，并且服务器无法连接网络，导致配置环境耗费了好几天的时间。
首先从网址https://github.com/jwyang/faster-rcnn.pytorch/tree/pytorch-1.0上下载下来代码。

在代码根目录下创建data文件夹

切换到data文件夹下，创建下图所示的pretrained_model和VOCdevkit2007两个文件夹。

其中pretrained_model用于存放下载好的预训练模型。VOCdevkit2007用于存放数据集。
切换到代码根目录下的lib文件夹（faster-rcnn.pytorch-pytorch-1.0\lib）下，编译Cuda依赖项。
cd lib
python setup.py build develop
顺利的话一步到位编译完成，报错需要看具体的报错信息。

准备自己的数据集

本文使用的是Pascal VOC数据集格式。
首先将随机划分训练集、测试集和验证集

import os
import random

trainval_percent = 0.7
train_percent = 0.8

xmlfilepath = "VOC2007/Annotations"
txtsavepath = "VOC2007/ImageSets/Main"

# 历遍"Annotations"文件夹然后返回列表
total_xml = os.listdir(xmlfilepath)

# 获取列表的总数
num = len(total_xml)
list = range(num)

tv = int(num * trainval_percent)
tr = int(tv * train_percent)
#te = int(tv * test_percent)

trainval = random.sample(list, tv)
train = random.sample(trainval, tr)
#test = random.sample(list, te)

ftrainval = open(txtsavepath + "/trainval.txt", 'w')
ftest = open(txtsavepath + "/test.txt", 'w')
ftrain = open(txtsavepath + "/train.txt", 'w')
fval = open(txtsavepath + "/val.txt", 'w')

for i in list:
    # 使用切片方法获取文件名(去掉后缀".xml")
    name = total_xml[i][:-4] + '\n'
    if i in trainval:
        ftrainval.write(name)
        if i in train:
            ftrain.write(name)
        else:
            fval.write(name)
    else:
        ftest.write(name)

ftrainval.close()
ftrain.close()
fval.close()
ftest.close()

最终划分出来的训练集+验证集（trainval）占数据集的70%，训练集（train）占trainval的80%.
我使用的数据集已经进行了边框标注，但是使用的是csv格式存储的标注信息，我写了段代码将其转换成了xml格式。

import os
import pandas as pd
import cv2
from xml.dom import minidom

parent_path = "/Data/"
annotations = "annotations/"
image_sets = "sets/"
images = "images/"
class_name = "xxxx"   # 此处替换为类别名
XML = "XML/"
IMG = "IMG/"

def csv_to_xml(parent_path, current_folder, current_file_name, list_data, label_name, current_generated_xmlfile_savepath):
    img = cv2.imread(os.path.join(parent_path, current_folder, current_file_name + ".png"))
    h, w, d = img.shape
    
    cv2.imwrite(os.path.join(parent_path, IMG, current_file_name + ".jpg"), img)

    doc = minidom.Document()