Loading Custom Models
Relevant source files
The following files were used as context for generating this wiki page:
The custom model loading system in this project allows users to utilize specifically trained Optical Character Recognition (OCR) models to bypass CAPTCHA protections. While the package includes a default model, the architecture is designed to support the execution of any compatible ONNX-formatted model by providing a path to the file during initialization.
The system relies on the onnxruntime engine for high-performance inference and the OcrModel wrapper class to handle pre-processing and post-processing. Users can transition from training a model in PyTorch to deploying it as an ONNX file for use in production environments.
Sources: README.md:1-5, mb_capcha_ocr/predict.py:12-16
Model Initialization and Configuration
The OcrModel class serves as the primary interface for loading models. By default, it searches for a file named model.onnx within the package directory. However, it accepts an optional model_path argument to load a custom-trained model from any filesystem location.
The OcrModel Interface
The initialization process involves sorting the character set (0-9, a-z, A-Z) to ensure label mapping consistency and creating an onnxruntime.InferenceSession.
The diagram above illustrates the conditional logic used to determine which model file to load into the inference engine.
Sources: mb_capcha_ocr/predict.py:12-17, setup.py:22-24
Inference and Data Flow
When a custom model is loaded, the predict method handles the transformation of raw images into the specific format required by the model's input layer. This ensures that even custom models benefit from the standardized pre-processing pipeline.
Input Pre-processing Requirements
Custom models must be trained to accept inputs with the following characteristics:
- Color Space: Greyscale ("L" mode).
- Dimensions: 160x50 pixels (Width x Height).
- Normalization: Pixel values scaled to a range of
0.0to1.0(divided by 255.0). - Tensor Shape: 4D array with shape
(1, 1, 50, 160).
Post-processing and Decoding
The output from the ONNX session is processed using a "Greedy Search" approach where argmax is applied to the character dimension to find the most likely character for each of the CAPTCHA positions.
| Step | Operation | Description |
|---|---|---|
| 1 | Image Conversion | Convert PIL image to Greyscale. |
| 2 | Resizing | Resize image to (160, 50). |
| 3 | Normalization | Convert to float32 and divide by 255. |
| 4 | Reshaping | Expand dimensions to match (Batch, Channel, Height, Width). |
| 5 | Execution | Run session.run with the processed image. |
| 6 | Decoding | Map indices back to string.digits + string.ascii_letters. |
Sources: mb_capcha_ocr/predict.py:19-32, train_model/core.py:72-78
Custom Model Training and Export
To create a custom model compatible with the loader, the project provides a training suite in the train_model/ directory. The architecture is based on a modified ResNet18.
Training Logic
Custom models are trained using MultiLabelSoftMarginLoss and the Adam optimizer. The model is trained to recognize a maximum of 6 characters per CAPTCHA.
The sequence diagram shows the progression from local datasets to a portable ONNX file.
Sources: train_model/train.py:35-55, train_model/core.py:72-83
Compatibility Specifications
Custom models must adhere to the following architecture constraints to work with the prediction script:
| Component | Specification | Source |
|---|---|---|
| Input Name | 'input' | train_model/train.py:53 |
| Output Name | 'output' | train_model/train.py:53 |
| Max Length | 6 characters | train_model/core.py:19 |
| Charset Size | 62 (0-9, a-z, A-Z) | train_model/core.py:14-18 |
| Output Shape | (-1, 6, 62) | train_model/core.py:82 |
Summary
The custom model loading feature provides flexibility for users to improve OCR accuracy by training on specific datasets while maintaining a simple API. By wrapping the complex onnxruntime logic and pre-processing steps into the OcrModel class, the system allows for seamless swapping of models via a simple file path parameter.