在 PyTorch 中使用 Detectron2 进行对象检测的指南

目标检测是深度学习的流行应用之一。让我们首先考虑一个现实生活中的例子。大多数人会在手机中使用 Google 相册，它会根据“事物”选项下的照片自动将照片分组。我在下面附上一个片段。

你可以观察到该应用程序能够从图片中识别对象并使用它们将它们分类为更广泛的类别。这是一个涉及对象检测的示例。

在本文中，我将使用名为 Detectron2 的最新稳健模型执行对象检测，使用 PyTorch 作为代码。

介绍 Detectron2

Facebook AI Research (FAIR) 提出了这个高级库，它在对象检测和分割问题上取得了惊人的结果。Detectron2 基于 maskrcnn 基准。它的实现是在 PyTorch 中。由于涉及大量计算，它需要 CUDA。

它支持边界框检测、实例分割、关键点检测、密集姿态检测等多项任务。它提供了预先训练的模型，你可以轻松地将其加载并用于新图像。

我将在下一节中介绍一个示例。

安装

第一步是安装detectron2库和需要的依赖

import torch
torch.__version__
import torchvision
#torchvision.__version__
!pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu102/torch1.7/index.html

现在，你必须导入detectron2 及其模块。

import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()
%matplotlib inline
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.structures import BoxMode

让我们也导入我们需要的公共库。

import numpy as np
import os, json, cv2, random
import matplotlib.pyplot as plt

使用预训练模型进行推理：代码

Detectron2 的许多预训练模型可以在MODEL_ZOO（https://github.com/facebookresearch/detectron2/blob/master/MODEL_ZOO.md）访问。这些模型已经在不同的数据集上进行了训练，可以随时使用。

即使人们在训练他们的自定义数据集时，他们也会使用这些预训练的权重来初始化他们的模型。事实证明，它可以减少训练时间并提高性能。我们将使用的模型是在 COCO 数据集上预训练的。

首先，我们必须定义对象检测模型的完整配置。我们从detectron2.config 模块中导入了'get_cfg' 函数，我们现在将使用它。

我选择了 Coco 实例分段配置（YAML 文件）。还有其他选项可用。你还必须设置模型的阈值分数（通常设置在 0.4 到 0.6 之间）。你可以从检查点加载配置的预训练权重。

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file('COCOInstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml'))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url('COCOInstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml')

完成配置部分后，我们使用配置初始化 DefaultPredictor。

predictor = DefaultPredictor(cfg)

现在可以开始预测图像了。

让我们在示例图像上使用它。下面的代码使用 OpenCV 库加载和读取图像。

!wget http://images.cocodataset.org/val2017/000000439715.jpg -O input.jpg
im = cv2.imread('./input.jpg')
print(im.shape)
plt.figure(figsize=(15,7.5))
plt.imshow(im[..., ::-1])

怎么做检测？

将输入图像传递给我们初始化的预测器

outputs = predictor(im[..., ::-1])

这个输出是一个字典。字典有实例（预测框）、分数、预测标签，我附上了代码片段的输出。

接下来，使用 Visualizer 类查看检测是如何执行的。可视化类具有绘制实例预测的功能。

v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
out = v.draw_instance_predictions(outputs['instances'].to('cpu'))
plt.figure(figsize=(20,10))
plt.imshow(out.get_image()[..., ::-1][..., ::-1])

你可以观察到模型检测到了所有的人和马。

我在照片上附加了另一个示例输出。

背景中的汽车也有97% 的准确率被检测到 。

自定义数据集上的 Detectron2

到目前为止，我们只是使用预训练的模型进行推理。但在某些情况下，你可能需要单独检测汽车、人等特定物体。你可能想从头开始在数据集上训练模型。

Detectron2 也为此提供了一种简单的方法。让我们看看如何操作。

准备数据集

我将使用气球数据集，目的是检测图像中的气球。这是一个比较简单的例子。

!wget https://github.com/matterport/Mask_RCNN/releases/download/v2.1/balloon_dataset.zip
!unzip balloon_dataset.zip > /dev/null

请注意，Detectron2 需要特定格式的数据。为了将我们的气球数据集转换成这样的格式，让我们定义一些辅助函数。

我们的函数将输入图像目录/文件夹路径作为输入。然后打开并加载 JSON 文件。我们通过JSON文件的记录枚举，得到图片路径。从路径中读取每张图像，并将其高度、权重、文件名和图像 ID 存储在字典“record”中。

接下来，我们通读注释，并将边界框详细信息存储在另一个字典“obj”中。在每个循环结束时，记录会附加到名为“dataset_dicts”的列表中。

类似地，边界框字典也附加到列表“objs”。该列表将依次被分配为记录字典中“annotations”键的值。

from detectron2.structures import BoxMode
def get_balloon_images(img_folder):
    json_file = os.path.join(img_folder, 'via_region_data.json')
    with open(json_file) as f:
        imgs_anns = json.load(f)
    dataset_dicts = []
    for idx, v in enumerate(imgs_anns.values()):
        record = {}
        filename = os.path.join(img_dir, v['filename'])
        height, width = cv2.imread(filename).shape[:2]
        record['file_name'] = filename
        record['image_id'] = idx
        record['height'] = height
        record['width'] = width
        annos = v['regions']
        objs = []
        for _, anno in annos.items():
            assert not anno['region_attributes']
            anno = anno['shape_attributes']
            px = anno['all_points_x']
            py = anno['all_points_y']
            poly = [(x + 0.5, y + 0.5) for x, y in zip(px, py)]
            poly = [p for x in poly for p in x]
            obj = {
                'bbox': [np.min(px), np.min(py), np.max(px), np.max(py)],
                'bbox_mode': BoxMode.XYXY_ABS,
                'segmentation': [poly],
                'category_id': 0,
            }
            objs.append(obj)
        record['annotations'] = objs
        dataset_dicts.append(record)
    return dataset_dicts

最后，此辅助函数返回具有注释的字典列表。下一步是注册这些训练和验证数据集。要注册数据集，你必须使用 DatasetCatalog.register 和 MetadataCatalog 方法。

for d in ['train', 'val']:
    DatasetCatalog.register('balloon_' + d, lambda d=d: get_balloon_images('balloon/' + d))
    MetadataCatalog.get('balloon_' + d).set(thing_classes=['balloon'])
balloon_metadata = MetadataCatalog.get('balloon_train')

训练数据可视化

我们已经注册了数据集。现在让我们看一下训练数据。下面的代码从气球训练数据集中随机抽取一个样本。

为了绘制实例检测，我们再次使用可视化类

dataset_dicts = get_balloon_images('balloon/train')
for d in random.sample(dataset_dicts, 3):
    img = cv2.imread(d['file_name'])
    visualizer = Visualizer(img[:, :, ::-1], metadata=balloon_metadata, scale=0.5)
    out = visualizer.draw_dataset_dict(d)
    plt.figure(figsize=(15,7))
    plt.imshow(out.get_image()[:, :, ::-1][..., ::-1])

自定义数据训练

让我们进入训练部分。为此，首先从 Detectron 的引擎模块中导入 DefaultTrainer。定义数据集和其他参数，如worker数、批次大小、类数（在本例中为 1）。

我们用预训练的权重初始化模型并进一步训练。最大迭代次数参数将根据数据集的大小和任务的复杂性而变化。

from detectron2.engine import DefaultTrainer
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file('COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml'))
cfg.DATASETS.TRAIN = ('balloon_train',)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 6
# Let training initialize from model zoo
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url('COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml')
cfg.SOLVER.IMS_PER_BATCH = 8
cfg.SOLVER.BASE_LR = 0.001
cfg.SOLVER.MAX_ITER = 500
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # faster, enough for this dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  #only one class (balloon)
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=False)
trainer.train()

这可能需要一段时间来训练！

结果

请注意，无论何时训练深度学习模型，都要保存其最终检查点。你可以轻松加载它以执行预测并获得推论。

下面的代码片段加载模型并初始化预测器。我们从验证数据集中抽取一些随机样本并将它们传递给预测器。

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, 'model_final.pth')  # path to the model we trained
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5   # set a testing threshold
predictor = DefaultPredictor(cfg)
from detectron2.utils.visualizer import ColorMode
dataset_dicts = get_balloon_images('balloon/val')
for d in random.sample(dataset_dicts, 2):    
    im = cv2.imread(d['file_name'])
    outputs = predictor(im) 
    v = Visualizer(im[:, :, ::-1],
                   metadata=balloon_metadata, 
                   scale=0.5, 
                   instance_mode=ColorMode.IMAGE_BW  
    )
    out = v.draw_instance_predictions(outputs['instances'].to('cpu'))
    plt.figure(figsize=(15,7))
    plt.imshow(out.get_image()[:, :, ::-1][..., ::-1])

这是输出！

☆ END ☆