Arxiv-tables数据集

下载并解析代码:

from datasets import load_dataset
import os
from PIL import Image, UnidentifiedImageError
import base64
import io

# Step 1: 加载数据集
dataset = load_dataset("staghado/ArXiv-tables")
data = dataset["train"]

# Step 2: 打印表头
print("表头（字段名）：")
print(data.column_names)

# Step 3: 创建保存目录
save_dir = "./saved_tables"
os.makedirs(save_dir, exist_ok=True)

# Step 4: 遍历每一行，保存图片
saved = 0
for idx, row in enumerate(data):
    img_bytes = row.get("table_image", None)
    if img_bytes:
        try:
            # 如果是str，先解码；如果是bytes，直接用
            if isinstance(img_bytes, str):
                img_bytes = base64.b64decode(img_bytes.strip('"').strip("'"))
            
            img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
            
            # 保存
            paper_id = row.get("arxiv_id", f"unknown_{idx}")
            page_num = row.get("page", "0")
            filename = f"{paper_id}_p{page_num}_{idx:04d}.png"
            filepath = os.path.join(save_dir, filename)
            
            img.save(filepath)
            saved += 1
            if saved % 100 == 0:
                print(f"✅ 已保存 {saved} 张图片...")
        
        except (base64.binascii.Error, UnidentifiedImageError, Exception) as e:
            print(f"[跳过] 第 {idx} 行坏图: {e}")
    else:
        print(f"[跳过] 第 {idx} 行没有 table_image 字段")

print(f"🎯 全部完成，成功保存 {saved} 张图片到 {save_dir}/")

日志:

表头（字段名）：
['id', 'arxiv_id', 'page', 'bounding_box', 'latex_content', 'extracted_content', 'similarity_score', 'table_image', 'page_image']
✅ 已保存 100 张图片...
✅ 已保存 200 张图片...
✅ 已保存 300 张图片...
✅ 已保存 400 张图片...
✅ 已保存 500 张图片...
✅ 已保存 600 张图片...
✅ 已保存 700 张图片...
✅ 已保存 800 张图片...
✅ 已保存 900 张图片...
✅ 已保存 1000 张图片...
✅ 已保存 1100 张图片...
✅ 已保存 1200 张图片...
✅ 已保存 1300 张图片...