制作具有内存限制的 h5py 文件的有效方法

Efficient way to make h5py file with memory constraint

假设我有如下图片:

root
|___dog
|    |___img1.jpg
|    |___img2.jpg
|    |___...
|    
|___cat
|___...

我想把图片文件做成h5py文件。

首先,我尝试读取所有图像文件并将其制作成 h5 文件。

import os
import numpy as np
import h5py
import PIL.Image as Image



datafile = h5py.File(data_path, 'w')


label_list = os.listdir('root')
for i, label in enumerate(label_list):
    files = os.listdir(os.path.join('root', label_list))
    for filename in files:
        img = Image.open(os.path.join('root', label, filename))
        ow, oh = 128, 128
        img = img.resize((ow, oh), Image.BILINEAR)
        data_x.append(np.array(img).tolist())
        data_y.append(i)


datafile = h5py.File(data_path, 'w')
datafile.create_dataset("data_image", dtype='uint8', data=data_x)
datafile.create_dataset("data_label", dtype='int64', data=data_y)

但由于内存限制我无法做到(每个文件夹有超过 200,000 张 224x224 大小的图像)。

那么,将此图片制作成 h5 文件的最佳方法是什么?

HDF5/h5py 数据集对象的内存占用比相同大小的 NumPy 数组小得多。 (这是使用 HDF5 的一个优势。)您可以在开始循环处理图像文件之前创建 HDF5 文件并分配数据集。然后您可以一次对图像进行操作(读取、调整大小和写入图像 0,然后是图像 1,等等)。

下面的代码为 200,000 张图像创建了必要的数据集。代码逻辑被重新安排以按照我描述的方式工作。 img_cnt 用于在现有数据集中定位新图像数据的变量。 (注意:我认为这按写的方式工作。但是没有数据,我无法测试,因此可能需要稍作调整。)如果你以后想调整数据集大小,你可以添加 maxshape=() create_dataset() 函数的参数。

# Open HDF5 and create datasets in advance
datafile = h5py.File(data_path, 'w')
datafile.create_dataset("data_image", (200000,224,224), dtype='uint8')
datafile.create_dataset("data_label", (200000,), dtype='int64')

label_list = os.listdir('root')
img_cnt = 0
for i, label in enumerate(label_list):
    files = os.listdir(os.path.join('root', label_list))
    for filename in files:
        img = Image.open(os.path.join('root', label, filename))
        ow, oh = 128, 128
        img = img.resize((ow, oh), Image.BILINEAR)
        datafile["data_image"][img_cnt,:,:] = np.array(img).tolist())
        datafile["data_label"][img_cnt] = i
        img_cnt += 1

datafile.close()