fix train acc=0 with filter <pad>

.gitignore

@@ -24,4 +24,5 @@ __pycache__/*
 dataset/test.json
 dataset/train.json
 app/
+models/
 

README.md

@@ -19,13 +19,16 @@ docker run --gpus all -it -v /tmp/trocr-chinese:/trocr-chinese trocr-chinese:lat
 vocab.txt
 1
 2
-3
-4
-5
 ...
 a
 b
 c
+```
+```[python]
+python gen_vocab.py \
+       --dataset_dataset_path "dataset/cust-data/0/*.txt" \
+       --cust_vocab ./cust-data/vocab.txt
+
 ```
 ### 初始化自定义数据集模型
 #### 下载预训练模型trocr模型权重
@@ -62,28 +65,54 @@ python train.py \
        --CUDA_VISIBLE_DEVICES 1
 ```
 
+#### 评估模型
+##### 拷贝checkpoint/trocr-custdata训练完成的pytorch_model.bin 到 ./cust-data/weights 目录下
+
+```[python]
+python eval.py \
+    --dataset_path "./data/cust-data/test/*/*.jpg" \
+    --cust_data_init_weights_path ./cust-data/weights    
+```
+
 ## 测试模型
 ```
 ## 拷贝训练完成的pytorch_model.bin 到 ./cust-data/weights 目录下
-index = 2300
+index = 2300 ##选择最好的或者最后一个step模型
 cp ./checkpoint/trocr-custdata/checkpoint-$index/pytorch_model.bin ./cust-data/weights
-python app.py --test_img test/test.jpg
+python app.py --cust_data_init_weights_path ./cust-data/weights --test_img test/test.jpg
 ```
 
 ## 预训练模型
 | 模型        | cer(字符错误率)           | acc(文本行)  | 下载地址  |训练数据来源 |训练耗时(GPU:3090) | 
 | ------------- |:-------------:| -----:|-----:|-----:|-----:|
-| hand-write(中文手写)      |0.011 | 0.940 |链接: https://pan.baidu.com/s/19f7iu9tLHkcT_zpi3UfqLQ  密码: punl |https://aistudio.baidu.com/aistudio/datasetdetail/102884/0 |8.5h|
-| 印章识别      |- | - |- |- |
-| im2latex(数学公式识别)      |- | - |- |https://zenodo.org/record/56198#.YkniL25Bx_S |
-| 表格识别      |- | - |- |链接：https://pan.baidu.com/s/1V0NT2XmQDDb0mHQlw7V7_w 提取码：oo4a |
+| hand-write(中文手写)      |0.011 | 0.940 |[hand-write](https://pan.baidu.com/s/19f7iu9tLHkcT_zpi3UfqLQ)  密码: punl |[数据集地址](https://aistudio.baidu.com/aistudio/datasetdetail/102884/0) |8.5h(10epoch)|
+| seal(印章识别)      |- | - |- |- |
+| im2latex(数学公式识别)      |- | - |- |[im2latex](https://zenodo.org/record/56198#.YkniL25Bx_S) ||
+| TAL_OCR_TABLE(表格识别)     |- | - |- |[TAL_OCR_TABLE](https://ai.100tal.com/dataset) |
+| TAL_OCR_MATH(小学低年级算式数据集)|- | - |- | [TAL_OCR_MATH](https://ai.100tal.com/dataset) |
+| TAL_OCR_CHN(手写中文数据集)|- | - |- | [TAL_OCR_CHN](https://ai.100tal.com/dataset) ||
+| HME100K(手写公式)|- | - |- | [HME100K](https://ai.100tal.com/dataset) |
 
 备注:后续所有模型会开源在这个目录下链接,可以自由下载. https://pan.baidu.com/s/1uSdWQhJPEy2CYoEULoOhRA  密码: vwi2
 ### 模型调用 
+#### 手写识别
+![image](img/hand.png)
 ```
 unzip hand-write.zip 
 python app.py --cust_data_init_weights_path hand-write --test_img test/hand.png
+
+## output: '醒我的昏迷,偿还我的天真。'
 ```
+
+#### 打印公式识别
+![image](img/im2latex.png)
+```
+unzip im2latex.zip 
+python app.py --cust_data_init_weights_path im2latex --test_img test/im2latex.png
+
+```
+
+
 ## 捐助
 如果此项目给您的工作带来了帮忙，希望您能贡献自己微薄的爱心,
 该项目的每一份收入将用着福利事业，每一季度在issues上公布捐赠明细!   

app.py

@@ -32,4 +32,4 @@ if __name__ == '__main__':
         generated_ids = model.generate(pixel_values[:, :, :].cpu())
 
     generated_text = decode_text(generated_ids[0].cpu().numpy(), vocab, vocab_inp)
-    print('time take:', round(time.time() - t, 2), "s ocr:", [generated_text.replace(' ', '\n')])
+    print('time take:', round(time.time() - t, 2), "s ocr:", [generated_text])

dataset.py

@@ -29,7 +29,13 @@ class trocrDataset(Dataset):
             txt_file = os.path.splitext(image_file)[0]+'.txt'
 
             with open(txt_file) as f:
-                    text = f.read().strip().replace('xa0','')
+                text = f.read().strip().replace('xa0','')
+                if text.startswith('[') and text.endswith(']'):
+                    ##list
+                    try:
+                       text = json.loads(text)
+                    except:
+                         pass
 
             image = Image.open(image_file).convert("RGB")
             image = self.transformer(image) ##图像增强函数
@@ -45,10 +51,13 @@ class trocrDataset(Dataset):
 
 def encode_text(text, max_target_length=128, vocab=None):
     """
+    ##自持自定义 list: ['<td>',"3","3",'</td>',....]
     {'input_ids': [0, 1092, 2, 1, 1],
     'attention_mask': [1, 1, 1, 0, 0]}
     """
-    text = list(text)
+    if type(text) is not list:
+       text = list(text)
+
     text = text[:max_target_length - 2]
     tokens = [vocab.get('<s>')]
     unk = vocab.get('<unk>')
@@ -76,9 +85,10 @@ def decode_text(tokens, vocab, vocab_inp):
     s_start = vocab.get('<s>')
     s_end = vocab.get('</s>')
     unk = vocab.get('<unk>')
+    pad = vocab.get('<pad>')
     text = ''
     for tk in tokens:
-        if tk not in [s_end, s_start]:
+        if tk not in [s_end, s_start , pad, unk]:
            text += vocab_inp[tk]
 
     return text

eval.py

@@ -1,6 +1,4 @@
 import os
-
-os.environ["CUDA_VISIBLE_DEVICES"] = '0'
 from PIL import Image
 import numpy as np
 import time
@@ -29,17 +27,16 @@ def compute_metrics(pred_str, label_str):
 
 
 if __name__ == '__main__':
-    parser = argparse.ArgumentParser(description='trocr fine-tune训练')
+    parser = argparse.ArgumentParser(description='trocr 模型评估')
     parser.add_argument('--cust_data_init_weights_path', default='./cust-data/weights', type=str,
                         help="初始化训练权重，用于自己数据集上fine-tune权重")
     parser.add_argument('--CUDA_VISIBLE_DEVICES', default='-1', type=str, help="GPU设置")
-    parser.add_argument('--test_img', default='test/test.jpg', type=str, help="img path")
     parser.add_argument('--dataset_path', default='dataset/HW-hand-write/HW_Chinese/*/*.[j|p]*', type=str,
                         help="img path")
-    parser.add_argument('--random_state', default=10086, type=int, help="用于训练集划分的随机数")
+    parser.add_argument('--random_state', default=None, type=int, help="用于训练集划分的随机数")
 
     args = parser.parse_args()
-
+    os.environ["CUDA_VISIBLE_DEVICES"] = args.CUDA_VISIBLE_DEVICES
     paths = glob(args.dataset_path)
     if args.random_state is not None:
         train_paths, test_paths = train_test_split(paths, test_size=0.05, random_state=args.random_state)

gen_vocab.py

@@ -0,0 +1,25 @@
+import os.path
+from glob import glob
+from tqdm import tqdm
+import codecs
+import argparse
+if __name__=='__main__':
+    parser = argparse.ArgumentParser(description='trocr vocab生成')
+    parser.add_argument('--cust_vocab', default="./cust-data/vocab.txt", type=str, help="自定义vocab文件生成")
+    parser.add_argument('--dataset_path', default="./dataset/train/*/*.jpg", type=str, help="自定义训练数字符集")
+    args = parser.parse_args()
+    paths = glob(args.dataset_path)
+    vocab = set()
+    for p in tqdm(paths):
+        with codecs.open(p, encoding='utf-8') as f:
+            txt = f.read().strip()
+        vocab.update(txt)
+    root_path = os.path.split(args.cust_vocab)
+    os.makedirs(root_path, exist_ok=True)
+    with open(args.cust_vocab, 'w') as f:
+        f.write('\n'.join(list(vocab)))
+
+
+
+
+

train.py

@@ -22,8 +22,11 @@ def compute_metrics(pred):
     labels_ids[labels_ids == -100] = processor.tokenizer.pad_token_id
     label_str = [decode_text(labels_id, vocab, vocab_inp) for labels_id in labels_ids]
     cer = cer_metric.compute(predictions=pred_str, references=label_str)
+
     acc = [pred == label for pred, label in zip(pred_str, label_str)]
+    print([pred_str[0], label_str[0]])
     acc = sum(acc)/(len(acc)+0.000001)
+
     return {"cer": cer, "acc": acc}
 
 if __name__ == '__main__':
@@ -34,10 +37,11 @@ if __name__ == '__main__':
     parser.add_argument('--dataset_path', default='./dataset/cust-data/*/*.jpg', type=str, help="训练数据集")
     parser.add_argument('--per_device_train_batch_size', default=32, type=int, help="train batch size")
     parser.add_argument('--per_device_eval_batch_size', default=8, type=int, help="eval batch size")
+    parser.add_argument('--max_target_length', default=128, type=int, help="训练文字字符数")
 
     parser.add_argument('--num_train_epochs', default=10, type=int, help="训练epoch数")
-    parser.add_argument('--eval_steps', default=5000, type=int, help="模型评估间隔数")
-    parser.add_argument('--save_steps', default=500, type=int, help="模型保存间隔步数")
+    parser.add_argument('--eval_steps', default=1000, type=int, help="模型评估间隔数")
+    parser.add_argument('--save_steps', default=1000, type=int, help="模型保存间隔步数")
 
     parser.add_argument('--CUDA_VISIBLE_DEVICES', default='0,1', type=str, help="GPU设置")
 
@@ -56,8 +60,8 @@ if __name__ == '__main__':
     vocab = processor.tokenizer.get_vocab()
     vocab_inp = {vocab[key]: key for key in vocab}
 
-    train_dataset = trocrDataset(paths=train_paths, processor=processor)
-    eval_dataset = trocrDataset(paths=test_paths, processor=processor)
+    train_dataset = trocrDataset(paths=train_paths, processor=processor, max_target_length=args.max_target_length)
+    eval_dataset = trocrDataset(paths=test_paths, processor=processor, max_target_length=args.max_target_length)
 
     model = VisionEncoderDecoderModel.from_pretrained(args.cust_data_init_weights_path)
     model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
@@ -99,6 +103,8 @@ if __name__ == '__main__':
         data_collator=default_data_collator,
     )
     trainer.train()
+    trainer.save_model(os.path.join(args.checkpoint_path, 'last'))
+    processor.save_pretrained(os.path.join(args.checkpoint_path, 'last'))