如何在PyTorch中修改图像分割掩码的类别标签顺序

问题根因 #

你当前使用的mask_four为one-hot编码的4通道分割掩码，最后一维的索引对应原始类别ID，你之前使用查找表直接索引的方法仅适用于单通道类别索引掩码，不适用于one-hot格式，同时查找表的映射方向设置错误，因此转换结果不符合预期。

对应解决方案 #

场景1：当前为one-hot编码掩码（形状为[N, H, W, 类别数]） #

你需要的类别映射规则为：

• 原始ID 0（背景）→ 新ID 3
• 原始ID 1（磨玻璃影）→ 新ID 0
• 原始ID 2（实变影）→ 新ID 1
• 原始ID 3（其他肺组织）→ 新ID 2

one-hot格式直接调整最后一维的通道顺序即可实现类别顺序调整，参考代码如下：

# 定义通道顺序：新通道i对应原始通道的order[i]
order = torch.tensor([1, 2, 3, 0], dtype=torch.long)
# 沿最后一维（类别通道）重排
test_mask2 = torch.index_select(test_mask, dim=-1, index=order)

# 可视化验证
fig, axes = plt.subplots(2, 2, figsize=(10, 8))
axes[0, 0].imshow(test_mask2[:, :, 0])
axes[0, 0].set_title('Ground glass')
axes[0, 1].imshow(test_mask2[:, :, 1])
axes[0, 1].set_title('Consolidation')
axes[1, 0].imshow(test_mask2[:, :, 2])
axes[1, 0].set_title('Lung other')
axes[1, 1].imshow(test_mask2[:, :, 3])
axes[1, 1].set_title('Background')

场景2：当前为单通道类别索引掩码（形状为[N, H, W]，像素值为类别ID） #

该场景下可以使用查找表实现转换，查找表的定义规则为lut[原始类别ID] = 新类别ID，参考代码如下：

lut = torch.tensor([3, 0, 1, 2], dtype=torch.long)
# 取单张掩码并转long类型
single_channel_mask = mask_single[31].long()
converted_mask = lut[single_channel_mask]

内容的提问来源于stack exchange，提问作者ALiCe P.