| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from rscd.models.backbones.seaformer_vmanba import SeaFormer_L |
|
|
| from rscd.models.backbones.cdloma import SS2D |
|
|
| class cdlamba(nn.Module): |
| def __init__(self, channels): |
| super().__init__() |
| self.backbone = SeaFormer_L(pretrained=True) |
| self.channels = channels |
| |
| self.css = nn.ModuleList() |
|
|
| for i in range(20): |
| self.css.append(SS2D(d_model = self.channels[i // 5], channel_first=True, stage_num= i // 5, depth_num= i % 5).cuda()) |
| |
| input_proj_list = [] |
|
|
| for i in range(4): |
| in_channels = self.channels[i] |
| input_proj_list.append(nn.Sequential( |
| nn.Conv2d(in_channels * 2, in_channels * 2, kernel_size=1), |
| nn.GroupNorm(32, in_channels * 2), |
| )) |
| |
| self.input_proj = nn.ModuleList(input_proj_list) |
|
|
| for proj in self.input_proj: |
| nn.init.xavier_uniform_(proj[0].weight, gain=1) |
| nn.init.constant_(proj[0].bias, 0) |
|
|
| def forward(self, xA, xB): |
| inA, inB = xA, xB |
| css_out = [] |
| for i in range(4): |
| fA = self.backbone(inA, i) |
| fB = self.backbone(inB, i) |
| |
| f = torch.concat([fA, fB], 1) |
|
|
| f1 = self.css[i * 5](f) |
| f2 = self.css[i * 5 + 1](f) |
| f3 = self.css[i * 5 + 2](f) |
| f4 = self.css[i * 5 + 3](f) |
| f5 = self.css[i * 5 + 4](f) |
|
|
| f = self.input_proj[i](f1 + f2 + f3 + f4 + f5) |
|
|
| cdaA, cdaB = torch.split(f, self.channels[i], 1) |
| css_out.append(cdaA - cdaB) |
| inA, inB = fA + cdaA, fB + cdaB |
|
|
| for i in range(1, 4): |
| css_out[i] = F.interpolate( |
| css_out[i], |
| scale_factor=(2 ** i, 2 ** i), |
| mode="bilinear", |
| align_corners=False, |
| ) |
|
|
| extract_out = torch.concat(css_out, dim=1) |
| |
| return extract_out |
|
|
