High-Level Architecture

Relevant source files

The following files were used as context for generating this wiki page:

• README.md
• setup.py
• train_model/train.py
• mb_capcha_ocr/predict.py
• train_model/core.py
• mb_capcha_ocr/init.py

The MBBank CAPTCHA OCR project provides a complete pipeline for training and deploying an Optical Character Recognition (OCR) model specifically tuned for MBBank CAPTCHA images. The system is divided into two primary modules: a training module for developing the model using PyTorch and an inference module for predicting text from images using the ONNX runtime.

The architecture leverages a modified ResNet-18 backbone to extract features from grayscale images, which are then mapped to a character set comprising digits and letters. The project facilitates the transition from research to production by exporting trained PyTorch models to the ONNX format for efficient cross-platform execution.

Sources: README.md:1-12, train_model/core.py:65-71

System Components

The project is structured into two main packages that handle the lifecycle of the OCR model.

1. Training Module (train_model/)

This module handles data ingestion, model definition, and the training loop. It uses PyTorch for deep learning operations.

• core.py: Defines the neural network architecture and data loading utilities.
• train.py: The entry point for training, performing optimization, and exporting the final model to ONNX format.

2. Prediction Module (mb_capcha_ocr/)

This is the consumer-facing library designed for lightweight inference.

• predict.py: Contains the OcrModel class which utilizes onnxruntime to perform predictions without requiring the full PyTorch stack.
• __init__.py: Exposes the OcrModel class for easier integration by external applications.

Sources: README.md:7-10, mb_capcha_ocr/predict.py:11-14, train_model/train.py:16-19

Data Architecture and Flow

The system processes CAPTCHA images through a specific pipeline from raw files to predicted text.

Character Encoding

The model supports a character set consisting of digits and both uppercase and lowercase English letters. Characters are converted to one-hot encoded vectors during training and mapped back to strings during inference.

Parameter	Value	Description
chars	0-9, a-z, A-Z (Sorted)	The available character set for recognition.
max_capcha	6	The fixed length of the CAPTCHA string.
chars_len	62	The total number of unique characters in the set.

Sources: train_model/core.py:12-25

Training Data Flow

The training process involves loading images from a local directory where labels are derived from filenames.

Inference Sequence

The inference module uses ONNX for high-performance prediction, bypassing the PyTorch training infrastructure.

Model Architecture

The OCR model is based on the ResNet-18 architecture, modified to accept single-channel (grayscale) input and produce multi-label output for sequence recognition.

Network Structure

The OcrModel class modifies a standard ResNet-18 in two key areas:

1. Input Layer: The first convolutional layer (conv1) is replaced to accept 1 input channel instead of 3.
2. Output Layer: The fully connected layer (fc) is replaced with a linear layer that outputs a vector of size chars_len * max_capcha (62 * 6 = 372).

Training Hyperparameters

The training script utilizes the following configuration:

Component	Setting	File Reference
Backbone	ResNet-18	core.py
Loss Function	MultiLabelSoftMarginLoss	train.py
Optimizer	Adam (lr=0.001)	train.py
Batch Size	64	train.py
Epochs	100	train.py

Sources: train_model/core.py:65-71, train_model/train.py:40-45

Component Relationships

The project maintains a clear separation between the library used for production and the scripts used for research.

Summary

The high-level architecture of mbbank-capcha-ocr centers on a robust ResNet-18 feature extractor capable of processing grayscale CAPTCHA images. By separating the training logic (PyTorch) from the prediction logic (ONNX), the project ensures that the final library remains lightweight and easy to integrate, as seen in the setup.py which only requires onnxruntime, pillow, and numpy for the core package installation.

Sources: setup.py:22-26, README.md:12-18