我的PyTorch模型输出一个分段图像，每个类别的值为(0,1,2)。在准备数据集时，我将黑色映射为0，红色映射为1，白色映射为2。我有两个问题:
-
如何展示每个类别代表什么？例如，看看这张图片:
我目前使用以下方法展示每个类别:output = net(input) input = input.cpu().squeeze() input = transforms.ToPILImage()(input) probs = F.softmax(output, dim=1) probs = probs.squeeze(0) full_mask = probs.squeeze().cpu().numpy() fig, (ax0, ax1, ax2, ax3, ax4) = plt.subplots(1, 5, figsize=(20,10), sharey=True) ax0.set_title('输入图像') ax1.set_title('背景类') ax2.set_title('神经元类') ax3.set_title('树突类') ax4.set_title('预测掩码') ax0.imshow(input) ax1.imshow(full_mask[0, :, :].squeeze()) ax2.imshow(full_mask[1, :, :].squeeze()) ax3.imshow(full_mask[2, :, :].squeeze()) full_mask = np.argmax(full_mask, 0) img = mask_to_image(full_mask)
但似乎类别之间存在共享像素，有没有更好的展示方法（我希望第一张图像仅显示背景类，第二张仅显示神经元类，第三张仅显示树突类）?
2.我的第二个问题是关于从掩码生成黑色、红色和白色的图像，目前掩码的形状为(512,512)，其值如下:
[[0 0 0 ... 0 0 0] [0 0 0 ... 2 0 0] [0 0 0 ... 2 2 0] ... [2 1 2 ... 2 2 2] [2 1 2 ... 2 2 2] [0 2 0 ... 2 2 2]]结果看起来像这样:
因为我使用以下代码转换为图像:
def mask_to_image(mask): return Image.fromarray((mask).astype(np.uint8))回答：
但似乎类别之间存在共享像素，有没有更好的展示方法（我希望第一张图像仅显示背景类，第二张仅显示神经元类，第三张仅显示树突类）?
是的，你可以沿0th维度取argmax，这样具有最高logit（未归一化的概率）的将为1，其余为零:
output = net(input)binary_mask = torch.argmax(output, dim=0).cpu().numpy()ax.set_title('神经元类')ax.imshow(binary_mask == 0)我的第二个问题是关于从掩码生成黑色、红色和白色的图像，目前掩码的形状为(512,512)，其值如下
你可以将[0, 1, 2]值扩展到零轴，使其按通道排列。现在，单个像素的[0, 0, 0]值在所有通道上将为黑色，[255, 255, 255]将为白色，[255, 0, 0]将为红色（因为PIL使用RGB格式）:
