Update app.py

app.py

@@ -1,6 +1,9 @@
 import os
 os.system('pip install git+https://github.com/huggingface/transformers.git --upgrade')
+os.system('pip install gradio --upgrade')
+os.system('pip freeze')
 
+import os
 import gradio as gr
 from transformers import ViTFeatureExtractor, ViTModel
 import torch
@@ -8,7 +11,16 @@ import torch.nn as nn
 import torchvision
 import matplotlib.pyplot as plt
 
-def get_attention_maps(pixel_values, attentions, nh):
+import cv2
+import numpy as np
+from tqdm import tqdm
+import glob
+from PIL import Image
+
+feature_extractor = ViTFeatureExtractor.from_pretrained("facebook/dino-vits8", do_resize=True, padding=True)
+model = ViTModel.from_pretrained("facebook/dino-vits8", add_pooling_layer=False)
+
+def get_attention_maps(pixel_values, attentions, nh, out, img_path):
   threshold = 0.6
   w_featmap = pixel_values.shape[-2] // model.config.patch_size
   h_featmap = pixel_values.shape[-1] // model.config.patch_size
@@ -25,57 +37,120 @@ def get_attention_maps(pixel_values, attentions, nh):
 
   # 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()
 
-  # save attentions heatmaps and return list of filenames
-  output_dir = '.'
-  os.makedirs(output_dir, exist_ok=True)
-  attention_maps = []
-  for j in range(nh):
-      fname = os.path.join(output_dir, "attn-head" + str(j) + ".png")
+  # sum all attentions
+  fname = os.path.join(out, "attn-" + os.path.basename(img_path))
+  plt.imsave(
+      fname=fname,
+      arr=sum(
+          attentions[i] * 1 / attentions.shape[0]
+          for i in range(attentions.shape[0])
+      ),
+      cmap="inferno",
+      format="png",
+  )
 
-      # save the attention map
-      plt.imsave(fname=fname, arr=attentions[j], format='png')
 
-      # append file name
-      attention_maps.append(fname)
+def inference(inp: str, out: str):
+  print(f"Generating attention images to {out}")
 
-  return attention_maps
+  # I had to process one at a time since colab was crashing...
+  for img_path in tqdm(sorted(glob.glob(os.path.join(inp, "*.jpg")))):
+      with open(img_path, "rb") as f:
+          img = Image.open(f)
+          img = img.convert("RGB")
+        
+      # normalize channels
+      pixel_values = feature_extractor(images=img, return_tensors="pt").pixel_values 
 
+      # forward pass
+      outputs = model(pixel_values, output_attentions=True, interpolate_pos_encoding=True)
 
-def visualize_attention(video):
-  return video
-  """
-  # normalize channels
-  pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values 
+      # get attentions of last layer
+      attentions = outputs.attentions[-1] 
+      nh = attentions.shape[1] # number of heads
 
-  # forward pass
-  outputs = model(pixel_values, output_attentions=True, interpolate_pos_encoding=True)
+      # we keep only the output patch attention
+      attentions = attentions[0, :, 0, 1:].reshape(nh, -1)
 
-  # get attentions of last layer
-  attentions = outputs.attentions[-1] 
-  nh = attentions.shape[1] # number of heads
+      # sum and save attention maps 
+      get_attention_maps(pixel_values, attentions, nh, out, img_path)
 
-  # 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"""
+def extract_frames_from_video(inp: str, out: str):
+    vidcap = cv2.VideoCapture(inp)
+    fps = vidcap.get(cv2.CAP_PROP_FPS)
+
+    print(f"Video: {inp} ({fps} fps)")
+    print(f"Extracting frames to {out}")
+
+    success, image = vidcap.read()
+    count = 0
+    while success:
+        cv2.imwrite(
+            os.path.join(out, f"frame-{count:04}.jpg"),
+            image,
+        )
+        success, image = vidcap.read()
+        count += 1
+    return fps
+
+def generate_video_from_images(inp: str, out_name: str, fps: int):
+  img_array = []
+  attention_images_list = sorted(glob.glob(os.path.join(inp, "attn-*.jpg")))
+
+  # Get size of the first image
+  with open(attention_images_list[0], "rb") as f:
+      img = Image.open(f)
+      img = img.convert("RGB")
+      size = (400, 400)
+      img_array.append(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))
+
+  print(f"Generating video {size} to {out_name}")
+
+  for filename in tqdm(attention_images_list[1:]):
+      with open(filename, "rb") as f:
+          img = Image.open(f)
+          img = img.convert("RGB")
+          img_array.append(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR))
+
+  out = cv2.VideoWriter(
+      out_name,
+      cv2.VideoWriter_fourcc(*"MP4V"),
+      fps,
+      size,
+  )
+
+  for i in range(len(img_array)):
+      out.write(img_array[i])
+  out.release()
+  print("Done")
+  return 
+
+def func(video):
+  frames_folder = os.path.join("output", "frames")
+  attention_folder = os.path.join("output", "attention")
+
+  os.makedirs(frames_folder, exist_ok=True)
+  os.makedirs(attention_folder, exist_ok=True)
+
+  fps = extract_frames_from_video(video, frames_folder)
+
+  inference(frames_folder,attention_folder)
+  generate_video_from_images(attention_folder, "video.mp4", fps)
+
+  return "video.mp4"
 
-feature_extractor = ViTFeatureExtractor.from_pretrained("facebook/dino-vits8", do_resize=False)
-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 =[['video.mp4']]
-iface = gr.Interface(fn=visualize_attention, 
-                     inputs=gr.inputs.Video(gr.inputs.Video()), 
-                     outputs=[gr.outputs.Video(label=f'result_video')],
+iface = gr.Interface(fn=func, 
+                     inputs=gr.inputs.Video(type=None), 
+                     outputs="video",
                      title=title,
                      description=description,
-                     article=article,
-                     examples=examples)
-iface.launch()
+                     article=article)