Spaces:
Runtime error
Runtime error
| ''' | |
| 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') |