diff --git a/README.md b/README.md index 2db1448..605a2ef 100644 --- a/README.md +++ b/README.md @@ -7,9 +7,12 @@

English | 中文

-

+ + TexTeller is an end-to-end formula recognition model based on ViT, capable of converting images into corresponding LaTeX formulas. @@ -21,6 +24,8 @@ TexTeller was trained with ~~550K~~7.5M image-formula pairs (dataset available [ ## 🔄 Change Log * 📮[2024-03-24] TexTeller 2.0 released! The training data for TexTeller 2.0 has been increased to 7.5M (about **15 times more** than TexTeller 1.0 and also improved in data quality). The trained TexTeller 2.0 demonstrated **superior performance** in the test set, especially in recognizing rare symbols, complex multi-line formulas, and matrices. + > [!INFO] + > [There](./assets/test.pdf) are more test images here and a horizontal comparison of recognition models from different companies. ## 🔑 Prerequisites @@ -138,8 +143,6 @@ In `TexTeller/src/globals.py` and `TexTeller/src/models/ocr_model/train/train_ar ## 🚧 Limitations -* Some complex multi-line scenarios are not well handled (e.g., long formulas mixed with matrices) - * Does not support scanned images and PDF document recognition * Does not support handwritten formulas diff --git a/assets/README_zh.md b/assets/README_zh.md index e2b301d..cc126b7 100644 --- a/assets/README_zh.md +++ b/assets/README_zh.md @@ -7,9 +7,12 @@

English | 中文

-

+ + TexTeller是一个基于ViT的端到端公式识别模型，可以把图片转换为对应的latex公式 @@ -21,6 +24,8 @@ TexTeller用了~~550K~~7.5M的图片-公式对进行训练(数据集可以在[ ## 🔄 变更信息 * 📮[2024-03-24] TexTeller2.0发布！TexTeller2.0的训练数据增大到了7.5M(相较于TexTeller1.0**增加了~15倍**并且数据质量也有所改善)。训练后的TexTeller2.0在测试集中展现出了**更加优越的性能**，尤其在生僻符号、复杂多行、矩阵的识别场景中。 + > [!INFO] + > 在[这里](./test.pdf)有更多的测试图片以及各家识别模型的横向对比。 ## 🔑 前置条件 @@ -28,7 +33,8 @@ python=3.10 pytorch -> 注意: 只有CUDA版本>= 12.0被完全测试过，所以最好使用>= 12.0的CUDA版本 +> [!WARNING] +> 只有CUDA版本>= 12.0被完全测试过，所以最好使用>= 12.0的CUDA版本 ## 🖼 关于把latex渲染成图片 @@ -62,6 +68,7 @@ pytorch #+e.g. python inference.py -img "./img.jpg" -cuda ``` + > [!NOTE] > 第一次运行时会在hugging face上下载所需要的checkpoints ## ❓ 常见问题：无法连接到Hugging Face @@ -102,9 +109,11 @@ pytorch 然后在浏览器里输入`http://localhost:8501`就可以看到web demo +> [!TIP] > 你可以改变`start_web.sh`的默认配置， 例如使用GPU进行推理(e.g. `USE_CUDA=True`) 或者增加beams的数量(e.g. `NUM_BEAM=3`)来获得更高的精确度 -**NOTE:** 如果你想直接把预测结果在网页上渲染成图片（比如为了检查预测结果是否正确）你需要确保[xelatex被正确安装](https://github.com/OleehyO/TexTeller?tab=readme-ov-file#Rendering-Predicted-Results) +> [!IMPORTANT] +> 如果你想直接把预测结果在网页上渲染成图片（比如为了检查预测结果是否正确）你需要确保[xelatex被正确安装](https://github.com/OleehyO/TexTeller?tab=readme-ov-file#Rendering-Predicted-Results) ## 📡 API调用 @@ -127,6 +136,7 @@ python server.py # default settings | `--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可用) | +> [!NOTE] > 一个客户端demo可以在`TexTeller/client/demo.py`找到，你可以参考`demo.py`来给server发送请求 ## 🏋️‍♂️ 训练 @@ -142,6 +152,7 @@ python server.py # default settings 如果你使用了不一样的数据集，你可能需要重新训练tokenizer来得到一个不一样的字典。配置好数据集后，可以通过以下命令来训练自己的tokenizer： 1. 在`TexTeller/src/models/tokenizer/train.py`中，修改`new_tokenizer.save_pretrained('./your_dir_name')`为你自定义的输出目录 + > [!IMPORTANT] > 如果要用一个不一样大小的字典(默认1W个token)，你需要在 `TexTeller/src/models/globals.py`中修改`VOCAB_SIZE`变量 2. **在 `TexTeller/src` 目录下**运行以下命令: @@ -162,20 +173,11 @@ python -m models.ocr_model.train.train 在`TexTeller/src/globals.py`和`TexTeller/src/models/ocr_model/train/train_args.py`中，你可以改变模型的架构以及训练的超参数。 +> [!NOTE] > 我们的训练脚本使用了[Hugging Face Transformers](https://github.com/huggingface/transformers)库, 所以你可以参考他们提供的[文档](https://huggingface.co/docs/transformers/v4.32.1/main_classes/trainer#transformers.TrainingArguments)来获取更多训练参数的细节以及配置。 ## 🚧 不足 -* 部分细节很多的公式无法做到100%的准确率 - - - -* 部分复杂的大型多行公式识别效果不佳(例如长公式与矩阵混合) - - - - > 如果遇到这种情况，你可以尝试把大型的多行公式分成多个小的子公式来识别。 - * 不支持扫描图片以及PDF文档识别 * 不支持手写体公式 diff --git a/assets/test.mp4 b/assets/test.mp4 new file mode 100644 index 0000000..91ae16b Binary files /dev/null and b/assets/test.mp4 differ diff --git a/assets/test.pdf b/assets/test.pdf new file mode 100644 index 0000000..f587024 Binary files /dev/null and b/assets/test.pdf differ diff --git a/src/models/ocr_model/model/TexTeller.py b/src/models/ocr_model/model/TexTeller.py index de712ac..3158d06 100644 --- a/src/models/ocr_model/model/TexTeller.py +++ b/src/models/ocr_model/model/TexTeller.py @@ -17,7 +17,7 @@ from transformers import ( class TexTeller(VisionEncoderDecoderModel): - REPO_NAME = '/home/lhy/code/TexTeller/src/models/ocr_model/train/train_result/TexTellerv2/checkpoint-356000' + REPO_NAME = '/home/lhy/code/TexTeller/src/models/ocr_model/train/train_result/TexTellerv2/checkpoint-588000' def __init__(self, decoder_path=None, tokenizer_path=None): encoder = ViTModel(ViTConfig( image_size=FIXED_IMG_SIZE, diff --git a/src/start_web.sh b/src/start_web.sh index 450dff2..6daef19 100755 --- a/src/start_web.sh +++ b/src/start_web.sh @@ -3,7 +3,8 @@ set -exu export CHECKPOINT_DIR="default" export TOKENIZER_DIR="default" -export USE_CUDA=False # True or False (case-sensitive) -export NUM_BEAM=1 +# export USE_CUDA=False # True or False (case-sensitive) +export USE_CUDA=True # True or False (case-sensitive) +export NUM_BEAM=10 streamlit run web.py diff --git a/src/web.py b/src/web.py index 6758e82..af934a1 100644 --- a/src/web.py +++ b/src/web.py @@ -95,7 +95,7 @@ def rendering(formula: str, out_img_path: Path) -> bool: return p.returncode == 0 def pdf_to_pngbytes(pdf_path): - images = convert_from_path(pdf_path, first_page=1, last_page=1) + images = convert_from_path(pdf_path, dpi=400,first_page=1, last_page=1) trimmed_images = trim(images[0]) png_image_bytes = io.BytesIO() trimmed_images.save(png_image_bytes, format='PNG')