手把手教你训练自己的Mask R-CNN图像实例分割模型（PyTorch官方教程）

近来在学习图像分割的相关算法，准备试试看Mask R-CNN的效果。

关于Mask R-CNN的详细理论说明，可以参见原作论文https://arxiv.org/abs/1703.06870，网上也有大量解读的文章。本篇博客主要是参考了PyTorch官方给出的训练教程，将如何在自己的数据集上训练Mask R-CNN模型的过程记录下来，希望能为感兴趣的读者提供一些帮助。

PyTorch官方教程（Object Detection finetuning tutorial）：

https://github.com/pytorch/tutorials/blob/master/_static/torchvision_finetuning_instance_segmentation.ipynb

或：

https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html

需要注意的是，TorchVision需要0.3之后的版本才可以使用。

目录

准备工作

数据集

定义模型

训练模型

1. 准备工作

2. 数据增强/转换

3. 训练

测试模型

---

 

准备工作

安装coco的api，主要用到其中的IOU计算的库来评价模型的性能。

git clone https://github.com/cocodataset/cocoapi.git
cd cocoapi/PythonAPI
python setup.py build_ext install

API的安装也可以参考另一篇：

https://blog.csdn.net/u013685264/article/details/100331064

数据集

本教程使用Penn-Fudan的行人检测和分割数据集来训练Mask R-CNN实例分割模型。Penn-Fudan数据集中有170张图像，包含345个行人的实例。图像中场景主要是校园和城市街景，每张图中至少有一个行人，具体的介绍和下载地址如下：

https://www.cis.upenn.edu/~jshi/ped_html/

# 下载Penn-Fudan dataset
wget https://www.cis.upenn.edu/~jshi/ped_html/PennFudanPed.zip
# 解压到当前目录
unzip PennFudanPed.zip

解压后的目录结构如下：

先看看Penn-Fudan数据集中的图像和mask：

from PIL import Image

Image.open('PennFudanPed/PNGImages/FudanPed00001.png')

mask = Image.open('PennFudanPed/PedMasks/FudanPed00001_mask.png')

mask.putpalette([
    0, 0, 0, # black background
    255, 0, 0, # index 1 is red
    255, 255, 0, # index 2 is yellow
    255, 153, 0, # index 3 is orange
])

mask

    

每一张图像都有对应的mask标注，不同的颜色表示不同的实例。在训练模型之前，需要写好数据集的载入接口。

import os
import torch
import numpy as np
import torch.utils.data
from PIL import Image


class PennFudanDataset(torch.utils.data.Dataset):
    def __init__(self, root, transforms=None):
        self.root = root
        self.transforms = transforms
        # load all image files, sorting them to ensure that they are aligned
        self.imgs = list(sorted(os.listdir(os.path.join(root, "PNGImages"))))
        self.masks = list(sorted(os.listdir(os.path.join(root, "PedMasks"))))

    def __getitem__(self, idx):
        # load images ad masks
        img_path = os.path.join(self.root, "