以AnyEdit數(shù)據(jù)集為例，我們展示了批量讀取parquet圖片文件的代碼，相關(guān)數(shù)據(jù)集地址如下：
Huggingface-AnyEdit

代碼如下所示，請注意使用時應(yīng)查看文件格式來適應(yīng)性修改代碼

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')# 查看第一行數(shù)據(jù)
sample_row = df.iloc[0]  # 獲取第一行數(shù)據(jù)
print(sample_row)#--------------------------------#
# 獲取data文件夾下的所有parquet文件
#--------------------------------#
parquet_dir      = './data'
parquet_files    = [f for f in os.listdir(parquet_dir) if f.endswith('.parquet')]# 創(chuàng)建新的文件夾來保存圖片
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')# 創(chuàng)建子文件夾
os.makedirs(image_dir, exist_ok=True)
os.makedirs(edited_dir, exist_ok=True)
os.makedirs(visual_input_dir, exist_ok=True)# 創(chuàng)建一個字典保存除了圖像數(shù)據(jù)外的所有內(nèi)容
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}"):# 提取圖片字節(jié)數(shù)據(jù)image_file_data   = row['image_file']edited_file_data  = row['edited_file']visual_input_data = row['visual_input']# 提取除了圖像列以外的其他數(shù)據(jù)row_data = {'image_id': row['image_id'],'edit_instruction': row['edit_instruction'],'edit_type': row['edit_type'],'input': row['input'],'output': row['output']}# 添加該行的非圖像數(shù)據(jù)到字典列表中non_image_data.append(row_data)# 獲取 image_id 來命名圖像文件image_id = row['image_id']# 確保字節(jié)數(shù)據(jù)存在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:# 將字節(jié)數(shù)據(jù)轉(zhuǎn)換為圖片image = Image.open(io.BytesIO(image_data['bytes']))# 使用 image_id 作為圖像文件名image_name = f"{image_id}.jpg"# 根據(jù)標簽選擇目標子文件夾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}")# 將非圖像數(shù)據(jù)保存為 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}'")