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绕过, 目标检测, 目标检测模型评估, 表格识别, 计算机视觉, 逆向工具