models/model.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import os
import clip
import sys
import numpy as np
from models.seg_post_model.models import SegModel

from torchvision.ops import roi_align


class Counting_with_SD_features_loca(nn.Module):
    def __init__(self, scale_factor):
        super(Counting_with_SD_features_loca, self).__init__()
        self.adapter = adapter_roi_loca()
        self.regressor = regressor_with_SD_features()


class Counting_with_SD_features_dino_vit_c3(nn.Module):
    def __init__(self, scale_factor, vit=None):
        super(Counting_with_SD_features_dino_vit_c3, self).__init__()
        self.adapter = adapter_roi_loca()
        self.regressor = regressor_with_SD_features_seg_vit_c3()

class Counting_with_SD_features_track(nn.Module):
    def __init__(self, scale_factor, vit=None):
        super(Counting_with_SD_features_track, self).__init__()
        self.adapter = adapter_roi_loca()
        self.regressor = regressor_with_SD_features_tra()


class adapter_roi_loca(nn.Module):
    def __init__(self, pool_size=[3, 3]):
        super(adapter_roi_loca, self).__init__()
        self.pool_size = pool_size
        self.conv1 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(256 * 3 * 3, 768)
        self.initialize_weights()
    def forward(self, x, boxes):
            num_of_boxes = boxes.shape[1]
            rois = []
            bs, _, h, w = x.shape
            if h != 512 or w != 512:
                x = F.interpolate(x, size=(512, 512), mode='bilinear', align_corners=False)
            if bs == 1:
                boxes = torch.cat([
                    torch.arange(
                        bs, requires_grad=False
                    ).to(boxes.device).repeat_interleave(num_of_boxes).reshape(-1, 1),
                    boxes.flatten(0, 1),
                ], dim=1)
                rois = roi_align(
                    x,
                    boxes=boxes, output_size=3,
                    spatial_scale=1.0 / 8, aligned=True
                )
                rois = torch.mean(rois, dim=0, keepdim=True)
            else:
                boxes = torch.cat([
                    boxes.flatten(0, 1),
                ], dim=1).split(num_of_boxes, dim=0)
                rois = roi_align(
                    x,
                    boxes=boxes, output_size=3,
                    spatial_scale=1.0 / 8, aligned=True
                )
                rois = rois.split(num_of_boxes, dim=0)
                rois = torch.stack(rois, dim=0)
                rois = torch.mean(rois, dim=1, keepdim=False)
            x = self.conv1(rois)
            x = x.view(x.size(0), -1)
            x = self.fc(x)
            return x
    
    def forward_boxes(self, x, boxes):
            num_of_boxes = boxes.shape[1]
            rois = []
            bs, _, h, w = x.shape
            if h != 512 or w != 512:
                x = F.interpolate(x, size=(512, 512), mode='bilinear', align_corners=False)
            if bs == 1:
                boxes = torch.cat([
                    torch.arange(
                        bs, requires_grad=False
                    ).to(boxes.device).repeat_interleave(num_of_boxes).reshape(-1, 1),
                    boxes.flatten(0, 1),
                ], dim=1)
                rois = roi_align(
                    x,
                    boxes=boxes, output_size=3,
                    spatial_scale=1.0 / 8, aligned=True
                )
                # rois = torch.mean(rois, dim=0, keepdim=True)
            else:
                raise NotImplementedError
            x = self.conv1(rois)
            x = x.view(x.size(0), -1)
            x = self.fc(x)
            return x

    def initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)



class regressor_with_SD_features(nn.Module):
    def __init__(self):
        super(regressor_with_SD_features, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Conv2d(324, 256, kernel_size=1, stride=1),
            nn.LeakyReLU(),
            nn.LayerNorm((64, 64))
        )
        self.layer2 = nn.Sequential(
            nn.Conv2d(256, 128, kernel_size=3, padding=1),
            nn.LeakyReLU(),
            nn.ConvTranspose2d(in_channels=128, out_channels=128, kernel_size=4, stride=2, padding=1),
        )
        self.layer3 = nn.Sequential(
            nn.Conv2d(128, 64, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(in_channels=64, out_channels=64, kernel_size=4, stride=2, padding=1),
        )
        self.layer4 = nn.Sequential(
            nn.Conv2d(64, 32, kernel_size=3, padding=1),
            nn.LeakyReLU(),
            nn.ConvTranspose2d(in_channels=32, out_channels=32, kernel_size=4, stride=2, padding=1),
        )
        self.conv = nn.Sequential(
            nn.Conv2d(32, 1, kernel_size=1),
            nn.ReLU()
        )
        self.norm = nn.LayerNorm(normalized_shape=(64, 64))
        self.initialize_weights()

    def forward(self, attn_stack, feature_list):
        attn_stack = self.norm(attn_stack)
        unet_feature = feature_list[-1]
        attn_stack_mean = torch.mean(attn_stack, dim=1, keepdim=True)
        unet_feature = unet_feature * attn_stack_mean
        unet_feature = torch.cat([unet_feature, attn_stack], dim=1) # [1, 324, 64, 64]
        x = self.layer1(unet_feature)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        out = self.conv(x)
        return out / 100
    
    def initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

from models.enc_model.unet_parts import *


class regressor_with_SD_features_seg_vit_c3(nn.Module):
    def __init__(self, n_channels=3, n_classes=2, bilinear=False):
        super(regressor_with_SD_features_seg_vit_c3, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear
        self.norm = nn.LayerNorm(normalized_shape=(64, 64))
        self.inc_0 = nn.Conv2d(n_channels, 3, kernel_size=3, padding=1)
        self.vit_model = SegModel(gpu=True, nchan=3, pretrained_model="", use_bfloat16=False)
        self.vit = self.vit_model.net

    def forward(self, img, attn_stack, feature_list):
        attn_stack = attn_stack[:, [1,3], ...]
        attn_stack = self.norm(attn_stack)
        unet_feature = feature_list[-1]
        unet_feature_mean = torch.mean(unet_feature, dim=1, keepdim=True)

        x = torch.cat([unet_feature_mean, attn_stack], dim=1) # [1, 324, 64, 64]

        if x.shape[-1] != 512:
            x = F.interpolate(x, size=(512, 512), mode="bilinear")
        x = self.inc_0(x)


        
        out = self.vit_model.eval(img.squeeze().cpu().numpy(), feat=x.squeeze().cpu().numpy())
        if out.dtype == np.uint16:
            out = out.astype(np.int16)
        out = torch.from_numpy(out).unsqueeze(0).to(x.device)
        return out
    
    def initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

class regressor_with_SD_features_tra(nn.Module):
    def __init__(self, n_channels=2, n_classes=2, bilinear=False):
        super(regressor_with_SD_features_tra, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear
        self.norm = nn.LayerNorm(normalized_shape=(64, 64))

        # segmentation
        self.inc_0 = nn.Conv2d(3, 3, kernel_size=3, padding=1)
        self.vit_model = SegModel(gpu=True, nchan=3, pretrained_model="", use_bfloat16=False)
        self.vit = self.vit_model.net

        self.inc_1 = nn.Conv2d(n_channels, 1, kernel_size=3, padding=1)
        self.mlp = nn.Linear(64 * 64, 320)

    def forward_seg(self, img, attn_stack, feature_list, mask, training=False):
        attn_stack = attn_stack[:, [1,3], ...]
        attn_stack = self.norm(attn_stack)
        unet_feature = feature_list[-1]
        unet_feature_mean = torch.mean(unet_feature, dim=1, keepdim=True)
        x = torch.cat([unet_feature_mean, attn_stack], dim=1) # [1, 324, 64, 64]

        if x.shape[-1] != 512:
            x = F.interpolate(x, size=(512, 512), mode="bilinear")
        x = self.inc_0(x)
        feat = x
        
        out = self.vit_model.eval(img.squeeze().cpu().numpy(), feat=x.squeeze().cpu().numpy())
        if out.dtype == np.uint16:
            out = out.astype(np.int16)
        out = torch.from_numpy(out).unsqueeze(0).to(x.device)
        return out, 0., feat

    def forward(self, attn_prev, feature_list_prev, attn_after, feature_list_after):
        assert attn_prev.shape == attn_after.shape, "attn_prev and attn_after must have the same shape"
        n_instances = attn_prev.shape[0] 
        attn_prev = self.norm(attn_prev) # [n_instances, 1, 64, 64]
        attn_after = self.norm(attn_after)
        
        x = torch.cat([attn_prev, attn_after], dim=1)  # n_instances, 2, 64, 64
        
        x = self.inc_1(x)
        x = x.view(1, n_instances, -1)  # Flatten the tensor to [n_instances, 64*64*4]
        x = self.mlp(x)  # Apply the MLP to get the output
        
        return x  # Output shape will be [n_instances, 4]

        
    
    def initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear):
                nn.init.xavier_normal_(m.weight)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)