One hot encoding a segmented image using pytorch

Question

I have a segmented image as a tensor of size [1,1,256,256]. The image is a binary segmented image. I want to one hot encode it to get an image of size [1,2,256,256]. I tried torch.nn.functional.one_hot(img, 2). But it gave me an image of size [1,256,256,2]. How do I get the desired tensor?

Answer 1

Try to use transpose():

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

transpose(1, 4) - swaps 1st and 4th dimension, returning the tensor of the shape of [1, 2, 256, 256, 1], squeeze(-1) removes the last dim resulting in [1 , 2, 256, 256] shaped tensor.