python imshow, set certain value to defined color

Question

I have a RGB Image, which I plot with matplotlib.pyplot.imshow and it works fine. But now I want to change the plot, that where the value of the picture is e.g 1, the color of the plot should change to white at all this positions.

Is there a way to do this?

Answer 1

I will answer the general question of how to set a particular value to a particular color regardless of the color map.

In the code below for illustration purposes I supposed that is the value -1 that you want to map white. You will be wanting to do something different for your code.

This technique uses a masked array to set the parts where your data is equal to -1 (the value you wish to map) and then uses cmap.set_bad() to assign the color white to this value.

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

value = -1
data = np.arange(100).reshape((10, 10))
data[5, :] = -1  # Values to set -1

masked_array = np.ma.masked_where(data == value, data)

cmap = matplotlib.cm.spring  # Can be any colormap that you want after the cm
cmap.set_bad(color='white')

plt.imshow(masked_array, cmap=cmap)
plt.show()

Hope it helps.

Answer 2

Assuming that your image is a single-channel image rather than a three-channel image, the required task can be performed by defining a palette that maps indices (e.g. gray level intensities or picture values) into colors:

import numpy as np
import matplotlib.pyplot as plt

palette = np.array([[  0,   0,   0],   # black
                    [255,   0,   0],   # red
                    [  0, 255,   0],   # green
                    [  0,   0, 255],   # blue
                    [255, 255, 255]])  # white
I = np.array([[ 0, 1, 2, 0],    # 2 rows, 4 columns, 1 channel
              [ 0, 3, 4, 0]])

Image conversion is efficiently accomplished through NumPy's broadcasting:

RGB = palette[I]

And this is how the transformed image looks like:

>>> RGB
array([[[  0,  0,    0],  # 2 rows, 4 columns, 3 channels
        [255,  0,    0],
        [  0, 255,   0],
        [  0,   0,   0]],
       
       [[  0,   0,   0],
        [  0,   0, 255],
        [255, 255, 255],
        [  0,   0,   0]]])

plt.imshow(RGB)

Answer 3

I am going to present a solution to the original question, which is extendable to assigning several values to several different colors respectively.

Solution

The solution involves creating a new three dimensional NumPy ndarray that contains, at each i,j position, an NumPy array with RGB values. This new data3d array is then plotted with imshow (or matshow)

# import packages 
import numpy as np
import matplotlib
import matplotlib.pyplot as plt

# generate data 
np.random.seed(42)
data = np.random.randint(low=1, high=4, size=(10,10)) # possible values are 1,2,3

# define color map 
color_map = {1: np.array([255, 0, 0]), # red
             2: np.array([0, 255, 0]), # green
             3: np.array([0, 0, 255])} # blue 

# make a 3d numpy array that has a color channel dimension   
data_3d = np.ndarray(shape=(data.shape[0], data.shape[1], 3), dtype=int)
for i in range(0, data.shape[0]):
    for j in range(0, data.shape[1]):
        data_3d[i][j] = color_map[data[i][j]]

# display the plot 
fig, ax = plt.subplots(1,1)
ax.imshow(data_3d)

# add numbers to the plot 
# thanks to tmdavison answer here https://stackoverflow.com/a/40890587/7871710
for i in range(0, data.shape[0]):
    for j in range(0, data.shape[1]):
        c = data[j,i]
        ax.text(i, j, str(c), va='center', ha='center')

Background I encountered a problem recently where I needed to plot several one-channel matrices and assign a distinct color to the 1's 2's and 3's in each matrix. The ordering of the 1's, 2's and 3's changed depending on the matrix, which meant that using a defined color scheme often led to the assignment of different colors to the same value in different matrices. For example in the first matrix the 1's were assigned to red, whereas in the second matrix the 1's were assigned to blue.

I spent a lot of time searching around stackoverflow for a solution, but never found anything that worked. Eventually, was able to work one out on my own. This is solution to the problem in this question, that is extendable to multiple values and independent of the ordering of the values in the main matrix. It will also work with matshow instead of imshow.