如何加速pd.DataFrame()将图像数据集转换为CSV的过程？

问题描述 #

我当前通过以下代码将原始图像数据集转换为CSV：第一列保存类别标签（即文件夹名称），后续列保存图像像素值。具体做法是将标签与像素值堆叠后加入列表，再通过pd.DataFrame(v)创建数据框，但该步骤耗时极长。请问有何方法可加速此过程并将数据追加至CSV文件？

# directory where your raw dataset of images is stored
root = 'test_case_images/Test_1/'
width = 224
height = 224
image_pixels = []
v=[]
# folder here will be the considered as class of image
for folder in os.listdir(root):
    for filename in os.listdir(os.path.join(root, folder)):
        current_filepath = (os.path.join(root,folder))
        # reading image
        img = cv2.imread(os.path.join(current_filepath,filename))
        # fix image color
        # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        # changing current image to grayscale
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        # resizing image
        img = cv2.resize(img,(width,height))
        # flatten image
        value = img.flatten()
        # stacking the label and the pixel value
        stacked_value = np.hstack((folder,value))
        v.append(stacked_value)
df = pd.DataFrame(v)
# saving the dataframe
df.to_csv('data.csv', mode='a',header=False, index=False)
v.clear()

解决方案 #

这个问题我之前也碰到过，核心瓶颈主要在列表动态追加+一次性转DataFrame的操作，以及重复的文件路径拼接和图像处理开销上。下面给你几个实用的优化方案，亲测能大幅提升速度：

1. 用Numpy数组替代Python列表存储数据 #

Python列表的动态追加会频繁触发内存扩容，效率极低。我们可以预先统计总图像数量，直接创建固定大小的Numpy数组来存储数据，从根源上解决内存操作的低效问题：

import os
import cv2
import numpy as np
import pandas as pd

root = 'test_case_images/Test_1/'
width = 224
height = 224
pixel_count = width * height

# 先统计总图像数，避免动态扩容
total_images = 0
for folder in os.listdir(root):
    folder_path = os.path.join(root, folder)
    if os.path.isdir(folder_path):  # 跳过非文件夹的文件
        total_images += len(os.listdir(folder_path))

# 预先创建数组：行=总图像数，列=1（标签）+像素数
data_array = np.empty((total_images, 1 + pixel_count), dtype=object)  # 标签用object类型，像素用数值类型
current_idx = 0

for folder in os.listdir(root):
    folder_path = os.path.join(root, folder)
    if not os.path.isdir(folder_path):
        continue
    for filename in os.listdir(folder_path):
        img_path = os.path.join(folder_path, filename)
        # 直接读取为灰度图，减少一次色彩转换
        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
        if img is None:  # 跳过损坏的图像
            continue
        img_resized = cv2.resize(img, (width, height))
        img_flattened = img_resized.flatten()
        # 填充数组
        data_array[current_idx, 0] = folder
        data_array[current_idx, 1:] = img_flattened
        current_idx += 1

# 转DataFrame并写入
df = pd.DataFrame(data_array[:current_idx])  # 跳过可能的空行（如果有损坏图像）
df.to_csv('data.csv', mode='a', header=False, index=False)

2. 分批次写入CSV，避免内存过载 #

如果数据集非常大（比如几万张以上），一次性存储所有数据到数组还是会占用大量内存。可以设置一个批次大小，每处理完一批就写入CSV，然后清空当前批次的数据：

import os
import cv2
import numpy as np
import pandas as pd

root = 'test_case_images/Test_1/'
width = 224
height = 224
pixel_count = width * height
batch_size = 1000  # 每1000张写入一次
batch_data = []

for folder in os.listdir(root):
    folder_path = os.path.join(root, folder)
    if not os.path.isdir(folder_path):
        continue
    for filename in os.listdir(folder_path):
        img_path = os.path.join(folder_path, filename)
        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
        if img is None:
            continue
        img_resized = cv2.resize(img, (width, height))
        img_flattened = img_resized.flatten()
        stacked_value = np.hstack((folder, img_flattened))
        batch_data.append(stacked_value)
        
        # 达到批次大小就写入
        if len(batch_data) >= batch_size:
            df_batch = pd.DataFrame(batch_data)
            df_batch.to_csv('data.csv', mode='a', header=False, index=False)
            batch_data.clear()  # 清空批次数据

# 写入剩余的不足一批的数据
if batch_data:
    df_batch = pd.DataFrame(batch_data)
    df_batch.to_csv('data.csv', mode='a', header=False, index=False)

3. 其他细节优化 #

• 减少路径拼接次数：提前拼接好folder_path，避免在循环中重复调用os.path.join
• 直接读取灰度图：用cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)代替先读彩色再转灰度，减少一次色彩转换操作
• 跳过损坏图像：增加if img is None的判断，避免处理损坏的图像导致报错
• 使用更高效的CSV写入方式：如果追求极致速度，可以用numpy.savetxt代替DataFrame写入，比如直接把批次的numpy数组写入CSV：

  np.savetxt('data.csv', np.array(batch_data), delimiter=',', fmt='%s', mode='a')
  

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