📄 English | 中文

𝚃𝚎𝚡𝚃𝚎𝚕𝚕𝚎𝚛

https://github.com/OleehyO/TexTeller/assets/56267907/532d1471-a72e-4960-9677-ec6c19db289f

TexTeller是一个基于TrOCR的端到端公式识别模型，能够将图像转换为相应的LaTeX公式。

TexTeller使用8000万对图像-公式数据进行训练（之前的数据集可以在这里获取），相比于LaTeX-OCR使用的10万数据集，TexTeller具有更强的泛化能力和更高的准确性，覆盖了大多数使用场景。

[!NOTE] 如果您想为这个项目提供反馈或建议，欢迎在讨论区开启讨论。

另外，如果您觉得这个项目有帮助，请不要忘记给它一个星星⭐️🙏️

感谢 北京邮电大学超级计算平台 对这项工作的支持😘

🔄 更新日志

📮[2024-06-06] TexTeller3.0发布！ 训练数据已增加到8000万（比TexTeller2.0多10倍，数据多样性也有所提高）。TexTeller3.0的新特性：
- 支持扫描图像、手写公式、中英文混合公式。
- 具备中英文打印图像的OCR能力。
📮[2024-05-02] 支持段落识别。
📮[2024-04-12] 公式检测模型发布！
📮[2024-03-25] TexTeller2.0发布！TexTeller2.0的训练数据增加到750万（比TexTeller1.0多15倍，数据质量也有所提高）。经过训练的TexTeller2.0在测试集上表现出卓越的性能，尤其是在识别罕见符号、复杂多行公式和矩阵方面。

这里有更多测试图像和各种识别模型的横向比较。

🚀 快速开始

克隆仓库：

git clone https://github.com/OleehyO/TexTeller

安装项目依赖：
```
pip install texteller
```

进入src/目录，在终端中运行以下命令开始推理：

python inference.py -img "/path/to/image.{jpg,png}" 
# 使用--inference-mode选项启用GPU(cuda或mps)推理
#+例如：python inference.py -img "img.jpg" --inference-mode cuda

第一次运行时，将从Hugging Face下载所需的检查点。

段落识别

如视频所示，TexTeller也能够识别整段文本。虽然TexTeller具有通用文本OCR能力，但我们仍建议使用段落识别以获得更好的结果：

下载权重并将公式检测模型放入src/models/det_model/model/目录
在src/目录中运行inference.py并添加-mix选项，结果将以markdown格式输出。
```
python inference.py -img "/path/to/image.{jpg,png}" -mix
```

TexTeller默认使用轻量级的PaddleOCR模型来识别中英文文本。您可以尝试使用更大的模型来获得更好的中英文识别结果：

检查点	模型描述	大小
ch_PP-OCRv4_det.onnx	默认检测模型，支持中英文文本检测	4.70M
ch_PP-OCRv4_server_det.onnx	高精度模型，支持中英文文本检测	115M
ch_PP-OCRv4_rec.onnx	默认识别模型，支持中英文文本识别	10.80M
ch_PP-OCRv4_server_rec.onnx	高精度模型，支持中英文文本识别	90.60M

将识别/检测模型的权重放在src/models/third_party/paddleocr/checkpoints/下的det/或rec/目录中，并将其重命名为default_model.onnx。

[!NOTE] 段落识别无法还原文档的结构，它只能识别其内容。

🌐 网页演示

进入src/目录并运行以下命令：

./start_web.sh

在浏览器中输入http://localhost:8501即可查看网页演示。

[!NOTE]

对于Windows用户，请运行start_web.bat文件。

使用onnxruntime + GPU进行推理时，需要安装onnxruntime-gpu。

🔍 公式检测

TexTeller的公式检测模型在3,415张中文教育材料图像（包含130多种布局）和8,272张来自IBEM数据集的图像上进行训练，支持对整张图像进行公式检测。

下载模型权重并将其放置在src/models/det_model/model/目录下[链接]。
在src/目录下运行以下命令，结果将保存在src/subimages/目录中

进阶：批量公式识别

在公式检测之后，在src/目录下运行以下命令：

python rec_infer_from_crop_imgs.py

这将使用之前的公式检测结果对所有裁剪出的公式进行批量识别，并将识别结果以txt文件形式保存在src/results/目录中。

📡 API使用

我们使用ray serve为TexTeller提供API接口，使您能够将TexTeller集成到自己的项目中。要启动服务器，首先需要进入src/目录，然后运行以下命令：

python server.py

参数	描述
`-ckpt`	权重文件路径，默认为TexTeller的预训练权重。
`-tknz`	分词器路径，默认为TexTeller的分词器。
`-port`	服务器服务端口，默认为8000。
`--inference-mode`	是否使用"cuda"或"mps"进行推理，默认为"cpu"。
`--num_beams`	束搜索的束数，默认为1。
`--num_replicas`	在服务器上运行的服务副本数量，默认为1个副本。您可以使用更多副本来实现更高的吞吐量。
`--ncpu_per_replica`	每个服务副本使用的CPU核心数，默认为1。
`--ngpu_per_replica`	每个服务副本使用的GPU数量，默认为1。您可以将此值设置在0到1之间，以在一个GPU上运行多个服务副本来共享GPU，从而提高GPU利用率。（注意，如果--num_replicas为2，--ngpu_per_replica为0.7，则必须有2个可用的GPU）
`-onnx`	使用Onnx Runtime进行推理，默认禁用

[!注意] 客户端演示可以在src/client/demo.py中找到，您可以参考demo.py向服务器发送请求

🏋️‍♂️ 训练

数据集

我们在src/models/ocr_model/train/dataset/目录中提供了一个示例数据集，您可以将自己的图像放在images/目录中，并在formulas.jsonl中为每个图像标注相应的公式。

准备好数据集后，您需要在**/train/dataset/loader.py中将DIR_URL变量更改为您自己数据集的路径

重新训练分词器

如果您使用的是不同的数据集，可能需要重新训练分词器以获得不同的词汇表。配置好数据集后，您可以使用以下命令训练自己的分词器：

在src/models/tokenizer/train.py中，将new_tokenizer.save_pretrained('./your_dir_name')更改为您的自定义输出目录

如果您想使用不同的词汇量大小（默认15K），需要在src/models/globals.py中更改VOCAB_SIZE变量
在src/目录下，运行以下命令：
```
python -m models.tokenizer.train
```

训练模型

在src/train_config.yaml中修改num_processes以匹配可用于训练的GPU数量（默认为1）。

在src/目录下，运行以下命令：

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

您可以在src/models/ocr_model/train/train.py中设置自己的分词器和检查点路径（更多信息请参考train.py）。如果您使用的架构和词汇表与TexTeller相同，您还可以使用自己的数据集对TexTeller的默认权重进行微调。

在src/globals.py和src/models/ocr_model/train/train_args.py中，您可以更改模型的架构和训练超参数。

[!注意] 我们的训练脚本使用Hugging Face Transformers库，因此您可以参考他们的文档获取有关训练参数的更多详细信息和配置。

TexTeller