Attention-OCR: 基于视觉注意力机制的图像文本识别模型

attention-ocr

Attention-OCR简介

Attention-OCR是一个开源的图像文本识别模型,由Qi Guo和Yuntian Deng开发。该模型采用了卷积神经网络(CNN)和带有视觉注意力机制的序列到序列(seq2seq)架构,能够高效准确地从图像中识别和提取文本信息。

Attention-OCR的主要特点包括:

基于TensorFlow框架实现,易于使用和扩展
采用视觉注意力机制,提高了识别准确率
支持创建TFRecords数据集和导出训练好的模型
兼容Google Cloud ML Engine,可进行大规模训练

本文将深入介绍Attention-OCR的工作原理、模型架构、使用方法以及在实际应用中的表现。

模型原理

Attention-OCR的核心思想是结合CNN的图像特征提取能力和seq2seq模型的序列生成能力,再加上注意力机制来提高识别准确率。其工作流程如下:

首先,使用滑动CNN对输入图像进行特征提取。图像会被调整为32像素高,同时保持原始宽高比。
然后,在CNN的输出之上叠加一个LSTM网络,进一步处理序列信息。
最后,使用带有注意力机制的解码器生成最终的文本输出。

这种架构设计使得模型能够"关注"图像中的关键区域,从而更准确地识别文本内容。

OCR示例

上图展示了Attention-OCR的工作过程,可以看到模型能够准确地识别出图像中的文本"jungle"。

安装与使用

安装

Attention-OCR可以通过pip直接安装:

pip install aocr

注意:该项目依赖TensorFlow 1.x版本。目前正在计划升级到TensorFlow 2,欢迎贡献PR。

创建数据集

要训练Attention-OCR模型,首先需要准备一个TFRecords格式的数据集。可以使用以下命令创建数据集:

aocr dataset ./datasets/annotations-training.txt ./datasets/training.tfrecords
aocr dataset ./datasets/annotations-testing.txt ./datasets/testing.tfrecords

其中annotations文件是一个简单的文本文件,包含图像路径和对应的标签:

datasets/images/hello.jpg hello
datasets/images/world.jpg world

训练模型

使用以下命令开始训练:

aocr train ./datasets/training.tfrecords

训练过程可能需要较长时间,因为模型同时训练CNN和注意力模型。可以通过--steps-per-checkpoint参数设置检查点保存频率。

测试与可视化

训练完成后,可以使用以下命令测试模型:

aocr test ./datasets/testing.tfrecords

如果想可视化注意力结果,可以添加--visualize参数:

aocr test --visualize ./datasets/testing.tfrecords

这将在results/correct目录下生成可视化结果,如下图所示:

注意力可视化示例

模型导出与部署

训练完成后,可以将模型导出为SavedModel或frozen graph格式:

# SavedModel (默认)
aocr export ./exported-model

# Frozen graph
aocr export --format=frozengraph ./exported-model

导出的模型可以使用TensorFlow Serving部署为HTTP REST API:

tensorflow_model_server --port=9000 --rest_api_port=9001 --model_name=yourmodelname --model_base_path=./exported-model

然后就可以向运行的服务器发送预测请求:

curl -X POST 
  http://localhost:9001/v1/models/aocr:predict 
  -H 'cache-control: no-cache' 
  -H 'content-type: application/json' 
  -d '{
  "signature_name": "serving_default",
  "inputs": {
      "input": { "b64": "<your image encoded as base64>" }
  }
}'

Google Cloud ML Engine训练

Attention-OCR还支持在Google Cloud ML Engine上进行大规模训练。具体步骤如下:

设置环境变量:

export JOB_PREFIX="aocr"
export REGION="us-central1"
export GS_BUCKET="gs://aocr-bucket"
export DATASET_UPLOAD_PATH="training.tfrecords"

上传训练数据集:

gsutil cp ./datasets/training.tfrecords $GS_BUCKET/$DATASET_UPLOAD_PATH

启动ML Engine任务:

export NOW=$(date +"%Y%m%d_%H%M%S")
export JOB_NAME="$JOB_PREFIX$NOW"
export JOB_DIR="$GS_BUCKET/$JOB_NAME"

gcloud ml-engine jobs submit training $JOB_NAME \
    --job-dir=$JOB_DIR \
    --module-name=aocr \
    --package-path=aocr \
    --region=$REGION \
    --scale-tier=BASIC_GPU \
    --runtime-version=1.2 \
    -- \
    train $GS_BUCKET/$DATASET_UPLOAD_PATH \
    --steps-per-checkpoint=500 \
    --batch-size=512 \
    --num-epoch=20