pooya-mohammadi/crnn-pytorch

pytorch implementation of crnn. A sample training of license plate is provided.

CRNN-Pytorch

Pytorch implementation of the CRNN model. In this repository I explain how to train a license plate-recognition model with pytorch-lightning.

CRNN Architecture:

Installation:

pip install -r requirements.txt

Dataset

Before training the model, it's a good practice to calculate the mean and std of the input dataset and therefore normalize the model using proper values instead of merely normalizing with magical 0.5. Before diving into the code make sure that the dataset has the following characters:

├── data-dir
│   ├── train
│   │  ├──<index_01>_<text>.jpg
│   │  ├──<index_02>_<text>.jpg
│   │  ├──...
│   ├── val
│   │  ├──<index_01>_<text>.jpg
│   │  ├──<index_02>_<text>.jpg
...

NOTE: Only .jpg, .png, and .jpeg extensions are supported!

Then checkout the alphabets.py module. It contains the alphabets characters that are required for training. If the existing alphabets do not meet your requirements create a new dictionary containing your required alphabets and then add it to the ALPHABETS variable with a specific name. You can get your character set using the following command:

python get_character_sets.py --data_directory <path-to-dataset>

The output will be like the following:

[INFO] characters: +ابتثجدزسشصطعقلمنهوپگی۰۱۲۳۴۵۶۷۸۹

Run the following command to get the mean and std of your input dataset:

python dataset.py --dataset_dir <your-train-dir> --alphabets +ابتثجدزسشصطعقلمنهوپگی۰۱۲۳۴۵۶۷۸۹  --batch_size 128

The output should be like below:

[INFO] MEAN: [0.4845], STD: [0.1884]

Run the following command to get the optimal img_w. For longer label length the img_w should be longer than usual, unless the ctc-loss returns nan.

python get_optimum_img_w.py --alphabets ابپتشثجدزسصطعفقکگلمنوهی+۰۱۲۳۴۵۶۷۸۹ --data_directory <your-train-dir>

the output should be like below:

[INFO] max_length of this dataset is 8, optimal img_w is: 100

Or get all stats using the following command:

python get_all_stats.py --data_directory <your-dataset-dir>

Get the stats and replace them with values provided for img_w, MEAN, STD, and N_CLASSES in the settings.py module under the BasicConfig class, or simply pass them as input arguments.

Train:

After modifying the aforementioned configs, run the following command to train the model:

python train.py

or

python train.py --img_w 100 --n_workers 8 --train_directory <your-trian-dir> --val_directory <your-val-dir> --mean 0.4845 --std 0.1884 --alphabets ابپتشثجدزسصطعفقکگلمنوهی+۰۱۲۳۴۵۶۷۸۹ 

To see all the configs:

python train.py -h

Output

optional arguments:
  -h, --help            show this help message and exit
  --train_directory TRAIN_DIRECTORY
                        path to the dataset, default: ./dataset
  --val_directory VAL_DIRECTORY
                        path to the dataset, default: ./dataset
  --output_dir OUTPUT_DIR
                        path to the output directory, default: ./output
  --epochs EPOCHS       number of training epochs
  --device DEVICE       what should be the device for training, default is cuda
  --mean MEAN [MEAN ...]
                        dataset channel-wise mean
  --std STD [STD ...]   dataset channel-wise std
  --img_w IMG_W         dataset img width size
  --n_workers N_WORKERS
                        number of workers used for dataset collection
  --batch_size BATCH_SIZE
                        batch size number
  --alphabets ALPHABETS
                        alphabets used in the process

Inference

For inference run the following code:

python crnn_inference.py --model_path {path-to-your-output-dir}/best.ckpt --img_path sample_images/۱۴ق۹۱۸۱۱_7073.jpg

The output should be like the following:

۱۴ق۹۱۸۱۱

Image examples:

Sample Persian Dataset is avalable by Amirkabir University of Technology in the following link:

https://ceit.aut.ac.ir/~keyvanrad/download/ML971/project/

Password: Upon request to AUT

Labeled Dataset

https://github.com/Vargha-Kh/crnn-pytorch/releases/tag/Persian-licence-plate-recognition

Foot-Notes:

1. For labeling tool checkout my project: https://github.com/pooya-mohammadi/ocr-labeling-tool

Forked repositories:

1. https://github.com/Vargha-Kh/crnn-pytorch

References

1. https://github.com/pooya-mohammadi/deep_utils
2. https://github.com/AryanShekarlaban/
3. https://arxiv.org/pdf/1507.05717.pdf

crnn-pytorch's Issues

Exporting the trained model into ONNX format.

Hello Mr. Pooya,

First of all, thank you for sharing your knowledge, and build this project.

I was wondering if there is an option to export the trained model into ONNX, or any guideline to do so if it is possible.

My concern about the current setup is using PyTorch in the process of inference, which is a quite big library to install for just doing the OCR step.

download dataset

Hi

I have downloaded dataset from this link "https://ceit.aut.ac.ir/~keyvanrad/download/ML971/project/" , but I can't extract it, I need the password for it.