update

.gitattributes

@@ -1,3 +1,29 @@
-*.pt filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
 dataset/* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,27 @@
+#Python:
+.py[cod]
+*.pyc
+*.so
+*.egg
+*.egg-info
+*.pth
+*.h5
+*.pbtxt
+*.hdf5
+*.weights
+*.trainedata
+*.xlsx
+dist*
+build
+.DS_Store*
+.ipynb_checkpoints
+__pycahe__
+dump.rdb
+define
+.idea
+__pycache__/*
+*.bert_model.*
+dataset/test.json
+dataset/train.json
+app/
+

Dockerfile

@@ -0,0 +1,6 @@
+FROM pytorch/pytorch:1.7.1-cuda11.0-cudnn8-devel
+LABEL version="0.1"  description="trocr-chinese" by="chineseocr"
+ADD requirements.txt /trocr-chinese/requirements.txt
+RUN pip install -r /trocr-chinese/requirements.txt -i http://mirrors.aliyun.com/pypi/simple/ --trusted-host mirrors.aliyun.com
+WORKDIR /trocr-chinese
+CMD ["python", "app.py"]

LICENSE

@@ -1,21 +0,0 @@
-MIT License
-
-Copyright (c) 2022 chineseocr
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
-SOFTWARE.

README.md

@@ -0,0 +1,89 @@
+# 基于trocr(beit+roberta)实现对中文场景文字识别
+trocr原地址(https://github.com/microsoft/unilm/tree/master/trocr)
+## 实现功能
+- [x]  单行/多行文字/横竖排文字识别
+- [x]  不规则文字（印章，公式等）
+- [ ]  表格识别
+- [ ]  模型蒸馏/DML(协作学习)
+- [ ]  Prompt Learning
+## 环境编译
+```
+docker build --network=host -t trocr-chinese:latest .
+docker run --gpus all -it -v /tmp/trocr-chinese:/trocr-chinese trocr-chinese:latest bash
+
+```
+## 训练
+### 初始化模型到自定义训练数据集
+#### 字符集准备参考cust-data/vocab.txt
+```
+vocab.txt
+1
+2
+3
+4
+5
+...
+a
+b
+c
+```
+### 初始化自定义数据集模型
+#### 下载预训练模型trocr模型权重
+链接: https://pan.baidu.com/s/1rARdfadQlQGKGHa3de82BA  密码: 0o65
+```
+python init_custdata_model.py \   
+    --cust_vocab ./cust-data/vocab.txt \  
+    --pretrain_model ./weights \
+    --cust_data_init_weights_path ./cust-data/weights
+```
+## cust_vocab 词库文件   
+## pretrain_model 预训练模型权重   
+## cut_data_init_weights_path 自定义模型初始化模型权重保存位置   
+
+### 训练模型
+#### 数据准备,数据结构如下图所示
+```
+dataset/cust-data/0/0.jpg
+dataset/cust-data/0/0.txt
+...
+dataset/cust-data/100/10000.jpg
+dataset/cust-data/100/10000.txt
+```
+
+#### 训练模型
+```
+python train.py \
+       --cut_data_init_weights_path ./cust-data/weights \
+       --checkpoint_path ./checkpoint/trocr-custdata \
+       --dataset_path "./dataset/cust-data/*/*.jpg" \
+       --per_device_train_batch_size 8 \
+       --CUDA_VISIBLE_DEVICES 1
+```
+
+## 测试模型
+```
+## 拷贝训练完成的pytorch_model.bin 到 ./cust-data/weights 目录下
+index = 2300
+cp ./checkpoint/trocr-custdata/checkpoint-$index/pytorch_model.bin ./cust-data/weights
+python app.py --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 |
+
+备注:后续所有模型会开源在这个目录下链接,可以自由下载. https://pan.baidu.com/s/1uSdWQhJPEy2CYoEULoOhRA  密码: vwi2
+### 模型调用 
+```
+unzip hand-write.zip 
+python app.py --cust_data_init_weights_path hand-write --test_img test/hand.png
+```
+## 捐助
+如果此项目给您的工作带来了帮忙，希望您能贡献自己微薄的爱心,
+该项目的每一份收入将用着福利事业，每一季度在issues上公布捐赠明细!   
+![image](img/chat.jpg
+

app.py

@@ -0,0 +1,35 @@
+import os
+from PIL import Image
+import time
+import torch
+import argparse
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+from dataset import decode_text
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='trocr fine-tune训练')
+    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")
+
+    args = parser.parse_args()
+
+    processor = TrOCRProcessor.from_pretrained(args.cust_data_init_weights_path)
+    vocab = processor.tokenizer.get_vocab()
+
+    vocab_inp = {vocab[key]: key for key in vocab}
+    model = VisionEncoderDecoderModel.from_pretrained(args.cust_data_init_weights_path)
+    model.eval()
+
+    vocab = processor.tokenizer.get_vocab()
+    vocab_inp = {vocab[key]: key for key in vocab}
+
+    t = time.time()
+    img = Image.open(args.test_img).convert('RGB')
+    pixel_values = processor([img], return_tensors="pt").pixel_values
+    with torch.no_grad():
+        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')])

cer.py

@@ -0,0 +1,159 @@
+# coding=utf-8
+# Copyright 2021 The HuggingFace Datasets Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Character Error Ratio (CER) metric. """
+
+from typing import List
+
+import jiwer
+import jiwer.transforms as tr
+from packaging import version
+
+import datasets
+from datasets.config import PY_VERSION
+
+
+if PY_VERSION < version.parse("3.8"):
+    import importlib_metadata
+else:
+    import importlib.metadata as importlib_metadata
+
+
+SENTENCE_DELIMITER = ""
+
+
+if version.parse(importlib_metadata.version("jiwer")) < version.parse("2.3.0"):
+
+    class SentencesToListOfCharacters(tr.AbstractTransform):
+        def __init__(self, sentence_delimiter: str = " "):
+            self.sentence_delimiter = sentence_delimiter
+
+        def process_string(self, s: str):
+            return list(s)
+
+        def process_list(self, inp: List[str]):
+            chars = []
+            for sent_idx, sentence in enumerate(inp):
+                chars.extend(self.process_string(sentence))
+                if self.sentence_delimiter is not None and self.sentence_delimiter != "" and sent_idx < len(inp) - 1:
+                    chars.append(self.sentence_delimiter)
+            return chars
+
+    cer_transform = tr.Compose(
+        [tr.RemoveMultipleSpaces(), tr.Strip(), SentencesToListOfCharacters(SENTENCE_DELIMITER)]
+    )
+else:
+    cer_transform = tr.Compose(
+        [
+            tr.RemoveMultipleSpaces(),
+            tr.Strip(),
+            tr.ReduceToSingleSentence(SENTENCE_DELIMITER),
+            tr.ReduceToListOfListOfChars(),
+        ]
+    )
+
+
+_CITATION = """\
+@inproceedings{inproceedings,
+    author = {Morris, Andrew and Maier, Viktoria and Green, Phil},
+    year = {2004},
+    month = {01},
+    pages = {},
+    title = {From WER and RIL to MER and WIL: improved evaluation measures for connected speech recognition.}
+}
+"""
+
+_DESCRIPTION = """\
+Character error rate (CER) is a common metric of the performance of an automatic speech recognition system.
+
+CER is similar to Word Error Rate (WER), but operates on character instead of word. Please refer to docs of WER for further information.
+
+Character error rate can be computed as:
+
+CER = (S + D + I) / N = (S + D + I) / (S + D + C)
+
+where
+
+S is the number of substitutions,
+D is the number of deletions,
+I is the number of insertions,
+C is the number of correct characters,
+N is the number of characters in the reference (N=S+D+C).
+
+CER's output is not always a number between 0 and 1, in particular when there is a high number of insertions. This value is often associated to the percentage of characters that were incorrectly predicted. The lower the value, the better the
+performance of the ASR system with a CER of 0 being a perfect score.
+"""
+
+_KWARGS_DESCRIPTION = """
+Computes CER score of transcribed segments against references.
+Args:
+    references: list of references for each speech input.
+    predictions: list of transcribtions to score.
+    concatenate_texts: Whether or not to concatenate sentences before evaluation, set to True for more accurate result.
+Returns:
+    (float): the character error rate
+
+Examples:
+
+    >>> predictions = ["this is the prediction", "there is an other sample"]
+    >>> references = ["this is the reference", "there is another one"]
+    >>> cer = datasets.load_metric("cer")
+    >>> cer_score = cer.compute(predictions=predictions, references=references)
+    >>> print(cer_score)
+    0.34146341463414637
+"""
+
+
+@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
+class CER(datasets.Metric):
+    def _info(self):
+        return datasets.MetricInfo(
+            description=_DESCRIPTION,
+            citation=_CITATION,
+            inputs_description=_KWARGS_DESCRIPTION,
+            features=datasets.Features(
+                {
+                    "predictions": datasets.Value("string", id="sequence"),
+                    "references": datasets.Value("string", id="sequence"),
+                }
+            ),
+            codebase_urls=["https://github.com/jitsi/jiwer/"],
+            reference_urls=[
+                "https://en.wikipedia.org/wiki/Word_error_rate",
+                "https://sites.google.com/site/textdigitisation/qualitymeasures/computingerrorrates",
+            ],
+        )
+
+    def _compute(self, predictions, references, concatenate_texts=False):
+        if concatenate_texts:
+            return jiwer.compute_measures(
+                references,
+                predictions,
+                truth_transform=cer_transform,
+                hypothesis_transform=cer_transform,
+            )["wer"]
+
+        incorrect = 0
+        total = 0
+        for prediction, reference in zip(predictions, references):
+            measures = jiwer.compute_measures(
+                reference,
+                prediction,
+                truth_transform=cer_transform,
+                hypothesis_transform=cer_transform,
+            )
+            incorrect += measures["substitutions"] + measures["deletions"] + measures["insertions"]
+            total += measures["substitutions"] + measures["deletions"] + measures["hits"]
+
+        return incorrect / total

cust-data/vocab.txt

@@ -0,0 +1,3076 @@
+线
+处
+失
+守
+边
+缘
+同
+花
+业
+发
+车
+补
+贴
+两
+千
+余
+从
+年
+审
+计
+署
+公
+布
+的
+障
+和
+改
+善
+民
+生
+头
+开
+了
+一
+个
+洞
+当
+他
+工
+作
+规
+格
+般
+为
+副
+情
+况
+授
+予
+安
+监
+库
+房
+用
+绳
+索
+抬
+区
+复
+赛
+刚
+在
+上
+邦
+数
+今
+日
+走
+势
+来
+看
+成
+本
+已
+经
+于
+束
+机
+制
+不
+健
+全
+企
+下
+调
+煤
+炭
+价
+以
+应
+瓶
+可
+乐
+然
+后
+再
+找
+美
+国
+运
+至
+广
+州
+港
+关
+单
+据
+出
+口
+流
+中
+心
+天
+基
+间
+人
+通
+过
+百
+分
+点
+期
+贷
+内
+查
+王
+挺
+革
+放
+础
+带
+动
+东
+北
+振
+兴
+武
+事
+早
+完
+万
+并
+未
+其
+合
+我
+们
+实
+需
+要
+认
+大
+连
+低
+款
+利
+率
+却
+时
+是
+前
+搭
+着
+高
+油
+展
+科
+技
+产
+山
+西
+牌
+银
+行
+多
+赚
+元
+李
+力
+男
+岁
+新
+晋
+商
+形
+象
+境
+比
+例
+与
+方
+家
+说
+故
+原
+因
+目
+正
+效
+能
+存
+愿
+偿
+尚
+启
+爱
+好
+这
+包
+括
+管
+特
+别
+现
+代
+犯
+罪
+泰
+金
+表
+示
+昨
+有
+主
+学
+医
+疗
+照
+顾
+谢
+赔
+等
+横
+幅
+对
+起
+指
+控
+矢
+难
+耐
+寂
+寞
+站
+街
+女
+救
+援
+员
+戈
+壁
+腹
+职
+各
+岗
+述
+她
+就
+死
+给
+添
+博
+裕
+祥
+长
+创
+升
+欧
+四
+架
+气
+很
+袁
+龙
+立
+足
+司
+重
+资
+组
+瞻
+性
+进
+步
+增
+相
+位
+杀
+跌
+政
+部
+马
+尔
+二
+程
+收
+购
+快
+此
+定
+会
+记
+够
+水
+空
+样
+设
+法
+修
+案
+征
+求
+意
+状
+态
+汇
+丰
+采
+忘
+参
+加
+专
+队
+刘
+周
+预
+披
+露
+知
+道
+爆
+么
+物
+试
+退
+臣
+名
+被
+近
+拖
+去
+尼
+娃
+友
+少
+次
+风
+灾
+习
+惯
+庞
+思
+几
+排
+环
+保
+护
+挥
+才
+问
+初
+判
+结
+论
+如
+码
+仓
+堆
+浮
+准
+识
+系
+统
+核
+身
+堪
+入
+手
+短
+信
+送
+欢
+层
+活
+汉
+市
+容
+卫
+嫌
+疑
+缴
+获
+五
+世
+及
+罗
+密
+朱
+批
+省
+级
+文
+电
+者
+之
+希
+腊
+错
+绝
+歉
+但
+亚
+访
+往
+约
+旦
+料
+月
+衣
+衫
+踩
+雷
+叶
+檀
+央
+降
+息
+局
+华
+网
+海
+陆
+联
+追
+也
+渐
+警
+里
+踪
+乏
+神
+小
+张
+扩
+食
+品
+检
+门
+委
+火
+距
+离
+枪
+击
+鉴
+值
+托
+模
+式
+差
+戴
+呼
+吸
+器
+建
+夫
+球
+场
+类
+双
+将
+得
+财
+告
+讯
+号
+织
+乌
+恰
+县
+晚
+江
+苏
+顺
+十
+青
+画
+落
+河
+显
+撑
+住
+庆
+吁
+协
+研
+胶
+南
+持
+客
+平
+台
+仅
+提
+供
+浦
+劳
+务
+派
+遣
+圈
+而
+钟
+共
+服
+施
+左
+隶
+受
+到
+支
+报
+韩
+旭
+叙
+阵
+雨
+总
+理
+助
+度
+更
+尽
+超
+额
+益
+谣
+言
+何
+昌
+致
+词
+取
+措
+炒
+股
+话
+宏
+观
+济
+使
+怪
+异
+抢
+蛋
+糕
+滑
+坡
+造
+校
+置
+志
+地
+还
+英
+刀
+捅
+逃
+件
+须
+悲
+份
+雅
+虎
+侠
+良
+果
+阿
+酋
+赞
+阶
+梯
+富
+党
+消
+费
+债
+减
+负
+申
+请
+旁
+听
+析
+士
+忠
+德
+震
+执
+读
+卖
+闻
+户
+遥
+自
+宜
+条
+选
+吧
+背
+居
+融
+围
+米
+直
+径
+清
+晰
+税
+策
+货
+若
+社
+责
+任
+子
+城
+乡
+享
+变
+币
+拉
+伯
+板
+块
+辞
+战
+斗
+非
+称
+迪
+先
+宁
+让
+土
+闲
+见
+简
+面
+际
+脚
+链
+锁
+险
+奶
+牛
+养
+殖
+京
+蓝
+色
+湾
+累
+宝
+尾
+段
+陈
+桔
+砸
+吨
+秦
+皇
+岛
+卡
+投
+图
+治
+买
+所
+界
+划
+剂
+做
+承
+诺
+均
+衡
+媒
+体
+些
+向
+证
+午
+三
+菱
+俄
+斯
+外
+交
+拳
+打
+伤
+霍
+姆
+哈
+涨
+廉
+担
+随
+量
+谓
+锅
+铁
+游
+饮
+威
+胁
+绩
+传
+播
+都
+最
+诉
+书
+节
+反
+明
+凌
+晨
+捡
+钱
+诈
+骗
+艺
+解
+决
+患
+纠
+纷
+由
+拿
+回
+暴
+咬
+断
+鹿
+肉
+府
+秀
+星
+耶
+古
+仍
+驾
+艇
+命
+介
+绍
+营
+停
+渠
+究
+似
+害
+众
+康
+梁
+拒
+农
+村
+耳
+伊
+防
+潜
+无
+含
+泪
+档
+困
+群
+另
+伙
+则
+窗
+响
+声
+刻
+殊
+域
+捕
+淄
+底
+项
+亿
+术
+违
+整
+凤
+镇
+宽
+松
+牺
+牲
+官
+辜
+路
+丹
+荣
+景
+许
+唱
+石
+膏
+易
+评
+川
+胡
+楼
+冰
+缺
+朋
+欲
+付
+标
+驶
+限
+较
+络
+第
+销
+种
+沸
+扬
+贸
+签
+没
+哺
+化
+切
+历
+史
+悠
+穿
+白
+配
+稿
+军
+转
+型
+宣
+缓
+质
+依
+什
+验
+满
+痛
+哭
+您
+票
+常
+属
+陪
+缉
+私
+每
+掌
+握
+矛
+装
+污
+院
+脱
+麦
+续
+你
+亡
+考
+根
+集
+即
+返
+召
+攻
+始
+必
+季
+太
+强
+按
+询
+源
+洋
+注
+焦
+真
+敢
+啊
+硬
+只
+燃
+招
+券
+薄
+弱
+争
+夺
+纪
+录
+该
+紧
+鲜
+扶
+坚
+师
+垄
+倍
+夜
+谭
+达
+罢
+休
+具
+屏
+幕
+深
+宗
+丧
+题
+卑
+姐
+曲
+阜
+孔
+林
+附
+傅
+触
+碰
+胸
+夏
+餐
+洲
+隔
+教
+堂
+福
+渡
+举
+沙
+岸
+驳
+馆
+首
+套
+洗
+侨
+席
+便
+老
+戒
+义
+谋
+迅
+速
+锣
+鼓
+精
+混
+挑
+扎
+推
+维
+探
+确
+列
+红
+冲
+输
+适
+慰
+促
+耕
+垫
+租
+竟
+允
+导
+宾
+律
+暖
+字
+刺
+汤
+灿
+辗
+骨
+纯
+构
+玲
+极
+甚
+柱
+售
+厘
+筑
+育
+捷
+激
+佩
+云
+齐
+权
+察
+材
+肠
+杆
+菌
+散
+感
+摸
+脑
+优
+袭
+辆
+赣
+伟
+娘
+麻
+详
+细
+豆
+择
+兼
+荷
+兰
+讲
+凶
+塌
+净
+损
+巴
+瘤
+药
+植
+俯
+视
+远
+募
+领
+佐
+算
+珠
+办
+昏
+逮
+功
+典
+折
+秒
+饲
+拥
+疆
+董
+侵
+献
+烟
+帆
+彭
+吃
+片
+龄
+尊
+微
+盘
+趋
+临
+汾
+羊
+辖
+独
+堵
+普
+或
+换
+赢
+迷
+烈
+虽
+母
+雇
+峻
+且
+团
+木
+荐
+接
+航
+鹏
+泽
+抨
+蚀
+朝
+疯
+倒
+苦
+恐
+慌
+又
+垣
+绸
+除
+透
+坛
+旨
+符
+恩
+摄
+像
+扫
+描
+遂
+咨
+想
+吵
+待
+彻
+屿
+甲
+卷
+尘
+染
+乙
+肝
+旧
+辽
+占
+滩
+渔
+船
+曾
+积
+枢
+纽
+盈
+框
+桥
+渣
+店
+末
+聚
+宵
+哪
+厅
+辉
+玻
+璃
+望
+孩
+鬼
+脸
+父
+剧
+议
+陷
+绿
+灯
+拆
+迁
+迈
+严
+孤
+迫
+喝
+湖
+妨
+备
+庭
+秘
+章
+钓
+鱼
+端
+豪
+迟
+跃
+滥
+黄
+宇
+杰
+帮
+玩
+候
+镜
+跑
+抓
+温
+挣
+霜
+袋
+拧
+盖
+邻
+眼
+谨
+慎
+楚
+映
+跟
+途
+滨
+溪
+引
+影
+郭
+亲
+泊
+廊
+坊
+尤
+润
+弯
+凹
+努
+缩
+把
+范
+破
+序
+假
+病
+葡
+萄
+牙
+野
+稳
+释
+裹
+避
+免
+撞
+隐
+蔡
+桃
+突
+移
+肥
+皂
+莎
+遗
+址
+儿
+醒
+毒
+继
+讨
+页
+侧
+舞
+己
+冒
+逼
+汹
+纱
+溃
+败
+驱
+督
+谩
+骂
+侮
+辱
+逾
+涉
+贪
+荡
+婚
+族
+悄
+残
+留
+厨
+它
+既
+妻
+雪
+课
+角
+寻
+吗
+赶
+干
+递
+毁
+频
+塔
+矿
+唐
+津
+仇
+卜
+惠
+归
+洒
+辛
+勤
+九
+射
+炸
+狂
+邮
+嘉
+艾
+缅
+甸
+愈
+演
+胀
+压
+赴
+黑
+岳
+庄
+臭
+粗
+婴
+厂
+抵
+届
+拍
+屋
+暂
+裁
+骚
+令
+危
+腾
+挪
+挂
+汗
+莱
+虚
+拟
+擎
+遭
+冷
+氛
+否
+挖
+姻
+互
+编
+洪
+粮
+赎
+误
+味
+恒
+久
+昆
+仑
+著
+泄
+漏
+罐
+凑
+填
+饱
+肚
+斤
+抗
+香
+竞
+爹
+硫
+硝
+载
+贞
+掉
+阁
+截
+慈
+枚
+箭
+弹
+订
+音
+室
+菲
+甘
+肃
+敦
+叫
+哄
+卿
+克
+顿
+陶
+冬
+怕
+觉
+语
+缮
+玉
+裴
+佳
+廖
+鲁
+伦
+帝
+答
+键
+偏
+怎
+貌
+叉
+巡
+拓
+捐
+赠
+测
+浙
+羽
+越
+略
+驼
+靠
+锋
+旅
+桑
+禁
+印
+止
+纳
+账
+拢
+估
+瘦
+田
+蒲
+轴
+慢
+倾
+斜
+舟
+箱
+杠
+侦
+孝
+襄
+阳
+阴
+笼
+罩
+酒
+惑
+峰
+摆
+僵
+虑
+奔
+沟
+滞
+晓
+逐
+驻
+栋
+墅
+哀
+鸿
+偶
+荒
+彩
+井
+六
+喷
+涌
+永
+盒
+综
+梅
+隆
+跳
+那
+脂
+遍
+储
+急
+坠
+姓
+冯
+忧
+厉
+喜
+梦
+醉
+诚
+溯
+怀
+孕
+逻
+蹲
+幸
+汽
+烹
+掺
+扣
+杂
+树
+枝
+抱
+佛
+塞
+坍
+墙
+镑
+抑
+孜
+勒
+柯
+闽
+版
+逆
+梳
+淡
+骑
+册
+光
+斩
+耀
+戏
+针
+乱
+某
+狱
+戚
+铆
+呈
+澄
+励
+伞
+兵
+狼
+催
+谈
+零
+黎
+澳
+床
+乎
+波
+潮
+韦
+念
+嫩
+籍
+旬
+借
+晤
+延
+伸
+抽
+萌
+翻
+拨
+沪
+惊
+腐
+畴
+拾
+胜
+悬
+崖
+仔
+阅
+俱
+春
+薪
+酬
+粘
+颇
+奈
+血
+液
+癌
+症
+啥
+孙
+暗
+冻
+勇
+垮
+洛
+乃
+呢
+剐
+蹭
+杯
+葛
+阻
+扰
+莲
+悉
+谅
+尖
+斌
+挤
+翁
+辩
+榨
+菜
+邵
+毗
+晒
+乳
+胺
+窖
+丈
+搜
+狐
+右
+缝
+爷
+狠
+揍
+铭
+泼
+粉
+敏
+馈
+七
+悔
+诱
+固
+剩
+娱
+廷
+郊
+绪
+涵
+俩
+萨
+翔
+钢
+鸟
+吕
+刑
+迎
+锐
+诸
+淋
+轮
+概
+割
+履
+泉
+删
+恶
+劣
+揭
+萎
+誓
+妇
+衍
+畜
+牧
+蒙
+徒
+甩
+奢
+侈
+攀
+顶
+肖
+遇
+猎
+炙
+吴
+谷
+摩
+擦
+郎
+滚
+殴
+趴
+灵
+寿
+旗
+腕
+鸥
+八
+沉
+盛
+殷
+莉
+慧
+森
+赖
+汛
+捧
+童
+辐
+半
+亩
+罚
+雾
+稀
+座
+封
+礼
+榜
+曝
+胖
+盲
+坐
+筹
+操
+敌
+忆
+怖
+钞
+猜
+砖
+坑
+匪
+夷
+贵
+洁
+粤
+邀
+跨
+毕
+贝
+侯
+耗
+兄
+弟
+徐
+扮
+炉
+肯
+浪
+陨
+郑
+削
+终
+钻
+亏
+闷
+宫
+盼
+烫
+芋
+谁
+轻
+亨
+桶
+淫
+秽
+轨
+奥
+朗
+瓦
+沂
+逊
+妙
+奖
+厌
+鼠
+纵
+兑
+沿
+倘
+懂
+辣
+牵
+穷
+秃
+斥
+艳
+碍
+胆
+勿
+淘
+鹤
+鸣
+欣
+捏
+歪
+嫁
+汰
+膜
+剖
+炼
+坦
+蛇
+茶
+妹
+循
+繁
+辨
+挨
+涂
+迹
+骄
+傲
+鸡
+疲
+软
+湘
+潭
+蔓
+燕
+扯
+衅
+瑟
+惜
+充
+携
+宋
+乘
+幼
+仿
+潘
+蜜
+垂
+园
+恢
+枯
+猛
+夹
+盾
+嚼
+疾
+函
+宿
+舍
+覆
+盟
+素
+鼻
+蒋
+珍
+摔
+哗
+晴
+漫
+盗
+奎
+拘
+旺
+躺
+浑
+棚
+草
+蓄
+宴
+晕
+骆
+灌
+瘪
+杨
+登
+飞
+丽
+遏
+煮
+秩
+闸
+桂
+伴
+涛
+翼
+坏
+琴
+穆
+霞
+柳
+挫
+胃
+邢
+劝
+卓
+耸
+浓
+笑
+溢
+辅
+逗
+坪
+坝
+忍
+熟
+奋
+裂
+苹
+罕
+拐
+扇
+押
+藏
+膀
+埃
+培
+谱
+趁
+冀
+魏
+巨
+骇
+纸
+奇
+葬
+棺
+埋
+踢
+锈
+喀
+牢
+沈
+啸
+班
+铸
+俗
+摧
+亦
+贩
+皮
+辰
+瑶
+酸
+庚
+妥
+番
+毫
+蔽
+柴
+仪
+拎
+废
+茨
+窜
+岩
+颜
+聊
+勘
+绘
+峡
+霄
+抛
+弃
+帖
+烧
+忙
+插
+灰
+揽
+囊
+盐
+酱
+彰
+哥
+闭
+亢
+哲
+敲
+笔
+邱
+掀
+壳
+偷
+泥
+遮
+厦
+隙
+凯
+赫
+写
+墨
+摘
+仁
+冤
+昭
+丝
+藉
+妓
+裸
+瑞
+丑
+陋
+婆
+浩
+厕
+惩
+沾
+贡
+哑
+铃
+宪
+拼
+刊
+锦
+赵
+搬
+叹
+饭
+摊
+裔
+斡
+旋
+雏
+肤
+囤
+妮
+丛
+衔
+伐
+猪
+蹈
+榔
+捶
+炮
+轰
+篇
+薛
+练
+墩
+谐
+鲸
+喘
+乔
+挡
+刷
+缠
+稽
+吻
+劫
+喉
+览
+竭
+浴
+宅
+亥
+丁
+唯
+蔬
+贯
+训
+呕
+牟
+热
+赁
+壮
+赂
+闯
+渤
+跪
+泛
+柏
+彪
+悍
+稍
+胎
+趣
+妈
+掏
+腰
+仰
+爬
+恤
+亮
+渝
+默
+骤
+戮
+顽
+怨
+氨
+氏
+钠
+莫
+吉
+歌
+娜
+搏
+纺
+钩
+瞧
+敛
+辈
+惟
+凳
+替
+邯
+郸
+膨
+圆
+宙
+熊
+淹
+吹
+陇
+凸
+蒂
+虹
+尺
+摇
+漠
+毛
+艰
+苑
+倚
+梭
+婿
+耘
+姜
+樊
+疼
+惋
+筒
+栏
+朴
+掘
+耻
+扛
+霓
+菏
+惨
+句
+橱
+桌
+椅
+淀
+巢
+疫
+泻
+斑
+憾
+彼
+夕
+怒
+腔
+仙
+弄
+瘫
+霉
+狙
+袖
+巧
+棒
+皆
+遵
+弗
+贫
+糊
+掖
+毯
+疏
+弥
+嘲
+讽
+孟
+咱
+凭
+赋
+肿
+彦
+饼
+寄
+智
+搞
+淤
+芬
+痕
+痪
+陕
+抄
+挟
+堡
+啦
+撤
+悦
+眠
+梨
+苯
+棍
+姑
+恋
+芦
+瞬
+寸
+弘
+伍
+掩
+踏
+绕
+傍
+睁
+汕
+沃
+呛
+洼
+株
+绒
+撼
+歧
+寨
+贬
+甜
+茅
+械
+磺
+磋
+徽
+盆
+赌
+贺
+屡
+屁
+崩
+烂
+郡
+绑
+啤
+谴
+桩
+溜
+肩
+碎
+愤
+浆
+咳
+嗽
+喊
+擅
+酿
+堤
+橙
+坞
+徘
+徊
+凡
+池
+谊
+曹
+碌
+冉
+拦
+匿
+吊
+劲
+渺
+轿
+玄
+祖
+傻
+焕
+鲍
+睡
+狈
+忽
+卢
+肋
+讶
+灼
+柔
+嘴
+薯
+吐
+拔
+漳
+誉
+冈
+奠
+吓
+殿
+契
+脊
+椎
+肾
+衰
+尸
+钾
+芯
+粹
+浇
+呵
+焉
+熔
+瞩
+澈
+裤
+夯
+畏
+碗
+贿
+荆
+帅
+柜
+2
+0
+,
+。
+、
+!
+:
+;
+?
+讹
+“
+蕉
+”
+辟
+衙
+
+刹
+吏
+6
+1
+9
+暑
+5
+7
+8
+3
+《
+》
+蓬
+勃
+鳌
+垃
+圾
+绵
+4
+塑
+瓜
+漂
+·
+曼
+俊
+甄
+舒
+淇
+…
+—
+篮
+L
+E
+裙
+粱
+珞
+镂
+A
+.
+(
+)
+脆
+矮
+颗
+厄
+冥
+%
+灶
+°
+诞
+N
+S
+詹
+―
+沼
+湿
+禽
+阔
+磁
+灭
+翰
+杜
+轩
+辕
+厚
+苇
+栖
+苍
+穹
+赏
+呆
+践
+迄
+C
+e
+r
+b
+u
+s
+-
+驰
+々
+"
+渗
+妆
+棉
+镍
+钴
+铂
+氧
+铜
+谜
+镉
+R
+奏
+鼎
+溧
+/
+虾
+滋
+幻
+姿
+唤
+颓
+畔
+缭
+静
+羞
+颤
+绡
+帐
+瓯
+凄
+昂
+坟
+炬
+卧
+闪
+焰
+怜
+刃
+躯
+剑
+炎
+胞
+幄
+奸
+燎
+邬
+讳
+亟
+勋
+钧
+睦
+锴
+梗
+矣
+卅
+藻
+贤
+吾
+猖
+獗
+绊
+鞋
+芳
+荟
+萃
+佬
+皱
+纹
+瞄
+掠
+怔
+宛
+竹
+韵
+铺
+串
+泡
+磨
+茧
+臂
+抖
+筷
+鼾
+寒
+锤
+泞
+锹
+铲
+镐
+刨
+窝
+酷
+剥
+佃
+轧
+碾
+叛
+篡
+恬
+匠
+敞
+阐
+榷
+伎
+〈
+〉
+圣
+肆
+虐
+糟
+蹋
+妄
+邪
+欺
+剽
+窃
+勾
+诫
+芒
+摹
+怳
+淑
+兮
+怡
+曰
+陌
+颠
+譬
+鸽
+旷
+萧
+乍
+泓
+吟
+哦
+咏
+闹
+惝
+慨
+焚
+泳
+庙
+兆
+舜
+恳
+寡
+巫
+僧
+凉
+躬
+讼
+尝
+诘
+霆
+沮
+苟
+秉
+君
+痘
+痤
+疮
+腺
+B
+酶
+泌
+丙
+紊
+肪
+①
+锌
+萝
+②
+肺
+兔
+鸭
+鲫
+蘑
+菇
+芹
+菠
+忌
+苋
+莴
+笋
+柿
+芽
+藕
+椹
+柚
+楂
+燥
+熏
+蒸
+肌
+芝
+咖
+啡
+椒
+蒜
+韭
+狗
+雀
+紫
+兽
+扒
+饥
+饿
+猕
+猴
+豹
+羚
+砾
+烤
+茎
+禹
+茂
+汶
+寺
+迥
+嵌
+镶
+陡
+珑
+剔
+晶
+莹
+壑
+虬
+瞰
+巍
+峨
+枕
+莽
+粼
+犹
+鳞
+悟
+雄
+羡
+鄙
+诗
+W
+F
+i
+g
+h
+t
+凛
+敬
+舵
+D
+J
+o
+n
+蔼
+脾
+卵
+雌
+酗
+丢
+爸
+穴
+痊
+脏
+梢
+腑
+匹
+袜
+锥
+犁
+耧
+锄
+耙
+鞭
+醋
+糖
+碱
+葱
+芥
+嗓
+夸
+徕
+炕
+啪
+杖
+哧
+扑
+韧
+拽
+顷
+煎
+舅
+甥
+帽
+峭
+茫
+蟹
+獾
+厩
+咴
+嘶
+腿
+骋
+埝
+塄
+蕾
+翠
+捋
+辫
+汁
+醇
+撒
+窑
+槛
+凿
+晌
+荫
+棋
+秋
+瞅
+禾
+涩
+竖
+惧
+哆
+嗦
+鹰
+啄
+睛
+抹
+扔
+蹬
+墟
+蚌
+埠
+涡
+淮
+P
+d
+a
+p
+l
+f
+O
+c
+澎
+飨
+拜
+宦
+魄
+湛
+丘
+阉
+w
+虔
+T
+V
+m
+舆
+趟
+浏
+昶
+烛
+×
+铅
+勺
+垛
+垒
+衩
+涔
+卸
+榴
+眨
+峦
+迭
+嶂
+钎
+叮
+崇
+岭
+滔
+甭
+掂
+怯
+汀
+巷
+赤
+桨
+粥
+爽
+嘈
+舀
+齿
+拗
+絮
+扁
+吆
+叠
+簿
+懒
+缆
+尉
+碑
+吞
+聂
+寝
+拱
+宠
+魂
+H
+k
+y
+呎
+豚
+纲
+啮
+肢
+棕
+趾
+蹼
+爪
+钝
+犬
+拣
+擀
+襟
+娇
+脖
+栈
+翅
+吼
+朵
+缰
+萍
+祝
+瞥
+磅
+倔
+叔
+喂
+胳
+膊
+瞪
+盯
+晃
+惹
+侄
+懵
+呀
+撸
+葫
+湃
+捣
+峪
+艘
+俘
+虏
+粪
+亭
+砍
+镰
+叭
+磕
+’
+棱
+绣
+咧
+咋
+俺
+‘
+咙
+仗
+搽
+胭
+渴
+柠
+檬
+哼
+莺
+溷
+浊
+窄
+羁
+缕
+艄
+桅
+郁
+恨
+奚
+役
+硕
+嗡
+娟
+煌
+窟
+毅
+擂
+蓦
+逢
+俨
+孪
+媚
+姝
+瑜
+~
+钙
+敝
+○
+枉
+颖
+玛
+荃
+郝
+潟
+扉
+屈
+厥
+欠
+揣
+幹
+磷
+纤
+暨
+沫
+琼
+涸
+沛
+旱
+祭
+俢
+伶
+歹
+俭
+撩
+杏
+昱
+佣
+钦
+伿
+'
+_
+寅
+娄
+咸
+長
+時
+張
+阈
+荧
+蛙
+龟
+虫
+蚂
+蚁
+膘
+獭
+愁
+撕
+嬉
+圳
+奕
+雕
+睹
+I
+U
+濒
+帕
+乞
+礁
+蝴
+蝶
+矫
+唇
+饰
+御
+酉
+郜
+侍
+濑
+`
+>
+镕
+滕
+碳
+陵
+鄂
+扈
+逝
+耿
+冢
+奉
+祀
+佰
+嘎
+瞒
+嫖
+糜
+涣
+畅
+舰
+瓷
+舸
+葆
+氮
+蕴
+庐
+黛
+阙
+瀑
+沓
+窥
+苔
+芙
+蓉
+垓
+敖
+甫
+辚
+弓
+戍
+杞
+卒
+烦
+啾
+琵
+琶
+湓
+铮
+倡
+贾
+悯
+沦
+憔
+悴
+徙
+谪
+浔
+枫
+荻
+弦
+眉
+捻
+裳
+幺
+蟆
+妒
+钿
+篦
+姨
+暮
+啼
+阑
+唧
+涯
+僻
+鹃
+猿
+岂
+笛
+哳
+泣
+浸
+幽
+咽
+凝
+歇
+迸
+帛
+舫
+佝
+鞍
+鬲
+遡
+澡
+沧
+嘛
+谒
+俅

dataset.py

@@ -0,0 +1,87 @@
+import os
+from PIL import Image
+from torch.utils.data import Dataset
+import torch
+class trocrDataset(Dataset):
+    """
+    trocr 训练数据集处理
+    文件数据结构
+    /tmp/0/0.jpg #image
+    /tmp/0/0.txt #text label
+    ....
+    /tmp/100/10000.jpg #image
+    /tmp/100/10000.txt #text label
+    """
+    def __init__(self, paths, processor, max_target_length=128, transformer=lambda x:x):
+        self.paths = paths
+        self.processor = processor
+        self.transformer = transformer
+        self.max_target_length = max_target_length
+        self.nsamples = len(self.paths)
+        self.vocab = processor.tokenizer.get_vocab()
+
+    def __len__(self):
+        return self.nsamples
+
+    def __getitem__(self, idx):
+            inx = idx %  self.nsamples
+            image_file = self.paths[idx]
+            txt_file = os.path.splitext(image_file)[0]+'.txt'
+
+            with open(txt_file) as f:
+                    text = f.read().strip().replace('xa0','')
+
+            image = Image.open(image_file).convert("RGB")
+            image = self.transformer(image) ##图像增强函数
+            pixel_values = self.processor(image, return_tensors="pt").pixel_values
+            #labels = encode_text(text, max_target_length=self.max_target_length, vocab=self.vocab)["input_ids"]
+            labels = encode_text(text, max_target_length=self.max_target_length, vocab=self.vocab)
+            labels = [label if label != self.processor.tokenizer.pad_token_id else -100 for label in labels]
+
+            encoding = {"pixel_values": pixel_values.squeeze(), "labels": torch.tensor(labels)}
+
+            return encoding
+
+
+def encode_text(text, max_target_length=128, vocab=None):
+    """
+    {'input_ids': [0, 1092, 2, 1, 1],
+    'attention_mask': [1, 1, 1, 0, 0]}
+    """
+    text = list(text)
+    text = text[:max_target_length - 2]
+    tokens = [vocab.get('<s>')]
+    unk = vocab.get('<unk>')
+    pad = vocab.get('<pad>')
+    mask = []
+    for tk in text:
+        token = vocab.get(tk, unk)
+        tokens.append(token)
+        mask.append(1)
+
+    tokens.append(vocab.get('</s>'))
+    mask.append(1)
+
+    if len(tokens) < max_target_length:
+        for i in range(max_target_length - len(tokens)):
+            tokens.append(pad)
+            mask.append(0)
+
+    return tokens
+    #return {"input_ids": tokens, 'attention_mask': mask}
+
+
+def decode_text(tokens, vocab, vocab_inp):
+    ##decode trocr
+    s_start = vocab.get('<s>')
+    s_end = vocab.get('</s>')
+    unk = vocab.get('<unk>')
+    text = ''
+    for tk in tokens:
+        if tk not in [s_end, s_start]:
+           text += vocab_inp[tk]
+
+    return text
+
+
+

eval.py

@@ -0,0 +1,80 @@
+import os
+
+os.environ["CUDA_VISIBLE_DEVICES"] = '0'
+from PIL import Image
+import numpy as np
+import time
+import torch
+import argparse
+from glob import glob
+from sklearn.model_selection import train_test_split
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+from dataset import decode_text
+from tqdm import tqdm
+from datasets import load_metric
+
+cer_metric = load_metric("./cer.py")
+
+
+def compute_metrics(pred_str, label_str):
+    """
+    计算cer,acc
+    :param pred:
+    :return:
+    """
+    cer = cer_metric.compute(predictions=pred_str, references=label_str)
+    acc = [pred == label for pred, label in zip(pred_str, label_str)]
+    acc = sum(acc) / (len(acc) + 0.000001)
+    return {"cer": cer, "acc": acc}
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='trocr fine-tune训练')
+    parser.add_argument('--cut_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="用于训练集划分的随机数")
+
+    args = parser.parse_args()
+
+    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)
+
+    else:
+        train_paths = []
+        test_paths = paths
+
+    print("train num:", len(train_paths), "test num:", len(test_paths))
+
+    processor = TrOCRProcessor.from_pretrained(args.cut_data_init_weights_path)
+    vocab = processor.tokenizer.get_vocab()
+
+    vocab_inp = {vocab[key]: key for key in vocab}
+    model = VisionEncoderDecoderModel.from_pretrained(args.cut_data_init_weights_path)
+    model.eval()
+    model.cuda()
+
+    vocab = processor.tokenizer.get_vocab()
+    vocab_inp = {vocab[key]: key for key in vocab}
+
+    pred_str, label_str = [], []
+    for p in tqdm(test_paths):
+        img = Image.open(p).convert('RGB')
+        txt_p = os.path.splitext(p)[0] + '.txt'
+        with open(txt_p) as f:
+            label = f.read().strip()
+        pixel_values = processor([img], return_tensors="pt").pixel_values
+
+        with torch.no_grad():
+            generated_ids = model.generate(pixel_values[:, :, :].cuda())
+
+        generated_text = decode_text(generated_ids[0].cpu().numpy(), vocab, vocab_inp)
+        pred_str.append(generated_text)
+        label_str.append(label)
+
+    res = compute_metrics(pred_str, label_str)
+    print(res)

init_custdata_model.py

@@ -0,0 +1,96 @@
+"""
+转换trocr 模型到自己数据集上的字符进行fine-tune
+"""
+import os
+import json
+
+os.environ["CUDA_VISIBLE_DEVICES"] = '-1'
+import argparse
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+from transformers import AutoConfig
+
+
+def read_vocab(vocab_path):
+    """
+    读取自定义训练字符集
+    vocab_path format:
+    1\n
+    2\n
+    ...
+    我\n
+    """
+    other = ["<s>", "<pad>", "</s>", "<unk>", "<mask>"]
+    vocab = {}
+    for ot in other:
+        vocab[ot] = len(vocab)
+
+    with open(vocab_path) as f:
+        for line in f:
+            line = line.strip('\n')
+            if line not in vocab:
+                vocab[line] = len(vocab)
+    return vocab
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='trocr fine-tune训练')
+
+    parser.add_argument('--cust_vocab', default="./cust-data/vocab.txt", type=str, help="自定义训练数字符集")
+
+    parser.add_argument('--pretrain_model', default='./weights', type=str, help="预训练bert权重文件")
+
+    parser.add_argument('--cust_data_init_weights_path', default='./cust-data/weights', type=str,
+                        help="初始化训练权重，用于自己数据集上fine-tune权重")
+    args = parser.parse_args()
+
+    processor = TrOCRProcessor.from_pretrained(args.pretrain_model)
+    pre_model = VisionEncoderDecoderModel.from_pretrained(args.pretrain_model)
+
+    pre_vocab = processor.tokenizer.get_vocab()
+
+    cust_vocab = read_vocab(args.cust_vocab)
+
+    keep_tokens = []
+    unk_index = pre_vocab.get('<unk>')
+    for key in cust_vocab:
+        keep_tokens.append(pre_vocab.get(key, unk_index))
+
+    processor.save_pretrained(args.cust_data_init_weights_path)
+
+    pre_model.save_pretrained(args.cust_data_init_weights_path)
+    ## 替换词库
+    with open(os.path.join(args.cust_data_init_weights_path, "vocab.json"), "w") as f:
+        f.write(json.dumps(cust_vocab, ensure_ascii=False))
+
+    ##替换模型参数
+    with open(os.path.join(args.cust_data_init_weights_path, "config.json")) as f:
+        model_config = json.load(f)
+
+    ## 替换roberta embedding层词库
+    model_config["decoder"]['vocab_size'] = len(cust_vocab)
+
+    ## 替换 attetion 字库
+    model_config['vocab_size'] = len(cust_vocab)
+
+    with open(os.path.join(args.cust_data_init_weights_path, "config.json"), "w") as f:
+        f.write(json.dumps(model_config, ensure_ascii=False))
+
+    ##加载cust model
+    cust_config = AutoConfig.from_pretrained(args.cust_data_init_weights_path)
+    cust_model = VisionEncoderDecoderModel(cust_config)
+
+    pre_model_weigths = pre_model.state_dict()
+    cust_model_weigths = cust_model.state_dict()
+
+    ##权重初始化
+    print("loading init weights..................")
+    for key in pre_model_weigths:
+        print("name:", key)
+        if pre_model_weigths[key].shape != cust_model_weigths[key].shape:
+            wt = pre_model_weigths[key][keep_tokens, :]
+            cust_model_weigths[key] = wt
+        else:
+            cust_model_weigths[key] = pre_model_weigths[key]
+
+    cust_model.load_state_dict(cust_model_weigths)
+    cust_model.save_pretrained(args.cust_data_init_weights_path)

requirements.txt

@@ -0,0 +1,5 @@
+datasets
+transformers==4.15.0
+scikit-learn
+jiwer==2.3.0
+python-Levenshtein==0.12.2

sources.list

@@ -0,0 +1,19 @@
+deb http://mirrors.aliyun.com/ubuntu/ bionic main restricted universe multiverse
+
+deb-src http://mirrors.aliyun.com/ubuntu/ bionic main restricted universe multiverse
+
+deb http://mirrors.aliyun.com/ubuntu/ bionic-security main restricted universe multiverse
+
+deb-src http://mirrors.aliyun.com/ubuntu/ bionic-security main restricted universe multiverse
+
+deb http://mirrors.aliyun.com/ubuntu/ bionic-updates main restricted universe multiverse
+
+deb-src http://mirrors.aliyun.com/ubuntu/ bionic-updates main restricted universe multiverse
+
+deb http://mirrors.aliyun.com/ubuntu/ bionic-backports main restricted universe multiverse
+
+deb-src http://mirrors.aliyun.com/ubuntu/ bionic-backports main restricted universe multiverse
+
+deb http://mirrors.aliyun.com/ubuntu/ bionic-proposed main restricted universe multiverse
+
+deb-src http://mirrors.aliyun.com/ubuntu/ bionic-proposed main restricted universe multiverse

train.py

@@ -0,0 +1,106 @@
+import os
+import argparse
+from glob import glob
+from dataset import trocrDataset, decode_text
+from transformers import TrOCRProcessor
+from transformers import VisionEncoderDecoderModel
+from transformers import default_data_collator
+from sklearn.model_selection import train_test_split
+from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments
+from datasets import load_metric
+
+def compute_metrics(pred):
+    """
+    计算cer,acc
+    :param pred:
+    :return:
+    """
+    labels_ids = pred.label_ids
+    pred_ids = pred.predictions
+
+    pred_str = [decode_text(pred_id, vocab, vocab_inp) for pred_id in pred_ids]
+    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)]
+    acc = sum(acc)/(len(acc)+0.000001)
+    return {"cer": cer, "acc": acc}
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='trocr fine-tune训练')
+    parser.add_argument('--cust_data_init_weights_path', default='./cust-data/weights', type=str,
+                        help="初始化训练权重，用于自己数据集上fine-tune权重")
+    parser.add_argument('--checkpoint_path', default='./checkpoint/trocr', type=str, help="训练模型保存地址")
+    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('--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('--CUDA_VISIBLE_DEVICES', default='0,1', type=str, help="GPU设置")
+
+    args = parser.parse_args()
+    print("train param")
+    print(args)
+    os.environ["CUDA_VISIBLE_DEVICES"] = args.CUDA_VISIBLE_DEVICES
+    print("loading data .................")
+    paths = glob(args.dataset_path)
+
+    train_paths, test_paths = train_test_split(paths, test_size=0.05, random_state=10086)
+    print("train num:", len(train_paths), "test num:", len(test_paths))
+
+    ##图像预处理
+    processor = TrOCRProcessor.from_pretrained(args.cust_data_init_weights_path)
+    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)
+
+    model = VisionEncoderDecoderModel.from_pretrained(args.cust_data_init_weights_path)
+    model.config.decoder_start_token_id = processor.tokenizer.cls_token_id
+    model.config.pad_token_id = processor.tokenizer.pad_token_id
+
+    model.config.vocab_size = model.config.decoder.vocab_size
+
+    model.config.eos_token_id = processor.tokenizer.sep_token_id
+    model.config.max_length = 256
+    model.config.early_stopping = True
+    model.config.no_repeat_ngram_size = 3
+    model.config.length_penalty = 2.0
+    model.config.num_beams = 4
+
+    cer_metric = load_metric("./cer.py")
+
+    training_args = Seq2SeqTrainingArguments(
+        predict_with_generate=True,
+        evaluation_strategy="steps",
+        per_device_train_batch_size=args.per_device_train_batch_size,
+        per_device_eval_batch_size=8,
+        fp16=True,
+        output_dir=args.checkpoint_path,
+        logging_steps=10,
+        num_train_epochs=args.num_train_epochs,
+        save_steps=args.eval_steps,
+        eval_steps=args.eval_steps,
+        save_total_limit=5
+    )
+
+    # seq2seq trainer
+    trainer = Seq2SeqTrainer(
+        model=model,
+        tokenizer=processor.feature_extractor,
+        args=training_args,
+        compute_metrics=compute_metrics,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        data_collator=default_data_collator,
+    )
+    trainer.train()
+
+
+
+
+