| """Hybrid change detector: SiamCD (Satlas Swin-v2-B) + organizer-baseline UNet-R18 (6ch early fusion) as an |
| expert branch. The final class heads see both experts (features + probabilities) and learn whom to trust. |
| |
| Inputs/outputs as in model_v2.SiamCD (pre/post in [0, 1]). |
| """ |
| import segmentation_models_pytorch as smp |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from model_v2 import SiamCD, cbr |
|
|
| MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) |
| STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) |
|
|
|
|
| class HybridCD(SiamCD): |
| def __init__(self, satlas_ckpt=None): |
| super().__init__(satlas_ckpt) |
| self.unet = smp.Unet(encoder_name="resnet18", encoder_weights=None, in_channels=6, classes=3) |
| uc = 16 |
| self.head_b = nn.Sequential(cbr(32 + 2 + uc + 2, 48), cbr(48, 32), nn.Conv2d(32, 1, 1)) |
| self.head_t = nn.Sequential(cbr(32 + 2 + uc + 2, 48), cbr(48, 32), nn.Conv2d(32, 1, 1)) |
| self.pres = nn.Sequential(nn.Linear(2 * (64 + 192) + 4, 128), nn.ReLU(inplace=True), nn.Linear(128, 2)) |
| self.register_buffer("mean", MEAN, persistent=False) |
| self.register_buffer("std", STD, persistent=False) |
|
|
| def unet_branch(self, pre, post): |
| x = torch.cat([(pre - self.mean) / self.std, (post - self.mean) / self.std], 1) |
| feats = self.unet.encoder(x) |
| try: |
| d = self.unet.decoder(feats) |
| except TypeError: |
| d = self.unet.decoder(*feats) |
| logits = self.unet.segmentation_head(d) |
| return d, logits |
|
|
| def forward(self, pre, post): |
| B, _, H, W = pre.shape |
| f = self.encode(torch.cat([pre, post], 0)) |
| fp, fq = [t[:B] for t in f], [t[B:] for t in f] |
| s = 0 |
| for lat, t in zip(self.sem_lat, f): |
| s = s + F.interpolate(lat(t), size=f[0].shape[-2:], mode="bilinear", align_corners=False) |
| sem = F.interpolate(self.sem_out(s), size=(H, W), mode="bilinear", align_corners=False) |
| sem_pre, sem_post = sem[:B], sem[B:] |
| z = [fu(torch.cat([a, b, b - a, (b - a).abs()], 1)) for fu, a, b in zip(self.fuse, fp, fq)] |
| x = z[3] |
| for i in (2, 1, 0): |
| x = F.interpolate(x, size=z[i].shape[-2:], mode="bilinear", align_corners=False) |
| x = self.dec[i](torch.cat([x, z[i]], 1)) |
| s1 = self.stem1(torch.cat([pre, post], 1)) |
| s2 = self.stem2(s1) |
| y = self.up2(torch.cat([F.interpolate(x, size=s2.shape[-2:], mode="bilinear", align_corners=False), s2], 1)) |
| y = self.up1(torch.cat([F.interpolate(y, size=(H, W), mode="bilinear", align_corners=False), s1], 1)) |
|
|
| ud, ul = self.unet_branch(pre, post) |
| up = ul.float().softmax(1)[:, 1:3].to(y.dtype) |
| sp, sq = sem_pre.softmax(1), sem_post.softmax(1) |
| cb = self.head_b(torch.cat([y, sq[:, 1:2], sp[:, 1:2], ud, up], 1)) |
| ct = self.head_t(torch.cat([y, sp[:, 2:3], sq[:, 2:3], ud, up], 1)) |
| g = torch.cat([x.mean((2, 3)), x.amax((2, 3)), z[3].mean((2, 3)), z[3].amax((2, 3)), |
| up.mean((2, 3)), up.amax((2, 3))], 1) |
| return torch.cat([cb, ct], 1), self.pres(g), sem_pre, sem_post, ul |
|
|
