PyTorch深度学习实战（26）——多对象实例分割

0. 前言

我们已经学习了多种图像分割算法，在本节中，我们将学习如何使用 Detectron2 平台以及 Google 开放图像数据集执行多对象实例分割任务。Detectron2 是 Facebook 团队打造的人工智能框架，其中包括了高性能的对象检测算法实现，包括 Mask R-CNN 模型等。Detectron2 支持一系列与目标检测和人体姿态估计等相关的任务，此外，Detectron2 还增加了对语义分割和全景分割的支持。通过利用 Detectron2，我们能够通过使用少量代码构建目标检测、分割和姿态估计模型。在本节中，我们将介绍如何从 Google 开放图像数据集中获取数据，将数据转换为 Detectron2 可接受的 COCO 格式，并训练模型执行实例分割，最后，介绍如何使用训练后的模型对新图像进行推理。

1. 获取并准备数据

Google 开放图像数据集 (Google Open Images) 是由 Google 开发和维护的一个大规模图像数据集，用于计算机视觉领域的研究和开发。该数据集包含了数百万张标记和分类的图像，涵盖了广泛的主题和场景。在本节中，我们使用 Google 开放图像数据集，为了减少训练的时间，我们仅仅使用其中一部分数据集，而不使用整个数据集。

(1) 安装所需的库：

pip install openimages

(2) 下载所需的标注文件，包括分割标注文件 train-annotations-object-segmentation.csv 和类别文件 dict.csv。

(3) 指定希望模型预测的类别 required_classes：

import pandas as pd

from glob import glob
import os
import numpy as np
import cv2

required_classes = 'person,dog,bird,car,elephant,football,jug,laptop,Mushroom,Pizza,Rocket,Shirt,Traffic sign,Watermelon,Zebra'
required_classes = [c.lower() for c in required_classes.lower().split(',')]

classes = pd.read_csv('dict.csv', header=None)
classes.columns = ['class','class_name']
classes = classes[classes['class_name'].map(lambda x: x in required_classes)]

(4) 获取 required_classes 对应的图像 ID 和掩码：

df = pd.read_csv('train-annotations-object-segmentation.csv')

data = pd.merge(df, classes, left_on='LabelName', right_on='class')

subset_data = data.groupby('class_name').agg({
   'ImageID': lambda x: list(x)[:2000]})
subset_data = ([y for x in subset_data.ImageID.tolist() for y in x])
subset_data = data[data['ImageID'].map(lambda x: x in subset_data)]
subset_masks = subset_data['MaskPath'].tolist()

考虑到海量数据集会增加训练成本，在 subset_data 中每个类别只使用 2000 张图像，我们也可以修改每个类别所用数据样本数量。使用以上代码，可以获取图像对应的 ImageId 和 MaskPath 值。接下来，我们继续从 Google 开放图像数据集中下载实际图像和蒙版。

(5) 下载掩码数据子集，Google 开放图像数据集有 16 个用于掩码 ZIP 文件，包括 train-masks-0.zip、train-masks-1.zip、train-masks-2.zip、train-masks-3.zip、train-masks-4.zip、train-masks-5.zip、train-masks-6.zip、train-masks-7.zip、train-masks-8.zip、train-masks-9.zip、train-masks-a.zip、train-masks-b.zip、train-masks-c.zip、train-masks-d.zip、train-masks-e.zip 和 train-masks-f.zip，下载以上 ZIP 文件。每个 ZIP 文件只包含 subset_masks 中的一些掩码，因此下载之后我们只读取所需掩码文件：

if not os.path.exists('masks'):
    os.mkdir('masks')

def fname(file):
    return file.split('/')[-1]

for c in '0123456789abcdef':
    tmp_mask = 'train-masks-{}/*.png'.format(c)
    tmp_masks = glob(tmp_mask)
    items = [(m, fname(m)) for m in tmp_masks]
    items = [(i,j) for (i,j) in items if j in subset_masks]
    for i,j in items:
        os.rename(i, f'masks/{
     j}')

(6) 下载 ImageId 对应的图片：

masks = glob('masks/*.png')

masks = [fname(mask) for mask in masks]

subset_data = subset_data[subset_data['MaskPath'].map(lambda x: x in masks)]
subset_imageIds = subset_data['ImageID'].tolist()

from openimages.download import _download_images_by_id

if not os.path.exists('images'):
    os.makedirs('images')

_download_images_by_id(subset_imageIds, 'train', './images/')

(7) 压缩所有图像、掩码和标注文件并保存，用于以后模型训练：

import zipfile
files = glob('images/*.jpg') + glob('masks/*.png') + ['train-annotations-object-segmentation.csv', 'dict.csv']
with zipfile.ZipFile('data.zip','w') as zipme:
    for file in files:
        zipme.write(file, compress_type=zipfile.ZIP_DEFLATED)

if not os.path.exists('train'):
    os.makedirs('train')
os.rename('images', 'train/myData')
os.rename('masks', 'train/annotations')

由于目标检测数据集有多种不同标注格式，为了标准化，Detectron 需要使用一种严格的训练数据格式。虽然可以使用自定义数据集格式并将其提供给 Detectron，但是将训练数据保存为 COCO 格式会使训练更加容易。这样，我们就在不改变训练数据的情况下利用其他训练算法(如 Detectron Transformers (DETR))。

(8) 以 COCO 格式定义所需类别：

import datetime

INFO = {
   
    "description": "MyData2023",
    "url": "None",
    "version": "1.0",
    "year": 2023,
    "contributor": "sizhky",
    "date_created": datetime.datetime.utcnow().isoformat(' ')
}

LICENSES = [
    {
   
        "id": 1,
        "name": "MIT"
    }
]

CATEGORIES = [{
   'id': id+1, 'name': name.replace('/',''), 'supercategory': 'none'} for id,(_,(name, clss_name)) in enumerate(classes.iterrows())]

在以上代码中，在 CATEGORIES 的定义中，创建了一个名为 supercategory 的新键。例如： Man 和 Woman 类别都属于 Person 这个超类别。在本节中，由于我们对超类别不感兴趣，将其指定为 none。

导入相关库，并创建一个空字典，其中包含保存 COCO JSON 文件所需的键值：

from pycococreatortools import pycococreatortools
from os import listdir
from os.path import isfile, join
from PIL import Image

coco_output = {
   
    "info": INFO,
    "licenses": LICENSES,
    "categories": CATEGORIES,
    "images": [],
    "annotations": []
}

定义包含图像位置和标注文件位置信息的变量：

ROOT_DIR = "train"
IMAGE_DIR, ANNOTATION_DIR = 'train/myData/', 'train/annotations/'
image_files = [f for f in listdir(IMAGE_DIR) if isfile(join(IMAGE_DIR, f))]
annotation_files = [f for f in listdir(ANNOTATION_DIR) if isfile(join(ANNOTATION_DIR, f))]

循环遍历每个图像文件名，并将信息填充到 coco_output 字典中的 images 键中：

image_id = 1

for image_filename in image_files:
    image = Image.open(IMAGE_DIR + '/' + image_filename)
    image_info = pycococreatortools.create_image_info(image_id, os.path.basename(image_filename), image.size)
    coco_output["images"].append(image_info)
    image_id = image_id + 1

(9) 循环遍历每个分段标注，并将信息填充到 coco_output 字典中的 annotations 键中：

def stem(filename):
    return filename.split('.')[0]
    
segmentation_id = 1
for annotation_filename in annotation_files:
    image_id = [f for f in coco_output['images'] if stem(f['file_name'])==annotation_filename.split('_')[0]][0]['id']
    class_id = [x['id'] for x in CATEGORIES if x['name'] in annotation_filename][0]
    category_info = {
   'id': class_id, 'is_crowd': 'crowd' in image_filename}
    binary_mask = np.asarray(Image.open(f'{
     ANNOTATION_DIR}/{
     annotation_filename}').convert('1')).astype(np.uint8)
 
    annotation_info = pycococreatortools.create_annotation_info(segmentation_id, image_id, category_info, binary_mask, image.size, tolerance=2)

    if annotation_info is not None:
        coco_output["annotations"].append(annotation_info)
        segmentation_id = segmentation_id + 1

(10) 将 coco_output 保存在 JSON 文件中：

coco_output['categories'] = [{
   'id': id+1, 'name': clss_name, 'supercategory': 'none'} for id,(_,(name, clss_name)) in enumerate(classes.iterrows())]

import json
with open('images.json', 'w') as output_json_file:
    json.dump(coco_output, output_json_file)

获取 COCO 格式的文件后，就可以很容易地使用 Detectron2 框架来训练模型。

2. 使用 Detectron2 训练实例分割模型

(1) 安装 Detectron2，在安装正确的库之前，首先应该检查 CUDA 和 PyTorch 版本，以 PyTorch 1.9 和 CUDA 11.3 为例，使用相应的文件：

$ pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu113/torch1.9/index.html
$ pip install pyyaml pycocotools

(2) 导入相关 detectron2 库：

from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.engine import DefaultTrainer

重新获取所需的类：

required_classes = 'person,dog,bird,car,elephant,football,jug,laptop,Mushroom,Pizza,Rocket,Shirt,Traffic sign,Watermelon,Zebra'
required_classes = [c.lower() for c in required_classes.lower().split(',')]

cclasses = pd.read_csv('dict.csv', header=None)
classes.columns = ['class','class_name']
classes = classes[classes['class_name'].map(lambda x: x in required_classes)]

(3) 使用 register_coco_instances 注册创建的数据集：

from detectron2.data.datasets import register_coco_instances
register_coco_instances("dataset_train", {
   }, "images.json", "train/myData")

register_coco_instances 函数接受四个参数，第一个参数是数据集的名称，用于唯一标识该数据集；第二个参数是一个空字典，通常用于存储一些额外的元数据信息；第三个参数是标注文件的路径；最后一个参数是图像文件夹的路径。可以通过类似方式注册任意数量的数据集。

(4) 定义 Detectron2 配置文件 cfg 中的参数。

配置文件 (cfg) 是一个特殊的 Detectron 对象，它包含用于训练模型的所有相关信息：

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("dataset_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml") # pretrained weights
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025 # pick a good LR
cfg.SOLVER.MAX_ITER = 8000 # instead of epochs, we train on 5000 batches
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512
cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(classes)

使用配置文件可以设置训练模型所需的所有主要超参数，merge_from_file 方法用于从配置文件导入核心参数，mask_rcnn_R_50_FPN_3x.yaml 配置文件使用以 FPN 作为主干网络进行预训练的 mask_rcnn，此配置文件还包含有关预训练模型的其他信息，例如优化器和损失函数。

(5) 训练模型：

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=True)
trainer.train()

训练完成后，可以使用该模型预测类别、边界框以及自定义数据集中定义的类别对象的分割掩码。

模型训练
将模型保存在指定文件夹中：

import shutil
import random
from matplotlib import pyplot as plt
shutil.copy('output/model_final.pth', 'output/trained_model.pth')

模型训练完成后，在下一小节中，我们将对新图像进行推断。

3. 对新图像进行推断

为了对新图像进行推断，我们首先加载模型，设置概率阈值，并将其传递给 DefaultPredictor 方法。

(1) 加载训练好的模型权重，使用相同的 cfg 配置文件并加载模型权重：

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "trained_model.pth")

(2) 设置对象属于某一类别的概率阈值：

cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.25

(3) 定义 predictor 方法：

predictor = DefaultPredictor(cfg)

(4) 对新图像进行分割并将其可视化。

加载图像文件，随机绘制 40 张图像分割结果：

from detectron2.utils.visualizer import ColorMode
files = glob('train/myData/*.jpg')
for _ in range(40):
    im = cv2.imread(random.choice(files))
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1],
                    scale=0.5,
                    metadata=MetadataCatalog.get("dataset_train"),
                    instance_mode=ColorMode.IMAGE_BW 
    )

    out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    plt.imshow(out.get_image())
    plt.show()

Visualizer 是 Detectron2 绘制对象实例的方式，模型预测输出是张量字典，Visualizer 将它们转换为像素信息并将绘制在图像上。Visualizer 接受以下参数作为输入：

im：要可视化的图像
scale：绘制图像的大小，在以上代码中，将图像缩小 50%
metadata：数据集的类别信息，主要是索引到类别的映射，用于将原始张量输入解码为人类的可读实际类别
instance_mode：指定模型突出显示的分割像素

创建了 Visualizer 类对象后，就可以用其绘制模型的实例预测并显示图像，代码输出结果如下所示：

模型预测结果

从以上输出可以看出，模型能够准确地识别出人物对应的像素。

小结

Detectron2 是用于计算机视觉任务的深度学习库，构建与 PyTorch 框架之上，并提供了一套灵活且高性能的工具，用于对象检测、实例分割和人体关键点检测等任务。其使用了先进的技术和优化策略，如特征金字塔网络、RoI Align 操作等，以提高模型的感受野和对小目标的检测能力。本节中，介绍了如何安装和使用 Detectron2 库，并通过使用 COCO 数据格式提高模型训练效率。