Getting an RGBA array from a matplotlib image

Question

I was using imshow to plot an array with a custom colormap and boundarynorm. However, this is going to be an automated script and I want to save the image produced

Answer 1

matplotlib.colors.Colormap.__call__ performs the cmap and returns the RGBA array. https://matplotlib.org/3.3.1/api/_as_gen/matplotlib.colors.Colormap.html#matplotlib.colors.Colormap.__call__

# created using numpy 1.18.5 and matplotlib 3.2.2
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm, colors

x = np.linspace(0, 2, 1000)
X, Y = np.meshgrid(x, x)
data = np.sin((X-2)**3 + Y**4)
print(f"data.shape: {data.shape}")
print(f"data.dtype: {data.dtype}")

cmap: colors.Colormap = cm.get_cmap("rainbow")
norm: colors.Normalize = colors.Normalize()
# set min and max values from data
norm.autoscale(data)

# move scalar values to range [0, 1]
# can skip and pass directly to cmap if data already [0, 1]
normalised = norm(data)

# create a RBGA array
# bytes=True gives a uint8 array (Unsigned integer 0 to 255)
data_rgba = cmap(normalised, bytes=True)
print(f"data_rgba.shape: {data_rgba.shape}")
print(f"data_rgba.dtype: {data_rgba.dtype}")

# pass RBGA array to imsave and set origin to upper or lower
plt.imsave("my_data.png", data_rgba, origin="lower")

Should also be possible in one step using matplotlib.cm.ScalarMappable.to_rgba but I haven't tried that. https://matplotlib.org/3.3.1/api/cm_api.html#matplotlib.cm.ScalarMappable.to_rgba