将多文件夹图像转换为HDF5文件用于分类模型的技术求助

实现思路 #

• 遍历目标根目录下的所有子文件夹，收集每张图像的路径及其所属文件夹名称（作为分类标签）
• 在HDF5文件中创建label和data两个专属组
• 对每张图像，在data组内创建对应编号的数据集存储预处理后的图像数据，同时在label组内创建同编号的数据集存储标签，保证两者一一匹配

完整代码 #

import os
import cv2
import h5py

# 配置参数
IMG_WIDTH = 30
IMG_HEIGHT = 30
ROOT_DIR = "./"  # 替换为你的图像文件夹根路径
HDF5_FILE = "image_classification_dataset.h5"

# 收集所有图像的路径和对应标签
image_info_list = []
for folder_name in os.listdir(ROOT_DIR):
    folder_path = os.path.join(ROOT_DIR, folder_name)
    if not os.path.isdir(folder_path):
        continue
    # 筛选常见图像格式文件，可根据需求扩展后缀
    for img_filename in os.listdir(folder_path):
        if img_filename.lower().endswith(('.jpeg', '.jpg', '.png')):
            img_full_path = os.path.join(folder_path, img_filename)
            image_info_list.append((img_full_path, folder_name))

# 创建并写入HDF5文件
with h5py.File(HDF5_FILE, 'w') as h5f:
    # 创建指定组
    data_group = h5f.create_group('data')
    label_group = h5f.create_group('label')
    
    # 逐个处理图像并写入对应组
    for idx, (img_path, label_str) in enumerate(image_info_list):
        # 读取并预处理图像
        img = cv2.imread(img_path)
        img_resized = cv2.resize(img, (IMG_WIDTH, IMG_HEIGHT))
        # 转换为RGB格式（cv2默认读取为BGR）
        img_rgb = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB)
        
        # 在data组创建对应编号的数据集存储图像
        data_group.create_dataset(str(idx), data=img_rgb, dtype='uint8')
        # 在label组创建同编号的数据集存储标签字符串
        label_group.create_dataset(str(idx), data=label_str.encode('utf-8'), dtype=h5py.string_dtype(encoding='utf-8'))

print(f"转换完成，生成HDF5文件：{HDF5_FILE}")

关键细节说明 #

• 组的创建：通过h5f.create_group()方法创建指定组，后续所有数据集都挂载到对应组下，满足你对分组的要求
• 数据集命名：使用图像的索引编号作为数据集名称，确保data和label组内的数据集一一对应
• 标签存储：将字符串标签编码为UTF-8格式，通过h5py.string_dtype指定字符串类型，避免出现乱码问题
• 图像预处理：统一调整图像尺寸为指定的30x30，并转换为RGB格式适配多数分类模型的输入要求

可选优化 #

• 若需要将标签转为整数编码（适合多数分类模型训练），可提前建立标签-整数映射字典，存储编码后的数值
• 针对超大规模图像数据集，可使用HDF5的分块存储功能，分批次写入数据以降低内存占用

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