Update app.py

app.py

@@ -43,7 +43,7 @@ def get_attention_maps(pixel_values, attentions, nh, out, img_path):
   attentions = attentions.detach().numpy()
 
   # sum all attentions
-  fname = os.path.join(out, "attn-" + os.path.basename(img_path))
+  fname = os.path.join(out, os.path.basename(img_path))
   plt.imsave(
       fname=fname,
       arr=sum(
@@ -51,14 +51,15 @@ def get_attention_maps(pixel_values, attentions, nh, out, img_path):
           for i in range(attentions.shape[0])
       ),
       cmap="inferno",
-      format="png",
+      format="jpg",
   )
-
+  return fname
 
 def inference(inp: str, out: str):
   print(f"Generating attention images to {out}")
 
   # I had to process one at a time since colab was crashing...
+  fnames = []
   for img_path in tqdm(sorted(glob.glob(os.path.join(inp, "*.jpg")))):
       with open(img_path, "rb") as f:
           img = Image.open(f)
@@ -78,71 +79,47 @@ def inference(inp: str, out: str):
       attentions = attentions[0, :, 0, 1:].reshape(nh, -1)
 
       # sum and save attention maps 
-      get_attention_maps(pixel_values, attentions, nh, out, img_path)
-
-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,
-  )
+      fnames.append(get_attention_maps(pixel_values, attentions, nh, out, img_path))
+  return fnames
 
-  for i in range(len(img_array)):
-      out.write(img_array[i])
-  out.release()
-  print("Done")
-  return 
 
 def func(video):
+  clip = VideoFileClip(video)
+  if clip.duration > 10:
+      return 'trim.mp4'
+
   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)
+  vid = VideoFileClip(inp)
+  fps = vid.fps
+
+  print(f"Video: {inp} ({fps} fps)")
+  print(f"Extracting frames to {frames_folder}")
+
+  vid.write_images_sequence(
+      os.path.join(frames_folder, "frame-count%03d.jpg"),
+  )
+
+  output_frame_fnames = inference(frames_folder,attention_folder)
 
-  inference(frames_folder,attention_folder)
-  generate_video_from_images(attention_folder, "video.mp4", fps)
+  new_clip = ImageSequenceClip(output_frame_fnames, fps=fps)
+  new_clip.write_videofile("my_new_video.mp4")
 
-  return "video.mp4"
+  return "my_new_video.mp4"
+
+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>"
+iface = gr.Interface(fn=func, 
+                     inputs=gr.inputs.Video(type=None), 
+                     outputs="video",
+                     title=title,
+                     description=description,
+                     article=article)
 
 
 title = "Interactive demo: DINO"