Skip to content

Vasbi-Crast/Cserobot

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

4 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Cserobot — Automated Engineering Drawing Attribution

Desktop application for automatic extraction of metadata (drawing number, part title, sheet number) from scanned engineering drawings using YOLOv9 and TrOCR neural networks. Built as a diploma project for JSC "Corporation Strategic Control Points" to automate their paper archive digitization pipeline.

Problem context: The enterprise had thousands of paper engineering drawings stored in physical archives. Converting them to an electronic archive required three steps: scanning, attribution (assigning metadata to each file), and database upload. The attribution step was the most labor-intensive — engineers manually opened each PDF, read the title block, and typed the metadata. This tool automates that step entirely.


Features

  • Batch PDF processing — scans a folder (and subfolders) for drawing PDFs
  • Title block detection — YOLOv9 locates the drawing number, part name, and sheet number fields
  • Text recognition — a fine-tuned TrOCR model reads Cyrillic text from each detected field
  • Automatic renaming — files are renamed to <Drawing Number> <Part Name> <Sheet>.pdf
  • Error reporting — generates a warnings log for files that couldn't be processed (missing fields, unreadable text, permission errors)
  • Modern GUI — CustomTkinter interface with progress bar, real-time status, and stop button

Perfect for engineering departments, archives, and any organization dealing with large volumes of technical documentation.


Model Training

The application relies on two custom-trained neural networks, each optimized for its specific task.

YOLOv9 (Title Block Detection)

The object detection model was trained on real scanned drawings provided by the enterprise. The dataset was manually annotated and managed using Roboflow. This approach ensured high detection accuracy for the specific layout, fonts, and physical artifacts (stains, folds, low contrast) of the company's historical documents.

TrOCR (Text Recognition)

The OCR model was fine-tuned on a synthetic dataset of 5,000 images generated with TRDG and PIL, then evaluated on a synthetic test set.

Metric Value
CER (Character Error Rate) 0.90%
WER (Word Error Rate) 3.86%
Exact Match Accuracy 93.40%

Evolution of the OCR model

Model CER Improvement
Tesseract (baseline) ~15.00%
microsoft/trocr-base-printed 5.09% ×3
raxtemur/trocr-base-ru (v1) 12.50% ×1.2
raxtemur/trocr-base-ru (v2, fine-tuned) 1.87% ×8
raxtemur/trocr-base-ru (v3, final) 0.90% ×17

Key insight: Using a community model pre-trained on Cyrillic (raxtemur/trocr-base-ru) combined with a hybrid dataset generator (PIL for multi-line text + TRDG for single-line) achieved CER comparable to commercial OCR services. Manual testing on real scanned drawings from the enterprise confirmed satisfactory performance in production conditions.


Technical Stack

Layer Technology
Object Detection YOLOv9 (Ultralytics) — title block field detection
OCR TrOCR (HuggingFace Transformers) — text recognition
Base OCR model raxtemur/trocr-base-ru (fine-tuned)
Dataset Management Roboflow (for YOLO dataset annotation and versioning)
GUI CustomTkinter + Tkinter
Image Processing OpenCV, Pillow, pdf2image
Dataset Generation TRDG, PIL
Training PyTorch, HuggingFace Seq2SeqTrainer
Metrics jiwer (CER/WER)
Language Python 3.11

Requirements

Software

  • Python: 3.11+
  • Poppler (for PDF-to-image conversion):
  • Git
  • Python libraries (installed via requirements.txt):
    • Deep Learning: torch>=2.12.0, torchvision>=0.17.0, transformers>=4.36.0
    • Computer Vision: ultralytics>=8.1.0, opencv-python>=4.8.0, Pillow>=9.5.0, numpy>=1.24.0
    • PDF Processing: pdf2image>=1.16.0
    • GUI: customtkinter>=5.2.0
    • Dataset Generation: trdg>=1.8.0
    • Metrics: jiwer>=3.0.0
    • Utilities: tqdm>=4.66.0, python-dotenv>=1.0.0

Hardware

Component Minimum Recommended
CPU Intel Core i3 / AMD Ryzen 3 Intel Core i5+ / AMD Ryzen 5+
RAM 4 GB 8+ GB
GPU None (CPU inference) NVIDIA GPU with 6+ GB VRAM (CUDA 12)
Storage 15 GB free 25+ GB (for models + dataset)

Installation

1. Clone the repository

git clone https://github.com/Vasbi-Crast/Cserobot.git
cd Cserobot

2. Create a virtual environment

python -m venv venv
source venv/bin/activate        # Linux/macOS
venv\Scripts\activate           # Windows

3. Install dependencies

# PyTorch with CUDA support (for GPU inference)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

# Other dependencies
pip install -r requirements.txt

4. Prepare model weights

Place the trained model weights in the corresponding directories:

# YOLOv9 weights for title block detection
mkdir -p assets/models/yolo
# Place YOLO_weight.pt in assets/models/yolo/

# TrOCR fine-tuned model
mkdir -p assets/models/trocr
# Place the trocr-drawing-ocr-ru-v3/ folder in assets/models/trocr/

Note: To use the application with your own drawings, train the models on your dataset. The YOLOv9 model should be annotated via Roboflow, and the TrOCR model can be fine-tuned using tools/ml_pipeline/train_trocr.py.

5. Configure HuggingFace token (optional)

If you plan to download models from the Hugging Face Hub, create a .env file in the ml_pipeline directory:

HF_TOKEN=hf_XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX

Get your token at huggingface.co/settings/tokens.


Usage

Launch the application

python src/Main.py

Workflow

  1. Select a folder containing PDF drawings (or leave the default path and click "Browse…")
  2. Click "Start" — the app processes all PDFs in the folder and subfolders
  3. Watch the progress — real-time status shows current file and percentage
  4. View results — files are renamed, and a warnings log is generated for any issues

Example

Before processing:

scan_001.pdf
drawing_2024.pdf
image.pdf

After processing:

САПР.123456.012 Вал.pdf
КГГ1.304514СБ Корпус.pdf
ИТФТМ.4033.2 Втулка.pdf

Project Structure

Cserobot/
── src/                              # Application source code
│   ├── Main.py                       # Entry point
│   ├── ML_part.py                    # ML pipeline (YOLO + TrOCR inference)
│   ├── OS_part.py                    # OS operations (PDF → JPEG → rename)
│   └── Tkinter_Graphics.py           # GUI (CustomTkinter)
│
├── tools/                            # Development utilities
│   ├── gen_ds.py                     # Synthetic dataset generator
│   └── ml_pipeline/
│       ├── train_trocr.py            # TrOCR fine-tuning script
│       └── evaluate.py               # Model evaluation (CER/WER)
│
├── assets/                           # Application resources
│   ├── icons/
│   │   └── icons.ico                 # App icon
│   ├── texts/
│   │   └── help_text.txt             # In-app help text
│   └── models/                       # Model weights
│       ├── yolo/
│       │   └── YOLO_weight.pt
│       └── trocr/
│           └── trocr-drawing-ocr-ru-v3/
│
├── data/                             # Data
│   ├── ocr_dataset/                  # Generated synthetic dataset
│   └── fonts/                        # TTF fonts for dataset generation
│
├── res/                              # Runtime logs and warnings
├── requirements.txt                  # Python dependencies
├── .gitignore
└── README.md

Development

Generate a synthetic dataset (for TrOCR)

python tools/gen_ds.py --num-samples 5000 --pool-size 5000

Fine-tune the TrOCR model

python tools/ml_pipeline/train_trocr.py \
    --model-name raxtemur/trocr-base-ru \
    --epochs 10 \
    --batch-size 5 \
    --learning-rate 1e-5 \
    --max-length 128

Evaluate the model

python tools/ml_pipeline/evaluate.py \
    --model-path assets/models/trocr/trocr-drawing-ocr-ru-v3 \
    --test-dir data/ocr_dataset/test

License

This project was developed as a diploma thesis at the Moscow College of Space Instrumentation (MTKP) in collaboration with JSC "Corporation Strategic Control Points".

About

Automated engineering drawing attribution using YOLOv9 and TrOCR. Desktop app that extracts metadata (drawing number, title, sheet) from scanned PDFs and renames files accordingly.

Topics

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages