parquet批量读取图片文件示例

以AnyEdit数据集为例，我们展示了批量读取parquet图片文件的代码，相关数据集地址如下：
Huggingface-AnyEdit

代码如下所示，请注意使用时应查看文件格式来适应性修改代码

import os
import pandas as pd
import io
import json
from tqdm import tqdm
from PIL import Image# 读取 parquet 文件
df = pd.read_parquet('./data/train-00000-of-00383.parquet')# 查看第一行数据
sample_row = df.iloc[0]  # 获取第一行数据
print(sample_row)#--------------------------------#
# 获取data文件夹下的所有parquet文件
#--------------------------------#
parquet_dir      = './data'
parquet_files    = [f for f in os.listdir(parquet_dir) if f.endswith('.parquet')]# 创建新的文件夹来保存图片
base_output_dir  = './extracted_images'
os.makedirs(base_output_dir, exist_ok=True)image_dir        = os.path.join(base_output_dir, 'image_files')
edited_dir       = os.path.join(base_output_dir, 'edited_files')
visual_input_dir = os.path.join(base_output_dir, 'visual_input_files')# 创建子文件夹
os.makedirs(image_dir, exist_ok=True)
os.makedirs(edited_dir, exist_ok=True)
os.makedirs(visual_input_dir, exist_ok=True)# 创建一个字典保存除了图像数据外的所有内容
non_image_data = []#-------------------#
# 遍历所有parquet文件
#-------------------#
for parquet_file in parquet_files:file_path = os.path.join(parquet_dir, parquet_file)print(f"Processing file: {file_path}")# 读取 parquet 文件df = pd.read_parquet(file_path)# 使用 tqdm 包装 DataFrame 的迭代器以显示进度条for index, row in tqdm(df.iterrows(), total=len(df), desc=f"Processing {parquet_file}"):# 提取图片字节数据image_file_data   = row['image_file']edited_file_data  = row['edited_file']visual_input_data = row['visual_input']# 提取除了图像列以外的其他数据row_data = {'image_id': row['image_id'],'edit_instruction': row['edit_instruction'],'edit_type': row['edit_type'],'input': row['input'],'output': row['output']}# 添加该行的非图像数据到字典列表中non_image_data.append(row_data)# 获取 image_id 来命名图像文件image_id = row['image_id']# 确保字节数据存在for image_data, label in [(image_file_data, 'image_file'), (edited_file_data, 'edited_file'), (visual_input_data, 'visual_input')]:if isinstance(image_data, dict) and 'bytes' in image_data:# 将字节数据转换为图片image = Image.open(io.BytesIO(image_data['bytes']))# 使用 image_id 作为图像文件名image_name = f"{image_id}.jpg"# 根据标签选择目标子文件夹if label   == 'image_file':target_dir = image_direlif label == 'edited_file':target_dir = edited_direlif label == 'visual_input':target_dir = visual_input_direlse:continue  # 如果标签没有匹配的，跳过# 定义目标保存路径target_path = os.path.join(target_dir, image_name)# 保存图片image.save(target_path)else:print(f"Skipping invalid data for {label} in row {index}")# 将非图像数据保存为 JSON 文件
json_file_path = './data.json'
with open(json_file_path, 'w') as json_file:json.dump(non_image_data, json_file)print(f"Non-image data saved to '{json_file_path}'")