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

目录

Mask-RCNN概述

训练自己数据步骤

工具Labelme

标注数据

源码需要改动地方

训练之后的测试结果

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 和MaskR-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。

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

训练自己数据步骤

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

工具Labelme

1、安装labelme工具

pip install labelmepip install pyqt5pip 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 osimport jsonimport numpy as npimport globimport shutilimport cv2from sklearn.model_selection import train_test_splitnp.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 = 0self.ann_id = 0def 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文件构建COCOdef 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 += 1self.img_id += 1instance = {}instance['info'] = 'spytensor created'instance['license'] = ['license']instance['images'] = self.imagesinstance['annotations'] = self.annotationsinstance['categories'] = self.categoriesreturn instance# 构建类别def _init_categories(self):for k, v in classname_to_id.items():category = {}category['id'] = vcategory['name'] = kself.categories.append(category)# 构建COCO的image字段def _image(self, obj, path):image = {}from labelme import utilsimg_x = utils.img_b64_to_arr(obj['imageData'])h, w = img_x.shape[:-1]image['height'] = himage['width'] = wimage['id'] = self.img_idimage['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_idannotation['image_id'] = self.img_idannotation['category_id'] = int(classname_to_id[label])annotation['segmentation'] = [np.asarray(points).flatten().tolist()]annotation['bbox'] = self._get_box(points)annotation['iscrowd'] = 0annotation['area'] = 1.0return 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.infmax_x = max_y = 0for 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格式iftemp_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