Python中如何调整字典内列表结构以实现预期数组存储格式？

我用以下代码初始化数据集字典：

train_dataset, test_dataset, valid_dataset = ({'image':[], 'mask':[]} for i in range(3))

我期望执行print(train_dataset[0])后，得到单个数组作为image和mask的值的结构（仅关注层级）：

{'image': array([[[238, 238, 239, ..., 226, 227, 230],
    [234, 238, 240, ..., 226, 227, 225],
    [255, 255, 255, ..., 255, 255, 255]],
    ...,
    [255, 255, 255, ..., 255, 255, 255],
    [255, 255, 255, ..., 255, 255, 255],
    [255, 255, 255, ..., 255, 255, 255]]], dtype=uint8),  
'mask': array([[[0., 0., 0., ..., 0., 0., 0.],
    [0., 0., 0., ..., 0., 0., 0.],
    [0., 0., 0., ..., 0., 0., 0.],
    ...,
    [2, 2, 2, ..., 2, 2, 2],
    [2, 2, 2, ..., 2, 2, 2],
    [2, 2, 2, ..., 2, 2, 2]]], dtype=uint8)}

但实际得到的是数组列表作为image的值：

{'image': [array([[40, 40, 40, ..., 59, 59, 59],
    [40, 40, 40, ..., 59, 59, 59],
    [40, 40, 40, ..., 59, 59, 59],
    ...,
    [42, 42, 42, ..., 41, 41, 41],
    [42, 42, 42, ..., 41, 41, 41],
    [42, 42, 42, ..., 41, 41, 41]], dtype=uint8),  
array([[11, 11, 14, ..., 28, 34, 35],
    [11, 12, 14, ..., 27, 32, 33],
    [12, 12, 15, ..., 21, 21, 21],
    ...,
    [7, 7, 6, ..., 13, 19, 20],
    [11, 11, 8, ..., 12, 20, 21],
    [12, 12, 9, ..., 12, 20, 21]], dtype=uint8),  
array([[19, 19, 19, ..., 39, 39, 39],
    [19, 19, 19, ..., 39, 39, 39],
    [19, 19, 19, ..., 39, 39, 39],
    ...,
    [0, 0, 0, ..., 0, 0, 0],
    [0, 0, 0, ..., 0, 0, 0],
    [0, 0, 0, ..., 0, 0, 0]], dtype=uint8)]}

请问如何调整代码，使字典内的结构符合预期？

问题出在初始化字典时，'image'和'mask'被设为了空列表，后续你应该是把多个数组追加到了列表里，导致最终是数组的列表。要达到预期结构，有两种调整方式：

方式1：直接赋值单张图片的数组 #

如果每个数据集字典对应单张图片的掩码对，初始化时可以把值设为None，后续填充时直接替换为数组：

# 初始化时用None占位
train_dataset, test_dataset, valid_dataset = ({'image': None, 'mask': None} for i in range(3))

# 填充数据（示例）
train_dataset['image'] = img_array  # img_array是单张图片的数组
train_dataset['mask'] = mask_array  # mask_array是对应掩码的数组

方式2：合并多个数组为多维数组 #

如果你的需求是把所有图片数组合并成一个多维数组（比如形状为(样本数, 高度, 宽度, 通道数)），可以先收集数组到列表，再用numpy合并：

import numpy as np

# 初始化空列表收集数组
train_dataset = {'image':[], 'mask':[]}

# 循环添加数组到列表（示例）
train_dataset['image'].append(img_array1)
train_dataset['image'].append(img_array2)
train_dataset['mask'].append(mask_array1)
train_dataset['mask'].append(mask_array2)

# 合并列表为单个数组
train_dataset['image'] = np.stack(train_dataset['image'])
train_dataset['mask'] = np.stack(train_dataset['mask'])

两种方式处理后，train_dataset['image']都会是单个数组，而非数组的列表，符合你预期的结构。

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