maskrcnn-benchmark框架下使用Maskrcnn训练自己的数据集

大家好，这是我的第一篇博客。作为深度学习的初学者，从实战的角度，详细介绍基于maskrcnn-benchmark框架下，使用Maskrcnn网络训练自己的数据集的过程。这其中，踩的坑不少，但也受益匪浅。不对的地方，还麻烦大家指正，一起进步。

---

简介

maskrcnn-benchmark框架由Facebook AI Research提出，基于Pytorch1.0，具有训练速度快，占用GPU较少等优点。详情可参考：https://github.com/facebookresearch/maskrcnn-benchmark，这里不再过多描述。

基础环境

操作系统：Ubuntu 16.04.6   4.15.0-96-generic GNU

编译环境：python 3.7.7  CUDA 10.2 opencv4.3.0等

由于设备问题，实验计算资源为网上租用，详情参考：https://www.longway-gpu.com/home；相关配置，如下图所示：

数据准备

训练数据来源于输电塔螺栓数据，标注工具为Labelme，由于标注得到的.json文件不符合所需的COCO数据集格式，故需要使用脚本文件转换。脚本文件、数据集组织如下：

脚本文件

import os
import shutil
import numpy as np
# train/val/test比例6:2:2
def data_split(img_path, json_path):
    filelist = os.listdir(img_path)
    filelist.remove('train')
    filelist.remove('val')
    filelist.remove('test')
    jsonfile_path = os.path.join(json_path, 'train')
    jsonfile_path = json_path
    jsonlist = os.listdir(jsonfile_path)
    # jsonlist.remove('train')
    # jsonlist.remove('val')
    # jsonlist.remove('test')
    trainlist = np.random.choice(filelist, size=187, replace=False)
    for filename in trainlist:
        filepath = os.path.join(img_path, filename)
        shutil.move(filepath, os.path.join(img_path, 'train'))
        imgname = os.path.splitext(filename)[0]
        for json in jsonlist:
            if imgname == os.path.splitext(json)[0]:
                jsonfile = os.path.join(jsonfile_path, json)
                shutil.move(jsonfile, os.path.join(json_path, 'train'))
                jsonlist.remove(json)
        filelist.remove(filename)

    # reslist = []
    # for item in filelist:
    #     if item not in trainlist and not os.path.isdir(os.path.join(img_path, item)):
    #         reslist.append(item)
    vallist = np.random.choice(filelist, size=63, replace=False)
    for filename in vallist:
        filepath = os.path.join(img_path, filename)
        shutil.move(filepath, os.path.join(img_path, 'val'))
        imgname = os.path.splitext(filename)[0]
        for json in jsonlist:
            if imgname == os.path.splitext(json)[0]:
                jsonfile = os.path.join(jsonfile_path, json)
                shutil.move(jsonfile, os.path.join(json_path, 'val'))
                jsonlist.remove(json)
        filelist.remove(filename)

    # testlist = []
    # for item in reslist:
    #     if item not in vallist:
    #         testlist.append(item)
    for filename in filelist:
        filepath = os.path.join(img_path, filename)
        shutil.move(filepath, os.path.join(img_path, 'test'))
        imgname = os.path.splitext(filename)[0]
        for json in jsonlist:
            if imgname == os.path.splitext(json)[0]:
                jsonfile = os.path.join(jsonfile_path, json)
                shutil.move(jsonfile, os.path.join(json_path, 'test'))
        #         jsonlist.remove(json)
        # filelist.remove(filename)

# 只需要修改train/vallist中的size大小以及img_path和json_path的路径即可运行
# 图像的文件夹路径
img_path = 'E:\\data_luoshuan\\imgs'
# json文件的路径
json_path = 'E:\\data_luoshuan\\labels'
data_split(img_path, json_path)

数据集组织

在/maskrcnn-benchmark/datasets目录下，创建文件夹，命名为coco。将上述转换得到的train/val/test文件夹放入coco目录下，并且创建annotations文件夹，将转换得到的train/val/test.json放入。如下图所示：

至此，数据集准备完毕。

---

安装框架