Spaces:
Sleeping
Sleeping
# Import the necessary libraries and modules | |
import os | |
import gradio as gr | |
from transformers import ViTImageProcessor, ViTFeatureExtractor, FlaxViTForImageClassification, ViTModel | |
from PIL import Image | |
import requests | |
import os | |
import torch | |
import torch.nn as nn | |
import torchvision | |
import matplotlib.pyplot as plt | |
def visualize_attention(name): | |
model_name = name.split(";")[0] | |
if len(name.split(";"))>1: | |
url = name.split(";")[1] | |
else: | |
url = "http://images.cocodataset.org/val2017/000000039769.jpg" | |
feature_extractor = ViTImageProcessor.from_pretrained(model_name, size=480) | |
pil_image = Image.open(requests.get(url, stream=True).raw) | |
device = "cpu" | |
pixel_values = feature_extractor(images=pil_image, return_tensors="pt").pixel_values.to(device) | |
model = ViTModel.from_pretrained(model_name, add_pooling_layer=False) | |
model.to(device) | |
outputs = model(pixel_values, output_attentions=True, interpolate_pos_encoding=True) | |
attentions = outputs.attentions[-1] # we are only interested in the attention maps of the last layer | |
nh = attentions.shape[1] # number of head | |
# we keep only the output patch attention | |
attentions = attentions[0, :, 0, 1:].reshape(nh, -1) | |
threshold = 0.6 | |
w_featmap = pixel_values.shape[-2] // model.config.patch_size | |
h_featmap = pixel_values.shape[-1] // model.config.patch_size | |
# we keep only a certain percentage of the mass | |
val, idx = torch.sort(attentions) | |
val /= torch.sum(val, dim=1, keepdim=True) | |
cumval = torch.cumsum(val, dim=1) | |
th_attn = cumval > (1 - threshold) | |
idx2 = torch.argsort(idx) | |
for head in range(nh): | |
th_attn[head] = th_attn[head][idx2[head]] | |
th_attn = th_attn.reshape(nh, w_featmap, h_featmap).float() | |
# interpolate | |
th_attn = nn.functional.interpolate(th_attn.unsqueeze(0), scale_factor=model.config.patch_size, mode="nearest")[0].cpu().numpy() | |
attentions = attentions.reshape(nh, w_featmap, h_featmap) | |
attentions = nn.functional.interpolate(attentions.unsqueeze(0), scale_factor=model.config.patch_size, mode="nearest")[0].cpu() | |
attentions = attentions.detach().numpy() | |
# show and save attentions heatmaps | |
output_dir = '.' | |
os.makedirs(output_dir, exist_ok=True) | |
torchvision.utils.save_image(torchvision.utils.make_grid(pixel_values, normalize=True, scale_each=True), os.path.join(output_dir, "img.png")) | |
for j in range(nh): | |
fname = os.path.join(output_dir, "attn-head" + str(j) + ".png") | |
plt.figure() | |
plt.imshow(attentions[j]) | |
plt.imsave(fname=fname, arr=attentions[j], format='png') | |
images = [] | |
for j in range(nh): | |
images.append(Image.open(os.path.join(output_dir, "attn-head" + str(j) + ".png"))) | |
return images | |
text_input = gr.Textbox(label="Enter the name of the model to use and optionally add in your own image jpg url with ; as a separator try out this: facebook/dino-vits8; https://upload.wikimedia.org/wikipedia/commons/thumb/3/3a/Cat03.jpg/481px-Cat03.jpg", placeholder = "facebook/dino-vits8; optionalurl.jpg") | |
attention_output = gr.Gallery(label="Attention Map") | |
iface = gr.Interface( | |
fn=visualize_attention, | |
inputs=text_input, | |
outputs=attention_output, | |
live=True, | |
capture_session=True, | |
title="Visualize Attention Maps", | |
description="This app uses a Vision Transformer to visualize the attention maps of an image.", | |
) | |
iface.launch() |