Visualizing Attention Maps about vision_transformer HOT 2 CLOSED

# `img` is an input image
# `w` are the corresponding multihead attention weights after the softmax,
# averaged across heads and rolled out as described in https://github.com/google-research/vision_transformer/issues/18).
# The 1D weight vector has also been reshaped back into the 2D image grid shape.
# `patch_size` is the width of the input patches, e.g. 16 for ViT/16.

import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageFilter
 
num_patches = img.shape[0] / patch_size
w = w / np.max(w)  # Normalization
mask = Image.fromarray(w * 255)
mask = mask.resize(size=img.shape[:2], resample=Image.LANCZOS)
mask = mask.convert("L")
mask = mask.filter(ImageFilter.GaussianBlur(radius=patch_size))
mask = np.array(mask).astype(float)
mask -= np.min(mask)
mask /= np.max(mask)
plt.imshow(np.zeros_like(img), vmin=0, vmax=0.5)
img = np.concatenate([img, mask[..., np.newaxis]], axis=2)
plt.imshow(img)