Upload gradcam_utils.py

gradcam_utils.py

@@ -0,0 +1,141 @@
+'''
+Grad-CAM visualization utilities
+
+- Based on https://keras.io/examples/vision/grad_cam/
+
+---
+- 2021-12-18 jkang first created 
+- 2022-01-16 
+    - copied from https://huggingface.co/spaces/jkang/demo-gradcam-imagenet/blob/main/utils.py
+    - updated for artis/trend classifier
+'''
+import matplotlib.cm as cm
+
+import os
+import re
+from glob import glob
+import numpy as np
+import tensorflow as tf
+tfk = tf.keras
+K = tfk.backend
+
+# Disable GPU for testing
+# os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
+
+
+def get_imagenet_classes():
+    '''Retrieve all 1000 imagenet classes/labels as dictionaries'''
+    classes = tfk.applications.imagenet_utils.decode_predictions(
+        np.expand_dims(np.arange(1000), 0), top=1000
+    )
+    idx2lab = {cla[2]: cla[1] for cla in classes[0]}
+    lab2idx = {idx2lab[idx]: idx for idx in idx2lab}
+    return idx2lab, lab2idx
+
+
+def search_by_name(str_part):
+    '''Search imagenet class by partial matching string'''
+    results = [key for key in list(lab2idx.keys()) if re.search(str_part, key)]
+    if len(results) != 0:
+        return [(key, lab2idx[key]) for key in results]
+    else:
+        return []
+
+
+def get_xception_model():
+    '''Get model to use'''
+    base_model = tfk.applications.xception.Xception
+    preprocessor = tfk.applications.xception.preprocess_input
+    decode_predictions = tfk.applications.xception.decode_predictions
+    last_conv_layer_name = "block14_sepconv2_act"
+
+    model = base_model(weights='imagenet')
+    grad_model = tfk.models.Model(
+        inputs=[model.inputs],
+        outputs=[model.get_layer(last_conv_layer_name).output,
+                 model.output]
+    )
+    return model, grad_model, preprocessor, decode_predictions
+
+
+def get_img_4d_array(image_file, image_size=(299, 299)):
+    '''Load image as 4d array'''
+    img = tfk.preprocessing.image.load_img(
+        image_file, target_size=image_size)  # PIL obj
+    img_array = tfk.preprocessing.image.img_to_array(
+        img)  # float32 numpy array
+    img_array = np.expand_dims(img_array, axis=0)  # 3d -> 4d (1,299,299,3)
+    return img_array
+
+
+def make_gradcam_heatmap(grad_model, img_array, pred_idx=None):
+    '''Generate heatmap to overlay with
+    - img_array: 4d numpy array
+    - pred_idx: eg. index out of 1000 imagenet classes
+        if None, argmax is chosen from prediction
+    '''
+    # Get gradient of pred class w.r.t. last conv activation
+    with tf.GradientTape() as tape:
+        last_conv_act, predictions = grad_model(img_array)
+        if pred_idx == None:
+            pred_idx = tf.argmax(predictions[0])
+        class_channel = predictions[:, pred_idx]  # (1,1000) => (1,)
+
+    # d(class_channel/last_conv_act)
+    grads = tape.gradient(class_channel, last_conv_act)
+    pooled_grads = tf.reduce_mean(grads, axis=(
+        0, 1, 2))  # (1,10,10,2048) => (2048,)
+
+    # (10,10,2048) x (2048,1) => (10,10,1)
+    heatmap = last_conv_act[0] @ pooled_grads[..., tf.newaxis]
+    heatmap = tf.squeeze(heatmap)  # (10,10)
+
+    # Normalize heatmap between 0 and 1
+    heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
+    return heatmap, pred_idx.numpy(), predictions.numpy().squeeze()
+
+
+def align_image_with_heatmap(img_array, heatmap, alpha=0.3, cmap='jet'):
+    '''Align the image with gradcam heatmap
+    - img_array: 4d numpy array
+    - heatmap: output of `def make_gradcam_heatmap()` as 2d numpy array
+    '''
+    img_array = img_array.squeeze()  # 4d => 3d
+
+    # Rescale to 0-255 range
+    heatmap_scaled = np.uint8(255 * heatmap)
+    img_array_scaled = np.uint8(255 * img_array)
+
+    colormap = cm.get_cmap(cmap)
+    colors = colormap(np.arange(256))[:, :3]  # mapping RGB to heatmap
+    heatmap_colored = colors[heatmap_scaled]  # ? still unclear
+
+    # Make RGB colorized heatmap
+    heatmap_colored = (tfk.preprocessing.image.array_to_img(heatmap_colored)  # array => PIL
+                       .resize((img_array.shape[1], img_array.shape[0])))
+    heatmap_colored = tfk.preprocessing.image.img_to_array(
+        heatmap_colored)  # PIL => array
+
+    # Overlay image with heatmap
+    overlaid_img = heatmap_colored * alpha + img_array_scaled
+    overlaid_img = tfk.preprocessing.image.array_to_img(overlaid_img)
+    return overlaid_img
+
+
+if __name__ == '__main__':
+    # Test GradCAM
+    examples = sorted(glob(os.path.join('examples', '*.jpg')))
+    idx2lab, lab2idx = get_imagenet_classes()
+
+    model, grad_model, preprocessor, decode_predictions = get_xception_model()
+
+    img_4d_array = get_img_4d_array(examples[0])
+    img_4d_array = preprocessor(img_4d_array)
+
+    heatmap = make_gradcam_heatmap(grad_model, img_4d_array, pred_idx=None)
+
+    img_pil = align_image_with_heatmap(
+        img_4d_array, heatmap, alpha=0.3, cmap='jet')
+
+    img_pil.save('test.jpg')
+    print('done')