[SOLVED]Convert color of tensors

[Ren_Pang (Local State)]

Add one note, that might not be the case when you use a cmap in matplotlib.
If you have a manual colormap that is torch.Tensor with shape [256, 3], you could try this cool method:

color_map = #Tensor of shape(256,3)
gray_image = (gray_image * 255).long() # Tensor values between 0 and 255 and LongTensor and shape of (512,512)
output = color_map[gray_image] #Tensor of shape (512,512,3)

from

Pytorch colormap gather operation

You can use this code color_map = #Tensor of shape(256,3) gray_image = (gray_image * 255).long() # Tensor values between 0 and 255 and LongTensor and shape of (512,512) output = color_map[gray_image] #Tensor of shape (512,512,3) Why does this work? The gray_image tensor is used as an index in the 0th dimension of color_map. The color_map tensor is accessed with this index and the new tensor created has the shape of gray_image inserted in the 0th dimension of output.