OCR 通用端到端模型GOT，迈向OCR2.0的新时代

AI浩

发布于 2024-10-22 05:19:27

45710

代码可运行

文章被收录于专栏：AI智韵AI智韵

运行总次数：0

代码可运行

摘要

在人工智能领域，光学字符识别（OCR）技术已经取得了显著的进展。随着技术的不断进步，我们正迈向OCR 2.0时代。本文将介绍由Vary团队开发的通用端到端模型GOT，这一模型在OCR领域具有革命性的潜力。

论文概览

论文标题：GOT: Towards OCR-2.0
发布平台：arXiv
链接：arXiv.org

模型特点

GOT模型是首个迈向OCR 2.0时代的通用端到端模型，它在多个方面展现了其先进性：

多任务支持：GOT模型支持多种OCR任务，包括场景文本OCR、文档OCR、细粒度OCR以及更通用的OCR任务。
输出格式多样：除了支持纯文本输出，GOT还能输出格式化文本，如Markdown格式，增强了文本的可读性和可编辑性。
结构优化：采用vision encoder + input embedding layer + decoder的架构，其中encoder部分采用带local attention的VITDet架构，有效管理显存使用。

训练方法

GOT模型的训练分为三个阶段：

第一阶段：高效预训练encoder，使用小型OPT-125M作为decoder，快速引入大量数据。
第二阶段：联合训练encoder-decoder，使用Qwen团队预训练的Qwen0.5B，适当增大decoder以适应OCR-2.0的知识需求。
第三阶段：锁定encoder，加强decoder以适配更多OCR应用场景，如支持坐标或颜色引导的细粒度OCR，动态分辨率OCR技术，多页OCR技术。

数据工程

研究团队在数据工程方面投入巨大，学习并使用了多种数据渲染工具，包括Latex、Mathpix-markdown-it、Matplotlib、Tikz、Verovio、Pyecharts等，以构造多样化的数据。

项目地址

对GOT模型感兴趣的研究者和开发者可以通过以下链接访问项目代码：GitHub - Ucas-HaoranWei/GOT-OCR2.0

安装

基础环境cuda11.8+torch2.0.1

克隆仓库并导航到GOT文件夹

git clone https://github.com/Ucas-HaoranWei/GOT-OCR2.0.git 
cd 'the GOT folder'

安装包

conda create -n got python=3.10 -y
conda activate got
pip install -e .

安装 Flash-Attention

pip install ninja
pip install flash-attn --no-build-isolation

GOT 权重

Huggingface
Google Drive
百度云密码: OCR2

演示

普通文本OCR:

python3 GOT/demo/run_ocr_2.0.py  --model-name  /GOT_weights/  --image-file  /an/image/file.png  --type ocr

格式文本OCR:

python3 GOT/demo/run_ocr_2.0.py  --model-name  /GOT_weights/  --image-file  /an/image/file.png  --type format

细粒度OCR:

python3 GOT/demo/run_ocr_2.0.py  --model-name  /GOT_weights/  --image-file  /an/image/file.png  --type format/ocr --box [x1,y1,x2,y2]

python3 GOT/demo/run_ocr_2.0.py  --model-name  /GOT_weights/  --image-file  /an/image/file.png  --type format/ocr --color red/green/blue

多裁剪OCR:

python3 GOT/demo/run_ocr_2.0_crop.py  --model-name  /GOT_weights/ --image-file  /an/image/file.png

多页OCR (图像路径包含多个.png文件):

python3 GOT/demo/run_ocr_2.0_crop.py  --model-name  /GOT_weights/ --image-file  /images/path/  --multi-page

渲染格式化OCR结果:

python3 GOT/demo/run_ocr_2.0.py  --model-name  /GOT_weights/  --image-file  /an/image/file.png  --type format --render

注意: 渲染结果可以在/results/demo.html中找到。请打开demo.html查看结果。

训练

训练样本可以在此链接找到。注意，在'conversations'-'human'-'value'中的'<image>'是必要的！
此代码库仅支持在我们GOT权重上的后训练（第二/第三阶段）。
如果你想从我们论文中描述的第一阶段训练，你需要这个仓库。

deepspeed   /GOT-OCR-2.0-master/GOT/train/train_GOT.py \
 --deepspeed /GOT-OCR-2.0-master/zero_config/zero2.json    --model_name_or_path /GOT_weights/ \
 --use_im_start_end True   \
 --bf16 True   \
 --gradient_accumulation_steps 2    \
 --evaluation_strategy "no"   \
 --save_strategy "steps"  \
 --save_steps 200   \
 --save_total_limit 1   \
 --weight_decay 0.    \
 --warmup_ratio 0.001     \
 --lr_scheduler_type "cosine"    \
 --logging_steps 1    \
 --tf32 True     \
 --model_max_length 8192    \
 --gradient_checkpointing True   \
 --dataloader_num_workers 8    \
 --report_to none  \
 --per_device_train_batch_size 2    \
 --num_train_epochs 1 \
 --learning_rate 2e-5   \
 --datasets pdf-ocr+scence \
 --output_dir /your/output.path

注意:

更改constant.py中相应的数据信息。
更改conversation_dataset_qwen.py中第37行为你的数据名称。

评估

使用Fox和OneChart基准，其他基准可以在权重下载链接中找到。
评估代码可以在GOT/eval中找到。
你可以使用evaluate_GOT.py运行评估。如果你有8个GPU，--num-chunks可以设置为8。

python3 GOT/eval/evaluate_GOT.py --model-name /GOT_weights/ --gtfile_path xxxx.json --image_path  /image/path/ --out_path /data/eval_results/GOT_mathpix_test/ --num-chunks 8 --datatype OCR