TexTeller: 高精度图像转LaTeX公式解决方案

引言

在当今数字化时代,将图像中的数学公式转换为可编辑的LaTeX代码是一项非常重要的技术。无论是学术研究、教育教学,还是科技文档编写,都有大量需要处理数学公式图像的场景。TexTeller应运而生,为这一需求提供了强大而精确的解决方案。

TexTeller简介

TexTeller是一个基于TrOCR (Transformer OCR)技术的端到端公式识别模型。它能够将包含数学公式的图像直接转换为对应的LaTeX代码,无需中间步骤。这种端到端的方法不仅简化了处理流程,还显著提高了识别的准确性。

TexTeller示例图

核心特点

海量训练数据: TexTeller使用了多达8000万对图像-公式数据进行训练,这是目前同类模型中规模最大的训练集。相比之下,此前的LaTeX-OCR模型仅使用了10万数据集。
强大的泛化能力: 得益于庞大而多样的训练数据,TexTeller展现出了卓越的泛化能力。它几乎可以应对所有常见的数学公式识别场景,包括复杂的多行公式和矩阵等。
高精度识别: TexTeller在各种测试集上都表现出了优异的识别准确率,特别是在识别罕见符号和复杂结构方面有显著优势。
多语言支持: 除了识别纯数学公式,TexTeller还支持中英文混合公式的识别,大大扩展了其应用范围。
扫描图像支持: TexTeller能够处理扫描后的图像,这对于处理纸质文档中的公式非常有用。
手写公式识别: 模型还具备识别手写数学公式的能力,为教育和研究领域提供了便利。

技术原理

TexTeller的核心是基于Transformer架构的OCR模型。这种模型在处理序列数据方面表现出色,非常适合处理数学公式这种具有复杂结构和长距离依赖的内容。

模型架构

TexTeller的模型架构主要包括以下几个部分:

图像编码器: 使用卷积神经网络(CNN)提取图像的视觉特征。
Transformer编码器: 将视觉特征转换为上下文相关的表示。
Transformer解码器: 根据编码的特征生成LaTeX代码序列。
Beam Search: 在解码过程中使用Beam Search策略,提高生成结果的质量。

训练过程

TexTeller的训练过程充分利用了大规模数据集和先进的训练技术:

数据增强: 使用各种图像处理技术创建多样化的训练样本,提高模型的鲁棒性。
预训练: 利用通用文本数据进行预训练,然后在数学公式数据上进行微调。
多GPU训练: 采用分布式训练策略,加速大规模数据集的训练过程。
渐进式学习: 从简单到复杂逐步增加训练难度,帮助模型更好地学习。

使用指南

安装

要开始使用TexTeller,首先需要克隆项目仓库并安装依赖:

git clone https://github.com/OleehyO/TexTeller
cd TexTeller
pip install -r requirements.txt

基本使用

进入src/目录,然后运行以下命令进行推理:

python inference.py -img "/path/to/image.{jpg,png}"

如果想使用GPU加速,可以添加--inference-mode选项:

python inference.py -img "img.jpg" --inference-mode cuda

段落识别

TexTeller不仅可以识别单个公式,还支持整个段落的识别。这对于处理包含多个公式的文档非常有用。使用段落识别功能需要以下步骤:

下载公式检测模型的权重文件,并放置在src/models/det_model/model/目录下。
在src/目录下运行以下命令:

python inference.py -img "/path/to/image.{jpg,png}" -mix

结果将以Markdown格式输出,方便后续编辑和使用。

Web演示

TexTeller还提供了一个简单的Web界面,方便用户直观地体验公式识别功能:

进入src/目录
运行命令: streamlit run app.py
在浏览器中打开http://localhost:8501即可使用Web演示

高级功能

API使用

TexTeller提供了API接口,方便集成到其他项目中。使用ray serve框架搭建的API服务,可以通过以下命令启动:

python serve.py

API服务支持多种参数配置,如模型权重路径、推理模式、服务端口等,详细参数说明可以参考项目文档。

公式检测

对于包含多个公式的复杂文档,TexTeller提供了公式检测功能:

下载检测模型权重
运行检测命令:

python det_infer.py -img "/path/to/image.{jpg,png}"

检测结果会保存在src/subimages/目录下。

批量识别

结合公式检测功能,TexTeller还支持批量识别文档中的所有公式:

python rec_infer_from_crop_imgs.py

这将对检测到的所有公式进行批量识别,并将结果保存为txt文件。

模型训练

TexTeller支持用户使用自己的数据集进行训练或微调。主要步骤包括:

准备数据集: 将图像放在src/models/ocr_model/train/dataset/images/目录,并在formulas.jsonl中添加对应的标注。
重新训练分词器(可选):

python -m models.tokenizer.train

训练模型:

accelerate launch --config_file ./train_config.yaml -m models.ocr_model.train.train

用户可以在配置文件中调整模型架构和训练超参数。

未来计划

TexTeller团队有着雄心勃勃的发展计划,包括:

进一步扩大训练数据集规模
增强对扫描图像的识别能力
完善对英文和中文场景的支持
提升手写公式的识别准确率
添加PDF文档直接识别功能
持续优化推理速度

结语

TexTeller作为一个强大的图像转LaTeX公式工具,为数学公式的数字化处理提供了极大便利。无论是学术研究、教育教学,还是科技文档编写,TexTeller都能提供高质量的公式识别服务。随着持续的优化和功能扩展,TexTeller有望成为这一领域的领先解决方案。

欢迎更多开发者和研究者加入TexTeller社区,共同推动数学公式识别技术的进步。如果您对这个项目感兴趣,不妨给它一个star⭐️,关注项目的最新动态。让我们一起,为数学公式的数字化处理贡献一份力量!