chore: refactor src

app.py

@@ -1,16 +1,21 @@
 import gradio as gr
 import torch
 
-from augmentations import get_videomae_transform
-from models import load_model
-from utils import create_plot, get_frames, get_videomae_outputs, prepare_frames_masks
+from src.augmentations import get_videomae_transform
+from src.models import load_model
+from src.utils import (
+    create_plot,
+    get_frames,
+    get_videomae_outputs,
+    prepare_frames_masks,
+)
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-transform = get_videomae_transform()
-
 
 def get_visualisations(mask_ratio, video_path):
+    transform = get_videomae_transform()
+
     frames, ids = get_frames(path=video_path, transform=transform)
 
     model, masks, patch_size = load_model(
@@ -36,6 +41,12 @@ def get_visualisations(mask_ratio, video_path):
 
 
 with gr.Blocks() as app:
+    gr.Markdown(
+        """
+        # VideoMAE Reconstruction Demo
+        To read more about the Self-Supervised Learning techniques for video please refer to the [Lightly AI blogpost on Self-Supervised Learning for Videos](www.lightly.ai/post/self-supervised-learning-for-videos).
+        """  # noqa: E501
+    )
     video = gr.Video(
         value="assets/example.mp4",
     )

requirements.txt

@@ -1,6 +1,8 @@
+eva-decord
 einops
-decord
+gradio
 numpy
+Pillow
 timm
 torch
 torchvision

models.py → src/models.py

@@ -8,7 +8,7 @@ import torch.nn.functional as F
 import torch.utils.checkpoint as checkpoint
 from timm.models.layers import drop_path, to_2tuple, trunc_normal_
 
-from augmentations import TubeMaskingGenerator
+from src.augmentations import TubeMaskingGenerator
 
 __all__ = ["load_model"]
 

utils.py → src/utils.py

@@ -139,6 +139,5 @@ def create_plot(images):
                 axes[i, j].set_title(column_names[j], fontsize=16)
 
     plt.tight_layout()
-    plt.show()
 
     return fig