基于 DETR 的目标检测 — 深度学习课程设计

学生: 黄延卓 | 学号: 2023904690 | 专业: 人工智能

项目概述

本项目基于 PyTorch 和 HuggingFace Transformers 实现了 DETR (Detection Transformer) 目标检测模型，在 COCO2017 数据集上进行训练与评估，并集成 SPD-Conv 和 CBAM 改进模块。

核心特性

• 模型: DETR-R50 + SPD-Conv + CBAM 改进，HuggingFace 预训练权重
• 数据集: COCO2017 子集 (500训练/100验证)
• GUI: PyQt5 图形界面，支持图片拖拽、视频检测、摄像头实时检测、CSV导出
• 实验: 5组消融实验 + SPD-Conv 30轮训练
• 报告: DOCX 课程设计报告 + PPTX 答辩PPT

环境配置

依赖

pip install -r requirements.txt

主要依赖:

• torch>=2.0.0, torchvision>=0.15.0
• transformers>=4.35.0
• pycocotools, opencv-python
• PyQt5 (GUI)
• python-docx, python-pptx (报告生成)

硬件要求

• GPU: NVIDIA GPU with ≥8GB VRAM (推荐 RTX 4060+)
• CPU: 任意现代 CPU
• 磁盘: ≥5GB (COCO子集约800MB)

项目结构

基于DETR的目标检测/
├── dataset/          # 数据加载与预处理
│   ├── coco_dataset.py         # COCO 数据集类
│   ├── coco_subset_builder.py  # 子集构建器
│   ├── transforms.py           # 数据增强
│   └── coco_class_names.py     # 类别名称映射
├── models/           # 模型定义
│   ├── detr_model.py           # DETR 封装 (HuggingFace)
│   ├── detr_improved.py        # 改进版 (SPD-Conv + CBAM)
│   ├── spd_conv.py             # SPD-Conv 模块
│   ├── cbam.py                 # CBAM 注意力模块
│   └── losses.py               # 损失函数
├── engine/           # 训练/评估/推理
│   ├── train.py                # 训练循环
│   ├── evaluate.py             # mAP 评估
│   └── infer.py                # 推理工具
├── gui/              # PyQt5 图形界面
│   ├── main_window.py          # 主窗口
│   ├── detect_worker.py        # 检测线程
│   └── result_panel.py         # 结果面板
├── utils/            # 工具函数
│   ├── draw_boxes.py           # 检测框绘制
│   └── metrics.py              # 评估指标
├── experiments/      # 消融实验
├── configs/          # 配置文件
├── scripts/          # 工具脚本
├── outputs/          # 输出目录
├── main.py           # 主入口
└── requirements.txt  # 依赖清单

快速开始

1. 下载数据

python main.py download --num_train 500 --num_val 100

2. 快速演示 (使用预训练权重)

python main.py demo

3. 命令行推理

python main.py infer --image path/to/image.jpg --threshold 0.5

4. 启动 GUI

python main.py gui

5. 训练模型

python main.py train --epochs 30

6. 运行消融实验

python main.py ablation --quick   # 快速模式
python main.py ablation --full    # 完整模式

实验结果 (最终)

消融实验 (每组 5 轮, COCO 子集 500/100)

实验	模型	mAP@0.5	mAP@0.5:0.95	参数量
Exp A	预训练 DETR-R50	0.684	0.481	41.5M
Exp B	微调 DETR-R50	0.641	0.417	41.5M
Exp C	+ SPD-Conv	0.649	0.423	41.6M
Exp D	+ CBAM	0.628	0.420	42.2M
Exp E	+ SPD+CBAM	0.628	0.420	41.5M

SPD-Conv 延长训练

训练量	mAP@0.5	mAP@0.5:0.95
5 轮	0.649	0.423
30 轮	0.668	0.451

关键发现

• 预训练模型仍是性能最优基线
• SPD-Conv 相比纯微调提升 0.8pp, 方向有效
• 30 轮训练相比 5 轮提升 1.9pp, 收敛趋势明确
• DETR 需要 300+ epoch 才能完全收敛, 有限资源下结果符合预期

参考资料

• DETR 论文: End-to-End Object Detection with Transformers
• SPD-Conv: No More Strided Convolutions or Pooling
• CBAM: Convolutional Block Attention Module
• HuggingFace DETR: facebook/detr-resnet-50