Dataset Directory & Structure

Relevant source files

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

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

Introduction

The dataset system in the MBBank CAPCHA OCR project is designed to handle raw CAPCHA images for training and evaluating a deep learning model based on ResNet18. The system relies on a specific file-naming convention where the image label is embedded directly in the filename, allowing for a streamlined ingestion process without the need for separate metadata files.

The core logic for interacting with this data is contained within the DatasetLoader class, which manages file discovery, label extraction via regex, and the partitioning of data into training and testing subsets. This structure supports both the training phase using PyTorch and the inference phase using ONNX Runtime.

Sources: README.md:32-34, train_model/core.py:34-40

Directory Organization

The project expects a specific directory structure to facilitate the training pipeline. All images must be contained within a single directory, typically named dataset/ relative to the training script.

Directory/File	Description
dataset/	Root directory for all training and testing images.
{label}.png	Image file where the filename represents the ground truth text.
train_model/	Contains the core.py and train.py scripts that process the dataset.

Sources: README.md:32-34, train_model/train.py:23

Naming Convention

The system uses the filename as the primary source of truth for labels. The filename (excluding the extension) must exactly match the characters present in the CAPCHA image. Supported extensions include .png, .jpg, and .jpeg.

Sources: README.md:32-34, train_model/core.py:61-63

Data Flow and Processing

The following diagram illustrates how raw image files are transformed into tensors suitable for model training.

The workflow ensures that data is partitioned and pre-processed consistently.
Sources: train_model/core.py:34-75, train_model/train.py:25-30

Dataset Loader Logic

The DatasetLoader class implements the torch.utils.data.Dataset interface and includes specific logic for data management and validation.

Data Partitioning and Limits

The loader includes hardcoded constraints to prevent model overfitting and manage resource usage:

• Capacity Limit: The dataset is capped at a maximum of 62,500 images.
• Train/Test Split: Data is split into a 9:1 ratio (90% training, 10% testing). If the dataset is too small, a minimum of 1 image is assigned to the test set.
• Randomization: The file list is shuffled twice: once before the split and once after the split for the specific subset.

Sources: train_model/core.py:43-54

Image Transformations

Images are standardized before being fed into the model. The training script defines a specific transformation pipeline:

1. Grayscale Conversion: Images are converted to "L" (8-bit pixels, black and white) to reduce input complexity.
2. Resizing: Images are scaled to a fixed resolution of 160x50 pixels.
3. Normalization: During training, transforms.ToTensor() is used. During prediction, pixel values are manually divided by 255.0.

Sources: train_model/train.py:25-28, mb_capcha_ocr/predict.py:22-24, train_model/core.py:64

Label Encoding

Labels are converted from strings into one-hot encoded vectors based on a predefined character set.

Parameter	Value/Description
chars	Digits (0-9) + ASCII letters (a-z, A-Z), sorted alphabetically.
max_capcha	Fixed at 6 characters.
Encoding	Single one-hot vector per character.

def encode(a):
    onehot = [0] * chars_len
    idx = chars.index(a)
    onehot[idx] += 1
    return onehot

Sources: train_model/core.py:14-25, mb_capcha_ocr/predict.py:13-16

Conclusion

The dataset structure is intentionally minimalist, leveraging filesystem metadata (filenames) to drive the labeling process. Through the DatasetLoader, the project enforces strict image dimensions and character encoding standards, ensuring that the ResNet18-based OCR model receives consistent inputs during both the training and inference phases.

Sources: train_model/core.py:34-75, README.md:32-34