[SOLVED]Convert color of tensors

Hi, thank you always for your help.
I am trying to convert an image set stored in torch.Tensor of like [number_of_images, width, height],
to a set of color images of like [number_of_images, 3, width, height].
I would like to use the color scheme ‘magma’ that can be used in matplotlib.
Is there any good way to realize it?

Thank you in advance :slight_smile:

You could convert the tensors to numpy arrays and apply the colormap on them:

import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib import cm


x = np.random.randn(100, 100).clip(0, 1)

magma = cm.get_cmap('magma')
x_transformed = magma(x)
plt.imshow(x_transformed)

plt.imshow(x)
plt.imshow(x, cmap='magma')

Note that the colormap expects normalized inputs, so I clipped the data to [0, 1].
Once you have these transformed arrays, you could convert them back to tensors via torch.from_numpy.

PS: What is your use case that you want to artificially inflate your channels?

1 Like

Hi ptrblck,
Thank you always for your support!
It solved my problem.
Thank you very much!

1 Like

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