CRSNet适配单HDF5格式训练测试数据集的方法咨询

原load_data函数的核心逻辑是匹配“单张图像-对应独立单h5标注”的存储结构，不需要改动原有图像预处理、密度图缩放、随机数据增强的业务逻辑，只需要替换数据读取的入口层，就能快速适配新数据集，具体改造步骤如下：

第一步：先确认新h5文件的内部存储结构 #

新拿到的全量h5文件没有统一的存储规范，必须先打开文件确认键名、数组形状、索引对应关系，常见的存储格式有两类：

• 按类别分键：比如images键存所有图像的数组（形状为N, H, W, 3），density键存所有密度图的数组（形状为N, H, W），两个数组的索引一一对应（第i张图配第i个密度图）
• 按样本ID分键：比如img_001存单张图像，gt_001存对应密度图

可以直接用这段临时代码排查结构：

import h5py
# 替换成你的全量h5文件路径
with h5py.File('./new_dataset.h5', 'r') as f:
    print("h5内所有键名：", list(f.keys()))
    for key in f.keys():
        print(f"键 {key} 对应的数组形状：{f[key].shape}")

第二步：重构load_data的读取逻辑，保持接口兼容 #

改造的核心是两点：一是避免每次取样本都重复打开h5文件拖慢训练速度，二是保持load_data返回的(img, target)对象格式和原代码完全一致，不需要改动后续CRSNet的训练流程、模型输入定义。
改造后的完整代码如下，关键修改点都做了注释：

import random
import os
from PIL import Image
import numpy as np
import h5py
import cv2

# 全局数据缓存，避免每次读取样本重复打开h5文件
DATA_CACHE = {}

def init_full_h5(h5_path, train_flag=True):
    """训练启动前调用，一次性把全量h5数据加载到内存缓存"""
    split_name = "train" if train_flag else "test"
    if split_name not in DATA_CACHE:
        with h5py.File(h5_path, "r") as f:
            # 注意：这里的键名要替换成第一步排查出来的实际键名
            # 比如常见的键是'images'存图像，'density'存密度图
            all_imgs = np.asarray(f["images"])
            all_dmaps = np.asarray(f["density"])
            DATA_CACHE[split_name] = {
                "imgs": all_imgs,
                "dmaps": all_dmaps,
                "sample_num": len(all_imgs)
            }

def load_data(sample_idx, train_flag=True):
    """
    改造后的加载函数：入参从原来的图像路径换成样本索引
    完全保留原有的所有预处理逻辑，和原模型接口100%兼容
    """
    split_name = "train" if train_flag else "test"
    if split_name not in DATA_CACHE:
        raise RuntimeError("请先在训练初始化阶段调用init_full_h5加载对应拆分的h5数据！")
    
    # 从缓存取对应索引的样本
    data = DATA_CACHE[split_name]
    img_np = data["imgs"][sample_idx]
    # 把numpy数组转回PIL图像，和原代码Image.open返回的对象格式完全一致
    img = Image.fromarray(img_np.astype(np.uint8)).convert("RGB")
    target = data["dmaps"][sample_idx]

    # 完全保留原有的随机数据增强逻辑（原代码默认关闭增强，需要开启的话把if False改成if train_flag即可）
    if False:
        crop_size = (img.size[0]//2, img.size[1]//2)
        dx = int(random.random() * img.size[0] * 1./2)
        dy = int(random.random() * img.size[1] * 1./2)
        img = img.crop((dx, dy, crop_size[0]+dx, crop_size[1]+dy))
        target = target[dy:crop_size[1]+dy, dx:crop_size[0]+dx]
        if random.random() > 0.8:
            target = np.fliplr(target)
            img = img.transpose(Image.FLIP_LEFT_RIGHT)

    # 完全保留原代码的密度图8倍下采样、乘64的缩放逻辑，不做任何改动
    target = cv2.resize(
        target, 
        (target.shape[1]//8, target.shape[0]//8), 
        interpolation=cv2.INTER_CUBIC
    ) * 64

    return img, target

第三步：适配Dataset类的调用逻辑 #

原来的PyTorch Dataset类是根据图像路径列表传参给load_data，现在只要把路径列表的入参换成样本索引即可，参考实现：

from torch.utils.data import Dataset

class NewFullH5Dataset(Dataset):
    def __init__(self, h5_path, train_flag=True):
        # 初始化时一次性加载全量h5数据
        init_full_h5(h5_path, train_flag)
        self.train_flag = train_flag
        self.sample_num = DATA_CACHE["train" if train_flag else "test"]["sample_num"]
    
    def __getitem__(self, idx):
        # 原来传img_path的逻辑替换为传样本索引
        img, target = load_data(idx, self.train_flag)
        return img, target
    
    def __len__(self):
        return self.sample_num

关键注意事项 #

• 原代码里img_path.replace('.jpg','.h5').replace('images','ground_truth')的路径匹配逻辑在新数据集上完全失效，直接删掉即可
• 要确认新h5存储的图像通道顺序：原代码用PIL读取的是RGB格式，如果h5里存的是cv2默认的BGR格式，需要在转PIL之前加一行img_np = img_np[..., ::-1]做通道转换，否则模型训练精度会异常
• 如果h5文件大小超过内存承载上限，就不要一次性把所有数据读进内存，改成常驻一个h5文件句柄，在取样本时按索引从句柄里切分读取即可，不要每次取样本都重复h5py.File()打开文件，避免IO卡顿

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