采用Mask RCNN分割自己的数据集
- 下载源工程代码：https://github.com/matterport/Mask_RCNN；
- 解压，进入主文件夹，安装所需模块：
pip3 install -r requirements.txt;
-
下载Jupyter Notebook以及coco权重：mask_rcnn_coco.h5，了解Jupyter Notebook的使用方法，采用Jupyter Notebook打开文件夹samples下的balloon或者shapes子文件夹，在web界面左上角的file下面将其download为python格式的文件；
说明: samples文件夹下面的几个示例例程基本都有三个文件：inspect_data,inspect_model,shapes，其中第一个文件是生成数据，第二个用来检测模型，可查看分步计算的结果并展示，第三个是主文件，用来训练或者预测分割；
-
结合shapes和balloon两个例程，在其基础上进行修改使得代码可以处理自己的数据，就是继承已有的类（mrcnn类）进行函数覆盖，使其适用于自己的data。在继承的过程中，除了sample文件夹下面的代码在改变，其他如mrcnn文件夹下面的文件都不要动。具体来说需要修改的有以下几个类：
-
类1： class BrainwebConfig(Config), 在其下面定义自己的环境配置，例如每个GPU处理的图像数目、每步迭代的次数等；
-
类2：class BrainwebDataset(utils.Dataset),继承原数据类，定义自己的类，该类需要覆盖的函数主要有两个：load_brainweb和load_mask。
在load_brainweb函数中需要注意的主要有两个函数：add_class，add_image，前者添加类别，后者添加有关图像的信息，特别需要注意的是，在这一步并不需要导入所有图像的数据，只需要添加原始图像所在路径和有关图像的信息即可，原add_image中没有的参数都将被合并至self.imageinfo的一个dict下，以方便后续使用，比如mask数据的导入，具体代码如下图：
# 添加图像
for img_name in train_img_names:
image_path = os.path.join(dataset_dir, img_name)
image = skimage.io.imread(image_path)
height, width = image.shape[:2]
self.add_image("brainweb", image_id = img_name, path = image_path,width = width, height=height)
在load_mask函数中需要利用到self.image_info中的信息，所以load_brainweb和load_mask要配合使用；具体代码如下：
# 添加图像
# 以下是自己编写的读取mask的代码
image_info = self.image_info[image_id]
mask = np.zeros([image_info["height"], image_info["width"], 3], dtype=np.uint8) # 每个实例中只有3种mask
image_name = image_info['id']
root_path = '/home/mdd/work/MRCNN_UNet/Mask_RCNN/self_Mask_RCNN/datasets/brainweb/mask_train_val/label_0_255_'
for r_num in range(3):
mask_name = root_path + str(r_num+1) + '/' + image_name
mask_data = scipy.misc.imread(mask_name) # 这里时读取的，后面显示会一块显示所以得去缓存
mask[:,:,r_num] = mask_data
return mask.astype(np.bool), np.array([1,2,3], dtype=np.int32)
- 除了修改覆盖上面两个类之外，剩下就是定义自己的训练(train)和测试(test)主函数。
- 训练主函数格式比较固定，基本都是以下形式：
def train(model,config):
"""Train the model."""
# Training dataset.
dataset_train = BrainwebDataset()
dataset_train.load_brainweb(args.dataset, "train") # 如果用自己的数据集，主要修改在这一句
dataset_train.prepare()
# Validation dataset
dataset_val = BrainwebDataset()
dataset_val.load_brainweb(args.dataset, "val") # 如果用自己的数据集，主要修改在这一句
dataset_val.prepare()
# 测试显示
image_id = random.choice(dataset_val.image_ids)
original_image, image_meta, gt_class_id, gt_bbox, gt_mask = modellib.load_image_gt(dataset_val, config, image_id = image_id, use_mini_mask=False)
scipy.misc.imshow(original_image)
# *** This training schedule is an example. Update to your needs ***
# Since we're using a very small dataset, and starting from
# COCO trained weights, we don't need to train too long. Also,
# no need to train all layers, just the heads should do it.
# 这里需要注意的是： 以下两次训练并不是独立的，而是累计的epochs，就是两次训练的总次数等于all layers第二次的训练epochs
print("Training network heads")
model.train(dataset_train, dataset_val,learning_rate = config.LEARNING_RATE, epochs=10, layers='heads')
print("Training all layers")
model.train(dataset_train, dataset_val, learning_rate=config.LEARNING_RATE / 10, epochs=20, layers="all")
print("Training all layers end!")
- 测试函数的话，主要用于多幅图像分割和Video的分割，这里主要讨论多幅图像的分割。如果测试图像是灰度图，则必须将其扩展到三通道彩色图，具体就是原灰度图复制三次，叠加在一起。代码如下：
def segment(model, image_path = None, video_path = None):
# 判断是否存在图像路径或者视频路径
assert image_path or video_path
# Image or video?
if image_path:
# Run model detection and generate the color splash effect
print("Running on {}".format(args.image_path))
# Read image
test_img_names = next(os.walk(args.image_path))[2]
for i,img_name in enumerate(test_img_names):
test_img = skimage.io.imread(args.image_path + '/' + img_name)
expand_image = np.zeros([test_img.shape[0],test_img.shape[1],3])
expand_image[:, :, 0] = test_img
expand_image[:, :, 1] = test_img
expand_image[:, :, 2] = test_img
test_img = expand_image
results = model.detect([test_img], verbose=1)[0]
# visualize.display_instances(test_img, results['rois'], results['masks'], results['class_ids'], ['BG','Class1','Class2','Class3'],results['scores']) # 显示图
# 下面将生成的mask保存到对应文件夹中
predict_masks = results['masks']
predict_ids = results['class_ids']
for j, predict_id in enumerate(predict_ids):
predict_int_mask = np.zeros([predict_masks.shape[0], predict_masks.shape[1]])
predict_int_mask[predict_masks[:, :, j]] = 1
save_name = args.image_path + '/' + img_name[:-4] + '_predict_' + str(predict_id) + img_name[-4:]
skimage.io.imsave(save_name, predict_int_mask)
- 加载MaskRCNN模型时，如果时训练模式，应该使用：
model = modellib.MaskRCNN(mode=“training”, config=config,model_dir=args.logs)
如果时推理测试模式，应该使用：
model = modellib.MaskRCNN(mode=“inference”, config=config,model_dir=args.logs) - 主函数中的各种解析指令如下：
if __name__ == '__main__':
import argparse
# Parse command line arguments
parser = argparse.ArgumentParser(description='Train Mask R-CNN to segment brain images.')
parser.add_argument("--command",metavar="<command>", help="'train' or 'segment'", default = 'train')
parser.add_argument('--dataset', required=False, metavar="/path/to/balloon/dataset/",
default = '/home/mdd/work/MRCNN_UNet/Mask_RCNN/self_Mask_RCNN/datasets/brainweb/',
help = 'Directory of the Balloon dataset')
parser.add_argument('--weights', required=False, metavar="/path/to/weights.h5",
default = 'coco',
help = "Path to weights .h5 file or 'coco'")
parser.add_argument('--logs', required=False, default=DEFAULT_LOGS_DIR, metavar="/path/to/logs/",
help = 'Logs and checkpoints directory (default=logs/)')
parser.add_argument('--image_path', required=False, metavar="path or URL to image",
default = '/home/mdd/work/MRCNN_UNet/Mask_RCNN/self_Mask_RCNN/datasets/brainweb/test',
help = 'Image to apply the color splash effect on')
parser.add_argument('--video_path', required=False, metavar="path or URL to video",
help = 'Video to apply the color splash effect on')
args = parser.parse_args()
print (args)