Upload 5 files

extractor/__init__.py

@@ -0,0 +1,2 @@
+# __init__.py
+print("Initializing extractor")

extractor/vf_extract.py

@@ -0,0 +1,72 @@
+import math
+import pandas as pd
+import cv2
+import os
+
+def extract_frames(video_path, sampled_path, frame_interval, residual=False):
+    try:
+        video_name = os.path.splitext(os.path.basename(video_path))[0]
+        cap = cv2.VideoCapture(video_path)
+        frames = []
+
+        if not cap.isOpened():
+            print(f"Error: Could not open video file {video_path}")
+            return frames
+
+        frame_count = 0
+        saved_frame_count = 0
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            if (frame_count % frame_interval == 0 and not residual) or (
+                    (frame_count - 1) % frame_interval == 0 and residual):
+                # suffix = '_next' if residual else ''
+                # output_filename = os.path.join(sampled_path, f'{video_name}_{saved_frame_count + 1}{suffix}.png')
+                # cv2.imwrite(output_filename, frame)
+                frames.append(frame)
+                saved_frame_count += 1
+
+            frame_count += 1
+        cap.release()
+        frame_type = 'next frames' if residual else 'sampled frames'
+        print(f'Extraction of {frame_type} for {video_name} completed!')
+        return frames
+    except Exception as e:
+        print(f"An unexpected error occurred: {e}")
+
+
+def process_video_residual(video_type, video_name, framerate, video_path, sampled_path):
+    if not os.path.exists(sampled_path):
+        os.makedirs(sampled_path)
+    # cap = cv2.VideoCapture(video_path)
+    # framerate = cap.get(cv2.CAP_PROP_FPS)
+    # print(f'framerate: {framerate}')
+    frame_interval = math.ceil(framerate / 2) if framerate < 2 else int(framerate / 2)
+    # print(f'Frame interval: {frame_interval}')
+
+    frames = extract_frames(video_path, sampled_path, frame_interval, residual=False)
+    frames_next = extract_frames(video_path, sampled_path, frame_interval, residual=True)
+    return frames, frames_next
+
+
+if __name__ == '__main__':
+    video_type = 'test'
+
+    if video_type == 'test':
+        ugcdata = pd.read_csv("../../metadata/test_videos.csv")
+
+    for i in range(len(ugcdata)):
+        video_name = ugcdata['vid'][i]
+        framerate = ugcdata['framerate'][i]
+        print(f'Processing video: {video_name}, framerate: {framerate}')
+
+        video_path = f"../../ugc_original_videos/{video_name}.mp4"
+        sampled_path = f'../../video_sampled_frame/original_sampled_frame/{video_name}/'
+
+        if not os.path.exists(sampled_path):
+            os.makedirs(sampled_path)
+
+        print(f'{video_name}')
+        frames, frames_next = process_video_residual(video_type, video_name, framerate, video_path, sampled_path)
+        print(f'Extracted {len(frames)} frames and {len(frames_next)} residual frames for video: {video_name}')

extractor/visualise_resnet.py

@@ -0,0 +1,187 @@
+import warnings
+warnings.filterwarnings("ignore")
+import os
+import glob
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+import torch
+from torchvision import models, transforms
+from thop import profile
+is_flop_cal = False
+
+# get the activation
+def get_activation(model, layer, input_img_data):
+    model.eval()
+    activations = []
+    inputs = []
+
+    def hook(module, input, output):
+        activations.append(output)
+        inputs.append(input[0])
+
+    hook_handle = layer.register_forward_hook(hook)
+    with torch.no_grad():
+        model(input_img_data)
+    hook_handle.remove()
+    return activations, inputs
+
+def get_activation_map(frame, layer_name, resnet50, device):
+    # image pre-processing
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+
+    # Apply the transformations (resize and normalize)
+    frame_tensor = transform(frame)
+
+    # adding index 0 changes the original [C, H, W] shape to [1, C, H, W]
+    if frame_tensor.dim() == 3:
+        frame_tensor = frame_tensor.unsqueeze(0)
+    # print(f'Image dimension: {frame_tensor.shape}')
+
+    # getting the activation of a given layer
+    layer_obj = eval(layer_name)
+    activations, inputs = get_activation(resnet50, layer_obj, frame_tensor)
+    activated_img = activations[0][0]
+    activation_array = activated_img.cpu().numpy()
+
+    # calculate FLOPs for layer
+    if is_flop_cal == True:
+        flops, params = profile(layer_obj, inputs=(inputs[0],), verbose=False)
+        if params == 0 and isinstance(layer_obj, torch.nn.Conv2d):
+            params = layer_obj.in_channels * layer_obj.out_channels * layer_obj.kernel_size[0] * layer_obj.kernel_size[1]
+            if layer_obj.bias is not None:
+                params += layer_obj.out_channels
+        # print(f"FLOPs for {layer_name}: {flops}, Params: {params}")
+    else:
+        flops, params = None, None
+    return activated_img, activation_array, flops, params
+
+def process_video_frame(video_name, frame, frame_number, all_layers, resnet50, device):
+    # create a dictionary to store activation arrays for each layer
+    activations_dict = {}
+    total_flops = 0
+    total_params = 0
+    for layer_name in all_layers:
+        fig_name = f"resnet50_feature_map_layer_{layer_name}"
+        combined_name = f"resnet50_feature_map"
+
+        activated_img, activation_array, flops, params = get_activation_map(frame, layer_name, resnet50, device)
+        if is_flop_cal == True:
+            total_flops += flops
+            total_params += params
+
+        # save activation maps as png
+        # png_path = f'../visualisation/resnet50/{video_name}/frame_{frame_number}/'
+        # npy_path = f'../features/resnet50/{video_name}/frame_{frame_number}/'
+        # os.makedirs(png_path, exist_ok=True)
+        # os.makedirs(npy_path, exist_ok=True)
+        # get_activation_png(png_path, fig_name, activated_img)
+        # save activation features as npy
+        # get_activation_npy(npy_path, fig_name, activation_array)
+        # save to the dictionary
+        activations_dict[layer_name] = activated_img
+
+    # print(f"total FLOPs for Resnet50 layerstack: {total_flops}, Params: {total_params}")
+    frame_npy_path = f'../features/resnet50/{video_name}/frame_{frame_number}_{combined_name}.npy'
+    return activations_dict, frame_npy_path, total_flops, total_params
+
+def get_activation_png(png_path, fig_name, activated_img, n=8):
+    fig = plt.figure(figsize=(10, 10))
+
+    # visualise activation map for 64 channels
+    for i in range(n):
+        for j in range(n):
+            idx = (n * i) + j
+            if idx >= activated_img.shape[0]:
+                break
+            ax = fig.add_subplot(n, n, idx + 1)
+            ax.imshow(activated_img[idx].cpu().numpy(), cmap='viridis')
+            ax.axis('off')
+
+    # save figures
+    fig_path = f'{png_path}{fig_name}.png'
+    print(fig_path)
+    print("----------------" + '\n')
+    plt.savefig(fig_path)
+    plt.close()
+
+def get_activation_npy(npy_path, fig_name, activation_array):
+    np.save(f'{npy_path}{fig_name}.npy', activation_array)
+
+if __name__ == '__main__':
+    device_name = "gpu"
+    if device_name == "gpu":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device("cpu")
+    print(f"Running on {'GPU' if device.type == 'cuda' else 'CPU'}")
+    # pre-trained ResNet-50 model to device
+    resnet50 = models.resnet50(pretrained=True).to(device)
+
+    all_layers = ['resnet50.conv1',
+                  'resnet50.layer1[0]', 'resnet50.layer1[1]', 'resnet50.layer1[2]',
+                  'resnet50.layer2[0]', 'resnet50.layer2[1]', 'resnet50.layer2[2]', 'resnet50.layer2[3]',
+                  'resnet50.layer3[0]', 'resnet50.layer3[1]', 'resnet50.layer3[2]', 'resnet50.layer3[3]',
+                  'resnet50.layer4[0]', 'resnet50.layer4[1]', 'resnet50.layer4[2]']
+
+    video_type = 'test'
+    # Test
+    if video_type == 'test':
+        metadata_path = "../../metadata/test_videos.csv"
+    # NR:
+    elif video_type == 'resolution_ugc':
+        resolution = '360P'
+        metadata_path = f"../../metadata/YOUTUBE_UGC_{resolution}_metadata.csv"
+    else:
+        metadata_path = f'../../metadata/{video_type.upper()}_metadata.csv'
+
+    ugcdata = pd.read_csv(metadata_path)
+    for i in range(len(ugcdata)):
+        video_name = ugcdata['vid'][i]
+        sampled_frame_path = os.path.join('../..', 'video_sampled_frame', 'sampled_frame', f'{video_name}')
+
+        print(f"Processing video: {video_name}")
+        image_paths = glob.glob(os.path.join(sampled_frame_path, f'{video_name}_*.png'))
+        frame_number = 0
+        for image in image_paths:
+            print(f"{image}")
+            frame_number += 1
+            process_video_frame(video_name, image, frame_number, all_layers, resnet50, device)
+
+# # ResNet-50 layers to visualize
+# layers_to_visualize_resnet50 = {
+#     'conv1': 0,
+#     'layer1.0.conv1': 2,
+#     'layer1.0.conv2': 3,
+#     'layer1.1.conv1': 5,
+#     'layer1.1.conv2': 6,
+#     'layer1.2.conv1': 8,
+#     'layer1.2.conv2': 9,
+#     'layer2.0.conv1': 11,
+#     'layer2.0.conv2': 12,
+#     'layer2.1.conv1': 14,
+#     'layer2.1.conv2': 15,
+#     'layer2.2.conv1': 17,
+#     'layer2.2.conv2': 18,
+#     'layer2.3.conv1': 20,
+#     'layer2.3.conv2': 21,
+#     'layer3.0.conv1': 23,
+#     'layer3.0.conv2': 24,
+#     'layer3.0.downsample.0': 25,
+#     'layer3.1.conv1': 27,
+#     'layer3.1.conv2': 28,
+#     'layer3.2.conv1': 30,
+#     'layer3.2.conv2': 31,
+#     'layer3.3.conv1': 33,
+#     'layer3.3.conv2': 34,
+#     'layer4.0.conv1': 36,
+#     'layer4.0.conv2': 37,
+#     'layer4.0.downsample.0': 38,
+#     'layer4.1.conv1': 40,
+#     'layer4.1.conv2': 41,
+#     'layer4.2.conv1': 43,
+#     'layer4.2.conv2': 44,
+# }

extractor/visualise_resnet_layer.py

@@ -0,0 +1,194 @@
+import warnings
+warnings.filterwarnings("ignore")
+import os
+import glob
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+import torch
+from torchvision import models, transforms
+from thop import profile
+is_flop_cal = False
+
+# get the activation
+def get_activation(model, layer, input_img_data):
+    model.eval()
+    activations = []
+    inputs = []
+
+    def hook(module, input, output):
+        activations.append(output)
+        inputs.append(input[0])
+
+    hook_handle = layer.register_forward_hook(hook)
+    with torch.no_grad():
+        model(input_img_data)
+    hook_handle.remove()
+    return activations, inputs
+
+def get_activation_map(frame, layer_name, resnet50, device):
+    # image pre-processing
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+
+    # Apply the transformations (resize and normalize)
+    frame_tensor = transform(frame)
+
+    # adding index 0 changes the original [C, H, W] shape to [1, C, H, W]
+    if frame_tensor.dim() == 3:
+        frame_tensor = frame_tensor.unsqueeze(0)
+    # print(f'Image dimension: {frame_tensor.shape}')
+
+    # getting the activation of a given layer
+    conv_idx = layer_name
+    layer_obj = eval(conv_idx)
+    activations, inputs = get_activation(resnet50, layer_obj, frame_tensor)
+    activated_img = activations[0][0]
+    activation_array = activated_img.cpu().numpy()
+
+    # calculate FLOPs for layer
+    if is_flop_cal == True:
+        flops, params = profile(layer_obj, inputs=(inputs[0],), verbose=False)
+        if params == 0 and isinstance(layer_obj, torch.nn.Conv2d):
+            params = layer_obj.in_channels * layer_obj.out_channels * layer_obj.kernel_size[0] * layer_obj.kernel_size[1]
+            if layer_obj.bias is not None:
+                params += layer_obj.out_channels
+        # print(f"FLOPs for {layer_name}: {flops}, Params: {params}")
+    else:
+        flops, params = None, None
+    return activated_img, activation_array, flops, params
+
+def process_video_frame(video_name, frame, frame_number, layer_name, resnet50, device):
+    # create a dictionary to store activation arrays for each layer
+    activations_dict = {}
+    total_flops = 0
+    total_params = 0
+    fig_name = f"resnet50_feature_map_layer_{layer_name}"
+    combined_name = f"resnet50_feature_map"
+
+    activated_img, activation_array, flops, params = get_activation_map(frame, layer_name, resnet50, device)
+    if is_flop_cal == True:
+        total_flops += flops
+        total_params += params
+
+    # save activation maps as png
+    # png_path = f'../visualisation/resnet50/{video_name}/frame_{frame_number}/'
+    # npy_path = f'../features/resnet50/{video_name}/frame_{frame_number}/'
+    # os.makedirs(png_path, exist_ok=True)
+    # os.makedirs(npy_path, exist_ok=True)
+    # get_activation_png(png_path, fig_name, activated_img)
+    # save activation features as pny
+    # get_activation_npy(npy_path, fig_name, activation_array)
+
+    # print(f"total FLOPs for Resnet50 layerstack: {total_flops}, Params: {total_params}")
+    frame_npy_path = f'../features/resnet50/{video_name}/frame_{frame_number}_{combined_name}.npy'
+    return activated_img, frame_npy_path, total_flops, total_params
+
+def get_activation_png(png_path, fig_name, activated_img, n=8):
+    fig = plt.figure(figsize=(10, 10))
+
+    # visualise activation map for 64 channels
+    for i in range(n):
+        for j in range(n):
+            idx = (n * i) + j
+            if idx >= activated_img.shape[0]:
+                break
+            ax = fig.add_subplot(n, n, idx + 1)
+            ax.imshow(activated_img[idx].cpu().numpy(), cmap='viridis')
+            ax.axis('off')
+
+    # save figures
+    fig_path = f'{png_path}{fig_name}.png'
+    print(fig_path)
+    print("----------------" + '\n')
+    plt.savefig(fig_path)
+    plt.close()
+
+def get_activation_npy(npy_path, fig_name, activation_array):
+    np.save(f'{npy_path}{fig_name}.npy', activation_array)
+
+if __name__ == '__main__':
+    device_name = "gpu"
+    if device_name == "gpu":
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device("cpu")
+    print(f"Running on {'GPU' if device.type == 'cuda' else 'CPU'}")
+    # pre-trained ResNet-50 model to device
+    resnet50 = models.resnet50(pretrained=True).to(device)
+
+    for idx, (name, layer) in enumerate(resnet50.named_children()):
+        print(f"Index: {idx}, Layer Name: {name}, Layer Type: {type(layer)}")
+    layer_name = 'layer4.2.conv2'
+
+    video_type = 'test'
+    # Test
+    if video_type == 'test':
+        metadata_path = "../../metadata/test_videos.csv"
+    # NR:
+    elif video_type == 'resolution_ugc':
+        resolution = '360P'
+        metadata_path = f"../../metadata/YOUTUBE_UGC_{resolution}_metadata.csv"
+    else:
+        metadata_path = f'../../metadata/{video_type.upper()}_metadata.csv'
+
+    ugcdata = pd.read_csv(metadata_path)
+    for i in range(len(ugcdata)):
+        video_name = ugcdata['vid'][i]
+        sampled_frame_path = os.path.join('../..', 'video_sampled_frame', 'sampled_frame', f'{video_name}')
+
+        print(f"Processing video: {video_name}")
+        image_paths = glob.glob(os.path.join(sampled_frame_path, f'{video_name}_*.png'))
+        frame_number = 0
+        for image in image_paths:
+            print(f"{image}")
+            frame_number += 1
+            process_video_frame(video_name, image, frame_number, layer_name, resnet50, device)
+
+# # ResNet-50 layers to visualize
+# layers_to_visualize_resnet50 = {
+#     'conv1': 0,
+#     'layer1.0.conv1': 2,
+#     'layer1.0.conv2': 3,
+#     'layer1.1.conv1': 5,
+#     'layer1.1.conv2': 6,
+#     'layer1.2.conv1': 8,
+#     'layer1.2.conv2': 9,
+#     'layer2.0.conv1': 11,
+#     'layer2.0.conv2': 12,
+#     'layer2.1.conv1': 14,
+#     'layer2.1.conv2': 15,
+#     'layer2.2.conv1': 17,
+#     'layer2.2.conv2': 18,
+#     'layer2.3.conv1': 20,
+#     'layer2.3.conv2': 21,
+#     'layer3.0.conv1': 23,
+#     'layer3.0.conv2': 24,
+#     'layer3.0.downsample.0': 25,
+#     'layer3.1.conv1': 27,
+#     'layer3.1.conv2': 28,
+#     'layer3.2.conv1': 30,
+#     'layer3.2.conv2': 31,
+#     'layer3.3.conv1': 33,
+#     'layer3.3.conv2': 34,
+#     'layer4.0.conv1': 36,
+#     'layer4.0.conv2': 37,
+#     'layer4.0.downsample.0': 38,
+#     'layer4.1.conv1': 40,
+#     'layer4.1.conv2': 41,
+#     'layer4.2.conv1': 43,
+#     'layer4.2.conv2': 44,
+# }
+
+# Index: 0, Layer Name: conv1, Layer Type: <class 'torch.nn.modules.conv.Conv2d'>
+# Index: 1, Layer Name: bn1, Layer Type: <class 'torch.nn.modules.batchnorm.BatchNorm2d'>
+# Index: 2, Layer Name: relu, Layer Type: <class 'torch.nn.modules.activation.ReLU'>
+# Index: 3, Layer Name: maxpool, Layer Type: <class 'torch.nn.modules.pooling.MaxPool2d'>
+# Index: 4, Layer Name: layer1, Layer Type: <class 'torch.nn.modules.container.Sequential'>
+# Index: 5, Layer Name: layer2, Layer Type: <class 'torch.nn.modules.container.Sequential'>
+# Index: 6, Layer Name: layer3, Layer Type: <class 'torch.nn.modules.container.Sequential'>
+# Index: 7, Layer Name: layer4, Layer Type: <class 'torch.nn.modules.container.Sequential'>
+# Index: 8, Layer Name: avgpool, Layer Type: <class 'torch.nn.modules.pooling.AdaptiveAvgPool2d'>
+# Index: 9, Layer Name: fc, Layer Type: <class 'torch.nn.modules.linear.Linear'>

extractor/visualise_vit_layer.py

@@ -0,0 +1,516 @@
+import os
+import glob
+import math
+from functools import partial
+import torch
+
+import ipywidgets as widgets
+import io
+from PIL import Image
+from torchvision import transforms
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from torch import nn
+from thop import profile
+is_flop_cal = False
+
+import warnings
+warnings.filterwarnings("ignore")
+
+# Step 2: Creating a Vision Transformer
+# normalise the torch
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    # type: (Tensor, float, float, float, float) -> Tensor
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)
+
+#用于执行无梯度截断正态分布初始化。这两个函数在模型初始化中使用，确保权重被适当地初始化。
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    def norm_cdf(x):
+        # computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+#对输入进行随机丢弃一部分元素，实现随机深度（Stochastic Depth）。
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    # work with diff dim tensors, not just 2D ConvNets
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+    random_tensor = keep_prob + \
+        torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    output = x.div(keep_prob) * random_tensor
+    return output
+
+#用于在残差块的主路径上应用 drop_path 函数。
+class DropPath(nn.Module):
+    """
+    Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+#一个多层感知机（MLP）类，包含两个线性层和一个激活函数，用于在残差块中对特征进行非线性映射。
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+# 自注意力机制类，用于在残差块中计算注意力权重并应用它们。
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C //
+                                  self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x, attn
+
+# 一个残差块类，包含一个自注意力模块和一个MLP模块。
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., attn_drop=0.,
+                 drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(
+            drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim,
+                       act_layer=act_layer, drop=drop)
+
+    def forward(self, x, return_attention=False):
+        y, attn = self.attn(self.norm1(x))
+        if return_attention:
+            return attn
+        x = x + self.drop_path(y)
+        x = x + self.drop_path(self.mlp(self.norm2(x)))
+        return x
+
+# 图像到块嵌入类，将输入图像分割成块并将它们映射到嵌入空间
+class PatchEmbed(nn.Module):
+    """
+    Image to Patch Embedding
+    """
+
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        num_patches = (img_size // patch_size) * (img_size // patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+        self.proj = nn.Conv2d(in_chans, embed_dim,
+                              kernel_size=patch_size, stride=patch_size)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        x = self.proj(x).flatten(2).transpose(1, 2)
+        return x
+
+# Vision Transformer模型的主要实现。包含多个残差块、嵌入层等。（还需要学里面每一步代码具体在做什么）
+class VisionTransformer(nn.Module):
+    """
+    Vision Transformer
+    """
+    def __init__(self, img_size=[224], patch_size=16, in_chans=3, num_classes=0, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., norm_layer=nn.LayerNorm, **kwargs):
+        super().__init__()
+        self.num_features = self.embed_dim = embed_dim
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size[0], patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, num_patches + 1, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # stochastic depth decay rule
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer)
+            for i in range(depth)])
+        self.norm = norm_layer(embed_dim)
+
+        # classifier head
+        self.head = nn.Linear(
+            embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def interpolate_pos_encoding(self, x, w, h):
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        class_pos_embed = self.pos_embed[:, 0]
+        patch_pos_embed = self.pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_embed.patch_size
+        h0 = h // self.patch_embed.patch_size
+
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w0 + 0.1, h0 + 0.1
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(
+                math.sqrt(N)), dim).permute(0, 3, 1, 2),
+            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
+            mode='bicubic',
+        )
+        assert int(
+            w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+    def prepare_tokens(self, x):
+        B, nc, w, h = x.shape
+        x = self.patch_embed(x)  # patch linear embedding
+
+        # add the [CLS] token to the embed patch tokens
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # add positional encoding to each token
+        x = x + self.interpolate_pos_encoding(x, w, h)
+
+        return self.pos_drop(x)
+
+    def forward(self, x):
+        x = self.prepare_tokens(x)
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        return x[:, 0], x[:, 1:]  # return CLS token and attention_features maps
+
+    def get_last_selfattention(self, x):
+        x = self.prepare_tokens(x)
+        for i, blk in enumerate(self.blocks):
+            if i < len(self.blocks) - 1:
+                x = blk(x)
+            else:
+                # return attention of the last block
+                # print(f"return attention of the last block: {x.shape}")
+                # print(blk(x, return_attention=True).shape)
+                return blk(x, return_attention=True)
+
+    def get_intermediate_layers(self, x, n=1):
+        x = self.prepare_tokens(x)
+
+        output = []
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if len(self.blocks) - i <= n:
+                output.append(self.norm(x))
+        return output
+
+# Vision Transformer 模型的生成器类，用于实例化和配置特定模型。
+class VitGenerator(object):
+    def __init__(self, name_model, patch_size, device, evaluate=True, random=False, verbose=False):
+        self.name_model = name_model
+        self.patch_size = patch_size
+        self.evaluate = evaluate
+        self.device = device
+        self.verbose = verbose
+        self.model = self._getModel()
+        self._initializeModel()
+        if not random:
+            self._loadPretrainedWeights()
+
+    def _getModel(self):
+        if self.verbose:
+            pass
+            # print((f"[INFO] Initializing {self.name_model} with patch size of {self.patch_size}"))
+        if self.name_model == 'vit_tiny':
+            model = VisionTransformer(patch_size=self.patch_size, embed_dim=192, depth=12, num_heads=3, mlp_ratio=4,
+                                      qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6))
+
+        elif self.name_model == 'vit_small':
+            model = VisionTransformer(patch_size=self.patch_size, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4,
+                                      qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6))
+
+        elif self.name_model == 'vit_base':
+            model = VisionTransformer(patch_size=self.patch_size, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4,
+                                      qkv_bias=True, norm_layer=partial(nn.LayerNorm, eps=1e-6))
+        else:
+            raise f"No model found with {self.name_model}"
+
+        return model
+
+    def _initializeModel(self):
+        if self.evaluate:
+            for p in self.model.parameters():
+                p.requires_grad = False
+
+            self.model.eval()
+
+        self.model.to(self.device)
+
+    def _loadPretrainedWeights(self):
+        if self.verbose:
+            pass
+            # print(("[INFO] Loading weights"))
+        url = None
+        if self.name_model == 'vit_small' and self.patch_size == 16:
+            url = "dino_deitsmall16_pretrain/dino_deitsmall16_pretrain.pth"
+
+        elif self.name_model == 'vit_small' and self.patch_size == 8:
+            url = "dino_deitsmall8_300ep_pretrain/dino_deitsmall8_300ep_pretrain.pth"
+
+        elif self.name_model == 'vit_base' and self.patch_size == 16:
+            url = "dino_vitbase16_pretrain/dino_vitbase16_pretrain.pth"
+
+        elif self.name_model == 'vit_base' and self.patch_size == 8:
+            url = "dino_vitbase8_pretrain/dino_vitbase8_pretrain.pth"
+
+        if url is None:
+            pass
+            # print((f"Since no pretrained weights have been found with name {self.name_model} and patch size {self.patch_size}, random weights will be used"))
+
+        else:
+            state_dict = torch.hub.load_state_dict_from_url(
+                url="https://dl.fbaipublicfiles.com/dino/" + url)
+            self.model.load_state_dict(state_dict, strict=True)
+        # print(url)
+
+    def get_last_selfattention(self, img):
+        return self.model.get_last_selfattention(img.to(self.device))
+
+    def __call__(self, x):
+        return self.model(x)
+
+# Step 3: Creating Visualization Functions
+def transform(img, img_size):
+    img = transforms.Resize(img_size)(img)
+    img = transforms.ToTensor()(img)
+    return img
+
+def visualize_predict(model, img_tensor, patch_size, device, video_name, frame_number, fig_name, combined_name):
+    if img_tensor.dim() == 3:
+        img_tensor = img_tensor.unsqueeze(0)
+    attention = visualize_attention(model, img_tensor, patch_size, device)
+    # save activation maps as png
+    # png_path = f'../visualisation/resnet50/{video_name}/frame_{frame_number}/'
+    # os.makedirs(png_path, exist_ok=True)
+    # get_activation_png(img, png_path, fig_name, attention)
+    # save activation features as npy
+    activations_dict, frame_npy_path = get_activation_npy(video_name, frame_number, fig_name, combined_name, attention)
+    return activations_dict, frame_npy_path
+
+def visualize_attention(model, img_tensor, patch_size, device):
+    # img_tensor: format [1, C, H, W]
+    # Adjust the image dimensions to be divisible by the patch size
+    w, h = img_tensor.shape[2] - img_tensor.shape[2] % patch_size, img_tensor.shape[3] - img_tensor.shape[3] % patch_size
+    img_tensor = img_tensor[:, :, :w, :h]
+
+    w_featmap = img_tensor.shape[-2] // patch_size
+    h_featmap = img_tensor.shape[-1] // patch_size
+
+    attentions = model.get_last_selfattention(img_tensor.to(device))
+    nh = attentions.shape[1]  # number of heads
+
+    # keep only the output patch attention
+    attentions = attentions[0, :, 0, 1:].reshape(nh, -1)
+    attentions = attentions.reshape(nh, w_featmap, h_featmap)
+    attentions = nn.functional.interpolate(attentions.unsqueeze(0), scale_factor=patch_size, mode="nearest")[0].cpu().numpy()
+
+    return attentions
+
+def get_activation_png(img, png_path, fig_name, attention):
+    n_heads = attention.shape[0]
+
+    # attention maps
+    for i in range(n_heads):
+        plt.imshow(attention[i], cmap='viridis') #cmap='viridis', cmap='inferno'
+        plt.title(f"Head n: {i + 1}")
+        plt.axis('off')  # Turn off axis ticks and labels
+
+        # Save figures
+        fig_path = f'{png_path}{fig_name}_head_{i + 1}.png'
+        print(fig_path)
+        plt.savefig(fig_path)
+        plt.close()
+
+    # head mean map
+    plt.figure(figsize=(10, 10))
+    image_name = fig_name.replace('vit_feature_map_', '')
+    text = [f"{image_name}", "Head Mean"]
+    for i, fig in enumerate([img, np.mean(attention, 0)]):
+        plt.subplot(1, 2, i+1)
+        plt.imshow(fig, cmap='viridis')
+        plt.title(text[i])
+        plt.axis('off')  # Turn off axis ticks and labels
+    fig_path1 = f'{png_path}{fig_name}_head_mean.png'
+    print(fig_path1)
+    print("----------------" + '\n')
+    plt.savefig(fig_path1)
+    plt.close()
+
+    # combine
+    # plt.figure(figsize=(20, 20))
+    # for i in range(n_heads):
+    #     plt.subplot(n_heads//3, 3, i+1)
+    #     plt.imshow(attention[i], cmap='inferno')
+    #     plt.title(f"Head n: {i+1}")
+    # plt.tight_layout()
+    # fig_path2 = png_path + fig_name + '_heads.png'
+    # print(fig_path2 + '\n')
+    # plt.savefig(fig_path2)
+    # plt.close()
+
+def get_activation_npy(video_name, frame_number, fig_name, combined_name, attention):
+    # save activation features as pny
+    # npy_path = f'../features/vit/{video_name}/frame_{frame_number}/'
+    # os.makedirs(npy_path, exist_ok=True)
+
+    mean_attention = attention.mean(axis=0)
+    frame_npy_path = f'../features/vit/{video_name}/frame_{frame_number}_{combined_name}.npy'
+
+    return mean_attention, frame_npy_path
+
+
+class Loader(object):
+    def __init__(self):
+        self.uploader = widgets.FileUpload(accept='image/*', multiple=False)
+        self._start()
+
+    def _start(self):
+        display(self.uploader)
+
+    def getLastImage(self):
+        try:
+            for uploaded_filename in self.uploader.value:
+                uploaded_filename = uploaded_filename
+            img = Image.open(io.BytesIO(
+                bytes(self.uploader.value[uploaded_filename]['content'])))
+
+            return img
+        except:
+            return None
+
+    def saveImage(self, path):
+        with open(path, 'wb') as output_file:
+            for uploaded_filename in self.uploader.value:
+                content = self.uploader.value[uploaded_filename]['content']
+                output_file.write(content)
+
+def process_video_frame(video_name, frame, frame_number, model, patch_size, device):
+    # resize image
+    if frame.dim() == 3:
+        frame = frame.unsqueeze(0)
+    if frame.shape[2:] != (224, 224):
+        frame_tensor = torch.nn.functional.interpolate(frame, size=(224, 224), mode='bicubic', align_corners=False)
+    else:
+        frame_tensor = frame
+
+    # Calculate FLOPs and Params
+    if is_flop_cal == True:
+        total_flops, total_params = profile(model.model, inputs=(frame_tensor,), verbose=False)
+        print(f"total FLOPs for ViT layerstack: {total_flops}, Params: {total_params}")
+    else:
+        total_flops, total_params = None, None
+
+    fig_name = f"vit_feature_map"
+    combined_name = f"vit_feature_map"
+
+    # activations_dict, frame_npy_path = visualize_predict(model, frame_tensor, patch_size, device, video_name, frame_number, fig_name, combined_name)
+    attention_features, frame_feature_npy_path = extract_features(model, frame_tensor, video_name, frame_number, combined_name)
+    return attention_features, frame_feature_npy_path, total_flops, total_params
+
+def extract_features(model, img_tensor, video_name, frame_number, combined_name):
+    if img_tensor.dim() == 3:
+        img_tensor = img_tensor.unsqueeze(0)
+    cls_token, attention_features = model(img_tensor)
+
+    attention_features = attention_features.squeeze(0)
+    frame_feature_npy_path = f'../features/vit/{video_name}/frame_attention_{frame_number}_{combined_name}.npy'
+    return attention_features, frame_feature_npy_path
+
+if __name__ == '__main__':
+    # Step 4: Visualizing Images
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+    if device.type == "cuda":
+        torch.cuda.set_device(0)
+
+    name_model = 'vit_base'
+    patch_size = 16
+
+    model = VitGenerator(name_model, patch_size,
+                         device, evaluate=True, random=False, verbose=True)
+
+    video_type = 'test'
+    # Test
+    if video_type == 'test':
+        metadata_path = "../../metadata/test_videos.csv"
+    # NR:
+    elif video_type == 'resolution_ugc':
+        resolution = '360P'
+        metadata_path = f"../../metadata/YOUTUBE_UGC_{resolution}_metadata.csv"
+    else:
+        metadata_path = f'../../metadata/{video_type.upper()}_metadata.csv'
+
+    ugcdata = pd.read_csv(metadata_path)
+    for i in range(len(ugcdata)):
+        video_name = ugcdata['vid'][i]
+        sampled_frame_path = os.path.join('../..', 'video_sampled_frame', 'sampled_frame', f'{video_name}')
+
+        print(f"Processing video: {video_name}")
+        image_paths = glob.glob(os.path.join(sampled_frame_path, f'{video_name}_*.png'))
+        frame_number = 0
+        for image in image_paths:
+            print(f"{image}")
+            frame_number += 1
+            process_video_frame(video_name, image, frame_number, model, patch_size, device)