使用Mask R-cnn训练自己的数据集

1.数据集的制作

使用的是labelme的labelme_json_to_dataset函数。
可以写一个简单的脚本让它批量执行。在json文件的同一目录下创建一个json2dataset.py文件，里面是这些内容：

import os
files=os.listdir('./')
files.remove('json2dataset.py')   # 删除这个py文件本身
for i in range(len(files)):
    os.system('labelme_json_to_dataset '+files[i])

然后在当前目录开启命令行，用python执行它就好了。
转换完之后，你看到的应该是这个样子：

 

 注意这个1.png最初你看到的应该是label.png，但是接下来在2.3中我们要把它转换成对应文件夹的名字。
如果你生成的文件夹里没有.yaml文件，不要着急，跟着下面的操作来。
首先在你的电脑上找到json_to_dataset.py这个脚本，打开之后做如下修改：

# 最前面加入导包
import yaml

# 中间是代码的主体部分就不贴了
# 然后在最下面main函数之前加上这一部分：
    logger.warning('info.yaml is being replaced by label_names.txt')
    info = dict(label_names=label_names)
    with open(osp.join(out_dir, 'info.yaml'), 'w') as f:
        yaml.safe_dump(info, f, default_flow_style=False)
    logger.info('Saved to: {}'.format(out_dir))

# 下面就是main函数了
if __name__ == '__main__':
    main()

2.3 将数据整理成模型认可的形式

先别着急把图存进去，在此之前，为了适应模型内部默认的路径格式，需要对label.png进行简单的重命名（否则就要去代码里边改，比较麻烦）。
比如你的json文件夹叫1_json，那这个png的图就应该改成1.png，为了大家批量操作，我写了个简单的脚本。

# 把label.png改名为1.png
import os
for root, dirs, names in os.walk(r'******out_dir'):   # 改成你自己的json文件夹所在的目录
    for dr in dirs:
        file_dir = os.path.join(root, dr)
        # print(dr)
        file = os.path.join(file_dir, 'label.png')
        # print(file)
        new_name = dr.split('_')[0] + '.png'
        new_file_name = os.path.join(file_dir, new_name)
        os.rename(file, new_file_name)

文件夹	内容
cv2_mask	json_to_dataset生成文件夹中的png格式label文件
json	labelme生成的json文件
labelme_json	json_to_dataset生成的文件夹
pic	尺寸标准化之后的原图

其他三个文件夹都比较容易准备好，cv2_mask这个挨个复制比较麻烦，可以用下面的方法批量复制：

import os
from shutil import copyfile
for root, dirs, names in os.walk(r'******'):   # 改成你自己的json文件夹所在的目录
    for dr in dirs:
        file_dir = os.path.join(root, dr)
        print(dr)
        file = os.path.join(file_dir,'label.png')
        print(file)
        new_name = dr.split('_')[0] + '.png'
        new_file_name = os.path.join(file_dir, new_name)
        print(new_file_name)
        
        tar_root = r'******my_data/cv2_mask'      # 目标路径
        tar_file = os.path.join(tar_root, new_name)
        copyfile(new_file_name, tar_file)

训练脚本

# -*- coding: utf-8 -*-
 
import os
import sys
import random
import math
import re
import time
import numpy as np
import cv2
# import matplotlib
# import matplotlib.pyplot as plt
import tensorflow as tf
from mrcnn.config import Config
# import utils
from mrcnn import model as modellib, utils
from mrcnn import visualize
import yaml
from mrcnn.model import log
from PIL import Image
 
# os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# Root directory of the project
ROOT_DIR = os.getcwd()
 
# ROOT_DIR = os.path.