Spaces:
Running
Running
| 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() | |