VisionLanguageGroup/MicroscopyMatching

1import torch2import torch.nn as nn3import torch.nn.functional as F4import os5import clip6import sys7import numpy as np8from models.seg_post_model.models import SegModel9 10from torchvision.ops import roi_align11 12 13class Counting_with_SD_features_loca(nn.Module):14    def __init__(self, scale_factor):15        super(Counting_with_SD_features_loca, self).__init__()16        self.adapter = adapter_roi_loca()17        self.regressor = regressor_with_SD_features()18 19 20class Counting_with_SD_features_dino_vit_c3(nn.Module):21    def __init__(self, scale_factor, vit=None):22        super(Counting_with_SD_features_dino_vit_c3, self).__init__()23        self.adapter = adapter_roi_loca()24        self.regressor = regressor_with_SD_features_seg_vit_c3()25 26class Counting_with_SD_features_track(nn.Module):27    def __init__(self, scale_factor, vit=None):28        super(Counting_with_SD_features_track, self).__init__()29        self.adapter = adapter_roi_loca()30        self.regressor = regressor_with_SD_features_tra()31 32 33class adapter_roi_loca(nn.Module):34    def __init__(self, pool_size=[3, 3]):35        super(adapter_roi_loca, self).__init__()36        self.pool_size = pool_size37        self.conv1 = nn.Conv2d(256, 256, kernel_size=3, padding=1)38        self.pool = nn.MaxPool2d(2)39        self.fc = nn.Linear(256 * 3 * 3, 768)40        self.initialize_weights()41    def forward(self, x, boxes):42            num_of_boxes = boxes.shape[1]43            rois = []44            bs, _, h, w = x.shape45            if h != 512 or w != 512:46                x = F.interpolate(x, size=(512, 512), mode='bilinear', align_corners=False)47            if bs == 1:48                boxes = torch.cat([49                    torch.arange(50                        bs, requires_grad=False51                    ).to(boxes.device).repeat_interleave(num_of_boxes).reshape(-1, 1),52                    boxes.flatten(0, 1),53                ], dim=1)54                rois = roi_align(55                    x,56                    boxes=boxes, output_size=3,57                    spatial_scale=1.0 / 8, aligned=True58                )59                rois = torch.mean(rois, dim=0, keepdim=True)60            else:61                boxes = torch.cat([62                    boxes.flatten(0, 1),63                ], dim=1).split(num_of_boxes, dim=0)64                rois = roi_align(65                    x,66                    boxes=boxes, output_size=3,67                    spatial_scale=1.0 / 8, aligned=True68                )69                rois = rois.split(num_of_boxes, dim=0)70                rois = torch.stack(rois, dim=0)71                rois = torch.mean(rois, dim=1, keepdim=False)72            x = self.conv1(rois)73            x = x.view(x.size(0), -1)74            x = self.fc(x)75            return x76    77    def forward_boxes(self, x, boxes):78            num_of_boxes = boxes.shape[1]79            rois = []80            bs, _, h, w = x.shape81            if h != 512 or w != 512:82                x = F.interpolate(x, size=(512, 512), mode='bilinear', align_corners=False)83            if bs == 1:84                boxes = torch.cat([85                    torch.arange(86                        bs, requires_grad=False87                    ).to(boxes.device).repeat_interleave(num_of_boxes).reshape(-1, 1),88                    boxes.flatten(0, 1),89                ], dim=1)90                rois = roi_align(91                    x,92                    boxes=boxes, output_size=3,93                    spatial_scale=1.0 / 8, aligned=True94                )95                # rois = torch.mean(rois, dim=0, keepdim=True)96            else:97                raise NotImplementedError98            x = self.conv1(rois)99            x = x.view(x.size(0), -1)100            x = self.fc(x)101            return x102 103    def initialize_weights(self):104        for m in self.modules():105            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):106                nn.init.xavier_normal_(m.weight)107                if m.bias is not None:108                    nn.init.constant_(m.bias, 0)109 110 111 112class regressor_with_SD_features(nn.Module):113    def __init__(self):114        super(regressor_with_SD_features, self).__init__()115        self.layer1 = nn.Sequential(116            nn.Conv2d(324, 256, kernel_size=1, stride=1),117            nn.LeakyReLU(),118            nn.LayerNorm((64, 64))119        )120        self.layer2 = nn.Sequential(121            nn.Conv2d(256, 128, kernel_size=3, padding=1),122            nn.LeakyReLU(),123            nn.ConvTranspose2d(in_channels=128, out_channels=128, kernel_size=4, stride=2, padding=1),124        )125        self.layer3 = nn.Sequential(126            nn.Conv2d(128, 64, kernel_size=3, padding=1),127            nn.ReLU(),128            nn.ConvTranspose2d(in_channels=64, out_channels=64, kernel_size=4, stride=2, padding=1),129        )130        self.layer4 = nn.Sequential(131            nn.Conv2d(64, 32, kernel_size=3, padding=1),132            nn.LeakyReLU(),133            nn.ConvTranspose2d(in_channels=32, out_channels=32, kernel_size=4, stride=2, padding=1),134        )135        self.conv = nn.Sequential(136            nn.Conv2d(32, 1, kernel_size=1),137            nn.ReLU()138        )139        self.norm = nn.LayerNorm(normalized_shape=(64, 64))140        self.initialize_weights()141 142    def forward(self, attn_stack, feature_list):143        attn_stack = self.norm(attn_stack)144        unet_feature = feature_list[-1]145        attn_stack_mean = torch.mean(attn_stack, dim=1, keepdim=True)146        unet_feature = unet_feature * attn_stack_mean147        unet_feature = torch.cat([unet_feature, attn_stack], dim=1) # [1, 324, 64, 64]148        x = self.layer1(unet_feature)149        x = self.layer2(x)150        x = self.layer3(x)151        x = self.layer4(x)152        out = self.conv(x)153        return out / 100154    155    def initialize_weights(self):156        for m in self.modules():157            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):158                nn.init.xavier_normal_(m.weight)159                if m.bias is not None:160                    nn.init.constant_(m.bias, 0)161 162from models.enc_model.unet_parts import *163 164 165class regressor_with_SD_features_seg_vit_c3(nn.Module):166    def __init__(self, n_channels=3, n_classes=2, bilinear=False):167        super(regressor_with_SD_features_seg_vit_c3, self).__init__()168        self.n_channels = n_channels169        self.n_classes = n_classes170        self.bilinear = bilinear171        self.norm = nn.LayerNorm(normalized_shape=(64, 64))172        self.inc_0 = nn.Conv2d(n_channels, 3, kernel_size=3, padding=1)173        self.vit_model = SegModel(gpu=True, nchan=3, pretrained_model="", use_bfloat16=False)174        self.vit = self.vit_model.net175 176    def forward(self, img, attn_stack, feature_list):177        attn_stack = attn_stack[:, [1,3], ...]178        attn_stack = self.norm(attn_stack)179        unet_feature = feature_list[-1]180        unet_feature_mean = torch.mean(unet_feature, dim=1, keepdim=True)181 182        x = torch.cat([unet_feature_mean, attn_stack], dim=1) # [1, 324, 64, 64]183 184        if x.shape[-1] != 512:185            x = F.interpolate(x, size=(512, 512), mode="bilinear")186        x = self.inc_0(x)187 188 189        190        out = self.vit_model.eval(img.squeeze().cpu().numpy(), feat=x.squeeze().cpu().numpy())191        if out.dtype == np.uint16:192            out = out.astype(np.int16)193        out = torch.from_numpy(out).unsqueeze(0).to(x.device)194        return out195    196    def initialize_weights(self):197        for m in self.modules():198            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):199                nn.init.xavier_normal_(m.weight)200                if m.bias is not None:201                    nn.init.constant_(m.bias, 0)202 203class regressor_with_SD_features_tra(nn.Module):204    def __init__(self, n_channels=2, n_classes=2, bilinear=False):205        super(regressor_with_SD_features_tra, self).__init__()206        self.n_channels = n_channels207        self.n_classes = n_classes208        self.bilinear = bilinear209        self.norm = nn.LayerNorm(normalized_shape=(64, 64))210 211        # segmentation212        self.inc_0 = nn.Conv2d(3, 3, kernel_size=3, padding=1)213        self.vit_model = SegModel(gpu=True, nchan=3, pretrained_model="", use_bfloat16=False)214        self.vit = self.vit_model.net215 216        self.inc_1 = nn.Conv2d(n_channels, 1, kernel_size=3, padding=1)217        self.mlp = nn.Linear(64 * 64, 320)218 219    def forward_seg(self, img, attn_stack, feature_list, mask, training=False):220        attn_stack = attn_stack[:, [1,3], ...]221        attn_stack = self.norm(attn_stack)222        unet_feature = feature_list[-1]223        unet_feature_mean = torch.mean(unet_feature, dim=1, keepdim=True)224        x = torch.cat([unet_feature_mean, attn_stack], dim=1) # [1, 324, 64, 64]225 226        if x.shape[-1] != 512:227            x = F.interpolate(x, size=(512, 512), mode="bilinear")228        x = self.inc_0(x)229        feat = x230        231        out = self.vit_model.eval(img.squeeze().cpu().numpy(), feat=x.squeeze().cpu().numpy())232        if out.dtype == np.uint16:233            out = out.astype(np.int16)234        out = torch.from_numpy(out).unsqueeze(0).to(x.device)235        return out, 0., feat236 237    def forward(self, attn_prev, feature_list_prev, attn_after, feature_list_after):238        assert attn_prev.shape == attn_after.shape, "attn_prev and attn_after must have the same shape"239        n_instances = attn_prev.shape[0] 240        attn_prev = self.norm(attn_prev) # [n_instances, 1, 64, 64]241        attn_after = self.norm(attn_after)242        243        x = torch.cat([attn_prev, attn_after], dim=1)  # n_instances, 2, 64, 64244        245        x = self.inc_1(x)246        x = x.view(1, n_instances, -1)  # Flatten the tensor to [n_instances, 64*64*4]247        x = self.mlp(x)  # Apply the MLP to get the output248        249        return x  # Output shape will be [n_instances, 4]250 251        252    253    def initialize_weights(self):254        for m in self.modules():255            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):256                nn.init.xavier_normal_(m.weight)257                if m.bias is not None:258                    nn.init.constant_(m.bias, 0)259