Add video

app.py

@@ -8,7 +8,7 @@ import torch.nn as nn
 import torchvision
 import matplotlib.pyplot as plt
 
-torch.hub.download_url_to_file('http://images.cocodataset.org/val2017/000000039769.jpg', 'cats.jpg')
+torch.hub.download_url_to_file('https://viratdata.org/video/VIRAT_S_010204_05_000856_000890.mp4', 'video.mp4')
 
 def get_attention_maps(pixel_values, attentions, nh):
   threshold = 0.6
@@ -47,19 +47,24 @@ def get_attention_maps(pixel_values, attentions, nh):
   return attention_maps
 
 
-def visualize_attention(image):
+def visualize_attention(video):
+  return video
+  """
   # normalize channels
   pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values 
+
   # forward pass
   outputs = model(pixel_values, output_attentions=True, interpolate_pos_encoding=True)
+
   # get attentions of last layer
   attentions = outputs.attentions[-1] 
   nh = attentions.shape[1] # number of heads
+
   # we keep only the output patch attention
   attentions = attentions[0, :, 0, 1:].reshape(nh, -1)
   attention_maps = get_attention_maps(pixel_values, attentions, nh)
   
-  return attention_maps
+  return attention_maps"""
 
 feature_extractor = ViTFeatureExtractor.from_pretrained("facebook/dino-vits8", do_resize=False)
 model = ViTModel.from_pretrained("facebook/dino-vits8", add_pooling_layer=False)
@@ -67,10 +72,10 @@ model = ViTModel.from_pretrained("facebook/dino-vits8", add_pooling_layer=False)
 title = "Interactive demo: DINO"
 description = "Demo for Facebook AI's DINO, a new method for self-supervised training of Vision Transformers. Using this method, they are capable of segmenting objects within an image without having ever been trained to do so. This can be observed by displaying the self-attention of the heads from the last layer for the [CLS] token query. This demo uses a ViT-S/8 trained with DINO. To use it, simply upload an image or use the example image below. Results will show up in a few seconds."
 article = "<p style='text-align: center'><a href='https://arxiv.org/abs/2104.14294'>Emerging Properties in Self-Supervised Vision Transformers</a> | <a href='https://github.com/facebookresearch/dino'>Github Repo</a></p>"
-examples =[['cats.jpg']]
+examples =[['video.mp4']]
 iface = gr.Interface(fn=visualize_attention, 
-                     inputs=gr.inputs.Image(shape=(480, 480), type="pil"), 
-                     outputs=[gr.outputs.Image(type='file', label=f'attention_head_{i}') for i in range(6)],
+                     inputs=gr.inputs.Video(gr.inputs.Video()), 
+                     outputs=[gr.outputs.Video(label=f'result_video')],
                      title=title,
                      description=description,
                      article=article,