1. 准备
  2. 训练
  3. 测试

参考 https://github.com/tensorflow/models/tree/master/research/attention_ocr

使用 TensorFlow Attention OCR API 进行图像文本提取

修改后的代码：https://github.com/lijiancheng0614/tensorflow_attention_ocr

准备

  1. 文件结构

    为了方便查看文件，使用以下文件结构。

    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    tensorflow_attention_ocr
    +-- README.md
    +-- common_flags.py
    +-- data_provider.py
    +-- datasets
    |   +-- __init__.py
    |   +-- data
    |   |   +-- fsns
    |   |       +-- charset_size=134.txt
    |   |       +-- fsns-00000-of-00001
    |   |       +-- links.txt
    |   |       +-- testdata
    |   |           +-- fsns_train_00.png
    |   |           +-- ...
    |   |           +-- fsns_train_31.png
    |   +-- fsns.py
    +-- eval.py
    +-- inception_preprocessing.py
    +-- infer.py
    +-- metrics.py
    +-- model.py
    +-- sequence_layers.py
    +-- train.py
    +-- utils.py

  2. 安装 TensorFlow

    参考 https://www.tensorflow.org/install/

  3. 准备数据

    需要数据集的 charset 和 TFRecord，这里使用官方生成的 fsns 数据集的 TFRecord。

    生成其它数据集的 TFRecord 见 https://github.com/lijiancheng0614/tensorflow_attention_ocr/tree/develop

  4. （可选）下载模型

    官方提供了一个训练好的模型，可用于 finetune 和测试。

    1
    2
    wget http://download.tensorflow.org/models/attention_ocr_2017_08_09.tar.gz
    tar zxf attention_ocr_2017_08_09.tar.gz

训练

如果使用现有模型进行预测则不需要训练。

  1. 训练

    从头开始训练：

    1
    2
    python train.py \
        --train_log_dir=train_logs

    使用预训练的 Inception 权重初始化训练模型：

    1
    2
    python train.py \
        --checkpoint_inception=inception_v3.ckpt

    使用一个 checkpoint 进行 fine tune：

    1
    2
    python train.py \
        --checkpoint=model.ckpt-399731

  2. 验证

    1
    2
    3
    4
    python eval.py \
        --train_log_dir=train_logs \
        --eval_log_dir=eval_logs \
        --num_batches=1

  3. 可视化 log

    可一边训练一边可视化训练的 log，访问 http://localhost:8001/ 即可看到 loss 等的变化。

    1
    tensorboard --logdir . --port 8001

测试

测试图片为 datasets/data/fsns/testdata/fsns_train_%02d.png，运行以下命令

1
2
3
python infer.py \
    --checkpoint=model.ckpt-399731 \
    --image_path_pattern=datasets/data/fsns/testdata/fsns_train_%02d.png

控制台中将输出每张图片的结果：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
Boulevard de Lunel░░░░░░░░░░░░░░░░░░░
Rue de Provence░░░░░░░░░░░░░░░░░░░░░░
Rue de Port Maria░░░░░░░░░░░░░░░░░░░░
Avenue Charles Gounod░░░░░░░░░░░░░░░░
Rue de l‘Aurore░░░░░░░░░░░░░░░░░░░░░░
Rue de Beuzeville░░░░░░░░░░░░░░░░░░░░
Rue d‘Orbey░░░░░░░░░░░░░░░░░░░░░░░░░░
Rue Victor Schoulcher░░░░░░░░░░░░░░░░
Rue de la Gare░░░░░░░░░░░░░░░░░░░░░░░
Rue des Tulipes░░░░░░░░░░░░░░░░░░░░░░
Rue André Maginot░░░░░░░░░░░░░░░░░░░░
Route de Pringy░░░░░░░░░░░░░░░░░░░░░░
Rue des Landelles░░░░░░░░░░░░░░░░░░░░
Rue des Ilettes░░░░░░░░░░░░░░░░░░░░░░
Avenue de Maurin░░░░░░░░░░░░░░░░░░░░░
Rue Théresa░░░░░░░░░░░░░░░░░░░░░░░░░░
Route de la Balme░░░░░░░░░░░░░░░░░░░░
Rue Hélène Roederer░░░░░░░░░░░░░░░░░░
Rue Emile Bernard░░░░░░░░░░░░░░░░░░░░
Place de la Mairie░░░░░░░░░░░░░░░░░░░
Rue des Perrots░░░░░░░░░░░░░░░░░░░░░░
Rue de la Libération░░░░░░░░░░░░░░░░░
Impasse du Capcir░░░░░░░░░░░░░░░░░░░░
Avenue de la Grand Mare░░░░░░░░░░░░░░
Rue Pierre Brossolette░░░░░░░░░░░░░░░
Rue de Provence░░░░░░░░░░░░░░░░░░░░░░
Rue du Docteur Mourre░░░░░░░░░░░░░░░░
Rue d‘Ortheuil░░░░░░░░░░░░░░░░░░░░░░░
Rue des Sarments░░░░░░░░░░░░░░░░░░░░░
Rue du Centre░░░░░░░░░░░░░░░░░░░░░░░░
Impasse Pierre Mourgues░░░░░░░░░░░░░░
Rue Marcel Dassault░░░░░░░░░░░░░░░░░░