这篇张使用PyTorch预训练的Mask R-CNN在自定义的训练集中做训练。

1 dataset的定义

dataset需要继承 torch.utils.data.Dataset 类, 并实现 __len__ 和__getitem__。

要注意的是 __getitem__ s需要返回如下内容:

• image: size为 (H, W)的PIL.Image

• target:包含如下对象的字典

  • boxes (FloatTensor[N, 4]): N个 bounding box，类型为左上坐标和右下坐标，[x0, y0, x1, y1]
  • labels (Int64Tensor[N]): bounding box的label，一般用 0 代表background
  • image_id (Int64Tensor[1]): 图像的id，每个图像的id应该是唯一的，它会用于evaluation
  • area (Tensor[N]): bounding box的面积，就是长乘宽.在基于COCO metric的evaluation的时候回用到，用来分别计算small, medium 和 large box的metric scores
  • iscrowd (UInt8Tensor[N]): iscrowd=True那么bounding box不会在evaluation的时候被计算
  • (可选) masks (UInt8Tensor[N, H, W]): 每个object的mask
  • (可选) keypoints (FloatTensor[N, K, 3]): 对于N个对象，包含K个keypoints，keypoints的格式为[x, y, visibility] 如果visibility=0，那么呆板keypoint被遮挡住，不可见

只要实现上面的dataset，就可以用PyTorch的预训练Mask R-CNN训练和evaluate，evaluate需要用到 pycocotools 安装命令为 pip install pycocotools

对于label有一个要注意的地方，model把0作为background，如果dataset不包含background类，那么label里面不应该包含0类，例如把猫定为1，把狗定为2，labels应该类似[1,2]，而不应该出现[0,2]这种情况

另外，如果希望在训练的时候使用aspect ratio grouping，也就是每个batch的图片的宽高比都类似，那么就应该实现get_height_and_width方法，返回的是图片的高和宽，如果不提供这个方法，会通过__getitem__来query全部的数据集，再选择宽高比类似的图片，这样速度会比实现get_height_and_width方法慢得多

2 实现自定义dataset

接下来，使用Penn-Fudan数据集来构建一个dataset，数据集的格式大概是

PennFudanPed/
  PedMasks/
    FudanPed00001_mask.png
    FudanPed00002_mask.png
    FudanPed00003_mask.png
    FudanPed00004_mask.png
    ...
  PNGImages/
    FudanPed00001.png
    FudanPed00002.png
    FudanPed00003.png
    FudanPed00004.png

下图是图片和segmentation的效果

每个图片有自己对应的segmentation mask，每个instance对应不同的颜色，

先安装需要的python库

pip3 install torch torchvision torchaudio cython pycocotools python-opencv opencv-contrib-python

接着，下载并解压数据集

wget https://www.cis.upenn.edu/~jshi/ped_html/PennFudanPed.zip
unzip PennFudanPed.zip

导入需要的Python库，其中的辅助脚本包括

• coco_eval.py
• coco_utils.py
• engine.py
• transforms.py
• utils.py

import os
import numpy as np
import torch
from PIL import Image

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

from engine import train_one_epoch, evaluate
import utils
import transforms as T

构建dataset

class PennFudanDataset(object):
    def __init__(self, root, transforms):
        self.root = root
        self.transforms = transforms
        # 加载全部图像和mask的名字
        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):
        # 加载选定的图像和mask
        img_path = os.path.join(self.root, "PNGImages", self.imgs[idx])
        mask_path = os.path.join(self.root, "PedMasks", self.masks[idx])
        img = Image.open(img_path).convert("RGB")
        # # mask不用转换为RGB，因为mask图片有class_num+1个值，其中0代表的是background
        mask = Image.open(mask_path)

        mask = np.array(mask)
        # np.unique()是去重，得到的是图片里包含的所有类别，instance segmentation的每个对象会被分割为不同的值
        obj_ids = np.unique(mask)
        # # 第一个id肯定是背景0，所以不要它，得到的是全部的类别
        obj_ids = obj_ids[1:]

        # mask转换为 n_objects个True/False 矩阵，用来对每个object的mask做分类, 最终形态为[n_class, h, w]
        masks = mask == obj_ids[:, None, None]

        # 对于每个mask，获取bbox坐标
        num_objs = len(obj_ids)
        boxes = []
        for i in range(num_objs):
            # np.where(condition) 取出值为True的x和y，取min和max，得到bbox
            pos = np.where(masks[i])
            xmin = np.min(pos[1])
            xmax = np.max(pos[1])
            ymin = np.min(pos[0])
            ymax = np.max(pos[0])
            boxes.append([xmin, ymin, xmax, ymax])
        
        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        # n_classes个object所属的类别
        labels = torch.ones((num_objs,), dtype=torch.int64)
        masks = torch.as_tensor(masks, dtype=torch.uint8)
        
        image_id = torch.tensor([idx])
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        # 假设没有遮挡，会在evaluation中用到
        iscrowd = torch.zeros((num_objs,), dtype=torch.int64)

        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)

        return img, target

    def __len__(self):
        return len(self.imgs)

获取预训练模型，并根据分类的类别数替换分类器和mask分类器的

def get_model_instance_segmentation(num_classes):
    # 加载在coco上预训练过的模型
    model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)

    # 获取分类器的输入feature size
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    # 替换预训练模型的头部
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # 获取mask的输入feature size
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    
    # 替换mask分类器
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
                                                       hidden_layer,
                                                       num_classes)

    return model

3 模型训练

这里使用的是基于Faster R-CNN的 Mask R-CNN模型，Faster R-CNN可以预测bounding box和class score

Mask R-CNN 在Faster R-CNN增加了一个额外的分支，来预测mask

数据增强

def get_transform(train):
    transforms = []
    transforms.append(T.ToTensor())
    if train:
        transforms.append(T.RandomHorizontalFlip(0.5))
    return T.Compose(transforms)

构建数据集，把之前的dataset用torch.utils.data.Subset包起来

# 如果可以用GPU，则在GPU上训练
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

# 数据集只有两个类别：背景和人
num_classes = 2
# 构建用于放到 DataLoader 的 dataset
dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))
dataset_test = PennFudanDataset('PennFudanPed', get_transform(train=False))

# 分割数据集
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-50])
dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])

# 设定data loader
data_loader = torch.utils.data.DataLoader(
    dataset, batch_size=2, shuffle=True, num_workers=4,
    collate_fn=utils.collate_fn)

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

获取模型，以及设置优化器与scheduler学习器

# 获取模型
model = get_model_instance_segmentation(num_classes)

# 模型转到显存
model.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)
# 学习率scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                               step_size=3,
                                               gamma=0.1)

开始训练10个epoch

# 训练10个epoch
num_epochs = 10

for epoch in range(num_epochs):
    train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
    #修改learning rate
    lr_scheduler.step()
    # 在测试集上evaluate
    evaluate(model, data_loader_test, device=device)

print("That's it!")

Epoch: [0]  [ 0/60]  eta: 0:00:46  lr: 0.000090  loss: 5.1038 (5.1038)  loss_classifier: 0.7602 (0.7602)  loss_box_reg: 0.1499 (0.1499)  loss_mask: 4.1816 (4.1816)  loss_objectness: 0.0105 (0.0105)  
...
IoU metric: bbox
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.738
 ...
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.329
 ...
IoU metric: segm
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.711
 ...
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.318
 ...
 ...
Epoch: [9]  [ 0/60]  eta: 0:00:46  lr: 0.000005  loss: 0.1645 (0.1645)  loss_classifier: 0.0191 (0.0191)  loss_box_reg: 0.0273 (0.0273)  loss_mask: 0.1137 (0.1137)  loss_objectness: 0.0002 (0.0002)  
...
That's it!

4 推理部分

读取一张图片，用于推理

image = next(iter(data_loader))[0][0]
image = image.reshape(1, image.shape[0], image.shape[1], image.shape[2])
image = image.cuda()

执行推理结果

output = model(image)

输出推理mask

plt.imshow(output[0]['masks'].detach().cpu().numpy()[0, 0])

plt.imshow(output[0]['masks'].detach().cpu().numpy()[1, 0])

把图片转回numpy.array格式

image_show = image.cpu().numpy()
image_show.shape

(1, 3, 441, 465)

要把它转成H,W,C格式才可以

# 把第一维压缩
image_show = image_show.squeeze(0)
# 交换(C, H, W)为(H, W, C)
image_show = np.transpose(image_show, [1, 2, 0])
# 之前归一化了，现在乘回255
image_show *= 255
# 必须用 uint8
image_show = image_show.astype('uint8')
# transpose之后存储不连续了，所以要用ascontiguousarray修复
image_show = np.ascontiguousarray(image_show)

画bbox：

for bbox in output[0]['boxes'].cpu().detach().numpy():
    cv2.rectangle(image_show, (bbox[0], bbox[1]), (bbox[2], bbox[3]), (0, 255, 0), 5)

plt.imshow(image_show)