models/clip_tqn.py

import sys
# 加入父文件夹路径到sys.path中  
sys.path.append(sys.path[0].replace('models', ''))

import re
import logging
import math
import json
import pathlib
import numpy as np
from copy import deepcopy
from pathlib import Path
from einops import rearrange
from collections import OrderedDict
from dataclasses import dataclass
from typing import Tuple, Union, Callable, Optional


import torch
from torch import nn
import torch.nn.functional as F
import torchvision.models as models
from torch.utils.checkpoint import checkpoint

from transformers import AutoModel,BertConfig,AutoTokenizer
# from pytorch_pretrained_vit import ViT

# from visualizer import get_local
from models.transformer_decoder import *

# from io import BytesIO
# from petrel_client.client import Client

# conf_path = '~/petreloss.conf'
# client = Client(conf_path) 
from torch.autograd import Function
import timm

class ReverseLayerF(Function):

    @staticmethod
    def forward(ctx, x, alpha):
        ctx.alpha = alpha
        return x.view_as(x)

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output.neg() * ctx.alpha, None
    

class DomainClassifier(nn.Module):
    '''一个单层分类器 带梯度反转层'''
    def __init__(self, domain_nums=4, feature_dims=768):
        super().__init__()
        self.domain_nums = domain_nums
        self.feature_dims = feature_dims
        self.fc = nn.Linear(feature_dims, domain_nums)

    def forward(self, x):
        reverse_x = ReverseLayerF.apply(x, 1.0)
        return self.fc(reverse_x)
    
class CLP_clinical(nn.Module):
    def __init__(self,
                bert_model_name: str,
                embed_dim: int = 768,
                freeze_layers:Union[Tuple[int, int], int] = None):
        super().__init__()
        self.bert_model = self._get_bert_basemodel(bert_model_name=bert_model_name, freeze_layers=freeze_layers)
        self.mlp_embed = nn.Sequential(
            nn.Linear(embed_dim, embed_dim),
            nn.GELU(),
            nn.Linear(embed_dim, embed_dim)
        )
        self.embed_dim = embed_dim
        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
        self.init_parameters()
    
    def init_parameters(self):
        nn.init.constant_(self.logit_scale, np.log(1 / 0.07))
        for m in self.mlp_embed:
            if isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, std=self.embed_dim ** -0.5)

    def _get_bert_basemodel(self, bert_model_name, freeze_layers=None):#12
        try:
            print(bert_model_name)
            config = BertConfig.from_pretrained(bert_model_name, output_hidden_states=True)#bert-base-uncased
            model = AutoModel.from_pretrained(bert_model_name, config=config)#, return_dict=True)
            print("Text feature extractor:", bert_model_name)
            print("bert encoder layers:",len(model.encoder.layer))
        except:
            raise ("Invalid model name. Check the config file and pass a BERT model from transformers lybrary")

        if freeze_layers is not None:
            for layer_idx in freeze_layers:
                for param in list(model.encoder.layer[layer_idx].parameters()):
                    param.requires_grad = False
        return model

    def encode_text(self, text):
        #input batch_size,token, return batch_size,dim 
        output = self.bert_model(input_ids = text['input_ids'],attention_mask = text['attention_mask'] )
        last_hidden_state, pooler_output, hidden_states = output[0],output[1],output[2]
        encode_out = self.mlp_embed(pooler_output)
        # encode_out = pooler_output
        return encode_out
    
    def forward(self, text):
        #input batch_size,token, return batch_size,dim 
        output = self.bert_model(input_ids = text['input_ids'],attention_mask = text['attention_mask'] )
        last_hidden_state, pooler_output, hidden_states = output[0],output[1],output[2]
        encode_out = self.mlp_embed(pooler_output)
        # encode_out = pooler_output
        return encode_out
    
    # def forward(self,text1,text2):
    #     text1_features = self.encode_text(text1)
    #     text2_features = self.encode_text(text2)
    #     text1_features = F.normalize(text1_features, dim=-1)
    #     text2_features = F.normalize(text2_features, dim=-1)
    #     return text1_features, text2_features, self.logit_scale.exp()

class ModelRes(nn.Module):
    def __init__(self, res_base_model):
        super(ModelRes, self).__init__()
        self.resnet_dict = {
                            "resnet50": models.resnet50(pretrained=True),
                            "resnet101": models.resnet101(pretrained=True),
                            "resnet152": models.resnet152(pretrained=True),
                            "resnet50_openai": None,
                            'resnet101_openai': None,
                            'resnet50x4_openai': None,
                            }
                            # "resnet50": models.resnet50(pretrained=True)}
        self.resnet = self._get_res_basemodel(res_base_model)
        # num_ftrs = int(self.resnet.fc.in_features/2)
        # self.res_features = nn.Sequential(*list(self.resnet.children())[:-3]) 224
        if 'openai' in res_base_model:
            # 重新定义res_features
            num_ftrs = int(self.resnet.attnpool.v_proj.in_features)
            self.res_features = nn.Sequential(*list(self.resnet.children())[:-1])
        else:
            num_ftrs = int(self.resnet.fc.in_features)
            self.res_features = nn.Sequential(*list(self.resnet.children())[:-2])
        # here num_ftrs = 2048 
        self.res_l1 = nn.Linear(num_ftrs, num_ftrs)
        self.res_l2 = nn.Linear(num_ftrs, 768)

    def _get_res_basemodel(self, res_model_name):
        try:
            res_model = self.resnet_dict[res_model_name]
            print("Image feature extractor:", res_model_name)
            return res_model
        except:
            raise ("Invalid model name. Check the config file and pass one of: resnet18 or resnet50")

    def forward(self, img):
        #return (batchsize, patch_num, dim)
        batch_size = img.shape[0]
        res_fea = self.res_features(img)
        # return res_fea
        # res_fea = F.adaptive_avg_pool2d(res_fea, (1, 1))
        res_fea = rearrange(res_fea,'b d n1 n2 -> b (n1 n2) d')
        h = rearrange(res_fea,'b n d -> (b n) d')
        x = self.res_l1(h)
        x = F.relu(x)
        x = self.res_l2(x)
        out_emb = rearrange(x,'(b n) d -> b n d',b=batch_size)
        out_pool = torch.mean(out_emb,dim=1)
        return out_emb,out_pool


class ModelConvNeXt(nn.Module):
    def __init__(self, convnext_base_model):
        super(ModelConvNeXt, self).__init__()
        self.convnext_dict = {"convnext-tiny": timm.create_model('convnextv2_tiny.fcmae_ft_in22k_in1k_384', pretrained=True, num_classes=1000),
                              "convnext-base": timm.create_model('convnextv2_base.fcmae_ft_in22k_in1k_384', pretrained=True, num_classes=1000),
                              }
        convnext = self._get_convnext_basemodel(convnext_base_model)
        num_ftrs = int(convnext.head.in_features)
        self.conv_features = nn.Sequential(*list(convnext.children())[:-2])
        self.conv_l1 = nn.Linear(num_ftrs, num_ftrs)
        self.conv_l2 = nn.Linear(num_ftrs, 768)
        
        
    def _get_convnext_basemodel(self, convnext_model_name):
        try:
            convnext_model = self.convnext_dict[convnext_model_name]
            print("Image feature extractor:", convnext_model_name)
            return convnext_model
        except:
            raise ("Invalid model name. Check the config file and pass one of: convnext-tiny, convnext-small or convnext-base")

    def forward(self, img):
        #return (batchsize, patch_num, dim)
        batch_size = img.shape[0]
        conv_fea = self.conv_features(img)
        conv_fea = F.adaptive_avg_pool2d(conv_fea, (1, 1))
        conv_fea = rearrange(conv_fea,'b d n1 n2 -> b (n1 n2) d')
        h = rearrange(conv_fea,'b n d -> (b n) d')
        x = self.conv_l1(h)
        x = F.relu(x)
        x = self.conv_l2(x)
        out_emb = rearrange(x,'(b n) d -> b n d',b=batch_size)
        out_pool = torch.mean(out_emb,dim=1)
        return out_emb,out_pool
    

# class ModelConvNeXt(nn.Module):
#     def __init__(self, convnext_base_model):
#         super(ModelConvNeXt, self).__init__()
#         self.convnext_dict = {"convnext-tiny": models.convnext_tiny(weights='ConvNeXt_Tiny_Weights.DEFAULT'),
#                               "convnext-small": models.convnext_small(weights='ConvNeXt_Small_Weights.DEFAULT'),
#                               "convnext-base": models.convnext_base(weights='ConvNeXt_Base_Weights.DEFAULT'),
#                               }
#         convnext = self._get_convnext_basemodel(convnext_base_model)
#         num_ftrs = int(convnext.classifier[-1].in_features)
#         self.conv_features = nn.Sequential(*list(convnext.children())[:-2])
#         self.conv_l1 = nn.Linear(num_ftrs, num_ftrs)
#         self.conv_l2 = nn.Linear(num_ftrs, 768)
        
        
#     def _get_convnext_basemodel(self, convnext_model_name):
#         try:
#             convnext_model = self.convnext_dict[convnext_model_name]
#             print("Image feature extractor:", convnext_model_name)
#             return convnext_model
#         except:
#             raise ("Invalid model name. Check the config file and pass one of: convnext-tiny, convnext-small or convnext-base")

#     def forward(self, img):
#         #return (batchsize, patch_num, dim)
#         batch_size = img.shape[0]
#         conv_fea = self.conv_features(img)
#         conv_fea = F.adaptive_avg_pool2d(conv_fea, (1, 1))
#         conv_fea = rearrange(conv_fea,'b d n1 n2 -> b (n1 n2) d')
#         h = rearrange(conv_fea,'b n d -> (b n) d')
#         x = self.conv_l1(h)
#         x = F.relu(x)
#         x = self.conv_l2(x)
#         out_emb = rearrange(x,'(b n) d -> b n d',b=batch_size)
#         out_pool = torch.mean(out_emb,dim=1)
#         return out_emb,out_pool
    
# import open_clip
# class ModelCLIP(nn.Module):
#     def __init__(self, clip_base_model):
#         super(ModelCLIP, self).__init__()
#         # 根据clip_base_model加载不同的模型
#         if clip_base_model == 'openai_EVA02-B-16':
#             model, _, preprocess = open_clip.create_model_and_transforms('EVA02-B-16', pretrained='merged2b_s8b_b131k')
            
#         elif clip_base_model == 'openai_convnext_base_w':
#             model, _, preprocess = open_clip.create_model_and_transforms('convnext_base_w', pretrained='laion2b_s13b_b82k_augreg')
        
#         else:
#             raise ("Invalid model name. Check the config file and pass one of: EVA02-B-16 or convnext_base_w")

  
        
class ModelEfficientV2(nn.Module):
    def __init__(self, efficientv2_base_model):
        super(ModelEfficientV2, self).__init__()
        self.efficientv2_dict = {"efficientnet_v2_s": models.efficientnet_v2_s(weights='EfficientNet_V2_S_Weights.IMAGENET1K_V1'),}
        self.efficientv2_model = self._get_efficientv2_basemodel(efficientv2_base_model)
        num_ftrs = int(self.efficientv2_model.classifier[-1].in_features)
        self.efficientv2_features = nn.Sequential(*list(self.efficientv2_model.children())[:-2])
        self.efficientv2_l1 = nn.Linear(num_ftrs, num_ftrs)
        self.efficientv2_l2 = nn.Linear(num_ftrs, 768)
        
    
    def _get_efficientv2_basemodel(self, efficientv2_model_name):
        try:
            efficientv2_model = self.efficientv2_dict[efficientv2_model_name]
            print("Image feature extractor:", efficientv2_model_name)
            return efficientv2_model
        except:
            raise ("Invalid model name. Check the config file and pass one of: efficientnetv2_rw_s")
    
    def forward(self, img):
        batch_size = img.shape[0]
        efficientv2_fea = self.efficientv2_features(img)
        # efficientv2_fea = F.adaptive_avg_pool2d(efficientv2_fea, (1, 1))
        # print(efficientv2_fea.shape)
        efficientv2_fea = rearrange(efficientv2_fea,'b d n1 n2 -> b (n1 n2) d')
        # print(efficientv2_fea.shape)
        h = rearrange(efficientv2_fea,'b n d -> (b n) d')
        # print(h.shape)
        x = self.efficientv2_l1(h)
        x = F.relu(x)
        x = self.efficientv2_l2(x)
        # print(x.shape)
        out_emb = rearrange(x,'(b n) d -> b n d',b=batch_size)
        out_pool = torch.mean(out_emb,dim=1)
        return out_emb,out_pool
        
        
    
class ModelDense(nn.Module):
    def __init__(self, dense_base_model):
        super(ModelDense, self).__init__()
        
        self.densenet_dict = {"densenet121": models.densenet121(weights='DenseNet121_Weights.IMAGENET1K_V1'),
                            "densenet161": models.densenet161(weights='DenseNet161_Weights.IMAGENET1K_V1'),
                            "densenet201": models.densenet201(weights='DenseNet201_Weights.IMAGENET1K_V1'),}
        self.densenet = self._get_dense_basemodel(dense_base_model)
        num_ftrs = int(self.densenet.classifier.in_features)
        self.dense_features = self.densenet.features
        self.dense_l1 = nn.Linear(num_ftrs, num_ftrs)
        self.dense_l2 = nn.Linear(num_ftrs, 768)

    def _get_dense_basemodel(self, dense_base_model):
        try:
            dense_model = self.densenet_dict[dense_base_model]
            print("Image feature extractor:", dense_base_model)
            return dense_model
        except:
            raise ("Invalid model name. Check the config file and pass one of: densenet121 or densenet161")

    def forward(self, img):
        batch_size = img.shape[0]
        dense_fea = self.dense_features(img)#N, 1024, 7,7
        dense_fea = rearrange(dense_fea,'b d n1 n2 -> b (n1 n2) d')
        h = rearrange(dense_fea,'b n d -> (b n) d')
        x = self.dense_l1(h)
        x = F.relu(x)
        x = self.dense_l2(x)
        out_emb = rearrange(x,'(b n) d -> b n d',b=batch_size)
        out_pool = torch.mean(out_emb,dim=1)
        return out_emb,out_pool

class TQN_Model(nn.Module):
    def __init__(self, 
            embed_dim: int = 768, 
            class_num: int = 1, 
            lam: list = [1, 0]
            ):
        super().__init__()
        self.d_model = embed_dim
        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
        # decoder_layer = TransformerDecoderLayer(self.d_model, 4, 1024,
                                        # 0.1, 'relu',normalize_before=True)
        decoder_layerV1 = TransformerDecoderLayerV1(self.d_model, 4, 1024,
                                        0.1, 'relu', True, lam)
        self.decoder_norm = nn.LayerNorm(self.d_model)
        # self.decoder = TransformerDecoder(decoder_layer, 4, self.decoder_norm,
                                # return_intermediate=False)
        self.decoderV1 = TransformerDecoderV1(decoder_layerV1, 4, self.decoder_norm,
                                return_intermediate=False)
        
        self.dropout_feas = nn.Dropout(0.1)

        # class_num = 2
        self.mlp_head = nn.Sequential( # nn.LayerNorm(768),
            nn.Linear(embed_dim, class_num)
        )
        self.apply(self._init_weights)
    
    @staticmethod
    def _init_weights(module):
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=0.02)

        elif isinstance(module, nn.MultiheadAttention):
            module.in_proj_weight.data.normal_(mean=0.0, std=0.02)
            module.out_proj.weight.data.normal_(mean=0.0, std=0.02)

        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
    
    # def forward(self, image_features, text_features):
    #     #image_features (batch_size,patch_num,dim)
    #     #text_features (query_num,dim)
    #     batch_size = image_features.shape[0]
    #     image_features = image_features.transpose(0,1)
    #     text_features = text_features.unsqueeze(1).repeat(1, batch_size, 1)
    #     image_features = self.decoder_norm(image_features)
    #     text_features = self.decoder_norm(text_features)
        
    #     # features = self.decoder(text_features, image_features,
    #     #         memory_key_padding_mask=None, pos=None, query_pos=None)
        
    #     image_features_pool = torch.mean(image_features,dim=0).unsqueeze(0)
    #     features = self.decoderV1(text_features, image_features, image_features_pool,
    #             memory_key_padding_mask=None, pos=None, query_pos=None)  
        

 
    #     features = self.dropout_feas(features).transpose(0,1)  #b,embed_dim
    #     out = self.mlp_head(features)  #(batch_size, query_num)
    #     # out = out.squeeze(-1)
    #     return out
    
    def forward(self, image_features, text_features, return_atten = False):
        #image_features (batch_size,patch_num,dim)
        #text_features (query_num,dim)
        batch_size = image_features.shape[0]
        image_features = image_features.transpose(0,1)
        text_features = text_features.unsqueeze(1).repeat(1, batch_size, 1)
        image_features = self.decoder_norm(image_features)
        text_features = self.decoder_norm(text_features)
        
        image_features_pool = torch.mean(image_features,dim=0).unsqueeze(0)
        features,atten_map = self.decoderV1(text_features, image_features, image_features_pool, 
                memory_key_padding_mask=None, pos=None, query_pos=None) 
        features = self.dropout_feas(features).transpose(0,1)  #b,embed_dim
        out = self.mlp_head(features)  #(batch_size, query_num)
        if return_atten:
            return out, atten_map
        else:
            return out



class TQN_Model_Ensemble(nn.Module):
    def __init__(self, 
            embed_dim: int = 768, 
            class_num: int = 1, 
            lam: list = [1, 0]
            ):
        super().__init__()
        self.d_model = embed_dim
        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))
        decoder_layerV1 = TransformerDecoderLayerV1(self.d_model, 4, 1024,
                                        0.1, 'relu', True, lam)
        self.decoder_norm = nn.LayerNorm(self.d_model)
        self.decoder_norm_1 = nn.LayerNorm(self.d_model)
        self.decoder_norm_2 = nn.LayerNorm(self.d_model)
        self.decoderV1 = TransformerDecoderV1(decoder_layerV1, 4, self.decoder_norm,
                                return_intermediate=False)
        self.decoderV1_1 = TransformerDecoderV1(decoder_layerV1, 4, self.decoder_norm_1,
                                return_intermediate=False)
        self.decoderV1_2 = TransformerDecoderV1(decoder_layerV1, 4, self.decoder_norm_2,
                                return_intermediate=False)
        
        self.dropout_feas = nn.Dropout(0.1)

        # class_num = 2
        self.mlp_head = nn.Sequential(nn.Linear(embed_dim, class_num))
        self.mlp_head_1 = nn.Sequential(nn.Linear(embed_dim, class_num))
        self.mlp_head_2 = nn.Sequential(nn.Linear(embed_dim, class_num))
        self.apply(self._init_weights)
    
    @staticmethod
    def _init_weights(module):
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=0.02)

        elif isinstance(module, nn.MultiheadAttention):
            module.in_proj_weight.data.normal_(mean=0.0, std=0.02)
            module.out_proj.weight.data.normal_(mean=0.0, std=0.02)

        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=0.02)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
    
    def forward(self, image_features, text_features, return_atten = False):

        batch_size = image_features.shape[0]
        image_features = image_features.transpose(0,1)
        text_features = text_features.unsqueeze(1).repeat(1, batch_size, 1)
        image_features = self.decoder_norm(image_features)
        image_features_1 = self.decoder_norm_1(image_features)
        image_features_2 = self.decoder_norm_2(image_features)

        text_features = self.decoder_norm(text_features)
        text_features_1 = self.decoder_norm_1(text_features)
        text_features_2 = self.decoder_norm_2(text_features)
        
        image_features_pool = torch.mean(image_features,dim=0).unsqueeze(0)
        image_features_pool_1 = torch.mean(image_features_1,dim=0).unsqueeze(0)
        image_features_pool_2 = torch.mean(image_features_2,dim=0).unsqueeze(0)


        features,atten_map = self.decoderV1(text_features, image_features, image_features_pool, 
                memory_key_padding_mask=None, pos=None, query_pos=None) 
        features = self.dropout_feas(features).transpose(0,1)  
        out = self.mlp_head(features)

        features_1,atten_map_1 = self.decoderV1_1(text_features_1, image_features_1, image_features_pool_1, 
                memory_key_padding_mask=None, pos=None, query_pos=None) 
        features_1 = self.dropout_feas(features_1).transpose(0,1)  
        out_1 = self.mlp_head_1(features_1)

        features_2,atten_map_2 = self.decoderV1_2(text_features_2, image_features_2, image_features_pool_2, 
                memory_key_padding_mask=None, pos=None, query_pos=None) 
        features_2 = self.dropout_feas(features_2).transpose(0,1)  
        out_2 = self.mlp_head_2(features_2)     


        out_stack = torch.stack([out, out_1, out_2])
        out = torch.mean(out_stack, dim=0)               
        





        if return_atten:
            return out, atten_map
        else:
            return out



# MIMIC时，batch_size=32, query_num=41, patch_num=256, dim=768
# img 256, 32, 768
# txt   1, 32, 768
# query41, 32, 768
# fts 41, 32, 768
# out 41, 32, 1
# 未经过sigmoid！计算loss时sigmoid！


if __name__ == "__main__":
            
    #torch 1.10.2 to torch 1.12.1
    #torchvision-0.11.3 to torchvision-0.13.1

    # image = torch.randn(1, 3, 224, 224) 
    # image_encoder = ModelRes(res_base_model = 'resnet50')
    # # image_encoder = ModelDense(dense_base_model = 'densenet121')
    # # image_encoder = ModelViT(vit_base_model = 'vit_b_32')
    # image_encoder(image)

    # image = torch.randn(256, 1, 768)
    # query = torch.randn(41, 768)
    # model = TQN_Model()
    # out = model(image, query)
    
    
    # img = torch.randn(1,3,512,512)
    img = torch.randn(2,3,224,224)
    # model = ModelConvNeXt(convnext_base_model = 'convnext-base')
    # model = ModelEfficientV2(efficientv2_base_model = 'efficientnet_v2_s')
    model = ModelRes(res_base_model = 'resnet50_openai')
    out_emb, out_pool = model(img)
    
    print(out_emb.size(), out_pool.size())