首页 > 解决方案 > 制作具有内存限制的 h5py 文件的有效方法

问题描述

假设我有如下图像:

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 文件的最佳方法是什么?

标签: pythonpython-3.xh5py

解决方案


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()

推荐阅读