Search code examples
pythonimagematplotlibcolorsrenderer

Extracting unique colors of an overlay AxesImage in matplotlib

This is a question about color information in matplotlib images.

I plot two arrays with the following code:

import numpy as np
import matplotlib.pyplot as plt

M1 = ([1,      2,      3,      np.nan], 
      [4,      5,      np.nan, np.nan], 
      [6,      7,      8,      9])

M2 = ([np.nan, np.nan, np.nan, np.nan],
      [np.nan, 1,      2,      3], 
      [np.nan, 4,      5,      6])

M1arr = ~np.isnan(M1)
M2arr = ~np.isnan(M2)

fig, ax = plt.subplots()

im1 = ax.imshow(M1arr, cmap="Reds",  alpha=0.5)
im2 = ax.imshow(M2arr, cmap="Blues", alpha=0.5)

#color_array = mystery_function(im1, im2, ax, fig) 

plt.show()

Output:
![![enter image description here

Is there a way to extract the colors from the plotted compound image that we finally see (to create a colorbar, for instance)? I have seen this amazing answer of how to reverse-engineer the colors from im1 or im2. BUT im2 is not the compound overlay image that we will finally see, im1 and im2 are seemingly only combined by plt.show(). I also tried to force matplotlib to pre-emptively generate the final image with plt.draw() and extract the image with ax.get_images(), alas the two images were still separated.

I am not interested in how to solve this differently - after useless attempts, I changed my strategy and plotted the combined matrix instead. My question is specifically, if we can extract the four colors from the AxesImage shortly before it is displayed.
Equally helpful would be information on how matplotlib combines the colours in the overlay. I tried summation of the colours in im1 and im2 (obviously wrong, because it can exceed 1) and mean value of each color channel (also wrong).

Solution

  • After some digging, I found the answer of how to get the rendered array. It is based on drawing the figure, as I tried before, and retrieving a memoryview of the renderer's buffer with fig.canvas.buffer_rgba():

    import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colors

M1 = ([1,      2,      3,      np.nan], 
      [4,      5,      np.nan, np.nan], 
      [6,      7,      8,      9])

M2 = ([np.nan, np.nan, np.nan, np.nan],
      [np.nan, 1,      2,      3], 
      [np.nan, 4,      5,      6])

M1arr = ~np.isnan(M1)
M2arr = ~np.isnan(M2)

fig, ax = plt.subplots()

ax.imshow(M1arr, cmap="Reds",  alpha=0.5)
ax.imshow(M2arr, cmap="Blues", alpha=0.5)


#get image without axis, so only the colors plotted in the overlay image are considered
ax.axis("off")
#get rgba values from image predrawn by the renderer
fig.canvas.draw()
im = fig.canvas.buffer_rgba()
#identify unique colors, containing white background
all_colors = np.unique(np.asarray(im).reshape(-1, 4), axis=0)
#remove white background color
colors_image = all_colors[:-1].reshape(4, -1)

#turn axes back on
ax.axis("on")
#determine cmap and norm for colorbar
cmapM1M2 = colors.ListedColormap(colors_image[::-1]/255)
normM1M2 = colors.BoundaryNorm(np.arange(-0.5,4), 4) 

#temporarily draw image of zeroes to get scalar mappable for colorbar
temp_im = ax.imshow(np.zeros(M1arr.shape), cmap=cmapM1M2, norm=normM1M2)
cbt = plt.colorbar(temp_im, ticks=np.arange(4), fraction=0.035)
#and remove this temporary image
temp_im.remove()

#label colorbar
cbt.ax.set_yticklabels(["M1 & M2 NaN", "only M1 values", "only M2 values", "M1 & M2 values"])
#and overlay image
ax.set_xticks(np.arange(M1arr.shape[1]))
ax.set_yticks(np.arange(M1arr.shape[0]))

plt.tight_layout()
plt.show()

    Output:
    enter image description here

    To use the identified unique colors in a colorbar, I draw a temporary image that later is removed. One could also draw a colorbar from scratch instead but it turned out rather annoying to control all parameters for this by hand.

    Update
    I just noticed that a colorbar is not required, it was just convenient in the previous question. So, we could alternatively just create a figure legend with patches:

    import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

M1 = ([1,      2,      3,      np.nan], 
      [4,      5,      np.nan, np.nan], 
      [6,      7,      8,      9])

M2 = ([np.nan, np.nan, np.nan, np.nan],
      [np.nan, 1,      2,      3], 
      [np.nan, 4,      5,      6])

M1arr = ~np.isnan(M1)
M2arr = ~np.isnan(M2)

fig, ax = plt.subplots()

ax.imshow(M1arr, cmap="Reds",  alpha=0.5)
ax.imshow(M2arr, cmap="Blues", alpha=0.5)

#after this, I want to extract the colors of the overlay image

#get image without axis, so only the colors plotted in the overlay image are considered
ax.axis("off")
#get rgba values from image predrawn by the renderer
fig.canvas.draw()
im = fig.canvas.buffer_rgba()
#identify unique colors, containing white background
all_colors = np.unique(np.asarray(im).reshape(-1, 4), axis=0)
#remove white background color
colors_image = all_colors[:-1].reshape(4, -1)


#turn axes back on
ax.axis("on")
#and overlay image
ax.set_xticks(np.arange(M1arr.shape[1]))
ax.set_yticks(np.arange(M1arr.shape[0]))

#create legend with patches of the four colors
categories = ["M1 & M2 NaN", "only M1 values", "only M2 values", "M1 & M2 values"]
fig.legend(handles=[mpatches.Patch(facecolor=col, edgecolor="k", label=categories[3-i]) for i, col in enumerate(colors_image/255)],
           loc="upper center", ncol = 2)

plt.show()

    Output:
    enter image description here

    Take home message: It is easier to define parameters before plotting than extracting information retrospectively from the backend's memory.