多类别3D Nii掩码转2D PNG图像输出黑白问题求助

我来帮你排查下问题根源，以及给出针对性的解决办法：

问题分析 #

你当前的代码把掩码值乘以255后转成float32，然后保存为单通道PNG。这里有两个核心问题：

• OpenCV的imwrite处理单通道float32图像时，会默认将其当作灰度图渲染，0对应黑色、255对应白色，多类别掩码的离散类别值被转换成了连续的灰度值，自然看不到类别区分。
• 多类别掩码的核心是离散的整数类别标签（比如0=背景、1=左心室、2=心肌等），直接乘以255会破坏原始标签的语义，反而让数值失去了类别区分的意义。

解决方案 #

根据你的需求，分两种场景给出修改后的代码：

场景1：保存原始整数掩码（适合后续模型训练）

如果只是需要把掩码保存为可读取的格式，供后续模型训练使用，不需要即时可视化彩色效果，直接保留原始类别整数即可：

import nibabel as nib
import numpy as np
import cv2
import os

# 加载掩码文件
input_image_mask = nib.load('masks/mask0A0S9V9_sa_gt.nii.gz')
input_mask = input_image_mask.get_fdata()  # 替代已弃用的get_data()，nibabel官方推荐用法

for i in range(input_mask.shape[2]):
    # 提取当前Z轴切片，转成uint8类型（只要类别数<=255就适用）
    input_slice_mask = input_mask[:, :, i].astype(np.uint8)
    # 构造保存路径
    path_mask = os.path.join('Desktop/Ilkay/M/masks/2d/', f"slicem{i}.png")
    # 保存单通道整数PNG
    cv2.imwrite(path_mask, input_slice_mask)

后续需要可视化时，再读取PNG并给不同类别分配颜色即可。

场景2：直接保存彩色可视化的多类别掩码

如果需要直接得到带颜色区分的可视化结果，可以手动给每个类别分配颜色：

import nibabel as nib
import numpy as np
import cv2
import os

# 自定义类别-颜色映射，根据你的实际类别数和需求调整
color_map = {
    0: (0, 0, 0),       # 背景：黑色
    1: (0, 0, 255),     # 类别1：红色
    2: (0, 255, 0),     # 类别2：绿色
    3: (255, 0, 0)      # 类别3：蓝色
}

input_image_mask = nib.load('masks/mask0A0S9V9_sa_gt.nii.gz')
input_mask = input_image_mask.get_fdata()

for i in range(input_mask.shape[2]):
    input_slice_mask = input_mask[:, :, i].astype(np.uint8)
    # 创建空白彩色图像（H, W, 3）
    color_mask = np.zeros((input_slice_mask.shape[0], input_slice_mask.shape[1], 3), dtype=np.uint8)
    # 给每个类别填充对应颜色
    for class_val, color in color_map.items():
        color_mask[input_slice_mask == class_val] = color
    # 保存彩色PNG
    path_mask = os.path.join('Desktop/Ilkay/M/masks/2d/', f"slicem_color{i}.png")
    cv2.imwrite(path_mask, color_mask)

额外注意事项 #

• 如果你的类别数量超过255，需要改用uint16类型保存，但大部分可视化工具对16位PNG支持有限，建议先将类别值归一化到0-255范围，或者优先使用彩色映射方案。
• 代码中替换了get_data()为get_fdata()，因为前者在nibabel的新版本中已经被弃用，后者是官方推荐的获取图像数据的方式。

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