Ludovic Gardy
Ludovic Gardy

Reputation: 41

Simple matrix to RGB

Let us say I want tu use the color map "jet" on a very simple matrix to display the graphic in matplotlib :

import numpy as np
import matplotlib.pyplot as plt

arr = np.array( [ [4,5,12,3], [2,4,5,4], [1,3,5,7], [3,9,10,12], [17,1,9,22] ] )
plt.imshow(arr, cmap="jet")

Now I would like to get back a matrix with the exact same shape (5,4) but in RGB dimensions (5,4,3), with the RGB values being the one presented on the matplotlib image. Any idea on how to do it please ?

Thank you !

Upvotes: 1

Views: 301

Answers (1)

Corralien
Corralien

Reputation: 120549

You have to compute the colors yourself:

# normalization
norm = (arr - np.min(arr)) / (np.max(arr) - np.min(arr))
cmap = plt.cm.jet  # colormap, it can be inferno, rainbow, viridis, etc

# convert to image and remove alpha channel (the fourth dimension)
image = np.round(255 * cmap(norm)).astype(int)[:, :, :3]

# check the result
plt.imshow(image)
plt.show()

Output:

enter image description here

Upvotes: 3

Related Questions