<reponame>MichaelLynn1996/ast
# -*- coding: utf-8 -*-
# @Time    : 6/10/21 5:04 PM
# @Author  : <NAME>
# @Affiliation  : Massachusetts Institute of Technology
# @Email   : <EMAIL>
# @File    : ast_models.py
import os
from typing import Tuple

import timm
import torch
import torch.nn as nn
from timm.models.layers import to_2tuple
from torch.cuda.amp import autocast

os.environ['TORCH_HOME'] = '../../pretrained_models'

Size_ = Tuple[int, int]


# override the timm package to relax the input shape constraint.
class PatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """

    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768, fstride=10, tstride=10):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        self.img_size = img_size
        self.patch_size = to_2tuple(patch_size)
        assert img_size[0] % patch_size[0] == 0 and img_size[1] % patch_size[1] == 0, \
            f"img_size {img_size} should be divided by patch_size {patch_size}."
        if patch_size == (4, 4):
            self.H = (img_size[0] // 4) * 2 - 1
            self.W = (img_size[1] // 4) * 2 - 1
            # self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]
            # self.H, self.W = self.get_shape(fstride, tstride, img_size, patch_size, embed_dim)
            self.num_patches = self.H * self.W
        else:
            self.H, self.W = img_size[0] // patch_size[0], img_size[1] // patch_size[1]
            self.num_patches = self.H * self.W
        if patch_size == (4, 4):
            self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=(patch_size[0] // 2,
                                                                                       patch_size[1] // 2))
            print(self.proj)
        else:
            self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
        self.norm = nn.LayerNorm(embed_dim)

    def get_shape(self, stride, img_size, patch_size, embed_dim):
        test_input = torch.randn(1, 1, img_size[0], img_size[1])
        test_proj = nn.Conv2d(1, embed_dim, kernel_size=patch_size, stride=stride)
        test_out = test_proj(test_input)
        print(test_out.shape)
        f_dim = test_out.shape[2]
        t_dim = test_out.shape[3]
        return f_dim, t_dim

    def forward(self, x) -> Tuple[torch.Tensor, Size_]:
        B, C, H, W = x.shape

        x = self.proj(x).flatten(2).transpose(1, 2)
        x = self.norm(x)
        # out_size = (H // self.patch_size[0], W // self.patch_size[1])
        # out_size = self.H, self.W
        if self.patch_size == (4, 4):
            out_size = self.H, self.W
        else:
            out_size = (H // self.patch_size[0], W // self.patch_size[1])

        return x, out_size


class ASSVT(nn.Module):
    """
    The AST model.
    :param label_dim: the label dimension, i.e., the number of total classes, it is 527 for AudioSet, 50 for ESC-50, and 35 for speechcommands v2-35
    :param fstride: the stride of patch splitting on the frequency dimension, for 16*16 patches, fstride=16 means no overlap, fstride=10 means overlap of 6
    :param tstride: the stride of patch splitting on the time dimension, for 16*16 patches, tstride=16 means no overlap, tstride=10 means overlap of 6
    :param input_fdim: the number of frequency bins of the input spectrogram
    :param input_tdim: the number of time frames of the input spectrogram
    :param imagenet_pretrain: if use ImageNet pretrained model
    :param audioset_pretrain: if use full AudioSet and ImageNet pretrained model
    :param model_size: the model size of AST, should be in [tiny224, small224, base224, base384], base224 and base 384 are same model, but are trained differently during ImageNet pretraining.
    """

    def __init__(self, label_dim=527, input_fdim=128, input_tdim=1024
                 , imagenet_pretrain=True, audioset_pretrain=False, model_size='large', verbose=True, in_chans=1):

        super(ASSVT, self).__init__()

        if verbose:
            print('---------------AST Model Summary---------------')
            print('ImageNet pretraining: {:s}, AudioSet pretraining: {:s}'.format(str(imagenet_pretrain),
                                                                                  str(audioset_pretrain)))
            # override timm input shape restriction
            # timm.models.twins.PatchEmbed = PatchEmbed

            print((input_fdim, input_tdim))
            img_size_0, img_size_1 = input_fdim, input_tdim
            self.left, self.right, self.top, self.bottom = 0, 0, 0, 0
            if input_fdim % 32 != 0:
                img_size_0 = (input_fdim // 32 + 1) * 32
                self.top = (img_size_0 - input_fdim) // 2
                self.bottom = img_size_0 - input_fdim - self.left
            if input_tdim % 32 != 0:
                img_size_1 = (input_tdim // 32 + 1) * 32
                self.left = (img_size_1 - input_tdim) // 2
                self.right = img_size_1 - input_tdim - self.left
            self.img_size = (img_size_0, img_size_1)
            print(self.img_size)

            # if AudioSet pretraining is not used (but ImageNet pretraining may still apply)
            if model_size == 'small':
                self.v = timm.create_model('twins_svt_small', pretrained=imagenet_pretrain, in_chans=in_chans
                                           , num_classes=label_dim, img_size=self.img_size)
            elif model_size == 'base':
                self.v = timm.create_model('twins_svt_base', pretrained=imagenet_pretrain, in_chans=in_chans
                                           , num_classes=label_dim, img_size=self.img_size)
            elif model_size == 'large':
                self.v = timm.create_model('twins_svt_large', pretrained=imagenet_pretrain, in_chans=in_chans
                                           , num_classes=label_dim, img_size=self.img_size)
            else:
                raise Exception('Model size must be one of small, base, large.')

            self.v.head = nn.Sequential(nn.LayerNorm(self.v.num_features),
                                        nn.Linear(self.v.num_features, label_dim))


    @autocast()
    def forward(self, x):
        """
        :param x: the input spectrogram, expected shape: (batch_size, time_frame_num, frequency_bins), e.g., (12, 1024, 128)
        :return: prediction
        """
        # expect input x = (batch_size, time_frame_num, frequency_bins), e.g., (12, 1024, 128)
        # x = x.unsqueeze(1)
        # x = x.transpose(2, 3)
        #
        # x = self.forward_features(x)
        # x = self.mlp_head(x)

        x = x.unsqueeze(1)
        x = x.transpose(2, 3)
        pad = nn.ZeroPad2d(padding=(self.left, self.right, self.top, self.bottom))
        x = pad(x)
        return self.v(x)
