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 pdfrom glob import glob
import os
import numpy as np
import cv2required_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_idif 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 datetimeINFO = {"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 Imagecoco_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 = 1for 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 数据格式提高模型训练效率。

系列链接

PyTorch深度学习实战（1）——神经网络与模型训练过程详解
PyTorch深度学习实战（2）——PyTorch基础
PyTorch深度学习实战（3）——使用PyTorch构建神经网络
PyTorch深度学习实战（4）——常用激活函数和损失函数详解
PyTorch深度学习实战（5）——计算机视觉基础
PyTorch深度学习实战（6）——神经网络性能优化技术
PyTorch深度学习实战（7）——批大小对神经网络训练的影响
PyTorch深度学习实战（8）——批归一化
PyTorch深度学习实战（9）——学习率优化
PyTorch深度学习实战（10）——过拟合及其解决方法
PyTorch深度学习实战（11）——卷积神经网络
PyTorch深度学习实战（12）——数据增强
PyTorch深度学习实战（13）——可视化神经网络中间层输出
PyTorch深度学习实战（14）——类激活图
PyTorch深度学习实战（15）——迁移学习
PyTorch深度学习实战（16）——面部关键点检测
PyTorch深度学习实战（17）——多任务学习
PyTorch深度学习实战（18）——目标检测基础
PyTorch深度学习实战（19）——从零开始实现R-CNN目标检测
PyTorch深度学习实战（20）——从零开始实现Fast R-CNN目标检测
PyTorch深度学习实战（21）——从零开始实现Faster R-CNN目标检测
PyTorch深度学习实战（22）——从零开始实现YOLO目标检测
PyTorch深度学习实战（23）——使用U-Net架构进行图像分割
PyTorch深度学习实战（24）——从零开始实现Mask R-CNN实例分割
PyTorch深度学习实战（25）——自编码器(Autoencoder)
PyTorch深度学习实战（26）——卷积自编码器(Convolutional Autoencoder)
PyTorch深度学习实战（27）——变分自编码器(Variational Autoencoder, VAE)
PyTorch深度学习实战（28）——对抗攻击(Adversarial Attack)
PyTorch深度学习实战（29）——神经风格迁移
PyTorch深度学习实战（30）——Deepfakes
PyTorch深度学习实战（31）——生成对抗网络(Generative Adversarial Network, GAN)
PyTorch深度学习实战（32）——DCGAN详解与实现
PyTorch深度学习实战（33）——条件生成对抗网络(Conditional Generative Adversarial Network, CGAN)
PyTorch深度学习实战（34）——Pix2Pix详解与实现