如何保存含PNG的ChartQA数据集并上传至S3用于VLM训练

JSON无法直接序列化PIL图像对象，针对你的场景（上传S3用于SmolVLM训练），提供三种实用解决方案：

方案1：保存图像文件，JSON存储路径（推荐大规模数据集） #

将PIL图像保存为单独文件，在JSON中替换为文件路径，上传时将图像目录和JSON文件一并传到S3，训练时从路径读取图像。

import os
import json
from PIL import Image
from datasets import load_dataset

# 加载数据集
dataset = load_dataset("HuggingFaceM4/ChartQA", split="train")

# 创建图像存储目录
img_dir = "chartqa_images"
os.makedirs(img_dir, exist_ok=True)

processed_samples = []
for idx, item in enumerate(dataset):
    user_content = item["user"]
    # 提取图像对象（假设user字段包含image键）
    img = user_content["image"]
    # 生成唯一文件名，避免重复
    img_filename = f"sample_{idx}.png"
    img_path = os.path.join(img_dir, img_filename)
    # 保存图像
    img.save(img_path, format="PNG")
    # 替换图像对象为路径
    user_content["image"] = img_path
    # 整理样本
    processed_samples.append({
        "system": item["system"],
        "user": user_content,
        "assistant": item["assistant"]
    })

# 保存处理后的JSON
with open("chartqa_train_processed.json", "w", encoding="utf-8") as f:
    json.dump(processed_samples, f, indent=2)

# 后续将img_dir和JSON文件一起上传到S3
# 示例AWS CLI命令：aws s3 cp --recursive chartqa_images s3://your-bucket/chartqa/images/
# aws s3 cp chartqa_train_processed.json s3://your-bucket/chartqa/

方案2：将图像编码为Base64字符串（适合小数据集） #

把PIL图像转为Base64编码字符串，直接存入JSON，无需额外图像文件，但会让JSON体积增大约30%，训练时需解码还原图像。

import json
import base64
from io import BytesIO
from PIL import Image
from datasets import load_dataset

dataset = load_dataset("HuggingFaceM4/ChartQA", split="train")

processed_samples = []
for item in dataset:
    user_content = item["user"]
    img = user_content["image"]
    # 将图像转为字节流再编码为Base64
    buffer = BytesIO()
    img.save(buffer, format="PNG")
    img_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
    # 替换图像对象为Base64字符串
    user_content["image"] = img_base64
    processed_samples.append({
        "system": item["system"],
        "user": user_content,
        "assistant": item["assistant"]
    })

# 保存JSON
with open("chartqa_train_base64.json", "w", encoding="utf-8") as f:
    json.dump(processed_samples, f, indent=2)

方案3：直接保存为Hugging Face Dataset格式（适配HF训练生态） #

利用Hugging Face Dataset的序列化机制，直接保存数据集对象，无需手动处理图像，上传到S3后可直接用load_from_disk加载，完美适配SmolVLM的HF微调流程。

from datasets import load_dataset

# 加载数据集
dataset = load_dataset("HuggingFaceM4/ChartQA", split="train")
# 保存到本地目录
dataset.save_to_disk("chartqa_hf_dataset")

# 上传到S3（示例AWS CLI命令）
# aws s3 cp --recursive chartqa_hf_dataset s3://your-bucket/chartqa/hf_dataset/

# 训练时加载数据集
# from datasets import load_from_disk
# dataset = load_from_disk("s3://your-bucket/chartqa/hf_dataset/")

方案选型建议 #

• 若数据集规模大，优先选方案1，文件体积小，训练加载高效；
• 若需单文件传输，可选方案2，但注意体积膨胀问题；
• 若基于Hugging Face生态微调SmolVLM，优先选方案3，最省心且兼容性最好。

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