Skip to content

This repository implements the the encoder and decoder model with attention model for OCR

Notifications You must be signed in to change notification settings

siyisan/Attention_ocr.pytorch

 
 

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

46 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

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

This repository implements the the encoder and decoder model with attention model for OCR, the encoder uses CNN+Bi-LSTM, the decoder uses GRU. This repository is modified from https://github.com/meijieru/crnn.pytorch
Earlier I had an open source version, but had some problems identifying images of fixed width. Recently I modified the model to support image recognition with variable width. The function is the same as CRNN. Due to the time problem, there is no pre-training model this time, which will be updated later.

requirements

pytorch 0.4.1
opencv_python

cd Attention_ocr.pytorch
pip install -r requirements.txt

Test

pretrained model coming soon

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 there uses the decoderV2 model for decoder.

The previous version

git checkout AttentionOcrV1

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
  • change the cnn bone model with inception net, densenet
  • realize the decoder with transformer model

About

This repository implements the the encoder and decoder model with attention model for OCR

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages

  • Python 100.0%