David310/Detect_AI-generated_Image

1from .resnet import resnet18, resnet34, resnet50, resnet101, resnet1522from .vision_transformer import vit_b_16, vit_b_32, vit_l_16, vit_l_323 4from torchvision import transforms5from PIL import Image6import torch 7import torch.nn as nn 8 9 10model_dict = {11    'resnet18': resnet18,12    'resnet34': resnet34,13    'resnet50': resnet50,14    'resnet101': resnet101,15    'resnet152': resnet152,16    'vit_b_16': vit_b_16,17    'vit_b_32': vit_b_32,18    'vit_l_16': vit_l_16,19    'vit_l_32': vit_l_3220}21 22 23CHANNELS = {24    "resnet50" : 2048,25    "vit_b_16" : 768,26}27 28 29 30class ImagenetModel(nn.Module):31    def __init__(self, name, num_classes=1):32        super(ImagenetModel, self).__init__()33 34        self.model = model_dict[name](pretrained=True)35        self.fc = nn.Linear(CHANNELS[name], num_classes) #manually define a fc layer here36        37 38    def forward(self, x):39        feature = self.model(x)["penultimate"]40        return self.fc(feature)41