import torch.nn as nn
import torch
#from pytorch_pretrained_vit import ViT,VoT,VVT
from vgg import vgg16 as backbone
from  common import global_para as gp



class Multi_transform_image(nn.Module):
    def __init__(self, load_pretrain=True):
        super(Multi_transform_image, self).__init__()
        #frontend_feat
        #self.frontend = backbone(pretrained=load_pretrain)
        self.pixel_model = ViT('B_16_imagenet1k', pretrained=True,num_classes=gp.FineClassNum,image_size=448)
    def forward(self, x):
        x = x.float()
        output_pixel = self.pixel_model(x)
        return output_pixel









if __name__ == '__main__':

    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print(device)
    net = Multi_transform().to(device)
    input1 = torch.randn(1, 3, 448,448).to(device)
    input2 = torch.randn(1, 3, 448, 448).to(device)
    output = net(input1,input2)
    print(output.shape)
    # torch.Size([1, 1, 96, 128]) torch.Size([1, 5]) torch.Size([1, 13])