integratedGradients.py

def plotAttributions(baseline,attributions,
                          image,
                          cmap=None,
                          overlay_alpha=0.4):


  #Sum of the abs value of the attriutions determining feature importance
  mask = tf.reduce_sum(tf.math.abs(attributions), axis=-1)


  fig, ax = plt.subplots(nrows=2, ncols=2, squeeze=False, figsize=(8, 8))

  #Original image
  ax[0, 1].set_title('Original')
  ax[0, 1].imshow(image)
  ax[0, 1].axis('off')

  #Baseline image
  ax[0, 0].set_title('Baseline')
  ax[0, 0].imshow(baseline)
  ax[0, 0].axis('off')

  #Just the feature attributions
  ax[1, 0].set_title('Attribution mask')
  ax[1, 0].imshow(mask, cmap=cmap)
  ax[1, 0].axis('off')

  #Attributaions overlayed over the image - overlay_alpha is interpolation used
  ax[1, 1].set_title('Overlay')
  ax[1, 1].imshow(mask, cmap=cmap)
  ax[1, 1].imshow(image, alpha=overlay_alpha)
  ax[1, 1].axis('off')

  plt.tight_layout()
  return fig