src/extractor/visualise_vit_layer.py

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)