This repository implements the the encoder and decoder model with attention model for OCR
Clone or download
Latest commit ab92bb0 Jan 8, 2019
Permalink
Type Name Latest commit message Commit time
Failed to load latest commit information.
data attention ocr Jan 4, 2019
expr/attentioncnn attention ocr Jan 4, 2019
models attention ocr Jan 4, 2019
test_img attention ocr Jan 4, 2019
.gitignore attention ocr Jan 4, 2019
README.md Update README.md Jan 8, 2019
class_attention.py attention ocr Jan 4, 2019
dataset.py attention ocr Jan 4, 2019
demo.py attention ocr Jan 4, 2019
requirements.txt add requirements.txt Dec 22, 2018
train.py attention ocr Jan 4, 2019
utils.py attention ocr Jan 4, 2019

README.md

attention-ocr.pytorch:Encoder+Decoder+attention model

This repository implements the the encoder and decoder model with attention model for OCR, and this repository is modified from https://github.com/meijieru/crnn.pytorch
There are many attention ocr repository which is finished with tensorflow, but they don't give the inference.py, besides i'm not good at tensorflow, i can't finish the inference.py by myself

requirements

pytorch 0.4.1
opencv_python

cd Attention_ocr.pytorch
pip install -r requirements.txt

Test

  1. download the pretrained model from attentioncnn.zip ,and put the pth files into the folder:expr/attentioncnn/

  2. change the test image's name in demo.py,the test images are in the test_img folder

python demo.py
  1. results
>>>predict_str:比赛,尽管很尽心尽力           => prob:0.6112725138664246
  1. some examples 结果.jpg
picture predict reading confidence
美国人不愿意与制鲜 0.33920
现之间的一个分享和 0.81095
中国通信学会主办、《 0.90660
此在战术上应大胆、勇 0.57111
年同期俱点83.50 0.14481
。留言无恶意心态成 0.31054
  1. the accuracy in the test data just stay around 88%, there is much thing to do

Train

  1. Here i choose a small dataset from Synthetic_Chinese_String_Dataset, about 270000+ images for training, 20000 images for testing. download the image data from Baidu
  2. the train_list.txt and test_list.txt are created as the follow form:
# path/to/image_name.jpg label
path/AttentionData/50843500_2726670787.jpg 情笼罩在他们满是沧桑
path/AttentionData/57724421_3902051606.jpg 心态的松弛决定了比赛
path/AttentionData/52041437_3766953320.jpg 虾的鲜美自是不可待言
  1. change the trainlist and vallist parameter in train.py, and start train
cd Attention_ocr.pytorch
python train.py --trainlist ./data/ch_train.txt --vallist ./data/ch_test.txt

then you can see in the terminel as follow: attentionocr

Reference

  1. crnn.pytorch
  2. Attention-OCR
  3. Seq2Seq-PyTorch
  4. caffe_ocr

TO DO

  • change LSTM to Conv1D, it can greatly accelerate the inference
  • to support images of different widths
  • change the cnn bone model with inception net, densenet