nielsr
HF staff
commited on
Commit
9a21097
1 Parent(s): ebb2c9b

Remove torchvision import

Browse files
Files changed (1) hide show
  1. app.py +2 -4
app.py CHANGED
@@ -4,8 +4,6 @@ os.system('pip install git+https://github.com/huggingface/transformers.git --upg
4
  import gradio as gr
5
  from transformers import ViTFeatureExtractor, ViTModel
6
  import torch
7
- import torch.nn as nn
8
- import torchvision
9
  import matplotlib.pyplot as plt
10
 
11
  torch.hub.download_url_to_file('http://images.cocodataset.org/val2017/000000039769.jpg', 'cats.jpg')
@@ -25,10 +23,10 @@ def get_attention_maps(pixel_values, attentions, nh):
25
  th_attn[head] = th_attn[head][idx2[head]]
26
  th_attn = th_attn.reshape(nh, w_featmap, h_featmap).float()
27
  # interpolate
28
- th_attn = nn.functional.interpolate(th_attn.unsqueeze(0), scale_factor=model.config.patch_size, mode="nearest")[0].cpu().numpy()
29
 
30
  attentions = attentions.reshape(nh, w_featmap, h_featmap)
31
- attentions = nn.functional.interpolate(attentions.unsqueeze(0), scale_factor=model.config.patch_size, mode="nearest")[0].cpu()
32
  attentions = attentions.detach().numpy()
33
 
34
  # save attentions heatmaps and return list of filenames
4
  import gradio as gr
5
  from transformers import ViTFeatureExtractor, ViTModel
6
  import torch
 
 
7
  import matplotlib.pyplot as plt
8
 
9
  torch.hub.download_url_to_file('http://images.cocodataset.org/val2017/000000039769.jpg', 'cats.jpg')
23
  th_attn[head] = th_attn[head][idx2[head]]
24
  th_attn = th_attn.reshape(nh, w_featmap, h_featmap).float()
25
  # interpolate
26
+ th_attn = torch.nn.functional.interpolate(th_attn.unsqueeze(0), scale_factor=model.config.patch_size, mode="nearest")[0].cpu().numpy()
27
 
28
  attentions = attentions.reshape(nh, w_featmap, h_featmap)
29
+ attentions = torch.nn.functional.interpolate(attentions.unsqueeze(0), scale_factor=model.config.patch_size, mode="nearest")[0].cpu()
30
  attentions = attentions.detach().numpy()
31
 
32
  # save attentions heatmaps and return list of filenames