查询网站死链接,唐山建设集团下岗职工网站,大连哪家网站技术开发公司好,商业网站设计与制作导 读 本文主要介绍如何通过微调训练TrOCR实现弯曲文本识别。
背景介绍 TrOCR#xff08;基于 Transformer 的光学字符识别#xff09;模型是性能最佳的 OCR 模型之一。在我们之前的文章中#xff0c;我们分析了它们在单行打印和手写文本上的表现。 TrOCR—基于Transforme…导 读 本文主要介绍如何通过微调训练TrOCR实现弯曲文本识别。
背景介绍 TrOCR基于 Transformer 的光学字符识别模型是性能最佳的 OCR 模型之一。在我们之前的文章中我们分析了它们在单行打印和手写文本上的表现。 TrOCR—基于Transformer的OCR入门 然而与任何其他深度学习模型一样它们也有其局限性。TrOCR 在处理开箱即用的弯曲文本时表现不佳。本文将通过在弯曲文本数据集上微调 TrOCR 模型使 TrOCR 系列更进一步。 实现步骤 从上一篇文章中我们知道TrOCR无法识别弯曲和垂直图像上的文本。这些图像是 SCUT-CTW1500 数据集的一部分。我们将在此数据集上训练 TrOCR 模型并再次运行推理来分析结果。这将使我们全面了解针对不同用例可以将 TrOCR 模型的边界拓展到什么程度。 我们将使用 Hugging Face Trainer API 来训练模型。要完成整个过程必须遵循以下步骤 准备并分析弯曲文本图像数据集。 从 Hugging Face 加载 TrOCR Small Printed 模型。 将拥抱脸部序列初始化为序列训练器 API。 定义评估指标 训练模型并运行推理。 弯曲文本数据集 SCUT-CTW1500 数据集以下简称 CTW1500包含数千张弯曲文本和野外文本的图像。 原始数据集可在官方 GitHub 存储库中获取。
https://github.com/Yuliang-Liu/Curve-Text-Detector 这包括训练集和测试集。只有训练集包含 XML 格式的标签。因此我们将训练集分为不同的训练和验证子集。 最终数据集包含 6052 个训练样本和 1651 个验证样本。每个图像的标签都存在于文本文件中并以换行符分隔。 让我们检查数据集中的一些图像及其文本标签。 从上图中可以看出一些事情。除了弯曲的文本图像之外数据集还包含模糊和朦胧的图像。这种现实世界的图像变化给深度学习模型带来了挑战。了解如此多样化的数据集中图像和文本的特征对于任何 OCR 模型的最先进性能至关重要。这对 TrOCR 模型提出了一个有趣的挑战自然地经过训练它在此类图像上的表现会明显更好。 微调弯曲文本上的 TrOCR 让我们进入本文的技术方面。从这里开始我们将详细讨论TrOCR训练过程的代码。 【1】安装并导入所需的库 第一步是安装所有必需的库。
!pip install -q transformers!pip install -q sentencepiece!pip install -q jiwer!pip install -q datasets!pip install -q evaluate!pip install -q -U accelerate !pip install -q matplotlib!pip install -q protobuf3.20.1!pip install -q tensorboard 其中一些关键概念 接下来我们导入所有必需的库和包。
import osimport osimport torchimport evaluateimport numpy as npimport pandas as pdimport glob as globimport torch.optim as optimimport matplotlib.pyplot as pltimport torchvision.transforms as transforms from PIL import Imagefrom zipfile import ZipFilefrom tqdm.notebook import tqdmfrom dataclasses import dataclassfrom torch.utils.data import Datasetfrom urllib.request import urlretrievefrom transformers import ( VisionEncoderDecoderModel, TrOCRProcessor, Seq2SeqTrainer, Seq2SeqTrainingArguments, default_data_collator) 上述代码块中的一些重要的导入语句是 现在设置种子以实现不同运行的可重复性并定义计算设备。
def seed_everything(seed_value): np.random.seed(seed_value) torch.manual_seed(seed_value) torch.cuda.manual_seed_all(seed_value) torch.backends.cudnn.deterministic True torch.backends.cudnn.benchmark False seed_everything(42) device torch.device(cuda if torch.cuda.is_available() else cpu) 【2】下载并提取数据集 下一个代码块包含一个帮助函数用于下载 CTW1500 数据并提取它。
def download_and_unzip(url, save_path): print(fDownloading and extracting assets...., end) # Downloading zip file using urllib package. urlretrieve(url, save_path) try: # Extracting zip file using the zipfile package. with ZipFile(save_path) as z: # Extract ZIP file contents in the same directory. z.extractall(os.path.split(save_path)[0]) print(Done) except Exception as e: print(\nInvalid file., e) URL rhttps://www.dropbox.com/scl/fi/vyvr7jbdvu8o174mbqgde/scut_data.zip?rlkeyfs8axkpxunwu6if9a2su71kxsdl1asset_zip_path os.path.join(os.getcwd(), scut_data.zip) # Download if asset ZIP does not exist.if not os.path.exists(asset_zip_path): download_and_unzip(URL, asset_zip_path) 提取模型后数据集结构将如下所示。
scut_data/├── scut_train├── scut_test├── scut_train.txt└── scut_test.txt 数据被提取到scut_data目录中。它包含保存训练和验证图像的 scut_train和子目录。scut_test这两个文本文件包含以下格式的注释。
006052.jpg ty Starts with Education006053.jpg Cardis006054.jpg YOU THE BUSINESS SIDE OF GREEN006055.jpg hat is... 每行包含一个图像文件名图像中的文本以空格分隔。文本文件中的行数与图像文件夹中的样本数相同。图像中的文本和图像文件名由第一个空格分隔。图像文件名不能包含任何空格否则将被视为文本的一部分。 【3】定义配置 在开始训练部分之前我们先定义训练、数据集和模型配置。
dataclass(frozenTrue)class TrainingConfig: BATCH_SIZE: int 48 EPOCHS: int 35 LEARNING_RATE: float 0.00005 dataclass(frozenTrue)class DatasetConfig: DATA_ROOT: str scut_data dataclass(frozenTrue)class ModelConfig: MODEL_NAME: str microsoft/trocr-small-printed 该模型将使用 48 的批量大小进行 35 个 epoch 的训练。优化器的学习率设置为 0.00005。较高的学习率会使训练过程不稳定从而从一开始就导致较高的损失。 此外我们还定义了根数据集目录和我们将使用的模型。TrOCR Small Printed 模型将进行微调因为它根据该数据集的实验展示了最佳性能。 【4】可视化一些样本 让我们可视化下载数据集中的一些图像及其文件名。
def visualize(dataset_path): plt.figure(figsize(15, 3)) for i in range(15): plt.subplot(3, 5, i1) all_images os.listdir(f{dataset_path}/scut_train) image plt.imread(f{dataset_path}/scut_train/{all_images[i]}) plt.imshow(image) plt.axis(off) plt.title(all_images[i].split(.)[0]) plt.show() visualize(DatasetConfig.DATA_ROOT) 【5】准备数据集 标签以文本文件格式存在。为了更顺利地准备数据加载器需要将它们修改为更简单的格式。让我们将训练和测试文本文件转换为 Pandas DataFrame。
train_df pd.read_fwf( os.path.join(DatasetConfig.DATA_ROOT, scut_train.txt), headerNone)train_df.rename(columns{0: file_name, 1: text}, inplaceTrue)test_df pd.read_fwf( os.path.join(DatasetConfig.DATA_ROOT, scut_test.txt), headerNone)test_df.rename(columns{0: file_name, 1: text}, inplaceTrue) 现在file_name 列包含与图像对应的所有文件名text 列包含图像中的文本。 下一步是定义增强/数据扩充。
# Augmentations.train_transforms transforms.Compose([ transforms.ColorJitter(brightness.5, hue.3), transforms.GaussianBlur(kernel_size(5, 9), sigma(0.1, 5)),]) 我们将ColorJitter和应用于GaussianBlur图像。无需对图像应用任何翻转旋转因为原始数据集中已经有足够的可变性。 准备数据集的最佳方法是编写自定义数据集类。这使我们能够更好地控制输入。以下代码块定义了一个CustomOCRDataset用于准备数据集的类。
class CustomOCRDataset(Dataset): def __init__(self, root_dir, df, processor, max_target_length128): self.root_dir root_dir self.df df self.processor processor self.max_target_length max_target_length def __len__(self): return len(self.df) def __getitem__(self, idx): # The image file name. file_name self.df[file_name][idx] # The text (label). text self.df[text][idx] # Read the image, apply augmentations, and get the transformed pixels. image Image.open(self.root_dir file_name).convert(RGB) image train_transforms(image) pixel_values self.processor(image, return_tensorspt).pixel_values # Pass the text through the tokenizer and get the labels, # i.e. tokenized labels. labels self.processor.tokenizer( text, paddingmax_length, max_lengthself.max_target_length ).input_ids # We are using -100 as the padding token. 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 该 __init__()方法接受根目录路径、DataFrame、TrOCR 处理器和最大标签长度作为参数。 该 __getitem__()方法首先从磁盘读取标签和图像。然后它通过变换传递图像以应用增强。TrOCRProcessor 以 PyTorch 张量格式返回标准化像素值。接下来文本标签通过分词器传递。如果标签短于 128 个字符则会用 -100 填充到长度 128。如果长于 128 个字符则会截断字符。最后它以字典的形式返回像素值和标签。 在创建训练集和验证集之前需要初始化 TrOCRProcessor以便将其传递给数据集类。
processor TrOCRProcessor.from_pretrained(ModelConfig.MODEL_NAME)train_dataset CustomOCRDataset( root_diros.path.join(DatasetConfig.DATA_ROOT, scut_train/), dftrain_df, processorprocessor)valid_dataset CustomOCRDataset( root_diros.path.join(DatasetConfig.DATA_ROOT, scut_test/), dftest_df, processorprocessor) 微调 TrOCR 模型的数据集准备过程到此结束。 【6】准备 TrOCR 小型打印模型 该类VisionEncoderDecoderModel使我们能够访问所有 TrOCR 模型。该from_pretrained()方法接受存储库名称来加载预训练模型。
model VisionEncoderDecoderModel.from_pretrained(ModelConfig.MODEL_NAME)model.to(device)print(model)# Total parameters and trainable parameters.total_params sum(p.numel() for p in model.parameters())print(f{total_params:,} total parameters.)total_trainable_params sum( p.numel() for p in model.parameters() if p.requires_grad)print(f{total_trainable_params:,} training parameters.) 该模型包含 6150 万个参数。将对所有参数进行微调以便对它们进行训练。 模型准备最重要的方面之一是模型配置。下面讨论这些配置。
# Set special tokens used for creating the decoder_input_ids from the labels.model.config.decoder_start_token_id processor.tokenizer.cls_token_idmodel.config.pad_token_id processor.tokenizer.pad_token_id# Set Correct vocab size.model.config.vocab_size model.config.decoder.vocab_sizemodel.config.eos_token_id processor.tokenizer.sep_token_id model.config.max_length 64model.config.early_stopping Truemodel.config.no_repeat_ngram_size 3model.config.length_penalty 2.0model.config.num_beams 4 预训练的 TrOCR 模型带有自己的一组预定义配置。然而为了微调模型我们将覆盖其中的一些内容其中包括标记 ID、词汇表大小以及序列结束标记。 此外提前停止设置为True。这确保了如果模型指标在连续几个时期没有改善则训练将停止。 【7】优化器和评估指标 为了优化模型权重我们选择权重衰减为 0.0005 的 AdamW 优化器。
optimizer optim.AdamW( model.parameters(), lrTrainingConfig.LEARNING_RATE, weight_decay0.0005) 评估指标将是 CER字符错误率。
cer_metric evaluate.load(cer) def compute_cer(pred): labels_ids pred.label_ids pred_ids pred.predictions pred_str processor.batch_decode(pred_ids, skip_special_tokensTrue) labels_ids[labels_ids -100] processor.tokenizer.pad_token_id label_str processor.batch_decode(labels_ids, skip_special_tokensTrue) cer cer_metric.compute(predictionspred_str, referenceslabel_str) return {cer: cer} 无需进一步详细说明CER 基本上是模型未正确预测的字符数。CER 越低模型的性能越好。 请注意我们在 CER 计算中跳过填充标记因为我们不希望填充标记影响模型的性能。 【8】TrOCR 的训练和验证 训练开始之前必须初始化训练参数。
training_args Seq2SeqTrainingArguments( predict_with_generateTrue, evaluation_strategyepoch, per_device_train_batch_sizeTrainingConfig.BATCH_SIZE, per_device_eval_batch_sizeTrainingConfig.BATCH_SIZE, fp16True, output_dirseq2seq_model_printed/, logging_strategyepoch, save_strategyepoch, save_total_limit5, report_totensorboard, num_train_epochsTrainingConfig.EPOCHS) 正在使用 FP16 训练因为它使用更少的 GPU 内存并且还允许我们使用更高的批量大小。此外日志记录和模型保存策略是基于纪元的。所有报告都将记录到张量板上。 这些训练参数将与其他所需参数一起传递给训练器 API。 # Initialize trainer.trainer Seq2SeqTrainer( modelmodel, tokenizerprocessor.feature_extractor, argstraining_args, compute_metricscompute_cer, train_datasettrain_dataset, eval_datasetvalid_dataset, data_collatordefault_data_collator) 训练过程可以通过调用训练器对象的train()方法来开始。 res trainer.train() Epoch Training Loss Validation Loss Cer1 3.822000 2.677871 0.6877392 2.497100 2.474666 0.6908003 2.180700 2.336284 0.627641...33 0.146800 2.130022 0.50420934 0.145800 2.167060 0.51109535 0.138300 2.120335 0.494496 训练结束时模型的CER 达到 49%考虑到所使用的小型 TrOCR 模型这是一个非常好的结果。 以下是 Tensorboard 日志中的 CER 图。 直到训练结束曲线呈下降趋势。尽管更长时间的训练可能会产生更好的结果但我们将继续使用现有的模型。 【9】使用微调 TrOCR 模型进行推理 训练完 TrOCR 模型后就可以对验证图像进行推理了。 第一步是从最后保存的检查点加载经过训练的模型。
processor TrOCRProcessor.from_pretrained(ModelConfig.MODEL_NAME)trained_model VisionEncoderDecoderModel.from_pretrained(seq2seq_model_printed/checkpoint-str(res.global_step)).to(device)res 对象包含一个 global_step 属性该属性保存模型训练的总步数。上面的代码块使用该属性来加载最终时期的权重。 接下来是一些辅助函数。第一个是读取图像。
def read_and_show(image_path): :param image_path: String, path to the input image. Returns: image: PIL Image. image Image.open(image_path).convert(RGB) return image 下一个辅助函数通过模型执行图像的前向传递。
def ocr(image, processor, model): :param image: PIL Image. :param processor: Huggingface OCR processor. :param model: Huggingface OCR model. Returns: generated_text: the OCRd text string. # We can directly perform OCR on cropped images. pixel_values processor(image, return_tensorspt).pixel_values.to(device) generated_ids model.generate(pixel_values) generated_text processor.batch_decode(generated_ids, skip_special_tokensTrue)[0] return generated_text 最后的辅助函数循环遍历目录中的所有图像并继续调用 ocr() 函数进行推理。
def eval_new_data( data_pathos.path.join(DatasetConfig.DATA_ROOT, scut_test, *), num_samples50): image_paths glob.glob(data_path) for i, image_path in tqdm(enumerate(image_paths), totallen(image_paths)): if i num_samples: break image read_and_show(image_path) text ocr(image, processor, trained_model) plt.figure(figsize(7, 4)) plt.imshow(image) plt.title(text) plt.axis(off) plt.show() eval_new_data( data_pathos.path.join(DatasetConfig.DATA_ROOT, scut_test, *), num_samples100) 我们正在对 100 个样本 (num_samples100) 进行推理。 以下是模型在训练前 OCR 错误的两个结果。 结果令人印象深刻。经过微调 TrOCR 模型它能够正确预测弯曲和垂直图像中的文本。 以下是模型表现良好的更多结果。 在这种情况下尽管最末端的文本被拉伸但模型仍然正确地预测它们。 在上述三种情况下即使文本模糊模型也能正确预测文本。 结论 在本文中我们在弯曲文本识别数据集上对 TrOCR 模型进行了微调。我们从数据集讨论开始。接下来是数据集准备和 TrOCR 模型的训练。训练结束后我们进行推理实验并分析结果。我们的结果表明即使在模糊或弯曲的文本图像上微调 TrOCR 模型也可以带来更好的性能。 OCR 不仅仅是识别场景中的文本还涉及使用 OCR 构建应用程序例如验证码识别器或将 TrOCR 识别器与车牌检测管道相结合。
