| |
| |
| |
| |
|
|
| import os, re, json, warnings |
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from einops import rearrange |
| from PIL import Image |
| import cv2 |
| from scipy import ndimage |
| from skimage.measure import label, regionprops |
| from scipy.ndimage import binary_fill_holes |
|
|
|
|
| |
| |
| |
|
|
| class ResidualBlock(nn.Module): |
| def __init__(self, in_c, out_c): |
| super().__init__() |
| self.conv = nn.Sequential( |
| nn.Conv3d(in_c, out_c, kernel_size=3, padding=1), |
| nn.InstanceNorm3d(out_c), |
| nn.LeakyReLU(0.1, inplace=True), |
| nn.Conv3d(out_c, out_c, kernel_size=3, padding=1), |
| nn.InstanceNorm3d(out_c), |
| ) |
| self.shortcut = (nn.Conv3d(in_c, out_c, kernel_size=1) |
| if in_c != out_c else nn.Identity()) |
| def forward(self, x): |
| return F.leaky_relu(self.conv(x) + self.shortcut(x), 0.1) |
|
|
|
|
| class AttentionGate(nn.Module): |
| def __init__(self, F_g, F_l, F_int): |
| super().__init__() |
| self.W_g = nn.Sequential(nn.Conv3d(F_g, F_int, kernel_size=1), nn.InstanceNorm3d(F_int)) |
| self.W_x = nn.Sequential(nn.Conv3d(F_l, F_int, kernel_size=1), nn.InstanceNorm3d(F_int)) |
| self.psi = nn.Sequential(nn.Conv3d(F_int, 1, kernel_size=1), nn.Sigmoid()) |
| def forward(self, g, x): |
| g1 = self.W_g(g) |
| x1 = self.W_x(x) |
| if g1.shape[2:] != x1.shape[2:]: |
| g1 = F.interpolate(g1, size=x1.shape[2:], mode="trilinear", align_corners=False) |
| return x * self.psi(F.leaky_relu(g1 + x1, 0.1)) |
|
|
|
|
| class ModalityFusionBlock(nn.Module): |
| def __init__(self, in_channels=4, fusion_dim=24, num_heads=4): |
| super().__init__() |
| self.M = in_channels |
| self.modal_proj = nn.ModuleList([nn.Linear(1, fusion_dim) for _ in range(in_channels)]) |
| self.norm_q = nn.LayerNorm(fusion_dim) |
| self.norm_kv = nn.LayerNorm(fusion_dim) |
| self.cross_attn = nn.MultiheadAttention(fusion_dim, num_heads, batch_first=True) |
| self.out_proj = nn.ModuleList([nn.Linear(fusion_dim, 1) for _ in range(in_channels)]) |
| def forward(self, x): |
| B = x.shape[0] |
| gap = x.mean(dim=[2, 3, 4]) |
| tokens = torch.cat([self.modal_proj[m](gap[:, m:m+1]).unsqueeze(1) for m in range(self.M)], dim=1) |
| Q = self.norm_q(tokens[:, 2:3, :]) |
| KV = self.norm_kv(tokens) |
| attn_out, _ = self.cross_attn(Q, KV, KV) |
| tokens = tokens.clone() |
| tokens[:, 2:3, :] += attn_out |
| bias = torch.cat([self.out_proj[m](tokens[:, m, :]).view(B, 1, 1, 1, 1) for m in range(self.M)], dim=1) |
| return x + bias |
|
|
|
|
| class RegionPrototypeModule(nn.Module): |
| def __init__(self, dim, num_regions=3, proto_dim=48): |
| super().__init__() |
| self.feat_proj = nn.Linear(dim, proto_dim) |
| self.prototypes = nn.Parameter(torch.randn(num_regions, proto_dim)) |
| self.score_proj = nn.Conv3d(num_regions, num_regions, kernel_size=1) |
| def forward(self, x): |
| B, C, H, W, D = x.shape |
| x_flat = rearrange(x, "b c h w d -> b (h w d) c") |
| x_proj = self.feat_proj(x_flat) |
| x_norm = F.normalize(x_proj, dim=-1) |
| p_norm = F.normalize(self.prototypes, dim=-1) |
| scores = torch.einsum("bnd,rd->bnr", x_norm, p_norm) |
| scores_3d = rearrange(scores, "b (h w d) r -> b r h w d", h=H, w=W, d=D) |
| return self.score_proj(scores_3d), scores[:, :, 2] |
|
|
|
|
| class ETBiasedSelfAttention(nn.Module): |
| def __init__(self, dim, num_heads=4): |
| super().__init__() |
| self.H, self.Dh = num_heads, dim // num_heads |
| self.scale = self.Dh ** -0.5 |
| self.qkv = nn.Linear(dim, dim * 3, bias=False) |
| self.out_lin = nn.Linear(dim, dim) |
| self.et_bias_scale = nn.Parameter(torch.tensor(0.05)) |
| def forward(self, x, et_bias=None): |
| B, N, D = x.shape |
| qkv = self.qkv(x).reshape(B, N, 3, self.H, self.Dh).permute(2, 0, 3, 1, 4) |
| Q, K, V = qkv[0], qkv[1], qkv[2] |
| attn = (Q @ K.transpose(-2, -1)) * self.scale |
| if et_bias is not None: |
| et_b = (self.et_bias_scale * torch.softmax(et_bias, dim=-1)).unsqueeze(1).unsqueeze(2) |
| attn += et_b |
| attn = torch.softmax(attn, dim=-1) |
| out = (attn @ V).transpose(1, 2).reshape(B, N, D) |
| return self.out_lin(out) |
|
|
|
|
| class ETGuidedTransformerBlock(nn.Module): |
| def __init__(self, dim): |
| super().__init__() |
| self.proto_module = RegionPrototypeModule(dim) |
| self.gate_feat = nn.Conv3d(dim, dim, kernel_size=1) |
| self.gate_region = nn.Conv3d(3, dim, kernel_size=1) |
| self.norm1 = nn.LayerNorm(dim) |
| self.attn = ETBiasedSelfAttention(dim) |
| self.norm2 = nn.LayerNorm(dim) |
| self.ffn = nn.Sequential(nn.Linear(dim, dim * 2), nn.GELU(), nn.Linear(dim * 2, dim)) |
| def forward(self, x): |
| B, C, H, W, D = x.shape |
| proto_scores, et_scores = self.proto_module(x) |
| gate = torch.sigmoid(self.gate_feat(x) + self.gate_region(proto_scores)) |
| x_gated = x * gate |
| x_flat = rearrange(x_gated, "b c h w d -> b (h w d) c") |
| x_flat = x_flat + self.attn(self.norm1(x_flat), et_bias=et_scores) |
| x_flat = x_flat + self.ffn(self.norm2(x_flat)) |
| return rearrange(x_flat, "b (h w d) c -> b c h w d", h=H, w=W, d=D), proto_scores |
|
|
|
|
| class RCMTUNetV4(nn.Module): |
| def __init__(self, in_channels=4, out_channels=4, features=(24, 48, 96, 192)): |
| super().__init__() |
| f = features |
| self.fusion = ModalityFusionBlock(in_channels=in_channels, fusion_dim=f[0], num_heads=4) |
| self.enc1 = ResidualBlock(in_channels, f[0]) |
| self.enc2 = ResidualBlock(f[0], f[1]) |
| self.enc3 = ResidualBlock(f[1], f[2]) |
| self.enc4 = ResidualBlock(f[2], f[3]) |
| self.pool = nn.MaxPool3d(2) |
| self.bottleneck = ETGuidedTransformerBlock(dim=f[3]) |
| self.up3 = nn.ConvTranspose3d(f[3], f[2], 2, stride=2) |
| self.ag3 = AttentionGate(f[2], f[2], f[2]//2) |
| self.dec3 = ResidualBlock(f[3], f[2]) |
| self.up2 = nn.ConvTranspose3d(f[2], f[1], 2, stride=2) |
| self.ag2 = AttentionGate(f[1], f[1], f[1]//2) |
| self.dec2 = ResidualBlock(f[2], f[1]) |
| self.up1 = nn.ConvTranspose3d(f[1], f[0], 2, stride=2) |
| self.ag1 = AttentionGate(f[0], f[0], f[0]//2) |
| self.dec1 = ResidualBlock(f[1], f[0]) |
| self.final_conv = nn.Conv3d(f[0], out_channels, 1) |
| self.deep_sup2 = nn.Conv3d(f[1], out_channels, 1) |
| self.deep_sup3 = nn.Conv3d(f[2], out_channels, 1) |
|
|
| def forward(self, x): |
| xf = self.fusion(x) |
| e1 = self.enc1(xf) |
| e2 = self.enc2(self.pool(e1)) |
| e3 = self.enc3(self.pool(e2)) |
| e4 = self.enc4(self.pool(e3)) |
| b, _ = self.bottleneck(e4) |
| def up_and_cat(up_layer, ag, dec, b_feat, enc_feat): |
| u = up_layer(b_feat) |
| if u.shape[2:] != enc_feat.shape[2:]: |
| u = F.interpolate(u, size=enc_feat.shape[2:], mode="trilinear", align_corners=False) |
| return dec(torch.cat([u, ag(g=u, x=enc_feat)], dim=1)) |
| d3 = up_and_cat(self.up3, self.ag3, self.dec3, b, e3) |
| d2 = up_and_cat(self.up2, self.ag2, self.dec2, d3, e2) |
| d1 = up_and_cat(self.up1, self.ag1, self.dec1, d2, e1) |
| return self.final_conv(d1) |
|
|
