pythonmatplotlibvisualizationcolorbarcolormap

Heatmap on top of image


I have Images of different objects (Pascal Voc) and I have a probability heatmap. I want to visualize it by plotting the image and somehow the heatmap on top of it. What would be the best way to do that?

I was thinking about using the alpha channel like this:

im_heat = np.zeros((image.shape[0],image.shape[1],4))
im_heat[:,:,:3] = image
im_heat[:,:,3] = np.rint(255/heatmap)
plt.imshow(im_heat, cmap='jet')
plt.colorbar()

How do I customize the colorbar to be from min(heatmap) to max(heatmap)? Or is there any better way to visualize the probabilities?


Solution

  • You can stack images and plots with matplotlib and then choose which handle to use for the colorbar. Using contourf the colorbar min and max values will be based on your heatmap (or you can pass vmin=min(heatmap) and vmax=max(heatmap) to contourf to be explicit about this range). The problem with this is the heatmap will cover you image (and setting transparency will make the entire thing transparent). The best option is to make a colormap which is transparent when near zero, as follows,

    import numpy as np
    import matplotlib.pyplot as plt
    import matplotlib.colors as mcolors
    import Image
    
    #2D Gaussian function
    def twoD_Gaussian((x, y), xo, yo, sigma_x, sigma_y):
        a = 1./(2*sigma_x**2) + 1./(2*sigma_y**2)
        c = 1./(2*sigma_x**2) + 1./(2*sigma_y**2)
        g = np.exp( - (a*((x-xo)**2) + c*((y-yo)**2)))
        return g.ravel()
    
    
    def transparent_cmap(cmap, N=255):
        "Copy colormap and set alpha values"
    
        mycmap = cmap
        mycmap._init()
        mycmap._lut[:,-1] = np.linspace(0, 0.8, N+4)
        return mycmap
    
    
    #Use base cmap to create transparent
    mycmap = transparent_cmap(plt.cm.Reds)
    
    
    # Import image and get x and y extents
    I = Image.open('./deerback.jpg')
    p = np.asarray(I).astype('float')
    w, h = I.size
    y, x = np.mgrid[0:h, 0:w]
    
    #Plot image and overlay colormap
    fig, ax = plt.subplots(1, 1)
    ax.imshow(I)
    Gauss = twoD_Gaussian((x, y), .5*x.max(), .4*y.max(), .1*x.max(), .1*y.max())
    cb = ax.contourf(x, y, Gauss.reshape(x.shape[0], y.shape[1]), 15, cmap=mycmap)
    plt.colorbar(cb)
    plt.show()
    

    which gives,

    enter image description here