Detectron2 目标检测和分割框架入门

1. Detectron2 简介

Detectron2 是由 Facebook AI Research (FAIR) 开发的目标检测和分割框架，基于 PyTorch 构建，是 Detectron 的 successor。它提供了丰富的目标检测、实例分割、语义分割等计算机视觉任务的模型和工具，是当前最流行的计算机视觉框架之一。

1.1 Detectron2 的主要特点

• 由 Facebook AI Research 开发：背靠顶级研究机构，持续更新和改进
• 支持多种目标检测和分割模型：包括 Faster R-CNN、Mask R-CNN、RetinaNet 等
• 模块化设计：易于扩展和定制
• 提供预训练模型：包含多种预训练模型，适用于不同场景
• 高性能实现：优化的训练和推理性能
• 丰富的工具和评估指标：提供完整的训练、评估和可视化工具

1.2 Detectron2 的应用场景

• 目标检测：识别图像中的物体并定位
• 实例分割：对每个物体进行像素级分割
• 语义分割：对图像进行像素级分类
• 关键点检测：检测人体关键点等
• 全景分割：同时进行实例分割和语义分割

2. 安装 Detectron2

2.1 环境要求

• Python 3.6 或更高版本
• PyTorch 1.6 或更高版本
• torchvision 0.7 或更高版本
• CUDA 10.1 或更高版本（推荐，用于 GPU 加速）

2.2 安装方法

1. 安装依赖：

# 安装 PyTorch 和 torchvision
pip install torch torchvision

# 安装其他依赖
pip install cython pyyaml==5.1
pip install -U 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

2. 安装 Detectron2：

# 从 GitHub 安装
pip install 'git+https://github.com/facebookresearch/detectron2.git'

3. Detectron2 核心概念

3.1 数据集 (Dataset)

Detectron2 使用 Dataset 类来表示训练、验证和测试数据。它支持多种数据格式，包括 COCO、PASCAL VOC 等。

3.2 数据加载器 (DataLoader)

DataLoader 负责将数据集转换为模型可以处理的批次数据，支持批处理和数据增强。

3.3 模型 (Model)

Model 是 Detectron2 的核心组件，定义了模型的前向传播逻辑和损失计算。

3.4 训练器 (Trainer)

Trainer 负责模型的训练过程，包括优化器选择、学习率调度、早停等。

3.5 评估器 (Evaluator)

Evaluator 负责评估模型在测试集上的性能，计算各种指标。

4. 基本使用

4.1 加载预训练模型

import detectron2
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
import cv2

# 配置
cfg = get_cfg()
cfg.merge_from_file("detectron2/configs/COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")
cfg.MODEL.WEIGHTS = "https://dl.fbaipublicfiles.com/detectron2/COCO-Detection/faster_rcnn_R_50_FPN_3x/137849458/model_final_280758.pkl"
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5

# 创建预测器
predictor = DefaultPredictor(cfg)

# 加载图像
img = cv2.imread("input.jpg")

# 预测
outputs = predictor(img)

# 可视化结果
v = Visualizer(img[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2.imshow("Output", v.get_image()[:, :, ::-1])
cv2.waitKey(0)

4.2 目标检测

import detectron2
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
import cv2

# 配置
cfg = get_cfg()
cfg.merge_from_file("detectron2/configs/COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")
cfg.MODEL.WEIGHTS = "https://dl.fbaipublicfiles.com/detectron2/COCO-Detection/faster_rcnn_R_50_FPN_3x/137849458/model_final_280758.pkl"
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5

# 创建预测器
predictor = DefaultPredictor(cfg)

# 加载图像
img = cv2.imread("street.jpg")

# 预测
outputs = predictor(img)

# 可视化结果
v = Visualizer(img[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2.imwrite("output.jpg", v.get_image()[:, :, ::-1])
print("检测完成，结果已保存到 output.jpg")

4.3 实例分割

import detectron2
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
import cv2

# 配置
cfg = get_cfg()
cfg.merge_from_file("detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.MODEL.WEIGHTS = "https://dl.fbaipublicfiles.com/detectron2/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl"
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5

# 创建预测器
predictor = DefaultPredictor(cfg)

# 加载图像
img = cv2.imread("people.jpg")

# 预测
outputs = predictor(img)

# 可视化结果
v = Visualizer(img[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2.imwrite("segmentation_output.jpg", v.get_image()[:, :, ::-1])
print("实例分割完成，结果已保存到 segmentation_output.jpg")

4.4 语义分割

import detectron2
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
import cv2

# 配置
cfg = get_cfg()
cfg.merge_from_file("detectron2/configs/COCO-SemanticSegmentation/panoptic_fpn_R_50_3x.yaml")
cfg.MODEL.WEIGHTS = "https://dl.fbaipublicfiles.com/detectron2/COCO-SemanticSegmentation/panoptic_fpn_R_50_3x/139514569/model_final_cafdb1.pkl"

# 创建预测器
predictor = DefaultPredictor(cfg)

# 加载图像
img = cv2.imread("scene.jpg")

# 预测
outputs = predictor(img)

# 可视化结果
v = Visualizer(img[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
v = v.draw_sem_seg(outputs["sem_seg"].argmax(dim=0).cpu())
cv2.imwrite("semantic_output.jpg", v.get_image()[:, :, ::-1])
print("语义分割完成，结果已保存到 semantic_output.jpg")

5. 自定义数据集

5.1 准备数据集

1. 创建数据集目录结构：

dataset/
├── annotations/
│   └── instances_train.json
└── images/
    └── train/

2. 标注数据：使用 COCO 格式标注数据。

5.2 注册数据集

from detectron2.data import DatasetCatalog, MetadataCatalog

# 注册训练数据集
def get_train_dicts():
    # 加载数据集
    # 返回 COCO 格式的数据集
    pass

DatasetCatalog.register("my_dataset_train", get_train_dicts)
MetadataCatalog.get("my_dataset_train").set(thing_classes=["person", "car", "bike"])

5.3 配置和训练

from detectron2.config import get_cfg
from detectron2.engine import DefaultTrainer

# 配置
cfg = get_cfg()
cfg.merge_from_file("detectron2/configs/COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")
cfg.DATASETS.TRAIN = ("my_dataset_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = "detectron2://COCO-Detection/faster_rcnn_R_50_FPN_3x/137849458/model_final_280758.pkl"
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 1000
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3  # 3 classes: person, car, bike

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

6. 模型评估和部署

6.1 评估模型

from detectron2.engine import DefaultTrainer
from detectron2.evaluation import COCOEvaluator

class CocoTrainer(DefaultTrainer):
    @classmethod
    def build_evaluator(cls, cfg, dataset_name, output_folder=None):
        if output_folder is None:
            output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
        return COCOEvaluator(dataset_name, output_dir=output_folder)

# 评估
trainer = CocoTrainer(cfg)
trainer.test()

6.2 导出模型

from detectron2.modeling import build_model
import torch

# 加载模型
model = build_model(cfg)
checkpoint = torch.load(cfg.MODEL.WEIGHTS)
model.load_state_dict(checkpoint["model"])

# 导出为 ONNX
dummy_input = torch.randn(1, 3, 640, 640)
torch.onnx.export(model, dummy_input, "model.onnx")

6.3 部署模型

使用 ONNX Runtime 部署模型：

import onnxruntime as rt
import cv2
import numpy as np

# 加载模型
session = rt.InferenceSession("model.onnx")
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 预处理图像
img = cv2.imread("input.jpg")
img = cv2.resize(img, (640, 640))
img = img.transpose(2, 0, 1)
img = img[np.newaxis, :, :, :]
img = img.astype(np.float32) / 255.0

# 推理
outputs = session.run([output_name], {input_name: img})

# 处理输出
# ...

7. 实用技巧

7.1 数据增强

Detectron2 提供了丰富的数据增强选项：

from detectron2.data import transforms as T

# 定义数据增强
transform_gen = T.AugmentationList([
    T.RandomFlip(prob=0.5, horizontal=True, vertical=False),
    T.RandomBrightness(0.8, 1.2),
    T.RandomContrast(0.8, 1.2),
    T.RandomSaturation(0.8, 1.2),
    T.ResizeShortestEdge([640, 672, 704, 736, 768, 800], sample_style="choice")
])

# 应用数据增强
aug_input = T.AugInput(img)
transforms = transform_gen(aug_input)
augmented_img = aug_input.image

7.2 模型调优

• 学习率调度：使用不同的学习率调度策略
• 批量大小：根据 GPU 内存调整批量大小
• 骨干网络：选择适合任务的骨干网络
• 超参数搜索：使用网格搜索或随机搜索寻找最佳超参数

7.3 性能优化

• 混合精度训练：使用 FP16 加速训练
• 分布式训练：使用多 GPU 加速训练
• 模型量化：使用 INT8 量化减少模型大小和加速推理
• 推理优化：使用 TensorRT 等工具加速推理

8. 应用案例

8.1 行人检测

import detectron2
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
import cv2

# 配置
cfg = get_cfg()
cfg.merge_from_file("detectron2/configs/COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")
cfg.MODEL.WEIGHTS = "https://dl.fbaipublicfiles.com/detectron2/COCO-Detection/faster_rcnn_R_50_FPN_3x/137849458/model_final_280758.pkl"
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5

# 创建预测器
predictor = DefaultPredictor(cfg)

# 加载图像
img = cv2.imread("street.jpg")

# 预测
outputs = predictor(img)

# 过滤出人体
person_instances = outputs["instances"][outputs["instances"].pred_classes == 0]  # 0 是 COCO 数据集中的人体类别

# 可视化结果
v = Visualizer(img[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
v = v.draw_instance_predictions(person_instances.to("cpu"))
cv2.imwrite("person_detection.jpg", v.get_image()[:, :, ::-1])
print("行人检测完成，结果已保存到 person_detection.jpg")

8.2 车辆检测

import detectron2
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
import cv2

# 配置
cfg = get_cfg()
cfg.merge_from_file("detectron2/configs/COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")
cfg.MODEL.WEIGHTS = "https://dl.fbaipublicfiles.com/detectron2/COCO-Detection/faster_rcnn_R_50_FPN_3x/137849458/model_final_280758.pkl"
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5

# 创建预测器
predictor = DefaultPredictor(cfg)

# 加载图像
img = cv2.imread("traffic.jpg")

# 预测
outputs = predictor(img)

# 过滤出车辆（汽车、卡车等）
vehicle_classes = [2, 3, 4, 6, 8]  # COCO 数据集中的车辆类别
vehicle_instances = outputs["instances"][np.isin(outputs["instances"].pred_classes, vehicle_classes)]

# 可视化结果
v = Visualizer(img[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
v = v.draw_instance_predictions(vehicle_instances.to("cpu"))
cv2.imwrite("vehicle_detection.jpg", v.get_image()[:, :, ::-1])
print("车辆检测完成，结果已保存到 vehicle_detection.jpg")

8.3 实例分割

import detectron2
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog
import cv2

# 配置
cfg = get_cfg()
cfg.merge_from_file("detectron2/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.MODEL.WEIGHTS = "https://dl.fbaipublicfiles.com/detectron2/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl"
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5

# 创建预测器
predictor = DefaultPredictor(cfg)

# 加载图像
img = cv2.imread("fruits.jpg")

# 预测
outputs = predictor(img)

# 可视化结果
v = Visualizer(img[:, :, ::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2.imwrite("instance_segmentation.jpg", v.get_image()[:, :, ::-1])
print("实例分割完成，结果已保存到 instance_segmentation.jpg")

9. 总结

Detectron2 是一个强大的目标检测和分割框架，它提供了丰富的模型和工具，使计算机视觉任务变得更加简单和高效。通过本教程的学习，你应该已经掌握了 Detectron2 的核心概念和基本使用方法，可以开始使用 Detectron2 进行自己的计算机视觉项目开发。

Detectron2 的模块化设计和丰富的预训练模型使其成为计算机视觉研究和应用的理想选择，而其基于 PyTorch 的实现则保证了灵活性和性能。无论是进行学术研究还是工业应用，Detectron2 都能为你提供强大的支持。

10. 进一步学习资源

• Detectron2 官方文档：https://detectron2.readthedocs.io/
• Detectron2 GitHub 仓库：https://github.com/facebookresearch/detectron2
• Facebook AI Research：https://ai.facebook.com/
• 计算机视觉相关论文：Detectron2 团队发表的研究论文
• COCO 数据集：https://cocodataset.org/