基于Qwen2-VL模型针对LaTeX OCR任务进行微调训练 - 数据处理

基于Qwen2-VL模型针对LaTeX OCR任务进行微调训练 - 数据处理

flyfish

基于Qwen2-VL模型针对LaTeX_OCR任务进行微调训练_-_LoRA配置如何写
基于Qwen2-VL模型针对LaTeX_OCR任务进行微调训练_-_单图推理
基于Qwen2-VL模型针对LaTeX_OCR任务进行微调训练_-_原模型_单图推理
基于Qwen2-VL模型针对LaTeX_OCR任务进行微调训练_-_原模型_多图推理
基于Qwen2-VL模型针对LaTeX_OCR任务进行微调训练_-_多图推理
基于Qwen2-VL模型针对LaTeX_OCR任务进行微调训练_-_数据处理
基于Qwen2-VL模型针对LaTeX_OCR任务进行微调训练_-_训练
基于Qwen2-VL模型针对LaTeX_OCR任务进行微调训练_-_训练过程
完整代码在文末

数据集介绍

数据集采用的是 开源的AI-ModelScope/LaTeX_OCR/

https://modelscope.cn/datasets/AI-ModelScope/LaTeX_OCR/summary

程序中下载的是full

本仓库有 5 个数据集

small 是小数据集，样本数 110 条，用于测试
full 是印刷体约 100k 的完整数据集。实际上样本数略小于 100k，因为用 LaTeX 的抽象语法树剔除了很多不能渲染的 LaTeX。
synthetic_handwrite 是手写体 100k 的完整数据集，基于 full 的公式，使用手写字体合成而来，可以视为人类在纸上的手写体。样本数实际上略小于 100k，理由同上。
human_handwrite 是手写体较小数据集，更符合人类在电子屏上的手写体。主要来源于 CROHME。我们用 LaTeX 的抽象语法树校验过了。
human_handwrite_print 是来自 human_handwrite 的印刷体数据集，公式部分和 human_handwrite 相同，图片部分由公式用 LaTeX 渲染而来。

代码实现

代码实现对LaTeX OCR相关数据集的下载、处理以及格式转换，使其能够适用于后续的模型训练等任务。先是从ModelScope平台下载包含图片与LaTeX文本的数据集，经过处理生成对应的CSV文件，然后再将CSV文件转换为特定格式的JSON文件，并按照一定比例划分出训练集和验证集。

LaTeX_OCR 存储的图像文件

latex_ocr_train.csv 文件里面是生成图像路径和文本，是一一对应的。
该文件是中间生成

latex_ocr_train.json和latex_ocr_val.json训练使用的数据

[{"id": "identity_9001","conversations": [{"role": "user","value": "/home/sss/datasets/LaTeX_OCR/9000.jpg"},{"role": "assistant","value": "\\ln { \\mathcal J } = - 2 i \\alpha \\mathrm { T r } \\left[ \\gamma _ { 5 } \\varphi ( \\frac { \\not \\! \\! D } { \\Lambda } ) \\right]"}]},{"id": "identity_9002","conversations": [{"role": "user","value": "/home/sss/datasets/LaTeX_OCR/9001.jpg"},{"role": "assistant","value": "{ \\cal P } _ { * } ( \\chi ) \\propto \\exp [ 3 N ( \\chi ) ] ."}]},{"id": "identity_9003","conversations": [{"role": "user","value": "/home/sss/datasets/LaTeX_OCR/9002.jpg"},{"role": "assistant","value": "M ( S , R ) _ { n } = \\left( \\begin{matrix} { 1 } & { 0 } \\\\ { 1 } & { 1 } \\\\ \\end{matrix} \\right) ^ { \\otimes ( n - 1 ) } \\ ,"}]},

import os
from modelscope.msdatasets import MsDataset
import pandas as pd
import json
import argparsedef download_and_process_dataset(max_data_number: int, dataset_id: str, subset_name: str, split: str, dataset_dir: str, csv_path: str) -> None:"""从ModelScope下载数据集并处理成图片和LaTeX文本对。Args:max_data_number (int): 最大处理图片数量。dataset_id (str): ModelScope上数据集的标识符。subset_name (str): 数据集中子集的名称。split (str): 数据集的划分（'train', 'test' 等）。dataset_dir (str): 图片保存目录。csv_path (str): 包含图片路径和LaTeX文本的CSV文件保存路径。"""# 从ModelScope加载数据集ds = MsDataset.load(dataset_id, subset_name=subset_name, split=split)print(f"加载的数据集包含 {len(ds)} 项")# 限制处理的总项数total = min(max_data_number, len(ds))# 如果不存在，则创建保存图片的目录os.makedirs(dataset_dir, exist_ok=True)# 初始化列表以存储图片路径及其对应的LaTeX文本image_paths = []texts = []for i in range(total):item = ds[i]text = item['text']image = item['image']# 保存图片并记录其路径image_path = os.path.abspath(os.path.join(dataset_dir, f'{i}.jpg'))image.save(image_path)# 将路径和文本添加到列表中image_paths.append(image_path)texts.append(text)# 每处理50张图片打印一次进度if (i + 1) % 50 == 0:print(f'正在处理第 {i+1}/{total} 张图片 ({(i+1)/total*100:.1f}%)')# 使用收集的数据创建DataFrame并保存为CSV文件df = pd.DataFrame({'image_path': image_paths,'text': texts,})df.to_csv(csv_path, index=False)print(f'数据处理完成，共处理了 {total} 张图片。')def convert_csv_to_json(csv_path: str, train_json_path: str, val_json_path: str, split_ratio: float) -> None:"""将CSV文件转换为JSON文件，并按照给定的比例划分训练集和验证集。Args:csv_path (str): CSV文件路径。train_json_path (str): 训练集JSON文件保存路径。val_json_path (str): 验证集JSON文件保存路径。split_ratio (float): 划分训练集和验证集的比例（例如：0.8 表示 80% 的数据作为训练集）。"""df = pd.read_csv(csv_path)conversations = []# 创建对话格式for i in range(len(df)):conversations.append({"id": f"identity_{i+1}","conversations": [{"role": "user","value": f"{df.iloc[i]['image_path']}"},{"role": "assistant", "value": str(df.iloc[i]['text'])}]})# 根据比例划分训练集和验证集split_index = int(len(conversations) * split_ratio)train_conversations = conversations[:split_index]val_conversations = conversations[split_index:]# 保存训练集with open(train_json_path, 'w', encoding='utf-8') as f:json.dump(train_conversations, f, ensure_ascii=False, indent=2)# 保存验证集with open(val_json_path, 'w', encoding='utf-8') as f:json.dump(val_conversations, f, ensure_ascii=False, indent=2)print(f'数据已成功转换并保存为 JSON 文件，训练集: {train_json_path}, 验证集: {val_json_path}')def main():"""主函数，用于处理命令行参数并调用相应的处理函数。"""# 设置参数解析器parser = argparse.ArgumentParser(description='下载并处理LaTeX OCR数据集，并将其转换为JSON格式')parser.add_argument('--max_data_number', type=int, default=10000, help='最大处理图片数量')parser.add_argument('--dataset_id', type=str, default='AI-ModelScope/LaTeX_OCR', help='ModelScope上的数据集ID')parser.add_argument('--subset_name', type=str, default='default', help='数据集中的子集名称')parser.add_argument('--split', type=str, default='train', help='数据集的划分（train/test等）')parser.add_argument('--dataset_dir', type=str, default='LaTeX_OCR', help='图片保存目录')parser.add_argument('--csv_path', type=str, default='./latex_ocr_train.csv', help='CSV文件保存路径')parser.add_argument('--train_json_path', type=str, default='./latex_ocr_train.json', help='训练集JSON文件保存路径')parser.add_argument('--val_json_path', type=str, default='./latex_ocr_val.json', help='验证集JSON文件保存路径')parser.add_argument('--split_ratio', type=float, default=0.9, help='训练集与验证集的划分比例（如0.9表示90%训练集，10%验证集）')args = parser.parse_args()# 检查数据集目录是否已存在if not os.path.exists(args.dataset_dir):download_and_process_dataset(args.max_data_number, args.dataset_id, args.subset_name, args.split, args.dataset_dir, args.csv_path)else:print(f'{args.dataset_dir} 目录已存在，跳过数据处理步骤。')# 转换CSV至JSON，并进行数据集划分convert_csv_to_json(args.csv_path, args.train_json_path, args.val_json_path, args.split_ratio)if __name__ == '__main__':main()

说明

1. 代码结构

LaTeXOCRDatasetProcessor类

• __init__方法：作为类的初始化函数，用于接收多个与数据集处理相关的参数，例如最大处理图片数量、数据集在ModelScope上的标识符、各种文件（图片、CSV、训练集JSON、验证集JSON）的保存路径以及训练集和验证集的划分比例等，将这些参数保存为类的属性，方便后续方法中使用。

• download_and_process_dataset方法：

  • 首先，从ModelScope平台加载指定的数据集，依据设定的最大处理图片数量，限制实际要处理的数据项数。
  • 接着，检查并创建用于保存图片的本地目录（若不存在的话）。
  • 然后，遍历数据集（在限制数量范围内），针对每条数据，获取其中的图片和LaTeX文本内容，将图片保存到本地指定目录，并记录其绝对路径，同时把图片路径与对应的LaTeX文本分别添加到对应的列表中。在处理过程中，每处理50张图片就打印一次处理进度，方便了解处理情况。
  • 最后，利用收集好的图片路径列表和LaTeX文本列表创建一个DataFrame，并将其保存为CSV文件，至此完成从原始数据集到CSV文件的处理流程，使得数据以一种更规整的格式存储，便于后续进一步操作。

• convert_csv_to_json方法：

  • 先读取之前生成的CSV文件内容到DataFrame中，然后按照特定的对话格式要求，遍历DataFrame中的每一行数据，将图片路径作为“用户”角色的内容，对应的LaTeX文本作为“助手”角色的内容，构建出一个个对话形式的字典，并添加到conversations列表中，以此将数据整理成符合后续训练要求的格式。
  • 之后，依据设定的划分比例，计算出训练集和验证集的分割索引，从而将conversations列表中的数据划分为训练集和验证集两部分。
  • 最后，分别将训练集和验证集的数据以JSON格式保存到对应的文件路径下，完成数据集格式的最终转换以及训练集、验证集的划分工作，生成的JSON文件可直接用于如模型训练等相关任务中。

main函数

• 首先，利用argparse库创建一个命令行参数解析器，定义了多个与数据集处理相关的参数，如各类文件路径、最大处理数量、划分比例等，并设置了相应的默认值和帮助信息，方便用户在运行代码时可以灵活指定不同的参数值，以满足不同的需求。
• 接着，解析命令行参数得到具体的参数值，然后使用这些参数值实例化LaTeXOCRDatasetProcessor类对象，创建一个专门用于处理数据集的实例。
• 之后，通过判断图片保存目录是否已存在，来决定是否调用download_and_process_dataset方法执行数据集下载及处理成CSV文件的步骤。如果目录已存在，就跳过这一步，避免重复处理。
• 最后，无论前面是否执行了下载处理步骤，都会调用convert_csv_to_json方法，将已有的（可能是刚生成的，也可能是之前就存在的）CSV文件转换为JSON文件，并划分出训练集和验证集，完成整个数据集处理及准备工作的完整流程。

2. 主程序

在if __name__ == '__main__'部分启动整个程序，调用main函数，按照上述main函数中的逻辑依次执行，先是解析命令行参数，接着实例化数据集处理类并根据实际情况处理数据集、转换格式并划分训练集和验证集，最终得到可用于后续任务