NyaNyav2/engineering-drawing-detection-ocr
GitHub: NyaNyav2/engineering-drawing-detection-ocr
基于 Detectron2 和 PaddleOCR 的工程图纸目标检测与 OCR 系统,能自动定位零件图、注释和表格区域并提取结构化文字内容。
Stars: 0 | Forks: 0
# 🔧 工程图纸目标检测与 OCR 系统
## 📞 联系方式
- **电话:** 0368820203
- **Email:** vansykien03@gmail.com
**Sotatek 计算机视觉评估**
自动检测和 OCR 工程图纸中 3 种类型对象的系统:
- **PartDrawing** — 技术详细图纸区域
- **Note** — 注释 / 标注区域
- **Table** — 技术数据表格区域
🌐 **Web Demo:** https://senators-richmond-montana-hist.trycloudflare.com
📦 **模型权重:** https://huggingface.co/siucapditbu/engineering-drawing-frcnn-x101
## 📐 方法与原理
### 检测
- **模型:** Faster R-CNN X-101-32x8d FPN 3x (Detectron2)
- **主干网络:** ResNeXt-101 32x8d — 在处理工程图纸(小细节、复杂布局)方面比 ResNet-50/101 更强大
- **按类别的置信度阈值:**
- PartDrawing ≥ 95%
- Note ≥ 50%
- Table ≥ 99%
- **数据增强:** RandomBrightness, RandomContrast, 多尺度调整大小 (640–960px)
### OCR
- **Note 区域:** PaddleOCR (角度分类 + 按行文本分组)
- **Table 区域:** PPStructure (单元格检测) → PaddleOCR (每个单元格的文本) → Markdown 表格
- **预处理:** 2倍放大 + 双边滤波器 + 自适应阈值 以提高清晰度
### 实验
| 模型 | mAP@50 (val) | 备注 |
|-------|-------------|-------|
| Faster R-CNN R-50 FPN 3x | ~0.72 | 基线 |
| Faster R-CNN R-101 FPN 3x | ~0.79 | 较好 |
| **Faster R-CNN X-101 FPN 3x** | **~0.85** | 最佳 — 用于生产环境 |
| RetinaNet R-101 FPN 3x | ~0.74 | 在小目标上表现较差 |
### 改进方向
- 尝试 **DINO / DINOv2** (ViT 主干网络) 作为 feature extractor 替代 ResNeXt — ViT 在工程图纸的复杂布局上能比 CNN 学习到更好的 global context
- **ViTDet** (plain ViT 主干网络 + Detectron2) — 移除 FPN,使用 window attention 处理高分辨率图像,而不会消耗过多 VRAM
- 微调 **Grounding DINO** (ViT + text-conditioned 检测) 以检测根据自定义类别名称的 open-vocabulary 目标 — 适合扩展到新类型的图纸
- 尝试 **RT-DETR-L/X** (混合 ViT 编码器) — 推理速度比 Faster R-CNN X-101 更快,同时保持高精度
- 使用 WBF 集成 **X-101 + ViTDet + Grounding DINO**,以提高对微小且被遮挡 Note 区域的召回率
- 集成 Meta 的 SAM 以自动为 object 区域提取详细的 segmentation mask,这在需要精确分析 PartDrawing 中复杂机械零件的边缘轮廓(而不是仅使用 bounding box 框选)时特别有用。
## 🛠️ 安装
### 要求
- Python 3.8+
- CUDA 11.8+ (推荐) 或 CPU
- Conda (推荐)
### 环境设置
```
# 1. 创建 conda 环境
conda create -n cv-engineering python=3.8 -y
conda activate cv-engineering
# 2. 安装 PyTorch (CUDA 12.1)
pip install torch==2.0.1+cu121 torchvision==0.15.2+cu121 \
--index-url https://download.pytorch.org/whl/cu121
# 3. 从源码安装 Detectron2
pip install git+https://github.com/facebookresearch/detectron2.git
# 4. 安装 PaddleOCR
pip install paddlepaddle-gpu==2.5.2 -i https://pypi.tuna.tsinghua.edu.cn/simple
pip install paddleocr>=2.7.0
# 5. 安装剩余的依赖
pip install opencv-python numpy Pillow gradio huggingface_hub
```
### 下载模型权重
```
# 选项 1: HuggingFace Hub (推荐)
python -c "
from huggingface_hub import hf_hub_download
hf_hub_download(
repo_id='siucapditbu/engineering-drawing-frcnn-x101',
filename='model_final.pth',
local_dir='detectron2_training/output_v4_ensemble/frcnn_x101/'
)
"
# 选项 2: 手动下载
# https://huggingface.co/siucapditbu/engineering-drawing-frcnn-x101/blob/main/model_final.pth
# 放置于: detectron2_training/output_v4_ensemble/frcnn_x101/model_final.pth
```
## 🚀 运行推理
### 完整 Pipeline (检测 + OCR)
```
cd detectron2_training
python pipeline_final.py
```
输出将保存至:
```
detectron2_training/output_v4_ensemble/frcnn_x101/ocr_pipeline_final/
├── crops/ ← ảnh crop từng object
├── json/ ← JSON output từng ảnh
├── markdown/ ← OCR text dạng .md
├── visualizations/ ← ảnh detection visualized
├── preprocessed/ ← ảnh preprocess OCR
├── table_structure_vis/ ← cell structure visualization
└── summary.json ← tổng hợp toàn bộ
```
### 仅测试检测 (无 OCR)
```
cd detectron2_training
python test_frcnn_x101_custom.py
```
### 本地 Web Demo
```
cd web_demo
python app.py
# 打开 http://localhost:7860
```
## 🏋️ 训练
```
cd detectron2_training/v4_ensemble
# 训练 frcnn_x101
python train_frcnn_x101.py
# 或选择特定的模型
python train_ensemble.py --model frcnn_x101
python train_ensemble.py --model frcnn_r101
python train_ensemble.py --model retinanet_r101
# 恢复训练
python train_ensemble.py --model frcnn_x101 --resume
```
**frcnn_x101 训练配置:**
- 最大迭代次数: 6000
- 基础 LR: 0.0002 (warmup 300 次迭代)
- LR 衰减: 0.1 于迭代 4000, 5200 次时
- Batch: 2 张图像/迭代
- 多尺度: 短边 640–960px,最大 1600px
- 数据增强: RandomBrightness(0.9–1.1), RandomContrast(0.9–1.1)
- 评估周期: 500 次迭代
## 📁 项目结构
```
├── detectron2_training/
│ ├── v4_ensemble/
│ │ ├── train_ensemble.py ← main training script
│ │ ├── train_frcnn_x101.py ← shortcut script
│ │ ├── train_frcnn_r101.py
│ │ ├── train_retinanet_r101.py
│ │ └── ensemble_infer.py ← ensemble inference (WBF)
│ ├── pipeline_final.py ← ★ main pipeline (detect + OCR)
│ ├── test_frcnn_x101_custom.py ← detection-only test
│ ├── train.json ← COCO format train annotations
│ └── val.json ← COCO format val annotations
├── web_demo/
│ ├── app.py ← Gradio local demo
│ ├── app_final_hf.py ← HuggingFace Spaces version
│ └── requirements.txt
└── README.md
```
## 📊 JSON 输出格式
```
{
"image": "drawing_001.jpg",
"objects": [
{
"id": 1,
"class": "Table",
"confidence": 0.9987,
"bbox": { "x1": 120, "y1": 340, "x2": 680, "y2": 520 },
"ocr_content": "| Col1 | Col2 |\n|------|------|\n| val1 | val2 |"
},
{
"id": 2,
"class": "Note",
"confidence": 0.9231,
"bbox": { "x1": 50, "y1": 600, "x2": 400, "y2": 650 },
"ocr_content": "GENERAL TOLERANCES AS PER ISO 2768"
},
{
"id": 3,
"class": "PartDrawing",
"confidence": 0.9999,
"bbox": { "x1": 10, "y1": 10, "x2": 500, "y2": 450 },
"ocr_content": ""
}
]
}
```
标签:Apex, Detectron2, Faster R-CNN, OCR, PaddleOCR, PPStructure, ResNet, ResNeXt, Vectored Exception Handling, Web演示, 人工智能, 光学字符识别, 凭据扫描, 图像分割, 图像处理, 图像预处理, 工程图纸处理, 技术图纸分析, 数据增强, 文档分析, 机器学习, 深度学习, 特征提取, 用户模式Hook绕过, 目标检测, 目标检测模型评估, 表格识别, 计算机视觉, 逆向工具