update README
This commit is contained in:
三洋三洋
@@ -25,10 +25,10 @@ TexTeller用了~~550K~~7.5M的图片-公式对进行训练(数据集可以在[
|
||||
## 🔄 变更信息
|
||||
|
||||
* 📮[2024-03-25] TexTeller2.0发布！TexTeller2.0的训练数据增大到了7.5M(相较于TexTeller1.0**增加了~15倍**并且数据质量也有所改善)。训练后的TexTeller2.0在测试集中展现出了**更加优越的性能**，尤其在生僻符号、复杂多行、矩阵的识别场景中。
|
||||
> 在[这里](./test.pdf)有更多的测试图片以及各家识别模型的横向对比。
|
||||
|
||||
* 📮[2024-04-11] 增加了整图推理的功能，只需额外安装onnxruntime库即可获取新功能！我们自行标注了3415张中文教材图片中的公式，并使用了8272张来自于IBEM英文论文公式检测数据集中的公式，基于RT-DETR-R50模型进行了公式目标检测的训练，并将训练好的模型导出为了onnx格式。以方便输入图片，一次性对图片中的所有公式进行识别。
|
||||
|
||||
> 在[这里](./test.pdf)有更多的测试图片以及各家识别模型的横向对比。
|
||||
>
|
||||
* 📮[2024-04-12] 训练了**公式检测模型**，从而增加了对整个文档进行公式检测+公式识别（整图推理）的功能！
|
||||
|
||||
## 🔑 前置条件
|
||||
|
||||
@@ -43,25 +43,22 @@ python=3.10
|
||||
|
||||
1. 克隆本仓库:
|
||||
|
||||
```bash
|
||||
git clone https://github.com/OleehyO/TexTeller
|
||||
```
|
||||
|
||||
```bash
|
||||
git clone https://github.com/OleehyO/TexTeller
|
||||
```
|
||||
2. [安装pytorch](https://pytorch.org/get-started/locally/#start-locally)
|
||||
|
||||
3. 安装本项目的依赖包:
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
4. 进入 `TexTeller/src`目录，在终端运行以下命令进行推理:
|
||||
|
||||
4. 进入`TexTeller/src`目录，在终端运行以下命令进行推理:
|
||||
|
||||
```bash
|
||||
python inference.py -img "/path/to/image.{jpg,png}"
|
||||
# use -cuda option to enable GPU inference
|
||||
#+e.g. python inference.py -img "./img.jpg" -cuda
|
||||
```
|
||||
```bash
|
||||
python inference.py -img "/path/to/image.{jpg,png}"
|
||||
# use -cuda option to enable GPU inference
|
||||
#+e.g. python inference.py -img "./img.jpg" -cuda
|
||||
```
|
||||
|
||||
> [!NOTE]
|
||||
> 第一次运行时会在hugging face上下载所需要的checkpoints
|
||||
@@ -72,27 +69,24 @@ python=3.10
|
||||
|
||||
1. 安装huggingface hub包
|
||||
|
||||
```bash
|
||||
pip install -U "huggingface_hub[cli]"
|
||||
```
|
||||
|
||||
```bash
|
||||
pip install -U "huggingface_hub[cli]"
|
||||
```
|
||||
2. 在能连接Hugging Face的机器上下载模型权重:
|
||||
|
||||
```bash
|
||||
huggingface-cli download OleehyO/TexTeller --include "*.json" "*.bin" "*.txt" --repo-type model --local-dir "your/dir/path"
|
||||
```
|
||||
|
||||
3. 把包含权重的目录上传远端服务器，然后把`TexTeller/src/models/ocr_model/model/TexTeller.py`中的`REPO_NAME = 'OleehyO/TexTeller'`修改为`REPO_NAME = 'your/dir/path'`
|
||||
```bash
|
||||
huggingface-cli download OleehyO/TexTeller --include "*.json" "*.bin" "*.txt" --repo-type model --local-dir "your/dir/path"
|
||||
```
|
||||
3. 把包含权重的目录上传远端服务器，然后把 `TexTeller/src/models/ocr_model/model/TexTeller.py`中的 `REPO_NAME = 'OleehyO/TexTeller'`修改为 `REPO_NAME = 'your/dir/path'`
|
||||
|
||||
如果你还想在训练模型时开启evaluate，你需要提前下载metric脚本并上传远端服务器：
|
||||
|
||||
1. 在能连接Hugging Face的机器上下载metric脚本
|
||||
|
||||
```bash
|
||||
huggingface-cli download evaluate-metric/google_bleu --repo-type space --local-dir "your/dir/path"
|
||||
```
|
||||
|
||||
2. 把这个目录上传远端服务器，并在`TexTeller/src/models/ocr_model/utils/metrics.py`中把`evaluate.load('google_bleu')`改为`evaluate.load('your/dir/path/google_bleu.py')`
|
||||
```bash
|
||||
huggingface-cli download evaluate-metric/google_bleu --repo-type space --local-dir "your/dir/path"
|
||||
```
|
||||
2. 把这个目录上传远端服务器，并在 `TexTeller/src/models/ocr_model/utils/metrics.py`中把 `evaluate.load('google_bleu')`改为 `evaluate.load('your/dir/path/google_bleu.py')`
|
||||
|
||||
## 🌐 网页演示
|
||||
|
||||
@@ -102,87 +96,105 @@ python=3.10
|
||||
./start_web.sh
|
||||
```
|
||||
|
||||
在浏览器里输入`http://localhost:8501`就可以看到web demo
|
||||
在浏览器里输入 `http://localhost:8501`就可以看到web demo
|
||||
|
||||
> [!TIP]
|
||||
> 你可以改变`start_web.sh`的默认配置， 例如使用GPU进行推理(e.g. `USE_CUDA=True`) 或者增加beams的数量(e.g. `NUM_BEAM=3`)来获得更高的精确度
|
||||
> 你可以改变 `start_web.sh`的默认配置， 例如使用GPU进行推理(e.g. `USE_CUDA=True`) 或者增加beams的数量(e.g. `NUM_BEAM=3`)来获得更高的精确度
|
||||
|
||||
> [!NOTE]
|
||||
> 对于Windows用户, 请运行 `start_web.bat`文件.
|
||||
|
||||
## 整图推理
|
||||
## 🧠 整图推理
|
||||
|
||||
TexTeller还支持对整张图片进行**公式检测+公式识别**，从而对整图公式进行检测，然后进行批公式识别。
|
||||
|
||||
### 下载权重
|
||||
在8272张IBEM数据集（https://zenodo.org/records/4757865）上训练，并导出的onnx模型：
|
||||
https://huggingface.co/TonyLee1256/texteller_det/resolve/main/rtdetr_r50vd_6x_coco_trained_on_IBEM_en_papers.onnx?download=true
|
||||
在2560张中文教材数据（100+版式）上训练，并导出的onnx模型：
|
||||
https://huggingface.co/TonyLee1256/texteller_det/blob/main/rtdetr_r50vd_6x_coco_trained_on_cn_textbook.onnx
|
||||
|
||||
英文文档公式检测 [[link](https://huggingface.co/TonyLee1256/texteller_det/resolve/main/rtdetr_r50vd_6x_coco_trained_on_IBEM_en_papers.onnx?download=true)]：在8272张[IBEM数据集](https://zenodo.org/records/4757865)上训练得到
|
||||
|
||||
中文文档公式检测 [[link](https://huggingface.co/TonyLee1256/texteller_det/blob/main/rtdetr_r50vd_6x_coco_trained_on_cn_textbook.onnx)]：在2560张中文教材数据(100+版式)上训练得到
|
||||
|
||||
### 公式检测
|
||||
cd TexTeller/src
|
||||
infer_det.py
|
||||
运行后，对整张图中的所有公式进行检测，绘制整图检测结果并保存，并将每一个检测出的目标单独裁剪并保存下来。
|
||||
|
||||
`TexTeller/src`目录下运行以下命令
|
||||
|
||||
```bash
|
||||
python infer_det.py
|
||||
```
|
||||
|
||||
对整张图中的所有公式进行检测，结果保存在 `TexTeller/src/subimages`
|
||||
|
||||
<div align="center">
|
||||
<img src="det_rec.png" width=400>
|
||||
</div>
|
||||
|
||||
### 公式批识别
|
||||
|
||||
在进行**公式检测后**， `TexTeller/src`目录下运行以下命令
|
||||
|
||||
```shell
|
||||
rec_infer_from_crop_imgs.py
|
||||
基于上一步公式检测的结果，对裁剪出的所有公式进行批量识别，将识别结果保存为txt文件。
|
||||
```
|
||||
|
||||
会基于上一步公式检测的结果，对裁剪出的所有公式进行批量识别，将识别结果在 `TexTeller/src/results`中保存为txt文件。
|
||||
|
||||
## 📡 API调用
|
||||
|
||||
我们使用[ray serve](https://github.com/ray-project/ray)来对外提供一个TexTeller的API接口，通过使用这个接口，你可以把TexTeller整合到自己的项目里。要想启动server，你需要先进入`TexTeller/src`目录然后运行以下命令:
|
||||
我们使用[ray serve](https://github.com/ray-project/ray)来对外提供一个TexTeller的API接口，通过使用这个接口，你可以把TexTeller整合到自己的项目里。要想启动server，你需要先进入 `TexTeller/src`目录然后运行以下命令:
|
||||
|
||||
```bash
|
||||
python server.py # default settings
|
||||
```
|
||||
|
||||
你可以给`server.py`传递以下参数来改变server的推理设置(e.g. `python server.py --use_gpu` 来启动GPU推理):
|
||||
你可以给 `server.py`传递以下参数来改变server的推理设置(e.g. `python server.py --use_gpu` 来启动GPU推理):
|
||||
|
||||
| 参数 | 描述 |
|
||||
| --- | --- |
|
||||
| `-ckpt` | 权重文件的路径，*默认为TexTeller的预训练权重*。|
|
||||
| `-tknz` | 分词器的路径， *默认为TexTeller的分词器*。|
|
||||
| `-port` | 服务器的服务端口， *默认是8000*。 |
|
||||
| `--use_gpu` | 是否使用GPU推理，*默认为CPU*。 |
|
||||
| `--num_beams` | beam search的beam数量， *默认是1*。 |
|
||||
| `--num_replicas` | 在服务器上运行的服务副本数量， *默认1个副本*。你可以使用更多的副本来获取更大的吞吐量。|
|
||||
| `--ncpu_per_replica` | 每个服务副本所用的CPU核心数，*默认为1*。 |
|
||||
| 参数 | 描述 |
|
||||
| ---------------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| `-ckpt` | 权重文件的路径，*默认为TexTeller的预训练权重*。 |
|
||||
| `-tknz` | 分词器的路径，*默认为TexTeller的分词器*。 |
|
||||
| `-port` | 服务器的服务端口，*默认是8000*。 |
|
||||
| `--use_gpu` | 是否使用GPU推理，*默认为CPU*。 |
|
||||
| `--num_beams` | beam search的beam数量，*默认是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可用) |
|
||||
|
||||
> [!NOTE]
|
||||
> 一个客户端demo可以在`TexTeller/client/demo.py`找到，你可以参考`demo.py`来给server发送请求
|
||||
> 一个客户端demo可以在 `TexTeller/client/demo.py`找到，你可以参考 `demo.py`来给server发送请求
|
||||
|
||||
## 🏋️♂️ 训练
|
||||
|
||||
### 数据集
|
||||
|
||||
我们在`TexTeller/src/models/ocr_model/train/dataset`目录中提供了一个数据集的例子，你可以把自己的图片放在`images`目录然后在`formulas.jsonl`中为每张图片标注对应的公式。
|
||||
我们在 `TexTeller/src/models/ocr_model/train/dataset`目录中提供了一个数据集的例子，你可以把自己的图片放在 `images`目录然后在 `formulas.jsonl`中为每张图片标注对应的公式。
|
||||
|
||||
准备好数据集后，你需要在`.../dataset/loader.py`中把 **`DIR_URL`变量改成你自己数据集的路径**
|
||||
准备好数据集后，你需要在 `.../dataset/loader.py`中把 **`DIR_URL`变量改成你自己数据集的路径**
|
||||
|
||||
### 重新训练分词器
|
||||
|
||||
如果你使用了不一样的数据集，你可能需要重新训练tokenizer来得到一个不一样的字典。配置好数据集后，可以通过以下命令来训练自己的tokenizer：
|
||||
|
||||
1. 在`TexTeller/src/models/tokenizer/train.py`中，修改`new_tokenizer.save_pretrained('./your_dir_name')`为你自定义的输出目录
|
||||
> 注意：如果要用一个不一样大小的字典(默认1W个token)，你需要在 `TexTeller/src/models/globals.py`中修改`VOCAB_SIZE`变量
|
||||
1. 在 `TexTeller/src/models/tokenizer/train.py`中，修改 `new_tokenizer.save_pretrained('./your_dir_name')`为你自定义的输出目录
|
||||
|
||||
> 注意：如果要用一个不一样大小的字典(默认1W个token)，你需要在 `TexTeller/src/models/globals.py`中修改 `VOCAB_SIZE`变量
|
||||
>
|
||||
2. **在 `TexTeller/src` 目录下**运行以下命令:
|
||||
|
||||
```bash
|
||||
python -m models.tokenizer.train
|
||||
```
|
||||
```bash
|
||||
python -m models.tokenizer.train
|
||||
```
|
||||
|
||||
### 训练模型
|
||||
|
||||
要想训练模型, 你需要在`TexTeller/src`目录下运行以下命令：
|
||||
要想训练模型, 你需要在 `TexTeller/src`目录下运行以下命令：
|
||||
|
||||
```bash
|
||||
python -m models.ocr_model.train.train
|
||||
```
|
||||
|
||||
你可以在`TexTeller/src/models/ocr_model/train/train.py`中设置自己的tokenizer和checkpoint路径（请参考`train.py`）。如果你使用了与TexTeller一样的架构和相同的字典，你还可以用自己的数据集来微调TexTeller的默认权重。
|
||||
你可以在 `TexTeller/src/models/ocr_model/train/train.py`中设置自己的tokenizer和checkpoint路径（请参考 `train.py`）。如果你使用了与TexTeller一样的架构和相同的字典，你还可以用自己的数据集来微调TexTeller的默认权重。
|
||||
|
||||
在`TexTeller/src/globals.py`和`TexTeller/src/models/ocr_model/train/train_args.py`中，你可以改变模型的架构以及训练的超参数。
|
||||
在 `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)来获取更多训练参数的细节以及配置。
|
||||
@@ -190,19 +202,14 @@ python -m models.ocr_model.train.train
|
||||
## 🚧 不足
|
||||
|
||||
* 不支持扫描图片以及PDF文档识别
|
||||
|
||||
* 不支持手写体公式
|
||||
|
||||
## 📅 计划
|
||||
|
||||
- [x] ~~使用更大的数据集来训练模型(7.5M样本，即将发布)~~
|
||||
|
||||
- [X] ~~使用更大的数据集来训练模型(7.5M样本，即将发布)~~
|
||||
- [ ] 扫描图片识别
|
||||
|
||||
- [ ] PDF文档识别 + 中英文场景支持
|
||||
|
||||
- [ ] 推理加速
|
||||
|
||||
- [ ] ...
|
||||
|
||||
## 💖 感谢
|
||||
|
||||
Reference in New Issue
Block a user