Spaces:
Runtime error
Runtime error
File size: 2,988 Bytes
2a9ec74 |
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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
from typing import Dict
import cv2
import numpy as np
import tensorflow as tf
from PIL import Image
from tensorflow import keras
RESOLUTION = 224
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],
)
rescale_layer = keras.layers.Rescaling(scale=1.0 / 127.5, offset=-1)
def preprocess_image(orig_image: Image, model_type: str, size=RESOLUTION):
"""Image preprocessing utility."""
# Turn the image into a numpy array and add batch dim.
image = np.array(orig_image)
image = tf.expand_dims(image, 0)
# If model type is vit rescale the image to [-1, 1].
if model_type == "original_vit":
image = rescale_layer(image)
# Resize the image using bicubic interpolation.
resize_size = int((256 / 224) * size)
image = tf.image.resize(image, (resize_size, resize_size), method="bicubic")
# Crop the image.
preprocessed_image = crop_layer(image)
# If model type is DeiT or DINO normalize the image.
if model_type != "original_vit":
image = norm_layer(preprocessed_image)
return orig_image, preprocessed_image.numpy()
def attention_rollout_map(
image: Image, attention_score_dict: Dict[str, np.ndarray], model_type: str
):
"""Computes attention rollout results.
Reference:
https://arxiv.org/abs/2005.00928
Code copied and modified from here:
https://github.com/jeonsworld/ViT-pytorch/blob/main/visualize_attention_map.ipynb
"""
num_cls_tokens = 2 if "distilled" in model_type else 1
# Stack the individual attention matrices from individual transformer blocks.
attn_mat = tf.stack(
[attention_score_dict[k] for k in attention_score_dict.keys()]
)
attn_mat = tf.squeeze(attn_mat, axis=1)
# Average the attention weights across all heads.
attn_mat = tf.reduce_mean(attn_mat, axis=1)
# To account for residual connections, we add an identity matrix to the
# attention matrix and re-normalize the weights.
residual_attn = tf.eye(attn_mat.shape[1])
aug_attn_mat = attn_mat + residual_attn
aug_attn_mat = (
aug_attn_mat / tf.reduce_sum(aug_attn_mat, axis=-1)[..., None]
)
aug_attn_mat = aug_attn_mat.numpy()
# Recursively multiply the weight matrices.
joint_attentions = np.zeros(aug_attn_mat.shape)
joint_attentions[0] = aug_attn_mat[0]
for n in range(1, aug_attn_mat.shape[0]):
joint_attentions[n] = np.matmul(
aug_attn_mat[n], joint_attentions[n - 1]
)
# Attention from the output token to the input space.
v = joint_attentions[-1]
grid_size = int(np.sqrt(aug_attn_mat.shape[-1]))
mask = v[0, num_cls_tokens:].reshape(grid_size, grid_size)
mask = cv2.resize(mask / mask.max(), image.size)[..., np.newaxis]
result = (mask * image).astype("uint8")
return result
|