基于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 argparse

def 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函数中的逻辑依次执行，先是解析命令行参数，接着实例化数据集处理类并根据实际情况处理数据集、转换格式并划分训练集和验证集，最终得到可用于后续任务