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() |