TrOCR识别后Pandas DataFrame仅存最后一行的问题及解决需求

问题原由 #

原代码存在两个核心问题导致无法正确保存所有识别结果：

1. 每次调用add_to_df时都基于初始空DataFramedf执行append，生成的新DataFrame未被赋值回原变量，导致之前的数据丢失，最终只保留最后一次的临时结果。
2. 向DataFrame传递的是image_path（Image对象）而非图片文件名字符串，导致image_name列存储为对象类型，不符合需求。

修正后的完整代码 #

from transformers import TrOCRProcessor, VisionEncoderDecoderModel
from PIL import Image
import pandas as pd
import json

# 初始化TrOCR处理器和模型
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")

# 定义图片识别函数
def process_image(image_path):
    pixel_values = processor(image_path, return_tensors="pt").pixel_values
    generated_ids = model.generate(pixel_values)
    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return generated_text

# 配置路径和图片列表（确保working_dir已定义）
working_dir = "./"  # 根据实际路径修改
images_1 = ["10.png", "11.png", "12.png"]  # 文件名需用字符串包裹

# 收集所有识别结果到列表（推荐方式，避免重复创建DataFrame）
results = []
for image_name in images_1:
    # 打开图片并转换为RGB
    image_path = Image.open(f'{working_dir}imgs/{image_name}').convert("RGB")
    # 执行识别
    generated_text = process_image(image_path)
    # 将文件名（字符串）和识别结果加入列表
    results.append({"image_name": image_name, "text": generated_text})

# 从列表生成DataFrame，确保image_name为字符串类型
df = pd.DataFrame(results, columns=["image_name", "text"])

# 将DataFrame转换为jsonl格式并保存
with open(f"{working_dir}labels.jsonl", "w", encoding="utf-8") as f:
    for record in df.to_dict(orient="records"):
        f.write(json.dumps(record, ensure_ascii=False) + "\n")

关键修改说明 #

• 数据收集方式优化：不再每次修改DataFrame，而是用列表results收集所有识别结果的字典，最后一次性生成DataFrame，既避免数据丢失，也提升效率。
• image_name类型修正：传递图片文件名（如"10.png"）而非Image对象到DataFrame，确保image_name列存储为字符串类型。
• 替换过时方法：放弃使用已被弃用的df.append方法，改用更高效的列表收集+批量生成DataFrame的方式。
• jsonl输出优化：写入文件时指定encoding="utf-8"，避免中文乱码问题。

内容的提问来源于stack exchange，提问作者Mohammed