Unet分割5通道输出图像获取主色及对应标签的实现问题

• Unet分割头输出的5通道是每个像素对应5个类别的预测概率，不是直接的one-hot编码或者RGB值，直接对概率矩阵做去重统计没有实际意义
• 你写的colors=counts.argmax()返回的是counts数组最大值的索引位置，不是对应的颜色值，输出0仅代表出现次数最多的概率向量排在去重结果的第0位，和实际标签、颜色没有关联

import numpy as np

# 1. 定义标签与RGB的映射关系，和你给出的规则完全对应
label_rgb_map = [
    [0, 255, 0],    # label 0 绿色
    [0, 0, 255],    # label 1 蓝色
    [255, 255, 0],  # label 2 黄色
    [255, 0, 0],    # label 3 红色
    [255, 255, 255] # label 4 白色
]

# 2. 处理模型预测结果，得到每个像素的分类标签
predicted_img = model.predict(image[:,:,:])[0]
# 对最后一维取argmax，得到每个像素概率最高的类别，输出维度为(128,128)
pred_label_mask = np.argmax(predicted_img, axis=-1)

# 3. 统计全图出现次数最多的主类别
unique_labels, label_counts = np.unique(pred_label_mask, return_counts=True)
main_label = unique_labels[label_counts.argmax()]

# 4. 查表得到主类别对应的RGB值
main_rgb = label_rgb_map[main_label]
print('主色RGB值：', main_rgb)

如果需要生成完整的(128,128,3)的彩色预测掩码，可以追加以下代码：

# 生成全图彩色预测掩码
pred_rgb_img = np.zeros((128, 128, 3), dtype=np.uint8)
for label, rgb in enumerate(label_rgb_map):
    pred_rgb_img[pred_label_mask == label] = rgb

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