David310/Detect_AI-generated_Image

1from .clip import clip 2from PIL import Image3import torch.nn as nn4 5 6CHANNELS = {7    "RN50" : 1024,8    "ViT-L/14" : 7689}10 11class CLIPModel(nn.Module):12    def __init__(self, name, num_classes=1):13        super(CLIPModel, self).__init__()14 15        self.model, self.preprocess = clip.load(name, device="cpu") # self.preprecess will not be used during training, which is handled in Dataset class 16        self.fc = nn.Linear( CHANNELS[name], num_classes )17 18 19    def forward(self, x, return_feature=False):20        features = self.model.encode_image(x) 21        # print(features.keys())22        """23        使用的是ViT-Large, 共24层24        选择第24、22、20层的[cls]feature做加权平均25        """26        if return_feature:27            return features['after_projection']28        # print(features['after_projection'].shape)29        # print(features['layer21'].shape)30        # print(features['layer19'].shape)31        # features = 0.5*features['after_projection'] + 0.3*features['layer21'] + 0.2*features['layer19']32        # print(features.shape)33        features = features['res_output']34        return self.fc(features)35 36