PyTorch深度学习实战（24）——从零开始实现Mask R-CNN实例分割

0. 前言

Mask R-CNN (Mask Region Convolutional Neural Network) 是基于深度学习的图像分割算法，它是在 Faster R-CNN 目标检测框架的基础上进行扩展和改进的。与传统目标检测方法相比，Mask R-CNN 不仅可以准确地检测图像中的对象，还可以为每个对象生成精确的像素级别的分割掩码。这意味着 Mask R-CNN 能够同时提供对象的边界框和具体的像素级别分割结果，从而更细粒度地理解图像中的结构和语义信息。在本节，将介绍 Mask R-CNN 架构的工作原理，并使用 PyTorch 实现 Mask R-CNN 进行实例分割。

1. Mask R-CNN

1.1 网络架构

Mask R-CNN 是一种用于目标检测和实例分割的深度学习算法，它扩展了 Faster R-CNN 算法，并增加了一个用于预测对象掩码 (mask) 的分支。
Mask R-CNN 架构可以用于在图像中识别/显示给定类别的对象实例，能够分割图像中类别相同的多个对象，Mask 表示由 Mask R-CNN 在像素级别完成的分割。掩码用于标注图像中的不同区域，使用图像分割模型可以将图像分成不同的区域，然后为每个区域分配一个掩码值。
Mask R-CNN 架构是对 Faster R-CNN 网络的扩展：

Mask R-CNN 架构修改了 Faster R-CNN 的 RoI Pooling 层，使用更加准确的 RoI Align 层
除了在最终层中预测对象的类别和边界框偏移量外，还增加了一个 mask head，用于预测对象的掩码
使用全卷积网络 (Fully Convolutional Network, FCN) 实现掩码预测。

Mask R-CNN 整体架构如下：

网络架构

除了用于获取类别和边界框信息的预测头外，在 Mask R-CNN 中还添加了 Mask 预测头 (Mask head) 获取掩码信息：

网络架构

接下来，我们介绍 Mask R-CNN 架构的基本组件。

1.2 RoI Align

在 Faster R-CNN 中，我们了解了 RoI Pooling 的缺点之一是在执行 RoI Pooling 操作时可能会丢失某些信息。例如，在下图 RoI Pooling 示例中：

ROI pooling

扫描二维码关注公众号，回复： 17147931 查看本文章

在上图中，区域提议形状为 5 x 7，需要将其转换为 2 x 2 的形状，将其转换为 2 x 2 形状时(这一过程也称为量化)，由于仅能保留最高值，因此会导致信息丢失，为了解决这一问题，提出了 RoI Align。
接下来，我们通过一个简单示例讲解 RoI Align 的工作原理，尝试将以下区域(以虚线表示)转换为 2 x 2 的形状：

RoI Align

该区域并非均匀分布在特征图中的所有单元格中。为了能够使用 2x2 的形状合理表示该区域，我们需要执行以下步骤。

(1) 首先，将该区域划分为形状相等的 2 x 2 网格：

RoI Align

(2) 在每个单元格中定义四个等距的点：

RoI Align

在上图中，两个连续点之间的距离为 0.75。

(3) 根据每个点到最近已知值的距离计算其加权平均值：

RoI Align

(4) 对单元格中的所有点重复以上加权均值计算过程：

RoI Align

(5) 对单元格中的点执行平均池化，并根据相同步骤计算所有单元格值：

RoI Align

可以看到，通过以上步骤 RoI Align 能够确保不会丢失信息。

1.3 Mask 检测头

使用 RoI Align，我们可以更准确地表示从区域提议网络 (Region Proposal Network, RPN) 获得的区域提议。然后，对于每个区域提议，根据 RoI Align 输出获取分割（掩码）输出。
在目标检测中，需要将 RoI Align 结果输入到展平层，用于预测目标物体的类别和边界框偏移量。在图像分割中，还需要预测包含目标对象的边界框内的像素。因此，除了类别和边界框偏移量外，还需要预测感兴趣区域内的掩码。
将预测的掩码叠加在原始图像上，将 RoI Align 的输出连接到卷积层上，以获得类似图像的结构，如下图所示：

Mask head
在上图中，使用特征金字塔网络 (feature pyramid network, FPN) 得到形状为 7 x 7 x 2048 的输出，该网络包含 2 个分支：

第一个分支在展平 FPN 输出后返回目标对象的类别和边界框
第二个分支在 FPN 的输出后执行卷积运算以获得掩码

形状为 14 x 14 输出对应的真实标签是区域提议对应的调整大小后图像。如果数据集中有 80 个不同的类别，则区域提议的真实标签形状为 80 x 14 x 14，每个像素值都是 1 或 0，表示该像素是否包含对象。因此，在预测像素类别的同时使用二元交叉熵损失函数。
模型训练完成后，可以用于检测区域，获取类别和边界框偏移量，并获取每个区域对应的掩码。在模型推理阶段时，首先检测图像中存在的目标对象并进行边界框校正，然后，将校正后的区域传递给 Mask 检测头，以预测该区域中不同像素对应的掩码。
了解了 Mask R-CNN 架构的工作原理后，我们将使用 PyTorch 实现 Mask R-CNN，以检测图像中的人物实例。

2. 使用 Mask R-CNN 实现实例分割

2.1 数据集分析

为了训练 Mask R-CNN 进行实例分割，我们将使用掩码标注的人物数据集，该数据集是根据 DE20K 数据集的子集创建的，可以在 ADE20K 数据集官方网站中下载该数据集，DE20K 数据集是一个大规模的图像和语义分割数据集，该数据集包含超过 `20,000 张高分辨率图像，覆盖了各种多样的场景和对象类别。本节中，我们只使用包含人物掩码的图像。

2.2 模型构建策略

为了训练 Mask R-CNN 图像分割模型，我们采用以下策略：

获取数据集，并据此创建数据集和数据加载器
创建符合 PyTorch 官方的 Mask R-CNN 实现所需格式的目标输出
下载预训练的 Faster R-CNN 模型，并为其附加一个 Mask R-CNN 预测头
训练模型 Mask R-CNN 模型
首先执行非极大值抑制，然后识别与图像中的人物对应的边界框和掩码

2.3 模型构建与训练

(1) 下载相关的数据集和模型训练工具。

为了训练模型，首先下载图像及掩码标注数据集，下载完成后解压文件。同时，为了快速训练模型，在 PyTorch GitHub 存储库中下载 engine.py、utils.py、transforms.py、coco_eval.py 和 coco_utils.py 文件。最后，使用 pip 安装 cocoapi：

$ pip install -q -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

(2) 导入所有必要的库并定义设备：

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
from glob import glob

from engine import train_one_epoch, evaluate
import utils
import transforms as T
import torch
device = 'cuda' if torch.cuda.is_available() else 'cpu'
import cv2
import numpy as np
from matplotlib import pyplot as plt
import random
from PIL import Image
from torch.utils.data import DataLoader, Dataset

(3) 获取包含人物掩码的图片。

遍历 images 和 annotations_instance 文件夹以获取文件名：

all_images = glob('images/training/*.jpg')
all_annots = glob('annotations_instance/training/*.png')

检查原始图像和人物实例的掩码表示：

f = 'ADE_train_00014184'

def find(item, original_list):
    results = []
    for o_i in original_list:
        if item in o_i:
            results.append(o_i)
    if len(results) == 1:
        print(results)
        return results[0]
    else:
        return results

im = cv2.imread(find(f, all_images), 1)
an = cv2.imread(find(f, all_annots), 1).transpose(2,0,1)
r,g,b = an
nzs = np.nonzero(b==4) # 4 stands for person
instances = np.unique(g[nzs])
masks = np.zeros((len(instances), *r.shape))
for ix,_id in enumerate(instances):
    masks[ix] = g ==_id

length = len(masks)+1
plt.subplot(1, length, 1)
plt.imshow(cv2.cvtColor(im, cv2.COLOR_BGR2RGB))
ix = 2
for m in masks:
    plt.subplot(1, length, ix)
    plt.imshow(m, cmap='gray')
    ix += 1
plt.show()

样本可视化
在该数据集中，对象实例以如下方式进行标注，RGB 中的红色通道对应于对象的类别，而绿色通道对应于实例编号(如果图像中有多个相同类别的对象)图像。此外，在数据集中，Person 类别的编码值为 4。

循环遍历标注文件并存储至少包含一个人物的文件：

annots = []
for ann in all_annots:
    _ann = cv2.imread(ann, 1).transpose(2,0,1)
    r,g,b = _ann
    if 4 not in np.unique(b):
        continue
    annots.append(ann)

读取图像的 R 通道获取掩码，然后遍历掩码并检查掩码中是否至少有一个像素值为 4 (人物类别)，如果存在值为 4 的像素，就将该图像的文件名添加到包含人物的文件列表中。

将数据文件拆分为训练和验证数据：

def stems(split):
    items_new = [item.split('/')[-1] for item in split]
    items = [item.split('.')[0] for item in items_new]
    return items

from sklearn.model_selection import train_test_split
_annots = stems(annots)

trn_items, val_items = train_test_split(_annots, random_state=2)

(4) 定义图像变换方法：

def get_transform(train):
    image_transforms = []
    image_transforms.append(T.PILToTensor())
    return T.Compose(image_transforms)

(5) 创建数据集类 MasksDataset。

定义 __init__ 方法，将图像名称 (items)、图像变换方法 (transforms) 和所用文件数 (N) 作为输入：

class MasksDataset(Dataset):
    def __init__(self, items, transforms, N):
        self.items = items
        self.transforms = transforms
        self.N = N

定义 get_mask 方法，获取与图像中存在的实例数量相同的多个掩码：

    def get_mask(self, path):
        an = cv2.imread(path, 1).transpose(2,0,1)
        r,g,b = an
        nzs = np.nonzero(b==4)
        instances = np.unique(g[nzs])
        masks = np.zeros((len(instances), *r.shape))
        for ix,_id in enumerate(instances):
            masks[ix] = g == _id
        return masks

定义 __getitem__ 方法，获取所需返回的图像和相应的目标值。每个人物(实例)都被视为不同的对象类别，也就是说，每个实例都是一个不同的类别。与 Faster R-CNN 模型类似，目标值以张量字典的形式返回：

    def __getitem__(self, ix):
        _id = self.items[ix]
        img_path = f'images/training/{
      
      _id}.jpg'
        mask_path = f'annotations_instance/training/{
      
      _id}.png'
        masks = self.get_mask(mask_path)
        obj_ids = np.arange(1, len(masks)+1)
        img = Image.open(img_path).convert("RGB")
        num_objs = len(obj_ids)

除了掩码本身，Mask R-CNN 还需要边界框信息：

        boxes = []
        for i in range(num_objs):
            obj_pixels = np.where(masks[i])
            xmin = np.min(obj_pixels[1])
            xmax = np.max(obj_pixels[1])
            ymin = np.min(obj_pixels[0])
            ymax = np.max(obj_pixels[0])
            if (((xmax-xmin)<=10) | (ymax-ymin)<=10):
                xmax = xmin+10
                ymax = ymin+10
            boxes.append([xmin, ymin, xmax, ymax])

在以上代码中，通过将边界框的 x 和 y 坐标的最小值加上 10 个像素来调整存在可疑标注信息的情况(即 Person 类别的高度或宽度小于 10 像素)。

将目标值转换为张量对象：

        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        labels = torch.ones((num_objs,), dtype=torch.int64)
        masks = torch.as_tensor(masks, dtype=torch.uint8)
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)
        image_id = torch.tensor([ix])

将目标值存储在字典中：

        target = {
    
    }
        target["boxes"] = boxes
        target["labels"] = labels
        target["masks"] = masks
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd

指定变换方法并返回图像和目标值：

        if self.transforms is not None:
            img, target = self.transforms(img, target)
        img = img/255.
        return img, target

定义 __len__ 方法：

    def __len__(self):
        return self.N

定义选择随机图像的函数：

    def choose(self):
        return self[random.randint(len(self))]

检查输入输出组合：

x = MasksDataset(trn_items, get_transform(train=True), N=100)
im, targ = x[3]

length = len(targ['masks'])+1
plt.subplot(1, length, 1)
print(type(im))
plt.imshow(im.permute(1,2,0).detach().cpu())
ix = 2
for m in targ['masks']:
    plt.subplot(1, length, ix)
    plt.imshow(m, cmap='gray')
    ix += 1
plt.show()

样本示例
在以上输出中，可以看到掩码的形状为 4 x 512 x 683，表示图像中有 4 个人物。
在 __getitem__ 方法中，对于图像中存在的每个对象(实例)，都有相应的掩码和边界框。此外，由于我们只有两个类别(背景类别和人类类别)，因此我们将人类类别指定为 1。
总体而言，在目标输出字典中包括对象类别、边界框、掩码、掩码的面积以及蒙版是否对应于人物，这些信息都可以在目标输出词典中找到。为了使用损失函数，需要将数据标准化为 torchvision.models.detection.maskrcnn_resnet50_fpn 类所要求的格式。

(5) 定义实例分割模型 (get_model_instance_segmentation)，使用预训练模型，仅重新初始化检测头来预测对象类别(背景类别和人物类别)。

首先，初始化预训练模型并替换 box_predictor 和 mask_predictor 检测头，以便学习最佳权重：

def get_model_instance_segmentation(num_classes):
    # 加载在COCO上预先训练的实例分割模型 
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)

    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # 获取掩码分类器的输入特征数
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    # 掩码预测器
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
                                                       hidden_layer,num_classes)
    return model

FastRCNNPredictor 需要两个输入——in_features (输入通道数)和 num_classes (类别数)。根据要预测的类别数，计算边界框预测数
MaskRCNNPredictor 需要三个输入——in_features_mask (输入通道数)、hidden_layer (输出通道数)和 num_classes (要预测的类别数)

获取模型的详细信息：

model = get_model_instance_segmentation(2).to(device)
print(model)

Faster R-CNN 网络和 Mask R-CNN 模型之间的主要区别在于 roi_heads 模块，该模块包含多个子模块：

box_head：对齐从 FPN 网络获取的输入并创建两个张量
box_predictor：使用 box_head 的输出预测每个 RoI 的类别和边界框偏移量
mask_roi_pool：对齐来自 FPN 网络的输出
mask_head：将 mask_roi_pool 的输出转换为可用于预测掩码的特征图
mask_predictor：从 mask_head 获取输出并预测最终的掩码

(6) 获取与训练和验证图像对应的数据集和数据加载器：

dataset = MasksDataset(trn_items, get_transform(train=True), N=len(trn_items))
dataset_test = MasksDataset(val_items, get_transform(train=False), N=len(val_items))

# 定义训练、测试数据管道
data_loader = torch.utils.data.DataLoader(
    dataset, batch_size=2, shuffle=True, num_workers=0,
    collate_fn=utils.collate_fn)

data_loader_test = torch.utils.data.DataLoader(
    dataset_test, batch_size=1, shuffle=False, num_workers=0,
    collate_fn=utils.collate_fn)

(7) 定义模型、超参数和优化器：

num_classes = 2
model = get_model_instance_segmentation(num_classes).to(device)
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005,
                            momentum=0.9, weight_decay=0.0005)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                                step_size=3,
                                                gamma=0.1)

模型将图像和目标字典作为输入，可以通过以下命令查看模型输出示例：

model.eval()
pred = model(dataset[0][0][None].to(device))

输出字典包含边界框 (BOXES)、与边界框相对应的类别 (LABELS)、与类别预测相对应的置信度分数 (SCORES) 以及掩码实例的位置 (MASKS)。该模型返回 100 个预测结果，因为通常图像中的对象数不超过 100 个。

获取检测到的实例数量：

print(pred[0]['masks'].shape)

一张图像通过以上模型最多可以获取 100 个掩码实例(非背景类别)。对于这 100 个实例，返回相应的边界框、类别标签和相应的置信度值。

(8) 训练模型：

trn_history = []
for epoch in range(num_epochs):
    # 模型训练
    res = train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
    trn_history.append(res)
    # 调整学习率
    lr_scheduler.step()

使用训练后的模型，就可以预测图像中的人物实例掩码。记录训练损失随时间的变化情况：

import matplotlib.pyplot as plt
plt.title('Training Loss') 
losses = [np.mean(list(trn_history[i].meters['loss'].deque)) for i in range(len(trn_history))]
plt.plot(losses)
plt.show()

模型监测
(9) 使用训练后的模型预测测试图像：

model.eval()
k=1
im = dataset_test[k][0]
plt.imshow(im.permute(1,2,0))
plt.show()
with torch.no_grad():
    prediction = model([im.to(device)])
    for i in range(len(prediction[0]['masks'])):
        plt.imshow(Image.fromarray(prediction[0]['masks'][i, 0].mul(255).byte().cpu().numpy()), cmap='gray')
        plt.title('Class: '+str(prediction[0]['labels'][i].cpu().numpy())+' Score:'+str(prediction[0]['scores'][i].cpu().numpy()))
        plt.show()

模型预测结果
从上图中可以看出，模型成功识别出图中的人物实例。此外，模型还预测图像中的多个置信度较低的其他分割实例。

3. 多类别实例分割

在上一小节中，我们学习了如何分割 Person 类别的多个实例。在本节中，我们将调整上一节中构建的模型，一次性分割图像中多个对象类别的多个实例。

(1) 获取包含感兴趣类别的图像，类别 ID 4、ID 5、ID 6 和 ID 6：

classes_list = [4,5,6,7]
annots = []
for ann in all_annots[:3000]:
    _ann = cv2.imread(ann, 1).transpose(2,0,1)
    r,g,b = _ann
    if np.array([num in np.unique(b) for num in classes_list]).sum()==0:
        continue
    annots.append(ann)

def stems(split):
    items_new = [item.split('/')[-1] for item in split]
    items = [item.split('.')[0] for item in items_new]
    return items

from sklearn.model_selection import train_test_split
_annots = stems(annots)

trn_items, val_items = train_test_split(_annots, random_state=2)

在以上代码中，获取至少包含一个感兴趣的类别 (classes_list) 的图像。

(2) 修改 MasksDataset 类中的 get_mask 方法，使其返回两个掩码以及与每个掩码对应的类别：

    def get_mask(self,path):
        an = cv2.imread(path, 1).transpose(2,0,1)
        r,g,b = an
        cls = list(set(np.unique(b)).intersection({
    
    4,5,6,7}))
        masks = []
        labels = []
        for _cls in cls:
          nzs = np.nonzero(b==_cls)
          instances = np.unique(g[nzs])
          for ix, _id in enumerate(instances):
              masks.append(g==_id)
              labels.append(classes_list.index(_cls)+1)
        return np.array(masks), np.array(labels)

在以上代码中，获取图像中感兴趣的类别，并存储在 cls 中，然后，循环遍历每个已识别的类别 (cls)，并将红色通道值对应于类别 (cls) 的位置存储在 nzs 中。接下来，获取这些位置上的实例 ID (instances)。此外，在返回掩码 masks 和标签 lables 数组之前，将实例 instances 追加到掩码数组中，并将实例对应的类别追加到标签数组中。

(3) 修改 __getitem__ 方法中的标签 (labels) 对象，使其包含从 get_mask 方法获得的标签：

    def __getitem__(self, ix):
        _id = self.items[ix]
        img_path = f'images/training/{
      
      _id}.jpg'
        mask_path = f'annotations_instance/training/{
      
      _id}.png'
        masks, labels = self.get_mask(mask_path)
        obj_ids = np.arange(1, len(masks)+1)
        img = Image.open(img_path).convert("RGB")
        num_objs = len(obj_ids)
        boxes = []
        for i in range(num_objs):
            obj_pixels = np.where(masks[i])
            xmin = np.min(obj_pixels[1])
            xmax = np.max(obj_pixels[1])
            ymin = np.min(obj_pixels[0])
            ymax = np.max(obj_pixels[0])
            if (((xmax-xmin)<=10) | (ymax-ymin)<=10):
                xmax = xmin+10
                ymax = ymin+10
            boxes.append([xmin, ymin, xmax, ymax])
        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        labels = torch.as_tensor(labels, dtype=torch.int64)
        masks = torch.as_tensor(masks, dtype=torch.uint8)
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)
        image_id = torch.tensor([ix])
        target = {
    
    }
        target["boxes"] = boxes
        target["labels"] = labels
        target["masks"] = masks
        target["image_id"] = image_id
        target["area"] = area
        target["iscrowd"] = iscrowd
        if self.transforms is not None:
            img, target = self.transforms(img, target)
        img = img/255.
        return img, target

4.定义模型时，指定 4 个类别：

num_classes = 5
model = get_model_instance_segmentation(num_classes).to(device)

绘制模型训练期间，训练损失随时间增加的变化情况：

模型监测

使用训练后的模型预测包含目标对象的样本图像：

模型预测

小结

Mask R-CNN 是一种在目标检测任务中引入了语义分割的强大框架，通过在 Faster R-CNN 基础上进行扩展，添加了额外的分支网络，不仅可以准确地检测对象的位置和类别，还可以生成每个实例的精确像素级别的语义分割掩码。其模块化的设计可以轻松地应用于不同的任务和数据集，并且可以通过添加更多的分支进行功能扩展，如实例关键点检测等。

系列链接