Search code examples
pythonmatplotlibcolorbar

Add color scale to matplotlib colorbar according to RGBA image channels


I am trying to plot a RGBA image with a colorbar representing color values.

The RGBA image is generated from raw data, transforming the 2d data array into a 6d-array with x, y, [R, G, B and A] according to the color input. E.g. 'green' will make it fill just the G channel with the values from the 2d-array, leaving R and B = 0 and A = 255. Like this:

data_to_color_equation

All solutions I found would apply a color map or limit the vmin and vmax of the colorbar but what I need is a colorbar that goes from pitch black to the brightest color present in the image. E.g. if I have an image in shades of purple, the color bar should go from 0 to 'full' purple with only shades of purple in it. The closest solution I found was this (https://pelson.github.io/2013/working_with_colors_in_matplotlib/), but it doesn't fit a "general" solution.
An image I'm getting is given below.

color_data_plot

import numpy as np
from ImgMath import colorize
import matplotlib.pyplot as plt
import Mapping

data = Mapping.getpeakmap('Au')
# data shape is (10,13) and len(data) is 10

norm_data = data/data.max()*255
color_data = colorize(norm_data,'green')
# color_data shape is (10,13,4) and len(color_data) is 10

fig, ax = plt.subplots()
im = plt.imshow(color_data)
fig.colorbar(im)
plt.show()

Solution

  • You could map your data with a custom, all-green, colormap

    import numpy as np
    import matplotlib.pyplot as plt
    from matplotlib.colors import ListedColormap
    
    # input 2D array
    data = np.random.randint(0,255, size=(10,13))
    
    z = np.zeros(256)
    colors = np.linspace(0,1,256)
    alpha = np.ones(256)
    
    #create colormap
    greencolors = np.c_[z,colors,z,alpha]
    cmap = ListedColormap(greencolors)
    
    
    im = plt.imshow(data/255., cmap=cmap, vmin=0, vmax=1)
    plt.colorbar(im)
    
    plt.show()
    

    enter image description here