TexTeller v2

README.md

@@ -7,9 +7,12 @@
     <p align="center">
         English | <a href="./assets/README_zh.md">中文</a>
     </p>
-    <p align="center">
+    <!-- <p align="center">
         <img src="./assets/web_demo.gif" alt="TexTeller_demo" width=800>
-    </p>
+    </p> -->
+    <video width="800" controls>
+        <source src="./assets/test.mp4" type="video/mp4">
+    </video>
 </div>
 
 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

assets/README_zh.md

@@ -7,9 +7,12 @@
     <p align="center">
         <a href="../README.md">English</a> | 中文
     </p>
-    <p align="center">
+    <!-- <p align="center">
         <img src="./web_demo.gif" alt="TexTeller_demo" width=800>
-    </p>
+    </p> -->
+    <video width="800" controls>
+        <source src="./test.mp4" type="video/mp4">
+    </video>
 </div>
 
 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%的准确率
-
-    <img src="" width=30, height=30> 
-
-* 部分复杂的大型多行公式识别效果不佳(例如长公式与矩阵混合)
-
-    <img src="" width=30, height=30> 
-
-    > 如果遇到这种情况，你可以尝试把大型的多行公式分成多个小的子公式来识别。
-
 * 不支持扫描图片以及PDF文档识别
 
 * 不支持手写体公式

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,

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

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')