Extending to New CAPTCHA Types
Relevant source files
The following files were used as context for generating this wiki page:
Extending the MBBank CAPTCHA OCR system to support new CAPTCHA types involves modifying the core model architecture, updating the character set, and providing a specialized dataset. The system is built on a ResNet-18 backbone and utilizes a multi-label classification approach to predict fixed-length character sequences from images.
The scope of extension covers adjustments to image preprocessing dimensions, character encoding logic, and the final fully connected layer of the neural network to accommodate different CAPTCHA lengths or symbol sets.
Data Preparation and Loading
To support a new CAPTCHA type, the primary requirement is a labeled dataset where filenames represent the ground truth labels. The DatasetLoader class expects images to be named in the format {label}.(png|jpg|jpeg).
Dataset Requirements
- Format: Images must be placed in a
dataset/directory. - Labeling: The filename (excluding extension) is used as the target label.
- Preprocessing: By default, the system converts images to grayscale ("L") and applies a resize transform.
Sources: README.md:27-29, train_model/core.py:48-67
The diagram above illustrates the data ingestion pipeline required for new CAPTCHA types. Sources: train_model/train.py:22-26, train_model/core.py:65-67
Modifying Model Constants
Extending the system requires updating specific global variables and class attributes located in train_model/core.py and mb_capcha_ocr/predict.py. These constants define the dimensions of the output layer and the mapping for character encoding/decoding.
Key Configuration Parameters
| Parameter | File | Purpose |
|---|---|---|
chars | core.py, predict.py | List of supported characters (alphanumeric by default). |
chars_len | core.py | Length of the character list; determines output vector size. |
max_capcha | core.py | Fixed length of the CAPTCHA string (default is 6). |
transform | train.py | Resize dimensions (default 160x50) for input tensors. |
Sources: train_model/core.py:12-15, mb_capcha_ocr/predict.py:12, train_model/train.py:22-25
Architectural Adjustments
The OcrModel class leverages a ResNet-18 architecture modified for single-channel (grayscale) input and a custom linear output head.
Input Layer
The first convolutional layer (conv1) is modified to accept a single input channel instead of three. If the new CAPTCHA type requires different input dimensions than the default 160x50, the kernel_size and training transforms must be updated.
Sources: train_model/core.py:73-75
Output Layer
The fully connected layer (fc) is recalculated based on the product of chars_len (the size of the character set) and max_capcha (the number of characters in each CAPTCHA). The forward pass then reshapes the flat output into a 3D tensor of shape (-1, max_capcha, chars_len).
Sources: train_model/core.py:76-80
The architecture modification ensures the model can output a multi-character prediction sequence. Sources: train_model/core.py:71-81
Training and Exporting
Once the code is adjusted for the new CAPTCHA type, the training process involves running the train.py script. The script uses MultiLabelSoftMarginLoss to optimize the model against one-hot encoded labels.
Training Workflow
- Initialization: Load the modified
OcrModeland move it to the target device (CUDA or CPU). - Optimization: Utilize the Adam optimizer with a default learning rate of 0.001.
- Export: The trained model is exported to ONNX format using
torch.onnx.exportfor use in the prediction module.
Sources: train_model/train.py:35-51, train_model/train.py:53-54
Inference for New Types
The mb_capcha_ocr/predict.py module must be synchronized with any changes made during training. Specifically, the chars list in the OcrModel class must match the list used during training to ensure correct decoding of the argmax results from the ONNX runtime session.
# From mb_capcha_ocr/predict.py:24-30
input_name = self.session.get_inputs()[0].name
pred = self.session.run(None, {input_name: img})[0]
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:24-30
Summary
Extending the MBBank CAPTCHA OCR to new formats requires a consistent update across the dataset labels, character mapping constants, and model output dimensions. By adjusting the chars list and max_capcha length in both the training core and the prediction logic, the system can be retrained for any fixed-length CAPTCHA variant using the provided ResNet-based transfer learning approach.