Pytorch版Mask-RCNN图像分割实战（自定义数据集）

Mask-RCNN概述

Mask R-CNN是一种广泛应用于目标检测和图像分割任务的深度学习模型，它是由Faster R-CNN（一种快速目标检测模型）和Mask R-CNN（一种实例分割模型）组成的。Mask R-CNN将Faster R-CNN中的RPN和RoI Pooling层替换成了RPN和RoI Align层，以实现像素级的图像分割，能够同时检测出多个对象，并对每个对象进行像素级的分割。

Mask R-CNN的主要思路是在Faster R-CNN的基础上增加一个分支网络，即Mask分支，该分支网络可以对检测出来的物体进行像素级的分割操作，得到每个物体实例的分割掩码。与 Faster R-CNN 相似， Mask R-CNN同样使用RPN产生候选框，在RoIPooling层和ROIAlign层中对候选框中的特征进行提取。在RoIAlign层中，Mask R-CNN通过双线性插值从特征图中提取出精确的特征，然后送入Mask分支中进行像素级别的分割，最终得到每个实例的精确掩码。

Facebook AI Research 开源了 Faster R-CNN 和 Mask R-CNN 的 PyTorch 1.0 实现基准：MaskRCNN-Benchmark。相比 Detectron 和 mmdetection，MaskRCNN-Benchmark 的性能相当，并拥有更快的训练速度和更低的 GPU 内存占用

优势如下：

• PyTorch 1.0：相当或者超越 Detectron 准确率的 RPN、Faster R-CNN、Mask R-CNN 实现；
• 非常快：训练速度是 Detectron 的两倍，是 mmdection 的 1.3 倍；
• 节省内存：在训练过程中使用的 GPU 内存比 mmdetection 少大约 500MB；
• 使用多 GPU 训练和推理；
• 批量化推理：可以在每 GPU 每批量上使用多张图像进行推理；
• 支持 CPU 推理：可以在推理时间内于 CPU 上运行。
• 提供几乎所有参考 Mask R-CNN 和 Faster R-CNN 配置的预训练模型，具有 1x 的 schedule。

基于Mask RCNN开源项目源码地址：https://github.com/facebookresearch/maskrcnn-benchmark

训练自己数据步骤

• 安装Labelme
• 标注数据
• 源码需要改动地方
• 训练之后测试结果

工具Labelme

1、安装labelme工具

pip install labelme
pip install pyqt5
pip install pillow==4.0.0

标注数据

1、使用labelme标注得到.json文件，将.json文件和原始图片放在一个文件夹里；

2、批量量转换：将labelme标注数据转换成coco数据集；

• 运行labelmetococo.py文件；
• 在我当前目录的dataset文件夹下，生成coco文件夹；
• coco文件夹下有annotations文件夹和images文件夹；
• annotations文件夹存放2个json文件；
• images文件夹存放train （存放：划分的用于训练的图像数据）和val （存放：划分的用于验证的图像数据） 两个文件夹；

下面是labelmetococo.py文件代码：

import os
import json
import numpy as np
import glob
import shutil
import cv2
from sklearn.model_selection import train_test_split

np.random.seed(41)

# 0为背景
classname_to_id = {
    "class1": 1

}
# 注意这里
# 需要从1开始把对应的Label名字写入：这里根据自己的Lable名字修改

class Lableme2CoCo:

    def __init__(self):
        self.images = []
        self.annotations = []
        self.categories = []
        self.img_id = 0
        self.ann_id = 0

    def save_coco_json(self, instance, save_path):
        json.dump(instance, open(save_path, 'w', encoding='utf-8'), ensure_ascii=False, indent=1)  # indent=2 更加美观显示

    # 由json文件构建COCO
    def to_coco(self, json_path_list):
        self._init_categories()
        for json_path in json_path_list:
            obj = self.read_jsonfile(json_path)
            self.images.append(self._image(obj, json_path))
            shapes = obj['shapes']
            for shape in shapes:
                annotation = self._annotation(shape)
                self.annotations.append(annotation)
                self.ann_id += 1
            self.img_id += 1
        instance = {}
        instance['info'] = 'spytensor created'
        instance['license'] = ['license']
        instance['images'] = self.images
        instance['annotations'] = self.annotations
        instance['categories'] = self.categories
        return instance

    # 构建类别
    def _init_categories(self):
        for k, v in classname_to_id.items():
            category = {}
            category['id'] = v
            category['name'] = k
            self.categories.append(category)

    # 构建COCO的image字段
    def _image(self, obj, path):
        image = {}
        from labelme import utils
        img_x = utils.img_b64_to_arr(obj['imageData'])
        h, w = img_x.shape[:-1]
        image['height'] = h
        image['width'] = w
        image['id'] = self.img_id
        image['file_name'] = os.path.basename(path).replace(".json", ".jpg")
        return image

    # 构建COCO的annotation字段
    def _annotation(self, shape):
        # print('shape', shape)
        label = shape['label']
        points = shape['points']
        annotation = {}
        annotation['id'] = self.ann_id
        annotation['image_id'] = self.img_id
        annotation['category_id'] = int(classname_to_id[label])
        annotation['segmentation'] = [np.asarray(points).flatten().tolist()]
        annotation['bbox'] = self._get_box(points)
        annotation['iscrowd'] = 0
        annotation['area'] = 1.0
        return annotation

    # 读取json文件，返回一个json对象
    def read_jsonfile(self, path):
        with open(path, "r", encoding='utf-8') as f:
            return json.load(f)

    # COCO的格式： [x1,y1,w,h] 对应COCO的bbox格式
    def _get_box(self, points):
        min_x = min_y = np.inf
        max_x = max_y = 0
        for x, y in points:
            min_x = min(min_x, x)
            min_y = min(min_y, y)
            max_x = max(max_x, x)
            max_y = max(max_y, y)
        return [min_x, min_y, max_x - min_x, max_y - min_y]

if __name__ == '__main__':

    # 需要把labelme_path修改为自己放images和json文件的路径
    labelme_path = "D:\\maskrcnn-benchmark-main\\dataset\\gps\\"
    # saved_coco_path = "../../../xianjin_data-3/"
    saved_coco_path = "D:\\maskrcnn-benchmark-main\\dataset\\"
    # saved_coco_path = "./"
    # 要把saved_coco_path修改为自己放生成COCO的路径，这里会在我当前COCO的文件夹下建立生成coco文件夹。
    print('reading...')
    # 创建文件
    if not os.path.exists("%scoco/annotations/" % saved_coco_path):
        os.makedirs("%scoco/annotations/" % saved_coco_path)
    if not os.path.exists("%scoco/images/train2017/" % saved_coco_path):
        os.makedirs("%scoco/images/train2017" % saved_coco_path)
    if not os.path.exists("%scoco/images/val2017/" % saved_coco_path):
        os.makedirs("%scoco/images/val2017" % saved_coco_path)
    # 获取images目录下所有的joson文件列表
    print(labelme_path + "\*.json")
    json_list_path = glob.glob(labelme_path + "\*.json")
    # json_list_path = glob.glob(labelme_path + "\*.png")
    print('json_list_path: ', len(json_list_path))
    # 数据划分,这里没有区分val2017和tran2017目录，所有图片都放在images目录下
    train_path, val_path = train_test_split(json_list_path, test_size=0.1, train_size=0.9)
    # 将训练集和验证集的比例是9：1，可以根据自己想要的比例修改。
    print("train_n:", len(train_path), 'val_n:', len(val_path))

    # 把训练集转化为COCO的json格式
    l2c_train = Lableme2CoCo()
    train_instance = l2c_train.to_coco(train_path)
    l2c_train.save_coco_json(train_instance, '%scoco/annotations/instances_train2017.json' % saved_coco_path)
    for file in train_path:
        # 存入png格式图片，原始图片有两种格式.png,.jpg
        # print("这里测试一下file："+file)
        img_name = file.replace('json', 'png')
        # print("这里测试一下img_name：" + img_name)
        temp_img = cv2.imread(img_name)
        # 图像为空说明为.jpg格式
        if  temp_img is None:
            img_name_jpg = img_name.replace('png', 'jpg')
            temp_img = cv2.imread(img_name_jpg)

        filenames = img_name.split("\\")[-1]
        cv2.imwrite("D:\\maskrcnn-benchmark-main\\dataset\\coco\\images\\train2017/{}".format(filenames), temp_img)
        # print(temp_img) #测试图像读取是否正确

    for file in val_path:
        # shutil.copy(file.replace("json", "jpg"), "%scoco/images/val2017/" % saved_coco_path)

        img_name = file.replace('json', 'png')
        temp_img = cv2.imread(img_name)
        if temp_img is None:
            img_name_jpg = img_name.replace('png', 'jpg')
            temp_img = cv2.imread(img_name_jpg)
        filenames = img_name.split("\\")[-1]
        cv2.imwrite("D:\\maskrcnn-benchmark-main\\dataset\\coco\\images\\val2017/{}".format(filenames), temp_img)
        

    # 把验证集转化为COCO的json格式
    l2c_val = Lableme2CoCo()
    val_instance = l2c_val.to_coco(val_path)
    l2c_val.save_coco_json(val_instance, '%scoco/annotations/instances_val2017.json' % saved_coco_path)

执行程序，生成的文件夹

源码需要改动地方：

假设你此时位于maskrcnn-benchmark/目录下,datasets的组织结构如下

datasets
-> coco
  -> annotations
    -> instances_train2014.json //训练标签
    -> instances_test2014.json  //验证标签
  -> train2014 //训练图片
  -> val2014  //验证图

为了方便建议仿照上面这样的coco标准命名，但上面的所有名字都不是一定要求这样写的，自己可以合理命名，只要程序中的各个与数据路径有关的地方一一对应上即可。

1. 以maskrcnn-benchmark/configse/2e_mask_rcnn_R_50_FPN_1x.yaml为例下面蓝色圈出来的地方必须修改：
2. 一个是类别，如果这里面没有类别，与default.py中的一致，那么程序会自动在defaults.py中寻找。
3. 特别还有一个是DATASETS，这个是用来指定本次训练所使用的数据集，为了方便我这里没有修改，就借用了coco_2014_train的名字,名字可以随意起，但是在maskrcnn-benchmark/maskrcnn_benchmark/config/paths_catalog.py中需要有对应的数据集的路径描述，且该路径要与datasets/coco/里面的文件对应起来
4. BASE_LR，迭代次数，保存间隔根据自己模型需要而定
5. 修改maskrcnn_benchmark/utils下的checkpoint，需要注释65- 68行（self.optimizer.load_state.. self.scheduler.load_...）

将myconfig下的paths_catalogs文件中的数据集路径改为自己的；

maskrcnn-benchmark/maskrcnn_benchmark/config/defaults.py

这是模型的一个总的默认配置文件，需要作出一定修改

_C.INPUT = CN()
# Size of the smallest side of the image during training
_C.INPUT.MIN_SIZE_TRAIN = (400,)  # (800,)
# Maximum size of the side of the image during training
_C.INPUT.MAX_SIZE_TRAIN = 667
# Size of the smallest side of the image during testing
_C.INPUT.MIN_SIZE_TEST = 400
# Maximum size of the side of the image during testing
_C.INPUT.MAX_SIZE_TEST = 667
#下面的两处修改也需要特别注意！！！必须和自己的类别相对应，如果没有分类，那么就为2
_C.MODEL.ROI_BOX_HEAD.NUM_CLASSES = 2 #类别数量加1
_C.MODEL.RETINANET.NUM_CLASSES = 2 #类别数量加1