first commit

.gitignore

@@ -0,0 +1,6 @@
+.DS_Store
+data*
+checkpoints
+*.pyc
+output_images
+*.pkl

README.md

@@ -10,4 +10,5 @@ pinned: false
 license: cc-by-nc-sa-4.0
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Demo for our ICASSP 2025 paper [Convolutional Prompting for Broad-Domain Retinal Vessel Segmentation](https://arxiv.org/abs/2412.18089).
+Please refer to [https://github.com/ruc-aimc-lab/dcp](https://github.com/ruc-aimc-lab/dcp) for more information.

app.py

@@ -0,0 +1,33 @@
+import gradio as gr
+from inference import Inference
+import os
+from huggingface_hub import snapshot_download
+
+#MODEL_ID = os.getenv("MODEL_ID", "your_username/your_model_name")  # 替换为你的模型ID
+
+model_path = snapshot_download(repo_id='AIMClab-RUC/UNet_DCP_1024')
+
+TEXT_OPTIONS = ["CFP", "UWF", "FFA", "SLO", "OCTA"]
+
+inference_engine = Inference(model_path=model_path)
+
+def main(image, text):
+    out = inference_engine.inference(image, text)
+    return out
+
+interface = gr.Interface(
+    fn=main,
+    inputs=[
+        gr.Image(type="numpy"),
+        gr.Dropdown(
+            choices=TEXT_OPTIONS,
+            label="Modality",
+            value=TEXT_OPTIONS[0]
+        )
+    ],
+    outputs=gr.Image(type="numpy"),
+    title="Broad domain retinal vessel segmentation",
+    description=""
+)
+
+interface.launch()

inference.py

@@ -0,0 +1,88 @@
+# Example code for running inference on a pre-trained model
+import os
+import json
+import numpy as np
+import cv2
+import torch
+from models import build_model
+
+
+# os.environ['CUDA_VISIBLE_DEVICES'] = "0,1"
+
+device = torch.device('cuda' if torch.cuda.is_available() else "cpu")
+
+def sigmoid(arr):
+    return 1. / (1 + np.exp(-arr))
+
+class Inference(object):
+    def __init__(self, model_path):
+        self.model_path = model_path
+        config_path = os.path.join(model_path, 'config.json')
+        with open(config_path) as fin:
+            params = json.load(fin)
+        self.model_params = params['model_params']
+        self.modality_mapping = params['modality_mapping']
+        self.model = self.load_model()
+        
+        
+    def inference(self, image, modality):
+        assert modality in self.modality_mapping, "Modality '{}' not supported".format(modality)
+        
+        image = self.load_image(image)
+        modality_idx = self.modality_mapping[modality]
+        modality_idx = torch.tensor([modality_idx])
+        with torch.no_grad():
+            output = self.model.predict(x=image, device=device, dataset_idx=modality_idx)
+        output = output.data.cpu().numpy()[0][0]
+        output = sigmoid(output) * 255
+        output = output.astype(np.uint8)
+        return output
+        
+    def load_image(self, image):
+        # Load the image and preprocess it
+        if isinstance(image, str):
+            image = cv2.imread(image)[:, :, [2, 1, 0]]
+        #image = image
+        image = cv2.resize(image, (self.model_params['size_w'], self.model_params['size_h']))
+        image = image.astype(np.float32) / 255.0
+        image = np.transpose(image, (2, 0, 1))
+        image = np.expand_dims(image, axis=0)
+        image = torch.tensor(image)
+        return image
+    
+    def load_model(self):
+        print('Loading model from {}'.format(self.model_path))
+        model = build_model(model_name=self.model_params['net'], 
+                            model_params=self.model_params, 
+                            training=False, 
+                            dataset_idx=list(self.modality_mapping.values()),
+                            pretrained=False)
+        #print(model.model.pos_promot3['0'])
+
+        model.set_device(device)
+        # model.requires_grad_false()
+        model.load_model(os.path.join(self.model_path, 'model.pkl'))
+        model.set_mode('eval')
+        
+        return model
+
+
+if __name__ == '__main__':
+    model_path = 'checkpoints/UNet_DCP_1024'
+    image_paths = [
+        'images/FFA.bmp',
+        'images/CFP.jpg',
+        'images/SLO.jpg',
+        'images/UWF.jpg',
+        'images/OCTA.png'
+        ]
+    modalities = ['FFA', 'CFP', 'SLO', 'UWF', 'OCTA']
+     
+    output_root = 'output_images'
+    os.makedirs(output_root, exist_ok=True)
+
+    inference = Inference(model_path)
+    
+    for image_path, modality in zip(image_paths, modalities):
+        output = inference.inference(image_path, modality)    
+        cv2.imwrite(os.path.join(output_root, '{}.png'.format(modality)), output)

models/UNet_p.py

@@ -0,0 +1,694 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Tuple
+
+
+def rand(size, val=0.01):
+    out = torch.zeros(size)
+    
+    nn.init.uniform_(out, -val, val)
+    return out
+
+#  from medsam
+def window_partition(x: torch.Tensor, window_size: int):
+    B, C, H, W = x.size()
+    pad_h = (window_size - H % window_size) % window_size
+    pad_w = (window_size - W % window_size) % window_size
+    if pad_h > 0 or pad_w > 0:
+        x = F.pad(x, (0, pad_w, 0, pad_h))
+    Hp, Wp = H + pad_h, W + pad_w
+    
+    x = x.view(B, C, Hp // window_size, window_size, Wp // window_size, window_size)
+    windows = x.permute(0, 2, 4, 1, 3, 5).contiguous().view(-1, C, window_size, window_size)
+    return windows, (Hp, Wp), (Hp // window_size, Wp // window_size)
+
+def prompt_partition(prompt: torch.Tensor, h_windows: int, w_windows: int):
+    # prompt: B, C, H, W
+    B, C, H, W = prompt.size()
+    prompt = prompt.view(B, 1, 1, C, H, W)
+    prompt = prompt.repeat((1, h_windows, w_windows, 1, 1, 1)).contiguous().view(-1, C, H, W)
+    return prompt
+    
+def window_unpartition(windows: torch.Tensor, window_size: int, pad_hw: Tuple[int, int], hw: Tuple[int, int]):
+    # windows: B * Hp // window_size * Wp // window_size, C, window_size, window_size
+    Hp, Wp = pad_hw
+    H, W = hw
+    B = (windows.shape[0] * window_size * window_size) // (Hp * Wp)
+    #                0          1                    2         3       4            5
+    x = windows.view(B, Hp // window_size, Wp // window_size, -1, window_size, window_size)
+    x = x.permute(0, 3, 1, 4, 2, 5).contiguous().view(B, -1, Hp, Wp)
+
+    if Hp > H or Wp > W:
+        x = x[:, :, :H, :W].contiguous()
+    return x
+
+
+class GELU(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        cdf = 0.5 * (1 + torch.erf(x / 2**0.5))
+        return x * cdf
+
+
+class OneLayerRes(nn.Module):
+    def __init__(self, in_features, out_features, kernel_size, padding) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, padding=padding)
+        self.weight = nn.Parameter(torch.zeros(1), requires_grad=True)
+
+    def forward(self, x):
+        x = x + self.weight * self.conv(x)
+        return x
+
+
+class MLP(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=GELU, drop=0.2):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        return x
+    
+
+class MultiHeadSelfAttention(nn.Module):
+    def __init__(self, dim, num_heads=8, drop_rate=0.2):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.norm = nn.LayerNorm(dim)
+
+        self.scale = head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=False)
+        self.drop = nn.Dropout(drop_rate)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x, heat=False):
+        B, N, C = x.shape
+        out = self.norm(x)
+        qkv = self.qkv(out).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]  
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.drop(attn)
+
+        out = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        out = self.proj(out)
+        out = self.drop(out)
+        out = x + out
+        if heat:
+            return out, attn
+        return out
+
+
+class MultiHeadAttention2D_POS(nn.Module):
+    def __init__(self, dim_q, dim_k, dim_v, embed_dim, num_heads=8, drop_rate=0.2, embed_dim_ratio=4, stride=1, slide=0):
+        super().__init__()
+        self.stride = stride
+        self.num_heads = num_heads
+
+        self.slide = slide
+
+        self.embed_dim_qk = embed_dim // embed_dim_ratio
+        
+        if self.embed_dim_qk % num_heads != 0:  
+            self.embed_dim_qk = (self.embed_dim_qk // num_heads + 1) * num_heads
+
+        self.embed_dim_v = embed_dim
+        if self.embed_dim_v % num_heads != 0:
+            self.embed_dim_v = (self.embed_dim_v // num_heads + 1) * num_heads
+        
+        head_dim = self.embed_dim_qk // num_heads
+
+        self.scale = head_dim ** -0.5
+
+        self.conv_q = nn.Conv2d(in_channels=dim_q, out_channels=self.embed_dim_qk, kernel_size=stride, padding=0, stride=stride)
+        self.conv_k = nn.Conv2d(in_channels=dim_k, out_channels=self.embed_dim_qk, kernel_size=stride, padding=0, stride=stride)
+        self.conv_v = nn.Conv2d(in_channels=dim_v, out_channels=self.embed_dim_v, kernel_size=stride, padding=0, stride=stride)
+                
+        self.drop = nn.Dropout(drop_rate)
+        self.proj_out = nn.Conv2d(in_channels=self.embed_dim_v, out_channels=dim_q, kernel_size=3, padding=1)
+        if self.stride > 1:
+            self.upsample = nn.Upsample(scale_factor=stride, mode='bilinear')
+        else:
+            self.upsample = nn.Identity()
+
+        self.gamma = nn.Parameter(torch.zeros(1), requires_grad=True)
+
+    def forward(self, q, k, v, heat=False):
+        B, _, H_q, W_q = q.size()
+        _, _, H_kv, W_kv = k.size()
+
+        H_q = H_q // self.stride
+        W_q = W_q // self.stride
+        H_kv = H_kv // self.stride
+        W_kv = W_kv // self.stride
+
+        proj_q = self.conv_q(q).reshape(B, self.num_heads, self.embed_dim_qk // self.num_heads, H_q * W_q).permute(0, 1, 3, 2).contiguous()
+        proj_k = self.conv_k(k).reshape(B, self.num_heads, self.embed_dim_qk // self.num_heads, H_kv * W_kv).permute(0, 1, 3, 2).contiguous()
+        proj_v = self.conv_v(v).reshape(B, self.num_heads, self.embed_dim_v // self.num_heads, H_kv * W_kv).permute(0, 1, 3, 2).contiguous()
+
+        attn = (proj_q @ proj_k.transpose(-2, -1)).contiguous() * self.scale  # B, self.num_heads, H_q * W_q, H_kv * W_kv
+        attn = attn.softmax(dim=-1)
+        attn = self.drop(attn)
+
+        out = (attn @ proj_v)  # B, self.num_heads, H_q * W_q, self.embed_dim // self.num_heads
+        out = out.transpose(2, 3).contiguous().reshape(B, self.embed_dim_v, H_q, W_q)
+
+        if self.slide > 0:
+            out = out[:, :,  self.slide // self.stride:]
+            q = q[:, :,  self.slide:]
+        
+        out = self.proj_out(out)
+        out = self.upsample(out)
+        out = self.drop(out)
+        out = q + out * self.gamma
+        return out
+
+
+class MultiHeadAttention2D_CHA(nn.Module):
+    def __init__(self, dim_q, dim_kv, stride, num_heads=8, drop_rate=0.2, slide=0):
+        super().__init__()
+        self.num_heads = num_heads
+        self.stride = stride
+        self.slide = slide
+        self.dim_q_out = dim_q - slide
+
+        self.conv_q = nn.Conv2d(in_channels=dim_q,  out_channels=dim_q * num_heads,  kernel_size=stride, stride=stride, groups=dim_q)
+        self.conv_k = nn.Conv2d(in_channels=dim_kv, out_channels=dim_kv * num_heads, kernel_size=stride, stride=stride, groups=dim_kv)
+        self.conv_v = nn.Conv2d(in_channels=dim_kv, out_channels=dim_kv * num_heads, kernel_size=stride, stride=stride, groups=dim_kv)
+        
+                
+        self.drop = nn.Dropout(drop_rate)
+        self.proj_out = nn.ConvTranspose2d(in_channels=self.dim_q_out * num_heads, out_channels=self.dim_q_out, kernel_size=stride,  stride=stride, groups=self.dim_q_out)
+        self.gamma = nn.Parameter(torch.zeros(1), requires_grad=True)
+
+    def forward(self, q, k, v, heat=False):
+        B, C_q, H_q, W_q = q.size()
+        _, C_kv, H_kv, W_kv = k.size()
+
+        proj_q = self.conv_q(q).reshape(B, self.num_heads, C_q,  -1)     # batch_size * num_heads * dim_q * (H * W)
+        proj_k = self.conv_k(k).reshape(B, self.num_heads, C_kv, -1) 
+        proj_v = self.conv_v(v).reshape(B, self.num_heads, C_kv, -1)     # batch_size * num_heads * dim_kv * (H * W)
+
+        scale = proj_q.size(3) ** -0.5
+        attn = (proj_q @ proj_k.transpose(-2, -1)).contiguous() * scale  # batch_size, num_heads, dim_q, dim_kv
+        attn = attn.softmax(dim=-1)
+        attn = self.drop(attn)
+
+        out = (attn @ proj_v)  #  batch_size, num_heads, dim_q, (H * W)
+        if self.slide > 0: 
+            out = out[:, :, :-self.slide]
+        out = out.reshape(B, self.num_heads * self.dim_q_out, H_q // self.stride, W_q // self.stride)
+        
+        out = self.proj_out(out)
+        out = self.drop(out)
+        out = q + out * self.gamma
+        return out
+
+
+class MultiHeadAttention2D_Dual2_2(nn.Module):
+    def __init__(self, dim_pos, dim_cha, embed_dim, att_fusion, num_heads=8, drop_rate=0.2, embed_dim_ratio=4, stride=1, cha_slide=0, pos_slide=0, use_conv=True):
+        super().__init__()
+        self.pos_att = MultiHeadAttention2D_POS(dim_q=dim_pos, dim_k=dim_pos, dim_v=dim_pos, embed_dim=embed_dim, num_heads=num_heads, drop_rate=drop_rate, embed_dim_ratio=embed_dim_ratio, stride=stride, slide=pos_slide)
+        self.cha_att = MultiHeadAttention2D_CHA(dim_q=dim_cha, dim_kv=dim_cha, num_heads=num_heads, drop_rate=drop_rate, slide=cha_slide, stride=stride)
+        self.att_fusion = att_fusion # concat, add
+
+        if att_fusion == 'concat':
+            channel_in  = 2 * (dim_pos - cha_slide)
+        if att_fusion == 'add':
+            channel_in  = (dim_pos - cha_slide) 
+        channel_out = dim_pos - cha_slide
+        
+        self.use_conv = use_conv
+        if use_conv:
+            self.conv_out = nn.Sequential(nn.Dropout2d(drop_rate, True), nn.Conv2d(channel_in, channel_out, kernel_size=3, padding=1))
+        else:
+            self.conv_out = nn.Identity()
+
+
+    def forward(self, qkv_pos, qkv_cha, heat=False):
+        if qkv_cha is None:
+            qkv_cha = qkv_pos
+        out_pos = self.pos_att(qkv_pos, qkv_pos, qkv_pos, heat)
+        out_cha = self.cha_att(qkv_cha, qkv_cha, qkv_cha, heat)
+
+        C = out_pos.size(1)
+        H = out_cha.size(2)
+        
+        if self.att_fusion == 'concat':
+            out  = torch.cat([out_pos[:, :, -H:], out_cha[:, :C, :]], dim=1)
+        if self.att_fusion == 'add':
+            out = (out_pos[:, :, -H:] + out_cha[:, :C, :]) / 2 
+        
+        out = self.conv_out(out)
+        return out
+
+
+
+class ResMLP(MLP):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=GELU, drop=0.2):
+        super().__init__(in_features=in_features, hidden_features=hidden_features, out_features=out_features, act_layer=act_layer, drop=drop)
+        self.norm = nn.LayerNorm(in_features)
+
+    def forward(self, x):
+        out = self.norm(x)
+        out = self.fc1(out)
+        out = self.act(out)
+        out = self.drop(out)
+        out = self.fc2(out)
+        out = out + x
+        return out
+
+
+class MHSABlock(nn.Module):
+    def __init__(self, dim, num_heads=8, drop_rate=0.2) -> None:
+        super().__init__()
+        self.mhsa = MultiHeadSelfAttention(dim=dim, num_heads=num_heads, drop_rate=drop_rate)
+        self.mlp = ResMLP(in_features=dim, hidden_features=dim*4, out_features=dim)
+
+    def forward(self, x, heat=False):
+        
+        if heat:
+            x, attn = self.mhsa(x, heat=True)
+        else:
+            x = self.mhsa(x)
+        x = self.mlp(x)
+        if heat:
+            return x, attn
+        return x
+
+
+class SelfAttentionBlocks(nn.Module):
+    def __init__(self, dim, block_num, num_heads=8, drop_rate=0.2):
+        super().__init__()
+        self.block_num = block_num
+        assert self.block_num >= 1
+
+        self.blocks = nn.ModuleList([MHSABlock(dim=dim, num_heads=num_heads, drop_rate=drop_rate)
+                                     for i in range(self.block_num)])
+
+    def forward(self, x, heat=False):
+        attns = []
+        for blk in self.blocks:
+            if heat:
+                x, attn = blk(x, heat=True)
+                attns.append(attn)
+            else:
+                x = blk(x)
+        if heat:
+            return x, attns
+        return x
+
+
+class conv_block(nn.Module):
+    def __init__(self,ch_in,ch_out):
+        super(conv_block,self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(ch_in, ch_out, kernel_size=3,stride=1,padding=1,bias=True),
+            nn.BatchNorm2d(ch_out),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(ch_out, ch_out, kernel_size=3,stride=1,padding=1,bias=True),
+            nn.BatchNorm2d(ch_out),
+            nn.ReLU(inplace=True)
+        )
+
+
+    def forward(self,x):
+        x = self.conv(x)
+        return x
+
+
+class up_conv(nn.Module):
+    def __init__(self,ch_in,ch_out):
+        super(up_conv,self).__init__()
+        self.up = nn.Sequential(
+            nn.Upsample(scale_factor=2),
+            nn.Conv2d(ch_in,ch_out,kernel_size=3,stride=1,padding=1,bias=True),
+		    nn.BatchNorm2d(ch_out),
+			nn.ReLU(inplace=True)
+        )
+
+    def forward(self,x):
+        x = self.up(x)
+        return x
+
+
+class Recurrent_block(nn.Module):
+    def __init__(self,ch_out,t=2):
+        super(Recurrent_block,self).__init__()
+        self.t = t
+        self.ch_out = ch_out
+        self.conv = nn.Sequential(
+            nn.Conv2d(ch_out,ch_out,kernel_size=3,stride=1,padding=1,bias=True),
+		    nn.BatchNorm2d(ch_out),
+			nn.ReLU(inplace=True)
+        )
+
+    def forward(self,x):
+        for i in range(self.t):
+
+            if i==0:
+                x1 = self.conv(x)
+            
+            x1 = self.conv(x+x1)
+        return x1
+
+
+class RRCNN_block(nn.Module):
+    def __init__(self,ch_in,ch_out,t=2):
+        super(RRCNN_block,self).__init__()
+        self.RCNN = nn.Sequential(
+            Recurrent_block(ch_out,t=t),
+            Recurrent_block(ch_out,t=t)
+        )
+        self.Conv_1x1 = nn.Conv2d(ch_in,ch_out,kernel_size=1,stride=1,padding=0)
+
+    def forward(self,x):
+        x = self.Conv_1x1(x)
+        x1 = self.RCNN(x)
+        return x+x1
+
+
+class single_conv(nn.Module):
+    def __init__(self,ch_in,ch_out):
+        super(single_conv,self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(ch_in, ch_out, kernel_size=3,stride=1,padding=1,bias=True),
+            nn.BatchNorm2d(ch_out),
+            nn.ReLU(inplace=True)
+        )
+
+    def forward(self,x):
+        x = self.conv(x)
+        return x
+
+
+class Attention_block(nn.Module):
+    def __init__(self,F_g, F_l, F_int):
+        super(Attention_block,self).__init__()
+        self.W_g = nn.Sequential(
+            nn.Conv2d(F_g, F_int, kernel_size=1,stride=1,padding=0,bias=True),
+            nn.BatchNorm2d(F_int)
+            )
+        
+        self.W_x = nn.Sequential(
+            nn.Conv2d(F_l, F_int, kernel_size=1,stride=1,padding=0,bias=True),
+            nn.BatchNorm2d(F_int)
+        )
+
+        self.psi = nn.Sequential(
+            nn.Conv2d(F_int, 1, kernel_size=1,stride=1,padding=0,bias=True),
+            nn.BatchNorm2d(1),
+            nn.Sigmoid()
+        )
+        
+        self.relu = nn.ReLU(inplace=True)
+        
+    def forward(self,g,x):
+        g1 = self.W_g(g)
+        x1 = self.W_x(x)
+        psi = self.relu(g1+x1)
+        psi = self.psi(psi)
+
+        return x*psi
+
+
+class R2AttUNetDecoder(nn.Module):
+    def __init__(self, channels, t=2):
+        super(R2AttUNetDecoder,self).__init__()
+        
+        self.Upsample = nn.Upsample(scale_factor=2, mode='bilinear')
+
+        self.Up5 = up_conv(ch_in=channels[4], ch_out=channels[3])
+        self.Att5 = Attention_block(F_g=channels[3], F_l=channels[3], F_int=channels[3]//2)
+        self.Up_RRCNN5 = RRCNN_block(ch_in=2 * channels[3], ch_out=channels[3], t=t)
+        
+        self.Up4 = up_conv(ch_in=channels[3], ch_out=channels[2])
+        self.Att4 = Attention_block(F_g=channels[2], F_l=channels[2], F_int=channels[2]//2)
+        self.Up_RRCNN4 = RRCNN_block(ch_in=2 * channels[2], ch_out=channels[2], t=t)
+        
+        self.Up3 = up_conv(ch_in=channels[2], ch_out=channels[1])
+        self.Att3 = Attention_block(F_g=channels[1], F_l=channels[1], F_int=channels[1]//2)
+        self.Up_RRCNN3 = RRCNN_block(ch_in=2 * channels[1], ch_out=channels[1], t=t)
+        
+        self.Up2 = up_conv(ch_in=channels[1], ch_out=channels[0])
+        self.Att2 = Attention_block(F_g=channels[0], F_l=channels[0], F_int=channels[0]//2)
+        self.Up_RRCNN2 = RRCNN_block(ch_in=2 * channels[0], ch_out=channels[0], t=t)
+
+    def forward(self, x1, x2, x3, x4, x5):
+        
+        out = self.Up5(x5)
+        x4_att = self.Att5(g=out, x=x4)
+        out = torch.cat((x4_att, out),dim=1)
+        out = self.Up_RRCNN5(out)
+        
+        out = self.Up4(out)
+        x3_att = self.Att4(g=out, x=x3)
+        out = torch.cat((x3_att, out),dim=1)
+        out = self.Up_RRCNN4(out)
+
+        out = self.Up3(out)
+        x2_att = self.Att3(g=out, x=x2)
+        out = torch.cat((x2_att, out),dim=1)
+        out = self.Up_RRCNN3(out)
+
+        out = self.Up2(out)
+        x1_att = self.Att2(g=out, x=x1)
+        out = torch.cat((x1_att, out),dim=1)
+        out = self.Up_RRCNN2(out)
+
+        out = self.Upsample(out)
+
+        return out
+
+
+class ConvBlock(nn.Module):
+    def __init__(self, ch_in, ch_out, kernel_size=3, stride=1, padding=0, bias=True):
+        super(ConvBlock, self).__init__()
+        self.conv1 = nn.Conv2d(ch_in, ch_out, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)
+        self.bn1 = nn.BatchNorm2d(ch_out)
+        self.conv2 = nn.Conv2d(ch_out, ch_out, kernel_size=kernel_size, stride=stride, padding=padding, bias=bias)
+        self.bn2 = nn.BatchNorm2d(ch_out)
+        self.activate = nn.LeakyReLU(negative_slope=0.01)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
+                nn.init.kaiming_normal_(m.weight)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.activate(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.activate(out)
+        return out
+
+
+class UNetDecoder(nn.Module):
+    def __init__(self, channels):
+        super(UNetDecoder,self).__init__()
+        
+        self.Upsample = nn.Upsample(scale_factor=2, mode='bilinear')
+
+        self.Up5 = up_conv(ch_in=channels[4], ch_out=channels[3])
+        self.conv5 = ConvBlock(ch_in=2 * channels[3], ch_out=channels[3], kernel_size=3, stride=1, padding=1)
+        
+        self.Up4 = up_conv(ch_in=channels[3], ch_out=channels[2])
+        self.conv4 = ConvBlock(ch_in=2 * channels[2], ch_out=channels[2], kernel_size=3, stride=1, padding=1)
+        
+        self.Up3 = up_conv(ch_in=channels[2], ch_out=channels[1])
+        self.conv3 = ConvBlock(ch_in=2 * channels[1], ch_out=channels[1], kernel_size=3, stride=1, padding=1)
+        
+        self.Up2 = up_conv(ch_in=channels[1], ch_out=channels[0])
+        self.conv2 = ConvBlock(ch_in=2 * channels[0], ch_out=channels[0], kernel_size=3, stride=1, padding=1)
+
+    def forward(self, x1, x2, x3, x4, x5):
+        
+        out = self.Up5(x5)
+        out = torch.cat((x4, out),dim=1)
+        out = self.conv5(out)
+        
+        out = self.Up4(out)
+        out = torch.cat((x3, out),dim=1)
+        out = self.conv4(out)
+
+        out = self.Up3(out)
+        out = torch.cat((x2, out),dim=1)
+        out = self.conv3(out)
+
+        out = self.Up2(out)
+        out = torch.cat((x1, out),dim=1)
+        out = self.conv2(out)
+
+        out = self.Upsample(out)
+
+        return out
+
+
+class U_Net_P(nn.Module):
+    def __init__(self, encoder, decoder, output_ch, num_classes):
+        super(U_Net_P, self).__init__()
+        
+        self.encoder = encoder
+        self.decoder = decoder
+    
+        self.Last_Conv = nn.Conv2d(output_ch, num_classes, kernel_size=3, stride=1, padding=1)
+
+
+    def forward(self, x):
+        # encoding path
+        x1, x2, x3, x4, x5 = self.encoder(x)
+        x = self.decoder(x1, x2, x3, x4, x5)
+        x = self.Last_Conv(x)
+
+        return x
+
+
+class Prompt_U_Net_P_DCP(nn.Module):
+    def __init__(self, encoder, decoder, output_ch, num_classes, dataset_idx, encoder_channels, prompt_init, pos_promot_channels, cha_promot_channels, embed_ratio, strides, local_window_sizes, att_fusion, use_conv):
+        super(Prompt_U_Net_P_DCP, self).__init__()
+        self.dataset_idx = dataset_idx
+        self.local_window_sizes = local_window_sizes
+
+        self.encoder = encoder
+        self.decoder = decoder
+    
+        self.Last_Conv = nn.Conv2d(output_ch, num_classes, kernel_size=3, stride=1, padding=1)
+        if prompt_init == 'zero':
+            p_init = torch.zeros
+        elif prompt_init == 'one':
+            p_init = torch.ones
+        elif prompt_init == 'rand':
+            p_init = rand        
+  
+        else:
+            raise Exception(prompt_init)
+        
+        self.pos_promot_channels = pos_promot_channels
+        pos_p1 = p_init((1, encoder_channels[0], pos_promot_channels[0], local_window_sizes[0]))
+        pos_p2 = p_init((1, encoder_channels[1], pos_promot_channels[1], local_window_sizes[1]))
+        pos_p3 = p_init((1, encoder_channels[2], pos_promot_channels[2], local_window_sizes[2]))
+        pos_p4 = p_init((1, encoder_channels[3], pos_promot_channels[3], local_window_sizes[3]))
+        pos_p5 = p_init((1, encoder_channels[4], pos_promot_channels[4], local_window_sizes[4]))
+        self.pos_promot1 = nn.ParameterDict({str(k): nn.Parameter(pos_p1.detach().clone(), requires_grad=True) for k in dataset_idx})
+        self.pos_promot2 = nn.ParameterDict({str(k): nn.Parameter(pos_p2.detach().clone(), requires_grad=True) for k in dataset_idx})
+        self.pos_promot3 = nn.ParameterDict({str(k): nn.Parameter(pos_p3.detach().clone(), requires_grad=True) for k in dataset_idx})
+        self.pos_promot4 = nn.ParameterDict({str(k): nn.Parameter(pos_p4.detach().clone(), requires_grad=True) for k in dataset_idx})
+        self.pos_promot5 = nn.ParameterDict({str(k): nn.Parameter(pos_p5.detach().clone(), requires_grad=True) for k in dataset_idx})
+        
+        self.cha_promot_channels = cha_promot_channels
+        cha_p1 = p_init((1, cha_promot_channels[0], local_window_sizes[0], local_window_sizes[0]))
+        cha_p2 = p_init((1, cha_promot_channels[1], local_window_sizes[1], local_window_sizes[1]))
+        cha_p3 = p_init((1, cha_promot_channels[2], local_window_sizes[2], local_window_sizes[2]))
+        cha_p4 = p_init((1, cha_promot_channels[3], local_window_sizes[3], local_window_sizes[3]))
+        cha_p5 = p_init((1, cha_promot_channels[4], local_window_sizes[4], local_window_sizes[4]))
+        self.cha_promot1 = nn.ParameterDict({str(k): nn.Parameter(cha_p1.detach().clone(), requires_grad=True) for k in dataset_idx})
+        self.cha_promot2 = nn.ParameterDict({str(k): nn.Parameter(cha_p2.detach().clone(), requires_grad=True) for k in dataset_idx})
+        self.cha_promot3 = nn.ParameterDict({str(k): nn.Parameter(cha_p3.detach().clone(), requires_grad=True) for k in dataset_idx})
+        self.cha_promot4 = nn.ParameterDict({str(k): nn.Parameter(cha_p4.detach().clone(), requires_grad=True) for k in dataset_idx})
+        self.cha_promot5 = nn.ParameterDict({str(k): nn.Parameter(cha_p5.detach().clone(), requires_grad=True) for k in dataset_idx})
+
+        self.strides = strides
+        
+        self.att1 = MultiHeadAttention2D_Dual2_2(dim_pos=encoder_channels[0], dim_cha=encoder_channels[0] + cha_promot_channels[0], embed_dim=encoder_channels[0], att_fusion=att_fusion, num_heads=8, embed_dim_ratio=embed_ratio, stride=strides[0], pos_slide=0, cha_slide=0, use_conv=use_conv)
+        self.att2 = MultiHeadAttention2D_Dual2_2(dim_pos=encoder_channels[1], dim_cha=encoder_channels[1] + cha_promot_channels[1], embed_dim=encoder_channels[1], att_fusion=att_fusion, num_heads=8, embed_dim_ratio=embed_ratio, stride=strides[1], pos_slide=0, cha_slide=0, use_conv=use_conv)
+        self.att3 = MultiHeadAttention2D_Dual2_2(dim_pos=encoder_channels[2], dim_cha=encoder_channels[2] + cha_promot_channels[2], embed_dim=encoder_channels[2], att_fusion=att_fusion, num_heads=8, embed_dim_ratio=embed_ratio, stride=strides[2], pos_slide=0, cha_slide=0, use_conv=use_conv)
+        self.att4 = MultiHeadAttention2D_Dual2_2(dim_pos=encoder_channels[3], dim_cha=encoder_channels[3] + cha_promot_channels[3], embed_dim=encoder_channels[3], att_fusion=att_fusion, num_heads=8, embed_dim_ratio=embed_ratio, stride=strides[3], pos_slide=0, cha_slide=0, use_conv=use_conv)
+        self.att5 = MultiHeadAttention2D_Dual2_2(dim_pos=encoder_channels[4], dim_cha=encoder_channels[4] + cha_promot_channels[4], embed_dim=encoder_channels[4], att_fusion=att_fusion, num_heads=8, embed_dim_ratio=embed_ratio, stride=strides[4], pos_slide=0, cha_slide=0, use_conv=use_conv)
+        
+    def get_cha_prompts(self, dataset_idx, batch_size):
+        if len(dataset_idx) != batch_size:
+            raise Exception(dataset_idx, self.dataset_idx, batch_size)
+        promots1 = torch.concatenate([self.cha_promot1[str(i)] for i in dataset_idx], dim=0)
+        promots2 = torch.concatenate([self.cha_promot2[str(i)] for i in dataset_idx], dim=0)
+        promots3 = torch.concatenate([self.cha_promot3[str(i)] for i in dataset_idx], dim=0)
+        promots4 = torch.concatenate([self.cha_promot4[str(i)] for i in dataset_idx], dim=0)
+        promots5 = torch.concatenate([self.cha_promot5[str(i)] for i in dataset_idx], dim=0)
+        return promots1, promots2, promots3, promots4, promots5
+    
+    def get_pos_prompts(self, dataset_idx, batch_size):
+        if len(dataset_idx) != batch_size:
+            raise Exception(dataset_idx, self.dataset_idx)
+        promots1 = torch.concatenate([self.pos_promot1[str(i)] for i in dataset_idx], dim=0)
+        promots2 = torch.concatenate([self.pos_promot2[str(i)] for i in dataset_idx], dim=0)
+        promots3 = torch.concatenate([self.pos_promot3[str(i)] for i in dataset_idx], dim=0)
+        promots4 = torch.concatenate([self.pos_promot4[str(i)] for i in dataset_idx], dim=0)
+        promots5 = torch.concatenate([self.pos_promot5[str(i)] for i in dataset_idx], dim=0)
+        return promots1, promots2, promots3, promots4, promots5
+
+    def forward(self, x, dataset_idx, return_features=False):
+        
+        if isinstance(dataset_idx, torch.Tensor):
+            dataset_idx = list(dataset_idx.cpu().numpy())
+        cha_promots1, cha_promots2, cha_promots3, cha_promots4, cha_promots5 = self.get_cha_prompts(dataset_idx=dataset_idx, batch_size=x.size(0))
+        pos_promots1, pos_promots2, pos_promots3, pos_promots4, pos_promots5 = self.get_pos_prompts(dataset_idx=dataset_idx, batch_size=x.size(0))
+        x1, x2, x3, x4, x5 = self.encoder(x)
+
+        if return_features:
+            pre_x1, pre_x2, pre_x3, pre_x4, pre_x5 = x1.detach().clone(), x2.detach().clone(), x3.detach().clone(), x4.detach().clone(), x5.detach().clone()
+
+        h1, w1 = x1.size()[2:]
+        h2, w2 = x2.size()[2:]
+        h3, w3 = x3.size()[2:]
+        h4, w4 = x4.size()[2:]
+        h5, w5 = x5.size()[2:]
+        x1, (Hp1, Wp1), (h_win1, w_win1) = window_partition(x1, self.local_window_sizes[0])
+        x2, (Hp2, Wp2), (h_win2, w_win2) = window_partition(x2, self.local_window_sizes[1])
+        x3, (Hp3, Wp3), (h_win3, w_win3) = window_partition(x3, self.local_window_sizes[2])
+        x4, (Hp4, Wp4), (h_win4, w_win4) = window_partition(x4, self.local_window_sizes[3])
+        x5, (Hp5, Wp5), (h_win5, w_win5) = window_partition(x5, self.local_window_sizes[4])
+
+        cha_promots1 = prompt_partition(cha_promots1, h_win1, w_win1)
+        cha_promots2 = prompt_partition(cha_promots2, h_win2, w_win2)
+        cha_promots3 = prompt_partition(cha_promots3, h_win3, w_win3)
+        cha_promots4 = prompt_partition(cha_promots4, h_win4, w_win4)
+        cha_promots5 = prompt_partition(cha_promots5, h_win5, w_win5)
+
+        pos_promots1 = prompt_partition(pos_promots1, h_win1, w_win1)
+        pos_promots2 = prompt_partition(pos_promots2, h_win2, w_win2)
+        pos_promots3 = prompt_partition(pos_promots3, h_win3, w_win3)
+        pos_promots4 = prompt_partition(pos_promots4, h_win4, w_win4)
+        pos_promots5 = prompt_partition(pos_promots5, h_win5, w_win5)
+
+        cha_x1, cha_x2, cha_x3, cha_x4, cha_x5 = torch.cat([x1, cha_promots1], dim=1), torch.cat([x2, cha_promots2], dim=1), torch.cat([x3, cha_promots3], dim=1), torch.cat([x4, cha_promots4], dim=1), torch.cat([x5, cha_promots5], dim=1)
+        pos_x1, pos_x2, pos_x3, pos_x4, pos_x5 = torch.cat([pos_promots1, x1], dim=2), torch.cat([pos_promots2, x2], dim=2), torch.cat([pos_promots3, x3], dim=2), torch.cat([pos_promots4, x4], dim=2), torch.cat([pos_promots5, x5], dim=2)
+        
+        x1, x2, x3, x4, x5 = self.att1(pos_x1, cha_x1), self.att2(pos_x2, cha_x2), self.att3(pos_x3, cha_x3), self.att4(pos_x4, cha_x4), self.att5(pos_x5, cha_x5)
+        
+        x1 = window_unpartition(x1, self.local_window_sizes[0], (Hp1, Wp1), (h1, w1))
+        x2 = window_unpartition(x2, self.local_window_sizes[1], (Hp2, Wp2), (h2, w2))
+        x3 = window_unpartition(x3, self.local_window_sizes[2], (Hp3, Wp3), (h3, w3))
+        x4 = window_unpartition(x4, self.local_window_sizes[3], (Hp4, Wp4), (h4, w4))
+        x5 = window_unpartition(x5, self.local_window_sizes[4], (Hp5, Wp5), (h5, w5))
+        
+        if return_features:
+            pro_x1, pro_x2, pro_x3, pro_x4, pro_x5 = x1.detach().clone(), x2.detach().clone(), x3.detach().clone(), x4.detach().clone(), x5.detach().clone()
+            
+            return (pre_x1, pre_x2, pre_x3, pre_x4, pre_x5), (pro_x1, pro_x2, pro_x3, pro_x4, pro_x5)
+
+        x = self.decoder(x1, x2, x3, x4, x5)
+        x = self.Last_Conv(x)
+
+        return x
+

models/__init__.py

@@ -0,0 +1 @@
+from .models import build_model

models/backbones/__init__.py

@@ -0,0 +1 @@
+from .backbones import build_backbone

models/backbones/backbones.py

@@ -0,0 +1,25 @@
+import os
+import torch
+from timm.models import efficientnet, convnext
+
+
+def build_backbone(model_name, pretrained):
+    model = getattr(Backbones, model_name)(pretrained=pretrained)
+    return model
+
+
+class Backbones(object):
+    @staticmethod
+    def efficientnet_b3_p(pretrained):
+        # channels: 24, 12, 40, 120, 384
+        # for test, pretrained can be set to False
+        model = efficientnet.efficientnet_b3_pruned(pretrained=pretrained, features_only=True)
+        
+        '''
+        # pre-downloaded weights
+        cp_path = os.path.join('checkpoints', 'effnetb3_pruned-59ecf72d.pth')
+        state_dict = torch.load(cp_path, map_location=torch.device('cpu'))
+        model.load_state_dict(state_dict=state_dict, strict=False)'''
+        return model
+    
+    

models/crit/__init__.py

@@ -0,0 +1,4 @@
+from .focal_loss import BFLoss
+from .mmd import MMDLinear
+from .dice import DiceLoss, DiceBCE
+from .get_bd import generate_BD

models/crit/dice.py

@@ -0,0 +1,36 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+e = 1-10
+
+
+def dice_loss(pred, target, need_sigmoid=True):
+    assert target.size() == pred.size()
+    if need_sigmoid:
+        pred = torch.sigmoid(pred)
+    intersect = 2 * (pred * target).sum() + e
+    union = (pred * pred).sum() + (target * target).sum() + e
+    return 1 - intersect / union
+
+
+class DiceLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, pred, target):
+        return dice_loss(pred=pred, target=target)
+    
+
+class DiceBCE(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, pred, target):
+        return  0.5 * dice_loss(pred=pred, target=target) + \
+              0.5 * F.binary_cross_entropy_with_logits(input=pred, target=target)
+              
+
+
+

models/crit/focal_loss.py

@@ -0,0 +1,23 @@
+from torch import nn
+import torch.nn.functional as F
+import torch
+
+def binary_focal_loss(pred, target, alpha=0.5, gamma=2):
+    assert pred.size() == target.size()
+    pred = torch.sigmoid(pred)
+    e = 1e-5
+    loss = alpha * target * (1 - pred) ** gamma * (pred + e).log() + (1 - alpha) * (1 - target) * pred ** gamma * (1 - pred + e).log()
+    loss = loss / (0.5 ** gamma)
+    return -loss.mean()
+
+
+class BFLoss(nn.Module):
+    def __init__(self, alpha=0.5, gamma=2):
+        super(BFLoss, self).__init__()
+        # alpha: the weight of fg
+        self.gamma = gamma
+        self.alpha = alpha
+
+    def forward(self, pred, target, *args, **kwargs):
+        return binary_focal_loss(pred, target, alpha=self.alpha, gamma=self.gamma)
+        

models/crit/get_bd.py

@@ -0,0 +1,12 @@
+import torch
+import torch.nn.functional as F
+
+
+def generate_BD(mask):
+    #print(mask.size())
+    # img = torch.from_numpy(img).unsqueeze(0).unsqueeze(0).float()
+    # mask = mask.float()
+    mask = torch.abs(mask - F.max_pool2d(mask, 3, 1, 1))
+    mask = mask.detach()
+    
+    return mask

models/crit/mmd.py

@@ -0,0 +1,28 @@
+import torch
+import torch.nn as nn
+from math import gcd
+
+
+def mmd_linear(f_of_X, f_of_Y):
+    delta = f_of_X - f_of_Y
+    loss = torch.mean(torch.mm(delta, torch.transpose(delta, 0, 1)))
+    return loss
+
+class MMDLinear(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, fea_source, fea_target):
+        n_s, d_s = fea_source.size()
+        n_t, d_t = fea_target.size()
+
+        assert d_s == d_t
+
+        if n_s != n_t:
+            n = int(n_s * n_t / gcd(n_s, n_t)) # 最小公倍数
+
+            fea_source = fea_source.repeat((int(n / n_s), 1))
+            fea_target = fea_target.repeat((int(n / n_t), 1))
+        return mmd_linear(fea_source, fea_target)
+
+

models/jtfn.py

@@ -0,0 +1,465 @@
+# ICCV2021, Joint Topology-preserving and Feature-refinement Network for Curvilinear Structure Segmentation
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .UNet_p import MultiHeadAttention2D_Dual2_2, rand, window_partition, window_unpartition, prompt_partition, OneLayerRes
+
+
+
+
+class SpatialAttention(nn.Module):
+    def __init__(self):
+        super(SpatialAttention, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(2, 1, kernel_size=(3, 3), padding=(1, 1)),
+            nn.Conv2d(1, 1, kernel_size=(5, 5), padding=(2, 2)),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        avg_out = torch.mean(x, dim=1, keepdim=True)
+        max_out, _ = torch.max(x, dim=1, keepdim=True)
+        x = torch.cat([avg_out, max_out], dim=1)
+        x = self.conv(x)
+        return x
+
+
+class ChannelAttention(nn.Module):
+    def __init__(self, channel, reduction=2):
+        super(ChannelAttention, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.max_pool = nn.AdaptiveMaxPool2d(1)
+
+        self.fc1 = nn.Conv2d(channel, channel // reduction, 1, bias=False)
+        self.fc2 = nn.Conv2d(channel // reduction, channel, 1, bias=False)
+        self.activate = nn.Sigmoid()
+
+    def forward(self, x):
+        avg_out = self.fc2(self.fc1(self.avg_pool(x)))
+        max_out = self.fc2(self.fc1(self.max_pool(x)))
+        out = avg_out + max_out
+        out = self.activate(out)
+        return out
+
+
+class GAU(nn.Module):
+    def __init__(self, in_channels, use_gau=True, reduce_dim=False, out_channels=None):
+        super(GAU, self).__init__()
+        self.use_gau = use_gau
+        self.reduce_dim = reduce_dim
+
+        if self.reduce_dim:
+            self.down_conv = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1),
+                nn.BatchNorm2d(out_channels),
+                nn.ReLU(inplace=True)
+            )
+            in_channels = out_channels
+
+        if self.use_gau:
+
+            self.sa = SpatialAttention()
+            self.ca = ChannelAttention(in_channels)
+
+            self.reset_gate = nn.Sequential(
+                nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=2, dilation=2),
+                nn.BatchNorm2d(out_channels),
+                nn.ReLU(inplace=True),
+            )
+
+    def forward(self, x, y):
+        if self.reduce_dim:
+            x = self.down_conv(x)
+
+        if self.use_gau:
+            y = F.interpolate(y, x.shape[-2:], mode='bilinear', align_corners=True)
+
+            comx = x * y
+            resx = x * (1 - y) # bs, c, h, w
+
+            x_sa = self.sa(resx) # bs, 1, h, w
+            x_ca = self.ca(resx) # bs, c, 1, 1
+
+            O = self.reset_gate(comx)
+            M = x_sa * x_ca
+
+            RF = M * x + (1 - M) * O
+        else:
+            RF = x
+        return RF
+
+
+class FIM(nn.Module):
+
+    def __init__(self, in_channels, out_channels, f_channels, use_topo=True, up=True, bottom=False):
+        super(FIM, self).__init__()
+        self.use_topo = use_topo
+        self.up = up
+        self.bottom = bottom
+
+        if self.up:
+            self.up_s = nn.Sequential(
+                nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2),
+                nn.BatchNorm2d(out_channels),
+                nn.ReLU(inplace=True)
+            )
+            self.up_t = nn.Sequential(
+                nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2),
+                nn.BatchNorm2d(out_channels),
+                nn.ReLU(inplace=True)
+            )
+        else:
+            self.up_s = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1),
+                nn.BatchNorm2d(out_channels),
+                nn.ReLU(inplace=True)
+            )
+            self.up_t = nn.Sequential(
+                nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1),
+                nn.BatchNorm2d(out_channels),
+                nn.ReLU(inplace=True)
+            )
+
+        self.decoder_s = nn.Sequential(
+            nn.Conv2d(out_channels + f_channels, out_channels, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+
+        '''self.inner_s = nn.Sequential(
+            nn.Conv2d(out_channels, 1, kernel_size=3, padding=1, bias=False),
+            nn.Sigmoid()
+        )'''
+
+        if self.bottom:
+            self.st = nn.Sequential(
+                nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1),
+                nn.BatchNorm2d(in_channels),
+                nn.ReLU(inplace=True)
+            )
+
+        if self.use_topo:
+            self.decoder_t = nn.Sequential(
+                nn.Conv2d(out_channels + out_channels, out_channels, kernel_size=3, stride=1, padding=1),
+                nn.BatchNorm2d(out_channels),
+                nn.ReLU(inplace=True)
+            )
+
+            self.s_to_t = nn.Sequential(
+                nn.Conv2d(out_channels, out_channels, kernel_size=1, stride=1),
+                nn.BatchNorm2d(out_channels),
+                nn.ReLU(inplace=True)
+            )
+
+            self.t_to_s = nn.Sequential(
+                nn.Conv2d(out_channels, out_channels, kernel_size=1, stride=1),
+                nn.BatchNorm2d(out_channels),
+                nn.ReLU(inplace=True)
+            )
+
+            self.res_s = nn.Sequential(
+                nn.Conv2d(out_channels * 2, out_channels, kernel_size=3, stride=1, padding=1),
+                nn.BatchNorm2d(out_channels),
+                nn.ReLU(inplace=True)
+            )
+
+            '''self.inner_t = nn.Sequential(
+                nn.Conv2d(out_channels, 1, kernel_size=3, padding=1, bias=False),
+                nn.Sigmoid()
+            )'''
+
+    def forward(self, x_s, x_t, rf):
+        if self.use_topo:
+            if self.bottom:
+                x_t = self.st(x_t)
+            #bs, c, h, w = x_s.shape
+            x_s = self.up_s(x_s)
+            x_t = self.up_t(x_t)
+
+            # padding
+            diffY = rf.size()[2] - x_s.size()[2]
+            diffX = rf.size()[3] - x_s.size()[3]
+
+            x_s = F.pad(x_s, [diffX // 2, diffX - diffX // 2,
+                          diffY // 2, diffY - diffY // 2])
+            x_t = F.pad(x_t, [diffX // 2, diffX - diffX // 2,
+                              diffY // 2, diffY - diffY // 2])
+
+            rf_s = torch.cat((x_s, rf), dim=1)
+            s = self.decoder_s(rf_s)
+            s_t = self.s_to_t(s)
+
+            t = torch.cat((x_t, s_t), dim=1)
+            x_t = self.decoder_t(t)
+            t_s = self.t_to_s(x_t)
+
+            s_res = self.res_s(torch.cat((s, t_s), dim=1))
+
+            x_s = s + s_res
+            # t_cls = self.inner_t(x_t)
+            # s_cls = self.inner_s(x_s)
+        else:
+            x_s = self.up_s(x_s)
+            #x_b = self.up_b(x_b)
+            # padding
+            diffY = rf.size()[2] - x_s.size()[2]
+            diffX = rf.size()[3] - x_s.size()[3]
+
+            x_s = F.pad(x_s, [diffX // 2, diffX - diffX // 2,
+                              diffY // 2, diffY - diffY // 2])
+
+            rf_s = torch.cat((x_s, rf), dim=1)
+            s = self.decoder_s(rf_s)
+            x_s = s
+            x_t = x_s
+            #t_cls = None
+            #s_cls = self.inner_s(x_s)
+        return x_s, x_t
+
+
+class JTFNDecoder(nn.Module):
+    def __init__(self, channels, use_topo) -> None:
+        super().__init__()
+        self.skip_blocks = []
+        for i in range(5):
+            self.skip_blocks.append(GAU(channels[i], use_gau=True, reduce_dim=False, out_channels=channels[i]))
+        self.fims = []
+        index = 3
+        for i in range(4):
+            if i == index:
+                self.fims.append(FIM(channels[i+1], channels[i], channels[i], use_topo=use_topo, up=True, bottom=True))
+            else:
+                self.fims.append(FIM(channels[i+1], channels[i], channels[i], use_topo=use_topo, up=True, bottom=False))
+        self.skip_blocks = nn.ModuleList(self.skip_blocks)
+        self.fims = nn.ModuleList(self.fims)
+
+    def forward(self, x1, x2, x3, x4, x5, y):
+        x1 = self.skip_blocks[0](x1, y)
+        x2 = self.skip_blocks[1](x2, y)
+        x3 = self.skip_blocks[2](x3, y)
+        x4 = self.skip_blocks[3](x4, y)
+        x5 = self.skip_blocks[4](x5, y)
+        
+        x5_seg, x5_bou = x5, x5
+        
+        x4_seg, x4_bou = self.fims[3](x5_seg, x5_bou, x4)
+        x3_seg, x3_bou = self.fims[2](x4_seg, x4_bou, x3)
+        x2_seg, x2_bou = self.fims[1](x3_seg, x3_bou, x2)
+        x1_seg, x1_bou = self.fims[0](x2_seg, x2_bou, x1)
+        
+        
+        return [x1_seg, x2_seg, x3_seg, x4_seg], [x1_bou, x2_bou, x3_bou, x4_bou]
+        
+
+class JTFN(nn.Module):
+    def __init__(self, encoder, decoder, channels, num_classes, steps) -> None:
+        super().__init__()
+        self.encoder = encoder
+        self.decoder = decoder
+        self.num_classes = num_classes
+        self.steps = steps
+        
+        self.conv_seg1_head = nn.Conv2d(channels[0], num_classes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.conv_seg2_head = nn.Conv2d(channels[1], num_classes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.conv_seg3_head = nn.Conv2d(channels[2], num_classes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.conv_seg4_head = nn.Conv2d(channels[3], num_classes, kernel_size=3, stride=1, padding=1, bias=False)
+
+        self.conv_bou1_head = nn.Conv2d(channels[0], num_classes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.conv_bou2_head = nn.Conv2d(channels[1], num_classes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.conv_bou3_head = nn.Conv2d(channels[2], num_classes, kernel_size=3, stride=1, padding=1, bias=False)
+        self.conv_bou4_head = nn.Conv2d(channels[3], num_classes, kernel_size=3, stride=1, padding=1, bias=False)
+
+        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        y = torch.zeros([B, self.num_classes, H, W], device=x.device)
+        
+        x1, x2, x3, x4, x5 = self.encoder(x)
+        
+        outputs = {}
+        for i in range(self.steps):
+            segs, bous = self.decoder(x1, x2, x3, x4, x5, y)
+            x1_seg, x2_seg, x3_seg, x4_seg = segs
+            x1_bou, x2_bou, x3_bou, x4_bou = bous
+            
+            x1_seg = self.conv_seg1_head(x1_seg)
+            x2_seg = self.conv_seg2_head(x2_seg)
+            x3_seg = self.conv_seg3_head(x3_seg)
+            x4_seg = self.conv_seg4_head(x4_seg)
+            
+            x1_bou = self.conv_bou1_head(x1_bou)
+            x2_bou = self.conv_bou2_head(x2_bou)
+            x3_bou = self.conv_bou3_head(x3_bou)
+            x4_bou = self.conv_bou4_head(x4_bou)
+            
+            y = x1_seg
+            outputs['step_{}_seg'.format(i)] = [x1_seg, x2_seg, x3_seg, x4_seg]
+            outputs['step_{}_bou'.format(i)] = [x1_bou, x2_bou, x3_bou, x4_bou]
+        y = self.upsample(y)
+        outputs['output'] = y
+        return outputs
+
+    def encoder_forward(self, x, dataset_idx):
+        # efficient net
+        x = self.encoder.conv_stem(x)
+        x = self.encoder.bn1(x)
+        features = []
+        if 0 in self.encoder._stage_out_idx:
+            features.append(x)  # add stem out
+        for i in range(len(self.encoder.blocks)):
+            for j, l in enumerate(self.encoder.blocks[i]):
+                if j == len(self.encoder.blocks[i]) - 1 and i + 1 in self.encoder._stage_out_idx:
+                    x = l(x, dataset_idx)
+                else:
+                    x = l(x)
+            if i + 1 in self.encoder._stage_out_idx:
+                features.append(x)
+        return features
+
+           
+
+class JTFN_DCP(JTFN):
+    def __init__(self, encoder, decoder, channels, num_classes, steps, dataset_idx, 
+                 local_window_sizes, encoder_channels, pos_promot_channels, cha_promot_channels, 
+                 embed_ratio, strides, att_fusion, use_conv) -> None:
+        super().__init__(encoder, decoder, channels, num_classes, steps)
+        self.dataset_idx = dataset_idx
+        self.local_window_sizes = local_window_sizes
+
+        self.pos_promot_channels = pos_promot_channels
+        pos_p1 = rand((1, encoder_channels[0], pos_promot_channels[0], local_window_sizes[0]), val=3. / encoder_channels[0] ** 0.5)
+        pos_p2 = rand((1, encoder_channels[1], pos_promot_channels[1], local_window_sizes[1]), val=3. / encoder_channels[1] ** 0.5)
+        pos_p3 = rand((1, encoder_channels[2], pos_promot_channels[2], local_window_sizes[2]), val=3. / encoder_channels[2] ** 0.5)
+        pos_p4 = rand((1, encoder_channels[3], pos_promot_channels[3], local_window_sizes[3]), val=3. / encoder_channels[3] ** 0.5)
+        pos_p5 = rand((1, encoder_channels[4], pos_promot_channels[4], local_window_sizes[4]), val=3. / encoder_channels[4] ** 0.5)
+        self.pos_promot1 = nn.ParameterDict({str(k): nn.Parameter(pos_p1.detach().clone(), requires_grad=True) for k in dataset_idx})
+        self.pos_promot2 = nn.ParameterDict({str(k): nn.Parameter(pos_p2.detach().clone(), requires_grad=True) for k in dataset_idx})
+        self.pos_promot3 = nn.ParameterDict({str(k): nn.Parameter(pos_p3.detach().clone(), requires_grad=True) for k in dataset_idx})
+        self.pos_promot4 = nn.ParameterDict({str(k): nn.Parameter(pos_p4.detach().clone(), requires_grad=True) for k in dataset_idx})
+        self.pos_promot5 = nn.ParameterDict({str(k): nn.Parameter(pos_p5.detach().clone(), requires_grad=True) for k in dataset_idx})
+        
+        self.cha_promot_channels = cha_promot_channels
+        cha_p1 = rand((1, cha_promot_channels[0], local_window_sizes[0], local_window_sizes[0]), val=3. / local_window_sizes[0])
+        cha_p2 = rand((1, cha_promot_channels[1], local_window_sizes[1], local_window_sizes[1]), val=3. / local_window_sizes[1])
+        cha_p3 = rand((1, cha_promot_channels[2], local_window_sizes[2], local_window_sizes[2]), val=3. / local_window_sizes[2])
+        cha_p4 = rand((1, cha_promot_channels[3], local_window_sizes[3], local_window_sizes[3]), val=3. / local_window_sizes[3])
+        cha_p5 = rand((1, cha_promot_channels[4], local_window_sizes[4], local_window_sizes[4]), val=3. / local_window_sizes[4])
+        self.cha_promot1 = nn.ParameterDict({str(k): nn.Parameter(cha_p1.detach().clone(), requires_grad=True) for k in dataset_idx})
+        self.cha_promot2 = nn.ParameterDict({str(k): nn.Parameter(cha_p2.detach().clone(), requires_grad=True) for k in dataset_idx})
+        self.cha_promot3 = nn.ParameterDict({str(k): nn.Parameter(cha_p3.detach().clone(), requires_grad=True) for k in dataset_idx})
+        self.cha_promot4 = nn.ParameterDict({str(k): nn.Parameter(cha_p4.detach().clone(), requires_grad=True) for k in dataset_idx})
+        self.cha_promot5 = nn.ParameterDict({str(k): nn.Parameter(cha_p5.detach().clone(), requires_grad=True) for k in dataset_idx})
+
+        self.strides = strides
+        
+        self.att1 = MultiHeadAttention2D_Dual2_2(dim_pos=encoder_channels[0], dim_cha=encoder_channels[0] + cha_promot_channels[0], embed_dim=encoder_channels[0], att_fusion=att_fusion, num_heads=8, embed_dim_ratio=embed_ratio, stride=strides[0], pos_slide=0, cha_slide=0, use_conv=use_conv)
+        self.att2 = MultiHeadAttention2D_Dual2_2(dim_pos=encoder_channels[1], dim_cha=encoder_channels[1] + cha_promot_channels[1], embed_dim=encoder_channels[1], att_fusion=att_fusion, num_heads=8, embed_dim_ratio=embed_ratio, stride=strides[1], pos_slide=0, cha_slide=0, use_conv=use_conv)
+        self.att3 = MultiHeadAttention2D_Dual2_2(dim_pos=encoder_channels[2], dim_cha=encoder_channels[2] + cha_promot_channels[2], embed_dim=encoder_channels[2], att_fusion=att_fusion, num_heads=8, embed_dim_ratio=embed_ratio, stride=strides[2], pos_slide=0, cha_slide=0, use_conv=use_conv)
+        self.att4 = MultiHeadAttention2D_Dual2_2(dim_pos=encoder_channels[3], dim_cha=encoder_channels[3] + cha_promot_channels[3], embed_dim=encoder_channels[3], att_fusion=att_fusion, num_heads=8, embed_dim_ratio=embed_ratio, stride=strides[3], pos_slide=0, cha_slide=0, use_conv=use_conv)
+        self.att5 = MultiHeadAttention2D_Dual2_2(dim_pos=encoder_channels[4], dim_cha=encoder_channels[4] + cha_promot_channels[4], embed_dim=encoder_channels[4], att_fusion=att_fusion, num_heads=8, embed_dim_ratio=embed_ratio, stride=strides[4], pos_slide=0, cha_slide=0, use_conv=use_conv)
+        
+    def get_cha_prompts(self, dataset_idx, batch_size):
+        if len(dataset_idx) != batch_size:
+            raise Exception(dataset_idx, self.dataset_idx, batch_size)
+        # print(dataset_idx, '***')
+        promots1 = torch.concatenate([self.cha_promot1[str(i)] for i in dataset_idx], dim=0)
+        promots2 = torch.concatenate([self.cha_promot2[str(i)] for i in dataset_idx], dim=0)
+        promots3 = torch.concatenate([self.cha_promot3[str(i)] for i in dataset_idx], dim=0)
+        promots4 = torch.concatenate([self.cha_promot4[str(i)] for i in dataset_idx], dim=0)
+        promots5 = torch.concatenate([self.cha_promot5[str(i)] for i in dataset_idx], dim=0)
+        return promots1, promots2, promots3, promots4, promots5
+    
+    def get_pos_prompts(self, dataset_idx, batch_size):
+        if len(dataset_idx) != batch_size:
+            raise Exception(dataset_idx, self.dataset_idx)
+        # print(dataset_idx, '***')
+        promots1 = torch.concatenate([self.pos_promot1[str(i)] for i in dataset_idx], dim=0)
+        promots2 = torch.concatenate([self.pos_promot2[str(i)] for i in dataset_idx], dim=0)
+        promots3 = torch.concatenate([self.pos_promot3[str(i)] for i in dataset_idx], dim=0)
+        promots4 = torch.concatenate([self.pos_promot4[str(i)] for i in dataset_idx], dim=0)
+        promots5 = torch.concatenate([self.pos_promot5[str(i)] for i in dataset_idx], dim=0)
+        return promots1, promots2, promots3, promots4, promots5      
+
+    def forward(self, x, dataset_idx, return_features=False):
+        if isinstance(dataset_idx, torch.Tensor):
+            dataset_idx = list(dataset_idx.cpu().numpy())
+        #print(dataset_idx)
+        cha_promots1, cha_promots2, cha_promots3, cha_promots4, cha_promots5 = self.get_cha_prompts(dataset_idx=dataset_idx, batch_size=x.size(0))
+        pos_promots1, pos_promots2, pos_promots3, pos_promots4, pos_promots5 = self.get_pos_prompts(dataset_idx=dataset_idx, batch_size=x.size(0))
+        
+        B, C, H, W = x.shape
+        y = torch.zeros([B, self.num_classes, H, W], device=x.device)
+        
+        x1, x2, x3, x4, x5 = self.encoder(x)
+        
+        if return_features:
+            pre_x1, pre_x2, pre_x3, pre_x4, pre_x5 = x1.detach().clone(), x2.detach().clone(), x3.detach().clone(), x4.detach().clone(), x5.detach().clone()
+        h1, w1 = x1.size()[2:]
+        h2, w2 = x2.size()[2:]
+        h3, w3 = x3.size()[2:]
+        h4, w4 = x4.size()[2:]
+        h5, w5 = x5.size()[2:]
+        x1, (Hp1, Wp1), (h_win1, w_win1) = window_partition(x1, self.local_window_sizes[0])
+        x2, (Hp2, Wp2), (h_win2, w_win2) = window_partition(x2, self.local_window_sizes[1])
+        x3, (Hp3, Wp3), (h_win3, w_win3) = window_partition(x3, self.local_window_sizes[2])
+        x4, (Hp4, Wp4), (h_win4, w_win4) = window_partition(x4, self.local_window_sizes[3])
+        x5, (Hp5, Wp5), (h_win5, w_win5) = window_partition(x5, self.local_window_sizes[4])
+
+        cha_promots1 = prompt_partition(cha_promots1, h_win1, w_win1)
+        cha_promots2 = prompt_partition(cha_promots2, h_win2, w_win2)
+        cha_promots3 = prompt_partition(cha_promots3, h_win3, w_win3)
+        cha_promots4 = prompt_partition(cha_promots4, h_win4, w_win4)
+        cha_promots5 = prompt_partition(cha_promots5, h_win5, w_win5)
+
+        pos_promots1 = prompt_partition(pos_promots1, h_win1, w_win1)
+        pos_promots2 = prompt_partition(pos_promots2, h_win2, w_win2)
+        pos_promots3 = prompt_partition(pos_promots3, h_win3, w_win3)
+        pos_promots4 = prompt_partition(pos_promots4, h_win4, w_win4)
+        pos_promots5 = prompt_partition(pos_promots5, h_win5, w_win5)
+
+        #print(x1.size(), x2.size(), x3.size(), x4.size(), x5.size())
+        cha_x1, cha_x2, cha_x3, cha_x4, cha_x5 = torch.cat([x1, cha_promots1], dim=1), torch.cat([x2, cha_promots2], dim=1), torch.cat([x3, cha_promots3], dim=1), torch.cat([x4, cha_promots4], dim=1), torch.cat([x5, cha_promots5], dim=1)
+        pos_x1, pos_x2, pos_x3, pos_x4, pos_x5 = torch.cat([pos_promots1, x1], dim=2), torch.cat([pos_promots2, x2], dim=2), torch.cat([pos_promots3, x3], dim=2), torch.cat([pos_promots4, x4], dim=2), torch.cat([pos_promots5, x5], dim=2)
+        
+        #print(x1.size(), x2.size(), x3.size(), x4.size(), x5.size())
+        x1, x2, x3, x4, x5 = self.att1(pos_x1, cha_x1), self.att2(pos_x2, cha_x2), self.att3(pos_x3, cha_x3), self.att4(pos_x4, cha_x4), self.att5(pos_x5, cha_x5)
+        
+        x1 = window_unpartition(x1, self.local_window_sizes[0], (Hp1, Wp1), (h1, w1))
+        x2 = window_unpartition(x2, self.local_window_sizes[1], (Hp2, Wp2), (h2, w2))
+        x3 = window_unpartition(x3, self.local_window_sizes[2], (Hp3, Wp3), (h3, w3))
+        x4 = window_unpartition(x4, self.local_window_sizes[3], (Hp4, Wp4), (h4, w4))
+        x5 = window_unpartition(x5, self.local_window_sizes[4], (Hp5, Wp5), (h5, w5))
+        
+        if return_features:
+            pro_x1, pro_x2, pro_x3, pro_x4, pro_x5 = x1.detach().clone(), x2.detach().clone(), x3.detach().clone(), x4.detach().clone(), x5.detach().clone()
+            
+            return (pre_x1, pre_x2, pre_x3, pre_x4, pre_x5), (pro_x1, pro_x2, pro_x3, pro_x4, pro_x5)
+        
+        outputs = {}
+        for i in range(self.steps):
+            segs, bous = self.decoder(x1, x2, x3, x4, x5, y)
+            x1_seg, x2_seg, x3_seg, x4_seg = segs
+            x1_bou, x2_bou, x3_bou, x4_bou = bous
+            
+            x1_seg = self.conv_seg1_head(x1_seg)
+            x2_seg = self.conv_seg2_head(x2_seg)
+            x3_seg = self.conv_seg3_head(x3_seg)
+            x4_seg = self.conv_seg4_head(x4_seg)
+            
+            x1_bou = self.conv_bou1_head(x1_bou)
+            x2_bou = self.conv_bou2_head(x2_bou)
+            x3_bou = self.conv_bou3_head(x3_bou)
+            x4_bou = self.conv_bou4_head(x4_bou)
+            
+            y = x1_seg
+            outputs['step_{}_seg'.format(i)] = [x1_seg, x2_seg, x3_seg, x4_seg]
+            outputs['step_{}_bou'.format(i)] = [x1_bou, x2_bou, x3_bou, x4_bou]
+        y = self.upsample(y)
+        outputs['output'] = y
+        return outputs
+
+

models/models.py

@@ -0,0 +1,131 @@
+from .processor import Processor, DCPProcessor, JTFNProcessor, JTFNDCPProcessor
+from .UNet_p import U_Net_P, R2AttUNetDecoder, UNetDecoder, Prompt_U_Net_P_DCP
+from .jtfn import JTFN, JTFNDecoder, JTFN_DCP
+from .backbones import build_backbone
+
+
+def build_model(model_name, model_params, training, dataset_idx, pretrained):
+    model = getattr(Models, model_name)(model_params=model_params, training=training, dataset_idx=dataset_idx, pretrained=pretrained)
+    return model
+
+
+class Models(object):
+    @staticmethod
+    def effi_b3_p_unet(model_params, training, dataset_idx, pretrained=True):
+        n_class = model_params['n_class']
+        channels = (24, 12, 40, 120, 384)
+
+        encoder = build_backbone('efficientnet_b3_p', pretrained=pretrained)
+        decoder = UNetDecoder(channels=channels)
+
+        seg_net = U_Net_P(encoder=encoder, decoder=decoder, output_ch=channels[0], num_classes=n_class)
+        model = Processor(model=seg_net, training_params=model_params, training=training)
+        return model
+    
+    
+    @staticmethod
+    def effi_b3_p_r2attunet(model_params, training, dataset_idx, pretrained=True):
+        n_class = model_params['n_class']
+        channels = (24, 12, 40, 120, 384)
+
+        encoder = build_backbone('efficientnet_b3_p', pretrained=pretrained)
+        decoder = R2AttUNetDecoder(channels=channels)
+
+        seg_net = U_Net_P(encoder=encoder, decoder=decoder, output_ch=channels[0], num_classes=n_class)
+        model = Processor(model=seg_net, training_params=model_params, training=training)
+        return model
+    
+    @staticmethod
+    def effi_b3_p_jtfn(model_params, training, dataset_idx, pretrained=True):
+        n_class = model_params['n_class']
+        channels = (24, 12, 40, 120, 384)
+        steps = model_params['steps']
+        
+        encoder = build_backbone('efficientnet_b3_p')
+        decoder = JTFNDecoder(channels=channels, use_topo=True)
+        
+        seg_net = JTFN(encoder=encoder, decoder=decoder, channels=channels, num_classes=n_class, steps=steps)
+        model = JTFNProcessor(model=seg_net, training_params=model_params, training=training)
+        return model
+
+    
+    @staticmethod
+    def prompt_effi_b3_p_unet_dcp(model_params, training, dataset_idx, pretrained=True):
+        n_class = model_params['n_class']
+        channels = [24, 12, 40, 120, 384]
+
+        cha_promot_channels = model_params['cha_promot_channels']
+        pos_promot_channels = model_params['pos_promot_channels']
+        local_window_sizes = model_params['local_window_sizes']
+        att_fusion = model_params['att_fusion']
+        prompt_init = model_params.get('prompt_init', 'rand')  # rand, zero, one
+        embed_ratio = model_params['embed_ratio']
+        strides = model_params['strides']
+        use_conv = model_params['use_conv']
+
+        encoder = build_backbone('efficientnet_b3_p', pretrained=pretrained)
+        decoder = UNetDecoder(channels=channels)
+
+        seg_net = Prompt_U_Net_P_DCP(encoder=encoder, decoder=decoder, output_ch=channels[0], num_classes=n_class, 
+                                         dataset_idx=dataset_idx, encoder_channels=channels, prompt_init=prompt_init, 
+                                         cha_promot_channels=cha_promot_channels, pos_promot_channels=pos_promot_channels,
+                                         embed_ratio=embed_ratio, strides=strides, local_window_sizes=local_window_sizes,
+                                         att_fusion=att_fusion, use_conv=use_conv)
+
+        model = DCPProcessor(model=seg_net, training_params=model_params, training=training)
+        return model
+    
+    @staticmethod
+    def prompt_effi_b3_p_r2attunet_dcp(model_params, training, dataset_idx, pretrained=True):
+        n_class = model_params['n_class']
+        channels = [24, 12, 40, 120, 384]
+
+        cha_promot_channels = model_params['cha_promot_channels']
+        pos_promot_channels = model_params['pos_promot_channels']
+        local_window_sizes = model_params['local_window_sizes']
+        att_fusion = model_params['att_fusion']
+        prompt_init = model_params.get('prompt_init', 'rand')  # rand, zero, one
+        embed_ratio = model_params['embed_ratio']
+        strides = model_params['strides']
+        use_conv = model_params['use_conv']
+
+        encoder = build_backbone('efficientnet_b3_p', pretrained=pretrained)
+        decoder = R2AttUNetDecoder(channels=channels)
+
+        seg_net = Prompt_U_Net_P_DCP(encoder=encoder, decoder=decoder, output_ch=channels[0], num_classes=n_class, 
+                                         dataset_idx=dataset_idx, encoder_channels=channels, prompt_init=prompt_init, 
+                                         cha_promot_channels=cha_promot_channels, pos_promot_channels=pos_promot_channels,
+                                         embed_ratio=embed_ratio, strides=strides, local_window_sizes=local_window_sizes,
+                                         att_fusion=att_fusion, use_conv=use_conv)
+
+        model = DCPProcessor(model=seg_net, training_params=model_params, training=training)
+        return model
+
+        
+    @staticmethod
+    def prompt_effi_b3_p_jtfn_dcp(model_params, training, dataset_idx, pretrained=True):
+        n_class = model_params['n_class']
+        steps = model_params['steps']
+        channels = [24, 12, 40, 120, 384]
+
+        cha_promot_channels = model_params['cha_promot_channels']
+        pos_promot_channels = model_params['pos_promot_channels']
+        local_window_sizes = model_params['local_window_sizes']
+        att_fusion = model_params['att_fusion']
+        embed_ratio = model_params['embed_ratio']
+        strides = model_params['strides']
+        use_conv = model_params['use_conv']
+
+        encoder = build_backbone('efficientnet_b3_p', pretrained=pretrained)
+        decoder = JTFNDecoder(channels=channels, use_topo=True)
+        seg_net = JTFN_DCP(encoder=encoder, decoder=decoder, channels=channels, num_classes=n_class, steps=steps,
+                                         dataset_idx=dataset_idx, local_window_sizes=local_window_sizes, 
+                                         encoder_channels=channels, 
+                                         cha_promot_channels=cha_promot_channels, pos_promot_channels=pos_promot_channels,
+                                         embed_ratio=embed_ratio, strides=strides, 
+                                         att_fusion=att_fusion, use_conv=use_conv)
+
+        model = JTFNDCPProcessor(model=seg_net, training_params=model_params, training=training)
+        return model
+    
+    

models/optimizer.py

@@ -0,0 +1,72 @@
+import torch.optim as optim
+import torch.nn as nn
+import torch
+import itertools
+
+
+def add_full_model_gradient_clipping(optim, clip_norm_val):
+
+    class FullModelGradientClippingOptimizer(optim):
+        def step(self, closure=None):
+            all_params = itertools.chain(*[x["params"] for x in self.param_groups])
+            torch.nn.utils.clip_grad_norm_(all_params, clip_norm_val)
+            super().step(closure=closure)
+
+    return FullModelGradientClippingOptimizer
+
+
+class Optimizer(object):
+    def __init__(self, models, training_params, sep_lr=None, sep_params=None, gradient_clip=0):
+        
+        params = []
+        for model in models:
+            if isinstance(model, nn.Parameter):
+                params += [model]
+            else:
+                params += list(model.parameters())
+        if sep_lr is not None:
+            print(sep_lr)
+            add_params = []
+            for model in sep_params:
+                if isinstance(model, nn.Parameter):
+                    add_params += [model]
+                else:
+                    add_params += list(model.parameters())
+            params = [{'params': params}, 
+                      {'params': add_params, 'lr': sep_lr}]
+
+
+        self.lr = training_params['lr']
+        self.weight_decay = training_params['weight_decay']
+        method = training_params['optimizer']
+        
+    
+        if method == 'SGD':
+            self.momentum = training_params['momentum']
+            if gradient_clip > 0:
+                self.optim = add_full_model_gradient_clipping(optim.SGD, gradient_clip)(params, lr=self.lr, momentum=self.momentum, weight_decay=self.weight_decay)
+            else:
+                self.optim = optim.SGD(params, lr=self.lr, momentum=self.momentum, weight_decay=self.weight_decay)
+        elif method == 'AdamW':
+            self.optim = optim.AdamW(params, lr=self.lr, weight_decay=self.weight_decay)
+        else:
+            raise Exception('{} is not supported'.format(method))
+
+        schedule_name = training_params['lr_schedule']
+        schedule_params = training_params['schedule_params']
+        if schedule_name == 'CosineAnnealingLR':
+            schedule_params['T_max'] = training_params['inter_val'] * 4
+        self.lr_schedule = getattr(optim.lr_scheduler, schedule_name)(self.optim, **schedule_params)
+        
+    def update_lr(self):
+        self.lr_schedule.step()
+    
+    def z_grad(self):
+        self.optim.zero_grad()
+
+    def g_step(self):
+        self.optim.step()
+
+    def get_lr(self):
+        for param_group in self.optim.param_groups:
+            return param_group['lr']

models/processor.py

@@ -0,0 +1,280 @@
+import torch
+import torch.nn as nn
+from .optimizer import Optimizer
+from .crit import DiceBCE, generate_BD
+from collections import OrderedDict
+
+import torch.nn.functional as F
+
+
+class BasicProcessor(object):
+    def __init__(self) -> None:
+        pass
+
+    def fit(self):
+        raise NotImplementedError
+
+    def predict(self):
+        raise NotImplementedError
+    
+    def set_mode(self, mode):
+        if mode == 'train':
+            self.model.train()
+        elif mode == 'eval':
+            self.model.eval()
+        else:
+            raise Exception('Invalid model mode {}'.format(mode))
+
+    def requires_grad_false(self):
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+    def set_device(self, device):
+        # print(device)
+        if isinstance(device, list):
+            if len(device) > 1:
+                self.model= nn.DataParallel(self.model, device_ids=device)
+                _device = 'cuda'
+            else:
+                _device = 'cuda:{}'.format(device[0])
+            self.model.to(_device)
+        else:
+            self.model.to(device)
+        
+    def save_model(self, path):
+        torch.save(self.model.state_dict(), path)
+
+    def load_model(self, path):
+        state_dict = torch.load(path, map_location='cpu')
+
+        remove_module = True
+        for k, v in state_dict.items():
+            if not k.startswith('module.'):
+                remove_module = False
+                break
+        if remove_module:
+        # create new OrderedDict that does not contain `module.`
+            new_state_dict = OrderedDict()
+            for k, v in state_dict.items():
+                name = k[7:] #remove 'module'
+                new_state_dict[name] = v
+
+            msg = self.model.load_state_dict(new_state_dict)
+        else:
+            msg = self.model.load_state_dict(state_dict)
+        print(msg)
+
+
+class Processor(BasicProcessor):
+    def __init__(self, model, training_params, training) -> None:
+        self.model = model
+
+        if training:
+            self.opt = Optimizer([self.model], training_params)
+            self.crit = DiceBCE()
+
+    def fit(self, xs, ys, device, **kwargs):
+        self.opt.z_grad()
+
+        if len(device) > 1:
+            _device = 'cuda'
+        else:
+            _device = 'cuda:{}'.format(device[0])
+        xs = xs.type(torch.FloatTensor).to(_device)
+        ys = ys.type(torch.FloatTensor).to(_device)
+
+        scores = self.model(xs)
+        loss = self.crit(scores, ys)
+
+        loss.backward()
+        self.opt.g_step()
+        self.opt.update_lr()
+
+        return scores, loss
+
+    def predict(self, x, device, **kwargs):
+        if len(device) > 1:
+            _device = 'cuda'
+        else:
+            _device = 'cuda:{}'.format(device[0])
+        x = x.type(torch.FloatTensor).to(_device)
+        return self.model(x)
+    
+    
+class DCPProcessor(BasicProcessor):
+    def __init__(self, model, training_params, training=True) -> None:
+        self.model = model
+        if training:
+            if 'prompt_lr' in training_params:
+                prompt_lr = training_params['prompt_lr']
+                self.opt = Optimizer([self.model.encoder, self.model.decoder, self.model.Last_Conv, self.model.att1, self.model.att2, self.model.att3, self.model.att4, self.model.att5], training_params, 
+                                     sep_lr=prompt_lr, sep_params=[self.model.cha_promot1, self.model.cha_promot2, self.model.cha_promot3, self.model.cha_promot4, self.model.cha_promot5, self.model.pos_promot1, self.model.pos_promot2, self.model.pos_promot3, self.model.pos_promot4, self.model.pos_promot5])
+            else:
+                self.opt = Optimizer([self.model], training_params)
+            self.crit = DiceBCE()
+
+    def fit(self, xs, ys, device, **kwargs):
+        dataset_idx = kwargs['dataset_idx']
+        self.opt.z_grad()
+        if len(device) > 1:
+            _device = 'cuda'
+        else:
+            _device = 'cuda:{}'.format(device[0])
+        
+        xs = xs.type(torch.FloatTensor).to(_device)
+        ys = ys.type(torch.FloatTensor).to(_device)
+        
+        scores = self.model(xs, dataset_idx)
+        loss = self.crit(scores, ys)
+
+        loss.backward()
+        
+        self.opt.g_step()
+        self.opt.update_lr()
+        
+        return scores, loss
+
+    def predict(self, x, device, **kwargs):
+        dataset_idx = kwargs['dataset_idx']
+        #print(dataset_idx)
+        if isinstance(device, list):
+            if len(device) > 1:
+                _device = 'cuda'
+            else:
+                _device = 'cuda:{}'.format(device[0])
+        else:
+            _device = device
+        
+        x = x.type(torch.FloatTensor).to(_device)
+    
+        return self.model(x, dataset_idx)
+   
+
+class JTFNProcessor(BasicProcessor):
+    def __init__(self, model, training_params, training=True) -> None:
+        # model_params = training_params['model_params']
+        # n_class = model_params['n_class']
+
+        self.model = model
+        self.steps = training_params['steps']
+
+        if training:
+            self.opt = Optimizer([self.model], training_params)
+            # self.crit = DiceLoss()
+            self.crit = DiceBCE()
+
+    def fit(self, xs, ys, device, **kwargs):
+        self.opt.z_grad()
+
+        #num_domains = len(xs)
+        batch_size = len(xs)
+
+        if len(device) > 1:
+            _device = 'cuda'
+        else:
+            _device = 'cuda:{}'.format(device[0])
+        #xs = torch.concatenate(xs, dim=0).type(torch.FloatTensor).to(_device)
+        #ys = torch.concatenate(ys, dim=0).type(torch.FloatTensor).to(_device)
+        xs = xs.type(torch.FloatTensor).to(_device)
+        ys = ys.type(torch.FloatTensor).to(_device)
+        
+        ys_boundary = generate_BD(ys)
+        _, _, h, w = ys.size()
+
+        outputs = self.model(xs)
+        loss = 0
+        for i in range(self.steps):
+            pred_seg = outputs['step_{}_seg'.format(i)]
+            pred_bou = outputs['step_{}_bou'.format(i)]
+            
+            for j in range(len(pred_seg)):
+                p_seg = F.interpolate(pred_seg[j], (h, w), mode='bilinear', align_corners=True)
+                p_bou = F.interpolate(pred_bou[j], (h, w), mode='bilinear', align_corners=True)
+                
+                loss += self.crit(p_seg, ys) + self.crit(p_bou, ys_boundary)
+        loss /= len(pred_seg)
+        loss.backward()
+        self.opt.g_step()
+        self.opt.update_lr()
+
+        
+        scores = outputs['output']
+        # _, C, H, W = scores.size()
+        
+        # scores = scores.view(num_domains, batch_size, C, H, W)
+        # scores = scores.cpu().numpy()
+        return scores, loss
+
+    def predict(self, x, device, **kwargs):
+        if len(device) > 1:
+            _device = 'cuda'
+        else:
+            _device = 'cuda:{}'.format(device[0])
+        x = x.type(torch.FloatTensor).to(_device)
+        outputs = self.model(x)
+        
+        return outputs['output']
+    
+
+class JTFNDCPProcessor(BasicProcessor):
+    def __init__(self, model, training_params, training=True) -> None:
+        # model_params = training_params['model_params']
+        # n_class = model_params['n_class']
+
+        self.model = model
+        self.steps = training_params['steps']
+
+        if training:
+            
+            self.opt = Optimizer([self.model], training_params)
+            # self.crit = DiceLoss()
+            self.crit = DiceBCE()
+
+    def fit(self, xs, ys, device, **kwargs):
+        dataset_idx = kwargs['dataset_idx']
+        self.opt.z_grad()
+
+        if len(device) > 1:
+            _device = 'cuda'
+        else:
+            _device = 'cuda:{}'.format(device[0])
+        xs = xs.type(torch.FloatTensor).to(_device)
+        ys = ys.type(torch.FloatTensor).to(_device)
+        
+        ys_boundary = generate_BD(ys)
+        _, _, h, w = ys.size()
+
+        outputs = self.model(xs, dataset_idx)
+        loss = 0
+        for i in range(self.steps):
+            pred_seg = outputs['step_{}_seg'.format(i)]
+            pred_bou = outputs['step_{}_bou'.format(i)]
+            
+            for j in range(len(pred_seg)):
+                p_seg = F.interpolate(pred_seg[j], (h, w), mode='bilinear', align_corners=True)
+                p_bou = F.interpolate(pred_bou[j], (h, w), mode='bilinear', align_corners=True)
+                
+                loss += self.crit(p_seg, ys) + self.crit(p_bou, ys_boundary)
+        loss /= len(pred_seg)
+        loss.backward()
+        self.opt.g_step()
+        self.opt.update_lr()
+
+        scores = outputs['output']
+
+        return scores, loss
+
+    def predict(self, x, device, **kwargs):
+        dataset_idx = kwargs['dataset_idx']
+        if len(device) > 1:
+            _device = 'cuda'
+        else:
+            _device = 'cuda:{}'.format(device[0])
+        x = x.type(torch.FloatTensor).to(_device)
+        
+        outputs = self.model(x, dataset_idx)
+        
+        return outputs['output']
+    
+    
+

requirements.txt

@@ -0,0 +1,4 @@
+numpy==1.26.4
+opencv-python==4.9.0.80
+torch==2.3.0
+timm==1.0.3