python - 使用 pytorch 对分割图像进行热编码

Question

我有一个分割图像作为大小为[1,1,256,256]的张量。该图像是二值分割图像。我想要对其进行热编码以获得大小为 [1,2,256,256] 的图像。我尝试了torch.nn.function.one_hot(img, 2)。但它给了我一个大小为 [1,256,256,2] 的图像。如何获得所需的张量？

Answer 1

尝试使用transpose():

img_one_hot = torch.nn.functional.one_hot(img, 2).transpose(1, 4).squeeze(-1)

transpose(1, 4) - 交换第一维和第四维，返回形状为 [1, 2, 256, 256, 1], squeeze(-1) 删除最后一个暗淡的结果，得到 [1 , 2, 256, 256] 形状的张量。