ONNX Runtime Integration
Relevant source files
The following files were used as context for generating this wiki page:
The ONNX Runtime integration in this project provides a high-performance cross-platform inference engine for the MBBank CAPTCHA OCR model. By utilizing the Open Neural Network Exchange (ONNX) format, the system decouples model training—which is performed using PyTorch—from model deployment, allowing for efficient execution without requiring a full PyTorch installation in production environments.
Sources: README.md:1-40, mb_capcha_ocr/predict.py:1-35
Architecture and Workflow
The integration follows a standard lifecycle: a PyTorch-based model is trained, exported to the .onnx format, and then loaded into an onnxruntime.InferenceSession for character prediction.
Model Export Flow
The training script utilizes the torch.onnx.export utility to convert the trained resnet18 based model into a static ONNX graph. This graph includes specific input and output names used later during inference.
The export process defines a fixed input shape of (1, 1, 50, 160) representing a single-channel (grayscale) CAPTCHA image.
Sources: train_model/train.py:56-58, train_model/core.py:75-84
Inference Architecture
The OcrModel class in the prediction module serves as the primary interface for the ONNX Runtime. It handles session initialization, image preprocessing to match the ONNX graph's expected input, and post-processing of the raw tensor output into human-readable text.
Sources: mb_capcha_ocr/predict.py:11-35
Components and Data Logic
Image Preprocessing
Before the ONNX session can process an image, it must be transformed to match the training data format. The predict.py script implements several manual transformations that mirror the torchvision.transforms used during training.
| Step | Operation | Logic |
|---|---|---|
| 1 | Color Conversion | img.convert("L") (Grayscale) |
| 2 | Resizing | img.resize((160, 50)) |
| 3 | Normalization | np.array / 255.0 |
| 4 | Dimensionality | np.expand_dims (Adds batch and channel dims) |
Sources: mb_capcha_ocr/predict.py:22-26, train_model/train.py:26-29
Post-Processing and Decoding
The ONNX model returns a raw tensor of shape (batch, max_capcha, chars_len). The integration uses numpy to decode this into characters.
- Argmax: Identification of the index with the highest probability for each of the 6 character positions.
- Mapping: Index translation using a sorted list of alphanumeric characters (
0-9a-zA-Z).
pred_labels = np.argmax(pred, axis=2)
pred_text = [''.join([self.chars[c] for c in pred_label]) for pred_label in pred_labels]
Sources: mb_capcha_ocr/predict.py:13-31, train_model/core.py:12-14
Configuration and Implementation Details
The integration is designed to be portable, with the model file bundled directly in the package via setup.py.
Class: OcrModel (Prediction)
Located in mb_capcha_ocr/predict.py, this class manages the ONNX session.
__init__(self, model_path: Optional[str]): Initializes theInferenceSession. If no path is provided, it defaults to looking formodel.onnxwithin the package directory.predict(self, img_data: Image): Executes the full inference pipeline from PIL image to string.
Sources: mb_capcha_ocr/predict.py:11-35, setup.py:21-23
Dependency Management
The runtime environment specifically requires onnxruntime for execution, whereas the training environment requires onnx and onnxscript for model generation.
| Environment | Key Dependencies |
|---|---|
| Inference | onnxruntime, numpy, Pillow |
| Training | torch, onnx, torchvision, matplotlib |
Sources: requirements.txt:1-5, train_model/train_requirements.txt:1-7
Summary
The ONNX Runtime integration allows the project to deliver a lightweight prediction library (mb_capcha_ocr) that is decoupled from the heavy PyTorch training dependencies. By standardizing on an 160x50 grayscale input and an argmax-based decoding strategy, the system ensures consistent CAPTCHA recognition performance between training and production environments.
Sources: README.md:42-55, mb_capcha_ocr/predict.py:1-35