README.md 5.69 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
## Attention-based Extraction of Structured Information from Street View Imagery

*A TensorFlow model for real-world image text extraction problems.*

This folder contains the code needed to train a new Attention OCR model on the
[FSNS dataset][FSNS] dataset to transcribe street names in France. You can
also use it to train it on your own data.

More details can be found in our paper:

["Attention-based Extraction of Structured Information from Street View
Imagery"](https://arxiv.org/abs/1704.03549)

## Contacts

Authors:
Zbigniew Wojna <zbigniewwojna@gmail.com>,
Alexander Gorban <gorban@google.com>

Pull requests:
[alexgorban](https://github.com/alexgorban)

## Requirements

25
1. Install the TensorFlow library ([instructions][TF]). For example:
26
27
28
29
30

```
virtualenv --system-site-packages ~/.tensorflow
source ~/.tensorflow/bin/activate
pip install --upgrade pip
31
pip install --upgrade tensorflow-gpu
32
33
```

34
2. At least 158GB of free disk space to download the FSNS dataset:
35

36
```
37
cd models/attention_ocr/python/datasets
38
aria2c -c -j 20 -i ../../../street/python/fsns_urls.txt
39
cd ..
40
41
```

42
43
3. 16GB of RAM or more; 32GB is recommended.
4. `train.py` works with both CPU and GPU, though using GPU is preferable. It has been tested with a Titan X and with a GTX980.
44
45
46
47
48
49
50
51
52

[TF]: https://www.tensorflow.org/install/
[FSNS]: https://github.com/tensorflow/models/tree/master/street

## How to use this code

To run all unit tests:

```
53
cd models/attention_ocr/python
54
55
56
57
58
59
60
61
62
python -m unittest discover -p  '*_test.py'
```

To train from scratch:

```
python train.py
```

63
64
To train a model using pre-trained Inception weights as initialization:

65
66
67
```
wget http://download.tensorflow.org/models/inception_v3_2016_08_28.tar.gz
tar xf inception_v3_2016_08_28.tar.gz
68
python train.py --checkpoint_inception=./inception_v3.ckpt
69
70
71
72
73
```

To fine tune the Attention OCR model using a checkpoint:

```
74
75
76
wget http://download.tensorflow.org/models/attention_ocr_2017_05_17.tar.gz
tar xf attention_ocr_2017_05_17.tar.gz
python train.py --checkpoint=model.ckpt-399731
77
78
```

Neal Wu's avatar
Neal Wu committed
79
## How to use your own image data to train the model
80
81
82
83
84

You need to define a new dataset. There are two options:

1. Store data in the same format as the FSNS dataset and just reuse the
[python/datasets/fsns.py](https://github.com/tensorflow/models/blob/master/attention_ocr/python/datasets/fsns.py)
Neal Wu's avatar
Neal Wu committed
85
module. E.g., create a file datasets/newtextdataset.py:
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
```
import fsns

DEFAULT_DATASET_DIR = 'path/to/the/dataset'

DEFAULT_CONFIG = {
    'name':
        'MYDATASET',
    'splits': {
        'train': {
            'size': 123,
            'pattern': 'tfexample_train*'
        },
        'test': {
            'size': 123,
            'pattern': 'tfexample_test*'
        }
    },
    'charset_filename':
        'charset_size.txt',
    'image_shape': (150, 600, 3),
    'num_of_views':
        4,
    'max_sequence_length':
        37,
    'null_code':
        42,
    'items_to_descriptions': {
        'image':
            'A [150 x 600 x 3] color image.',
        'label':
            'Characters codes.',
        'text':
            'A unicode string.',
        'length':
            'A length of the encoded text.',
        'num_of_views':
            'A number of different views stored within the image.'
    }
}


def get_split(split_name, dataset_dir=None, config=None):
  if not dataset_dir:
    dataset_dir = DEFAULT_DATASET_DIR
  if not config:
    config = DEFAULT_CONFIG

  return fsns.get_split(split_name, dataset_dir, config)
```
You will also need to include it into the `datasets/__init__.py` and specify the
dataset name in the command line.

```
python train.py --dataset_name=newtextdataset
```

Neal Wu's avatar
Neal Wu committed
143
Please note that eval.py will also require the same flag.
144

Alexander Gorban's avatar
Alexander Gorban committed
145
146
147
To learn how to store a data in the FSNS
 format please refer to the https://stackoverflow.com/a/44461910/743658.

148
149
150
151
152
153
2. Define a new dataset format. The model needs the following data to train:

- images: input images,  shape [batch_size x H x W x 3];
- labels: ground truth label ids,  shape=[batch_size x seq_length];
- labels_one_hot: labels in one-hot encoding,  shape [batch_size x seq_length x num_char_classes];

Neal Wu's avatar
Neal Wu committed
154
155
Refer to [python/data_provider.py](https://github.com/tensorflow/models/blob/master/attention_ocr/python/data_provider.py#L33)
for more details. You can use [python/datasets/fsns.py](https://github.com/tensorflow/models/blob/master/attention_ocr/python/datasets/fsns.py)
156
157
158
159
160
as the example.

## How to use a pre-trained model

The inference part was not released yet, but it is pretty straightforward to
Neal Wu's avatar
Neal Wu committed
161
implement one in Python or C++.
162
163
164
165
166
167
168
169
170
171
172
173

The recommended way is to use the [Serving infrastructure](https://tensorflow.github.io/serving/serving_basic).

Alternatively you can:
1. define a placeholder for images (or use directly an numpy array)
2. [create a graph ](https://github.com/tensorflow/models/blob/master/attention_ocr/python/eval.py#L60)
`endpoints = model.create_base(images_placeholder, labels_one_hot=None)`
3. [load a pretrained model](https://github.com/tensorflow/models/blob/master/attention_ocr/python/model.py#L494)
4. run computations through the graph:
`predictions = sess.run(endpoints.predicted_chars, feed_dict={images_placeholder:images_actual_data})`
5. Convert character IDs (predictions) to UTF8 using the provided charset file.

174
175
176
## Disclaimer

This code is a modified version of the internal model we used for our paper.
177
Currently it reaches 83.79% full sequence accuracy after 400k steps of training.
178
179
The main difference between this version and the version used in the paper - for
the paper we used a distributed training with 50 GPU (K80) workers (asynchronous
180
181
updates), the provided checkpoint was created using this code after ~6 days of
training on a single GPU (Titan X) (it reached 81% after 24 hours of training),
Alexander Gorban's avatar
Alexander Gorban committed
182
the coordinate encoding is disabled by default.