FCN模型预测分割图像可视化报错：Invalid shape (128,128,2) for image data
2026-7-9
解决imshow可视化FCN预测结果时的形状错误
报错的核心原因是plt.imshow()不支持通道数为2的图像输入。你的FCN模型输出preds_test[iix]形状为(128,128,2)，对应2个类别的预测概率图，np.squeeze()仅能移除维度为1的轴，无法处理通道数为2的情况。
解决方法：
- 从多通道预测结果中提取每个像素的类别索引：用
np.argmax()对最后一个通道维度取最大值，得到形状为(128,128)的单通道分割掩码 - 再用
imshow可视化这个掩码
修改后的代码示例：
print('Segmented_Test_Mask') # 对通道维度取argmax，得到每个像素对应的类别 segment_mask = np.argmax(preds_test[iix], axis=-1) plt.imshow(segment_mask) plt.show() # 加上show()确保图像正常显示
补充说明：
- 如果模型输出是经过softmax的概率图，
argmax会选出每个像素概率最高的类别，正好对应分割掩码 - 若需要可视化单个类别的概率分布，可以单独提取对应通道的结果（比如
preds_test[iix][..., 0]）再用imshow显示
内容的提问来源于stack exchange，提问作者MOHAMMAD
相关产品推荐
相关产品推荐
