DiffusionDet：用于目标检测的扩散模型

DiffusionDet 是首个将扩散模型应用于目标检测的工作。

DiffusionDet：用于目标检测的扩散模型
陈守法、孙培泽、宋一冰、罗平
arXiv 2211.09788

更新

• (2022年11月) 代码已发布。

模型

方法	边框 AP（1 步）	边框 AP（4 步）	下载
COCO-Res50	45.5	46.1	模型
COCO-Res101	46.6	46.9	模型
COCO-SwinBase	52.3	52.7	模型
LVIS-Res50	30.4	31.8	模型
LVIS-Res101	31.9	32.9	模型
LVIS-SwinBase	40.6	41.9	模型

快速入门

安装说明和使用方法请参阅 DiffusionDet 快速入门。

许可证

本项目采用 CC-BY-NC 4.0 许可协议。详情请见 LICENSE。

引用 DiffusionDet

如果您在研究中使用了 DiffusionDet，或希望引用此处发布的基准结果，请使用以下 BibTeX 条目。

@article{chen2022diffusiondet,
      title={DiffusionDet: Diffusion Model for Object Detection},
      author={Chen, Shoufa and Sun, Peize and Song, Yibing and Luo, Ping},
      journal={arXiv preprint arXiv:2211.09788},
      year={2022}
}

DiffusionDet 快速上手指南

DiffusionDet 是首个将扩散模型（Diffusion Model）应用于目标检测任务的工作。本指南帮助开发者快速搭建环境并运行预训练模型。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 18.04+)
• Python: 3.8 或更高版本
• GPU: 支持 CUDA 的 NVIDIA 显卡（建议显存 >= 16GB 以运行 SwinBase 等大模型）
• CUDA: 与您的 PyTorch 版本匹配的 CUDA toolkit

前置依赖: 本项目基于 Detectron2 框架开发，安装前需先正确安装 PyTorch 和 Detectron2。

安装步骤

1. 创建虚拟环境并安装 PyTorch

建议使用 Conda 管理环境。以下命令使用国内镜像源加速下载：

conda create -n diffdet python=3.8 -y
conda activate diffdet

# 安装 PyTorch (根据实际 CUDA 版本选择，此处以 CUDA 11.3 为例，使用清华源)
pip install torch==1.10.1+cu113 torchvision==0.11.2+cu113 torchaudio==0.10.1 -f https://pypi.tuna.tsinghua.edu.cn/simple

2. 安装 Detectron2

DiffusionDet 依赖特定版本的 Detectron2。

# 使用国内镜像安装 detectron2 依赖
pip install 'git+https://github.com/facebookresearch/fvcore.git' 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

# 编译安装 detectron2 (使用清华源加速)
pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/cu113/torch1.10/index.html --extra-index-url https://pypi.tuna.tsinghua.edu.cn/simple

3. 克隆代码并安装 DiffusionDet

git clone https://github.com/ShoufaChen/DiffusionDet.git
cd DiffusionDet

# 安装项目依赖
pip install -r requirements.txt --extra-index-url https://pypi.tuna.tsinghua.edu.cn/simple

# 以可编辑模式安装当前包
pip install -e .

基本使用

以下示例展示如何使用预训练的 COCO-Res50 模型进行单步推理（1 step）。

1. 下载预训练模型

从 GitHub Release 下载模型权重，或使用 wget：

mkdir models
wget -P models https://github.com/ShoufaChen/DiffusionDet/releases/download/v0.1/diffdet_coco_res50.pth

2. 运行推理

使用提供的配置文件和下载的权重运行评估或推理。以下命令在 COCO val2017 数据集上运行评估（需提前准备好 COCO 数据集）：

python train_net.py --eval-only \
  --config-file configs/diffdet.coco.res50.yaml \
  --num-gpus 1 \
  MODEL.WEIGHTS models/diffdet_coco_res50.pth \
  MODEL.DIFFUSION.NUM_SAMPLE_STEP 1

参数说明：

• --config-file: 指定模型配置文件。
• MODEL.WEIGHTS: 指定下载的预训练权重路径。
• MODEL.DIFFUSION.NUM_SAMPLE_STEP: 设置扩散采样步数（1 或 4，步数越多精度略高但速度越慢）。

若仅需对单张图片进行测试，可参考 GETTING_STARTED.md 中的 Demo 脚本编写简单的 Python 调用代码，加载配置与权重后即可输入图像获取检测框。