README.md 2.2 KB
Newer Older
农夫山泉2号's avatar
农夫山泉2号 已提交
1
attention-ocr.pytorch:Encoder+Decoder+attention model
M
meijieru 已提交
2 3
======================================

4 5
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.
X
xiaohang 已提交
6

农夫山泉2号's avatar
农夫山泉2号 已提交
7
# requirements
农夫山泉2号's avatar
农夫山泉2号 已提交
8 9 10 11 12
pytorch 0.4.1  
opencv_python
```bash
cd Attention_ocr.pytorch
pip install -r requirements.txt
农夫山泉2号's avatar
农夫山泉2号 已提交
13 14 15
```

# Test
16
pretrained model coming soon
农夫山泉2号's avatar
农夫山泉2号 已提交
17

农夫山泉2号's avatar
农夫山泉2号 已提交
18

农夫山泉2号's avatar
农夫山泉2号 已提交
19 20 21 22
# Train 
1. Here i choose a small dataset from [Synthetic_Chinese_String_Dataset](https://github.com/chenjun2hao/caffe_ocr), about 270000+ images for training, 20000 images for testing.
download the image data from [Baidu](https://pan.baidu.com/s/1hIurFJ73XbzL-QG4V-oe0w)
2. the train_list.txt and test_list.txt are created as the follow form:
农夫山泉2号's avatar
农夫山泉2号 已提交
23 24
```
# path/to/image_name.jpg label
农夫山泉2号's avatar
农夫山泉2号 已提交
25 26 27
path/AttentionData/50843500_2726670787.jpg 情笼罩在他们满是沧桑
path/AttentionData/57724421_3902051606.jpg 心态的松弛决定了比赛
path/AttentionData/52041437_3766953320.jpg 虾的鲜美自是不可待言
农夫山泉2号's avatar
农夫山泉2号 已提交
28
```
农夫山泉2号's avatar
农夫山泉2号 已提交
29 30 31 32 33 34
3. change the **trainlist** and **vallist** parameter in train.py, and start train
```bash
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:
35 36
![attentionocr](./test_img/md_img/attentionV2.png)
there uses the decoderV2 model for decoder.
农夫山泉2号's avatar
农夫山泉2号 已提交
37

C
chenjun2hao 已提交
38 39 40 41 42 43

# The previous version  

**_git checkout AttentionOcrV1_**


农夫山泉2号's avatar
农夫山泉2号 已提交
44 45 46 47
# Reference
1. [crnn.pytorch](https://github.com/meijieru/crnn.pytorch)
2. [Attention-OCR](https://github.com/da03/Attention-OCR)
3. [Seq2Seq-PyTorch](https://github.com/MaximumEntropy/Seq2Seq-PyTorch)
农夫山泉2号's avatar
农夫山泉2号 已提交
48
4. [caffe_ocr](https://github.com/senlinuc/caffe_ocr)
农夫山泉2号's avatar
农夫山泉2号 已提交
49

C
chenjun2hao 已提交
50

农夫山泉2号's avatar
农夫山泉2号 已提交
51 52
# TO DO
- [ ] change LSTM to Conv1D, it can greatly accelerate the inference
农夫山泉2号's avatar
农夫山泉2号 已提交
53
- [ ] change the cnn bone model with inception net, densenet
54
- [ ] realize the decoder with transformer model