Captum: Allow colour map selection for single channel images

[FL33TW00D]

I have some single channel numpy matricies that I am feeding to my network. The output of captum is great, but it only allows single channel "original images" to be displayed in black and white.

[bilalsal]

Thank you for this suggestion, @FL33TW00D. We understand your need to have a custom color map for your input.
As a quick solution, we suggest something like the following (assuming X is your 1-channel image)

cmap = "your map" # a matplotlib cmap
from captum.attr import visualization as viz
fig_and_axes = viz.visualize_image_attr_multiple(...)
fig_and_axes[1][0].imshow(X, cmap=cmap) # the subplot for the original image
fig_and_axes[0] # show the matplotlib Fig object

Hope this helps

[FL33TW00D]

Hi Bilal,
It works perfectly thank you.