Python：遍历4维ndarray并批量保存其中的单通道图像

直接遍历数组第0维的索引即可实现批量处理，无需重复编写代码，以下提供两种实现方案：

方案1：使用PIL直接存储数组（推荐，效率更高） #

该方案跳过matplotlib画布渲染步骤，速度更快，无内存泄漏风险：

import os
import numpy as np
from PIL import Image

generate_images_path = '/content/drive/My Drive/test/predicted image/'
# 自动创建保存路径，避免路径不存在报错
os.makedirs(generate_images_path, exist_ok=True)

# 遍历所有图像的索引
for img_idx in range(preds_test_t.shape[0]):
    # 提取单张图像的二维数组，去掉最后一维的单通道维度
    single_img = preds_test_t[img_idx, :, :, 0]
    # 若数组值为0-1范围的浮点数，转为0-255的uint8格式适配图像存储要求
    if single_img.max() <= 1.0:
        single_img = (single_img * 255).astype(np.uint8)
    # 生成单通道灰度图像对象
    img = Image.fromarray(single_img, mode='L')
    # 按序号命名保存
    save_path = os.path.join(generate_images_path, f'predicted image ({img_idx + 1}).tiff')
    img.save(save_path)

方案2：兼容原有matplotlib写法的循环版本 #

如果需要保留原有matplotlib的渲染参数，可以用循环改写原有代码，注意每次用完关闭画布释放内存：

import os
import matplotlib.pyplot as plt

generate_images_path = '/content/drive/My Drive/test/predicted image/'
os.makedirs(generate_images_path, exist_ok=True)

dpi = 96
# 修正原代码的高度参数，和数组实际高度464对齐
figsize = 512 / float(dpi), 464 / float(dpi)

for img_idx in range(preds_test_t.shape[0]):
    fig = plt.figure(figsize=figsize) 
    ax = fig.add_axes([0, 0, 1, 1])
    ax.imshow(preds_test_t[img_idx,:,:,0], cmap=plt.cm.gray, interpolation = "nearest")
    plt.axis('off')
    ax.set(xlim=[-0.5, 512 - 0.5], ylim=[464 - 0.5, -0.5], aspect=1)
    save_path = os.path.join(generate_images_path, f'predicted image ({img_idx + 1}).tiff')
    fig.savefig(save_path, dpi=dpi, bbox_inches='tight', pad_inches = 0)
    # 关闭画布释放内存，避免批量处理时卡顿崩溃
    plt.close(fig)

注意事项 #

• 若你的ndarray值为0_{1范围的归一化浮点数，必须转为0}255的uint8类型，否则会出现存储的图像全黑/全白的异常
• 原代码中使用的462高度为笔误，需和数组实际高度464对齐，避免图像被裁剪
• 批量处理数量大时优先使用PIL方案，渲染开销低，稳定性更高

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