Spaces:
Running
Running
File size: 2,106 Bytes
0f09377 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import matplotlib.cm as cm
import numpy as np
import tensorflow as tf
from tensorflow import keras
def make_gradcam_heatmap(img_array, grad_model, pred_index=None):
with tf.GradientTape(persistent=True) as tape:
preds, base_top, swin_top = grad_model(img_array)
if pred_index is None:
pred_index = tf.argmax(preds[0])
class_channel = preds[:, pred_index]
grads = tape.gradient(class_channel, base_top)
pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
base_top = base_top[0]
heatmap_a = base_top @ pooled_grads[..., tf.newaxis]
heatmap_a = tf.squeeze(heatmap_a)
heatmap_a = tf.maximum(heatmap_a, 0) / tf.math.reduce_max(heatmap_a)
heatmap_a = heatmap_a.numpy()
grads = tape.gradient(class_channel, swin_top)
pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
swin_top = swin_top[0]
heatmap_b = swin_top @ pooled_grads[..., tf.newaxis]
heatmap_b = tf.squeeze(heatmap_b)
heatmap_b = tf.maximum(heatmap_b, 0) / tf.math.reduce_max(heatmap_b)
heatmap_b = heatmap_b.numpy()
return heatmap_a, heatmap_b, preds
def save_and_display_gradcam(
img,
heatmap,
target=None,
pred=None,
cam_path="cam.jpg",
cmap="jet", # inferno, viridis
alpha=0.6,
plot=None,
image_shape=None,
):
# Rescale heatmap to a range 0-255
heatmap = np.uint8(255 * heatmap)
# Use jet colormap to colorize heatmap
jet = cm.get_cmap(cmap)
# Use RGB values of the colormap
jet_colors = jet(np.arange(256))[:, :3]
jet_heatmap = jet_colors[heatmap]
# Create an image with RGB colorized heatmap
jet_heatmap = keras.utils.array_to_img(jet_heatmap)
jet_heatmap = jet_heatmap.resize((img.shape[0], img.shape[1]))
jet_heatmap = keras.utils.img_to_array(jet_heatmap)
# Superimpose the heatmap on original image
superimposed_img = img + jet_heatmap * alpha
superimposed_img = keras.utils.array_to_img(superimposed_img)
size_w, size_h = image_shape[:2]
superimposed_img = superimposed_img.resize((size_h, size_w))
return superimposed_img
|