Getting Started Guide

Relevant source files

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

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

The mbbank-capcha-ocr project provides a specialized Optical Character Recognition (OCR) system designed to recognize characters in MBBank CAPTCHA images. It encompasses two primary modules: a training pipeline for developing the OCR model using PyTorch and an inference library for making predictions using ONNX Runtime.

This guide outlines the environment setup, dataset preparation, model training process, and subsequent deployment for prediction. The system utilizes a ResNet-18 backbone modified to output multi-label classifications corresponding to CAPTCHA characters.

Sources: README.md:1-8, setup.py:16-19, train_model/core.py:65-72

Project Architecture and Components

The repository is structured into two main directories to separate training logic from inference logic.

Component	Description	Relevant Files
Prediction Core	Contains the OcrModel class for ONNX inference and the pre-trained model file.	mb_capcha_ocr/
Training Module	Contains scripts for dataset loading, model definition (PyTorch), and training loops.	train_model/
Model File	The trained weights exported in ONNX format for efficient cross-platform execution.	mb_capcha_ocr/model.onnx

Sources: README.md:10-13, setup.py:22-24

Data Flow Overview

The following diagram illustrates the lifecycle of the OCR system, from raw image dataset to a functional prediction service.

The workflow transitions from a PyTorch-based training environment to an ONNX-based prediction environment.
Sources: README.md:33-54, train_model/train.py:53-54, mb_capcha_ocr/predict.py:15-18

Environment Setup

The project requires specific dependencies depending on whether you are training a new model or simply using the library for predictions.

Installation for Prediction

To use the pre-trained model for OCR, install the base requirements:

• onnxruntime
• pillow
• numpy

Sources: requirements.txt:1-3, setup.py:26

Installation for Training

Training requires a more extensive stack including deep learning frameworks:

1. Clone the repository and navigate to the training directory.
2. Install dependencies: torch, torchvision, onnx, and matplotlib.

Sources: README.md:15-31, train_model/train_requirements.txt:1-7

Training the Model

The training process involves preparing a labeled dataset and running the training script which outputs both a model file and a loss visualization.

Dataset Preparation

• Images must be placed in the train_model/dataset/ directory.
• Filename format: {label}.(png|jpg|jpeg) (e.g., A1B2C3.png).
• The DatasetLoader automatically handles a 90/10 train-test split and limits the dataset to 62,500 images to prevent overfitting.

Sources: README.md:33-35, train_model/core.py:32-46

Execution and Model Export

Running python train.py initiates the following sequence:

1. Loads images and applies transforms (Resize to 160x50, ToTensor).
2. Initializes a ResNet-18 based model with a modified first layer for grayscale input (1 channel) and a final linear layer outputting chars_len * max_capcha.
3. Optimizes using Adam with MultiLabelSoftMarginLoss.
4. Exports the resulting model to model.onnx.

Sources: train_model/train.py:24-54, train_model/core.py:65-72

Class hierarchy for the training components.
Sources: train_model/core.py:26-78

Using the Model for Prediction

Once model.onnx is available, the mb_capcha_ocr.OcrModel class provides a high-level API for character recognition.

Prediction Logic

The inference engine performs the following steps:

1. Preprocessing: Converts image to grayscale ("L"), resizes to 160x50, and normalizes pixel values to [0, 1].
2. Inference: Runs the processed image through the onnxruntime.InferenceSession.
3. Decoding: Extracts the character with the highest probability (argmax) for each of the 6 expected character positions.

Sources: mb_capcha_ocr/predict.py:20-33

Code Example

from PIL import Image
from mb_capcha_ocr import OcrModel

# Initializes with local model.onnx by default
model = OcrModel()  
img = Image.open("captcha.png")
predicted_text = model.predict(img)
print(predicted_text)

Sources: README.md:46-52, mb_capcha_ocr/predict.py:16-18

Technical Summary

The mbbank-capcha-ocr system effectively bridges the gap between complex neural network training and lightweight deployment. By leveraging PyTorch for training and ONNX for execution, it provides an efficient solution for automating MBBank CAPTCHA recognition.

Sources: README.md:1-5, setup.py:1-29