Upload utils.py

utils.py

@@ -0,0 +1,66 @@
+from typing import Dict
+
+import numpy as np
+import tensorflow as tf
+from PIL import Image
+from tensorflow import keras
+
+RESOLUTION = 224
+PATCH_SIZE = 16
+
+
+crop_layer = keras.layers.CenterCrop(RESOLUTION, RESOLUTION)
+norm_layer = keras.layers.Normalization(
+    mean=[0.485 * 255, 0.456 * 255, 0.406 * 255],
+    variance=[(0.229 * 255) ** 2, (0.224 * 255) ** 2, (0.225 * 255) ** 2],
+)
+
+
+def preprocess_image(orig_image: Image, size: int):
+    """Image preprocessing utility."""
+    image = np.array(orig_image)
+    image_resized = tf.expand_dims(image, 0)
+    resize_size = int((256 / 224) * size)
+    image_resized = tf.image.resize(
+        image_resized, (resize_size, resize_size), method="bicubic"
+    )
+    image_resized = crop_layer(image_resized)
+    return image_resized.numpy().squeeze(), norm_layer(image_resized).numpy()
+
+
+# Reference:
+# https://github.com/facebookresearch/dino/blob/main/visualize_attention.py
+
+
+def get_cls_attention_map(
+    preprocessed_image: np.ndarray,
+    attn_score_dict: Dict[str, np.ndarray],
+    block_key="ca_ffn_block_0_att",
+):
+    """Utility to generate class saliency map modeling spatial-class relationships."""
+    w_featmap = preprocessed_image.shape[2] // PATCH_SIZE
+    h_featmap = preprocessed_image.shape[1] // PATCH_SIZE
+
+    attention_scores = attn_score_dict[block_key]
+    nh = attention_scores.shape[1]  # Number of attention heads.
+
+    # Taking the representations from CLS token.
+    attentions = attention_scores[0, :, 0, 1:].reshape(nh, -1)
+
+    # Reshape the attention scores to resemble mini patches.
+    attentions = attentions.reshape(nh, w_featmap, h_featmap)
+
+    attentions = np.mean(attentions, axis=0)
+    attentions = (attentions - attentions.min()) / (
+        attentions.max() - attentions.min()
+    )
+    attentions = np.expand_dims(attentions, -1)
+
+    # Resize the attention patches to 224x224 (224: 14x16)
+    attentions = tf.image.resize(
+        attentions,
+        size=(h_featmap * PATCH_SIZE, w_featmap * PATCH_SIZE),
+        method="bicubic",
+    )
+
+    return attentions.numpy()