import torch import torch.nn as nn import torch.nn.functional as F # --------------------------------------------------------------------------- # 0. Normalization helpers # --------------------------------------------------------------------------- def get_norm_layer(norm_type: str, num_channels: int): if norm_type == "instance": return nn.InstanceNorm2d(num_channels, affine=True) elif norm_type == "batch": return nn.BatchNorm2d(num_channels) else: raise ValueError(f"Invalid norm_type: {norm_type}") # --------------------------------------------------------------------------- # 1. Basic Conv Block: Conv3x3 -> Norm -> ReLU (x2) # --------------------------------------------------------------------------- class BasicBlock(nn.Module): def __init__(self, in_ch: int, out_ch: int, norm_type: str = "instance", dilation: int = 1): super().__init__() pad = dilation self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=pad, dilation=dilation) self.norm1 = get_norm_layer(norm_type, out_ch) self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=pad, dilation=dilation) self.norm2 = get_norm_layer(norm_type, out_ch) self.relu = nn.ReLU(inplace=True) def forward(self, x): x = self.relu(self.norm1(self.conv1(x))) x = self.relu(self.norm2(self.conv2(x))) return x # --------------------------------------------------------------------------- # 2. DCSA Block (paper Sec. 3.2) # Channel: avg-pool + two 1x1 convs, ADD back to features # Spatial: avg+max concat → standard 3x3 path + dilated 3x3 path (×0.1) → sum # --------------------------------------------------------------------------- class ChannelAttention(nn.Module): """Paper: average pooling + two consecutive 1x1 convs; correlations added back.""" def __init__(self, channels: int, reduction: int = 8): super().__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) hidden = max(channels // reduction, 4) self.mlp = nn.Sequential( nn.Conv2d(channels, hidden, 1), nn.ReLU(inplace=True), nn.Conv2d(hidden, channels, 1), ) def forward(self, x): attn = self.mlp(self.avg_pool(x)) return x + attn class SpatialAttention(nn.Module): """Paper: avg+max concat, then standard and dilated 3x3-norm-relu paths; dilated ×0.1.""" def __init__(self, channels: int, dilation: int = 3, factor: float = 0.1, norm_type: str = "batch"): super().__init__() self.factor = factor # Input to both paths is 2-channel (avg || max); map to `channels` then back to attention map self.std_path = nn.Sequential( nn.Conv2d(2, channels, 3, padding=1, dilation=1), get_norm_layer(norm_type, channels), nn.ReLU(inplace=True), nn.Conv2d(channels, 1, 1), ) pad_d = dilation self.dil_path = nn.Sequential( nn.Conv2d(2, channels, 3, padding=pad_d, dilation=dilation), get_norm_layer(norm_type, channels), nn.ReLU(inplace=True), nn.Conv2d(channels, 1, 1), ) def forward(self, x): avg_out = torch.mean(x, dim=1, keepdim=True) max_out, _ = torch.max(x, dim=1, keepdim=True) concat = torch.cat([avg_out, max_out], dim=1) attn = self.std_path(concat) + self.factor * self.dil_path(concat) return x + attn class DCSABlock(nn.Module): def __init__(self, channels: int, factor: float = 0.1, norm_type: str = "batch"): super().__init__() self.channel_attn = ChannelAttention(channels) self.spatial_attn = SpatialAttention(channels, factor=factor, norm_type=norm_type) def forward(self, x): x = self.channel_attn(x) x = self.spatial_attn(x) return x # --------------------------------------------------------------------------- # 3. Reconstruction Sub-network (Encoder-Decoder + Instance Norm + DCSA) # --------------------------------------------------------------------------- class ReconstructionNet(nn.Module): def __init__(self, in_ch: int = 3, base_ch: int = 32, norm_type: str = "instance", dcsa_factor: float = 0.1): super().__init__() chs = [base_ch, base_ch * 2, base_ch * 4, base_ch * 8] self.level_channels = [chs[0], chs[1], chs[2]] self.enc1 = BasicBlock(in_ch, chs[0], norm_type) self.dcsa1 = DCSABlock(chs[0], factor=dcsa_factor, norm_type=norm_type) self.enc2 = BasicBlock(chs[0], chs[1], norm_type) self.dcsa2 = DCSABlock(chs[1], factor=dcsa_factor, norm_type=norm_type) self.enc3 = BasicBlock(chs[1], chs[2], norm_type) self.dcsa3 = DCSABlock(chs[2], factor=dcsa_factor, norm_type=norm_type) self.enc4 = BasicBlock(chs[2], chs[3], norm_type) self.dcsa4 = DCSABlock(chs[3], factor=dcsa_factor, norm_type=norm_type) self.pool = nn.MaxPool2d(2) self.up3 = nn.ConvTranspose2d(chs[3], chs[2], 2, stride=2) self.dec3 = BasicBlock(chs[2] * 2, chs[2], norm_type) self.dcsa_d3 = DCSABlock(chs[2], factor=dcsa_factor, norm_type=norm_type) self.up2 = nn.ConvTranspose2d(chs[2], chs[1], 2, stride=2) self.dec2 = BasicBlock(chs[1] * 2, chs[1], norm_type) self.dcsa_d2 = DCSABlock(chs[1], factor=dcsa_factor, norm_type=norm_type) self.up1 = nn.ConvTranspose2d(chs[1], chs[0], 2, stride=2) self.dec1 = BasicBlock(chs[0] * 2, chs[0], norm_type) self.dcsa_d1 = DCSABlock(chs[0], factor=dcsa_factor, norm_type=norm_type) self.out_img = nn.Conv2d(chs[0], in_ch, 1) self.up3_jnd = nn.ConvTranspose2d(chs[3], chs[2], 2, stride=2) self.dec3_jnd = BasicBlock(chs[2] * 2, chs[2], norm_type) self.dcsa_d3j = DCSABlock(chs[2], factor=dcsa_factor, norm_type=norm_type) self.up2_jnd = nn.ConvTranspose2d(chs[2], chs[1], 2, stride=2) self.dec2_jnd = BasicBlock(chs[1] * 2, chs[1], norm_type) self.dcsa_d2j = DCSABlock(chs[1], factor=dcsa_factor, norm_type=norm_type) self.up1_jnd = nn.ConvTranspose2d(chs[1], chs[0], 2, stride=2) self.dec1_jnd = BasicBlock(chs[0] * 2, chs[0], norm_type) self.dcsa_d1j = DCSABlock(chs[0], factor=dcsa_factor, norm_type=norm_type) self.out_jnd = nn.Conv2d(chs[0], 1, 1) def forward(self, x): e1 = self.dcsa1(self.enc1(x)) p1 = self.pool(e1) e2 = self.dcsa2(self.enc2(p1)) p2 = self.pool(e2) e3 = self.dcsa3(self.enc3(p2)) p3 = self.pool(e3) e4 = self.dcsa4(self.enc4(p3)) d3 = self.dcsa_d3(self.dec3(torch.cat([self.up3(e4), e3], dim=1))) d2 = self.dcsa_d2(self.dec2(torch.cat([self.up2(d3), e2], dim=1))) d1 = self.dcsa_d1(self.dec1(torch.cat([self.up1(d2), e1], dim=1))) rec_img = torch.sigmoid(self.out_img(d1)) d3j = self.dcsa_d3j(self.dec3_jnd(torch.cat([self.up3_jnd(e4), e3], dim=1))) d2j = self.dcsa_d2j(self.dec2_jnd(torch.cat([self.up2_jnd(d3j), e2], dim=1))) d1j = self.dcsa_d1j(self.dec1_jnd(torch.cat([self.up1_jnd(d2j), e1], dim=1))) rec_jnd = torch.sigmoid(self.out_jnd(d1j)) feats_inp = [e1, e2, e3] feats_rec = [d1, d2, d3] feats_jnd = [d1j, d2j, d3j] return rec_img, rec_jnd, feats_inp, feats_rec, feats_jnd # --------------------------------------------------------------------------- # 4. DiffNeck — (Rec,Inp), (JND,Inp), (Rec,JND,Inp); separate paths ×0.1 # --------------------------------------------------------------------------- class DiffNeckLevel(nn.Module): def __init__(self, in_ch: int, out_ch: int, factor: float = 0.1, norm_type: str = "batch"): super().__init__() self.factor = factor self.block_rec_inp = BasicBlock(in_ch * 2, out_ch, norm_type) # (Rec, Inp) self.block_jnd_inp = BasicBlock(in_ch * 2, out_ch, norm_type) # (JND, Inp) self.block_all = BasicBlock(in_ch * 3, out_ch, norm_type) # (Rec, JND, Inp) def forward(self, rec_feat, inp_feat, jnd_feat): d_rec_inp = self.block_rec_inp(torch.cat([rec_feat, inp_feat], dim=1)) d_jnd_inp = self.block_jnd_inp(torch.cat([jnd_feat, inp_feat], dim=1)) d_all = self.block_all(torch.cat([rec_feat, jnd_feat, inp_feat], dim=1)) # Separate paths balanced by factor 0.1; unified path kept at full weight return self.factor * (d_rec_inp + d_jnd_inp) + d_all class DiffNeckMultiScale(nn.Module): """Alias-friendly name; also exposed as DiffNeck.""" def __init__(self, channels_list, out_ch: int, factor: float = 0.1, norm_type: str = "batch"): super().__init__() self.levels = nn.ModuleList([ DiffNeckLevel(ch, out_ch, factor=factor, norm_type=norm_type) for ch in channels_list ]) self.out_ch = out_ch def forward(self, feats_rec, feats_inp, feats_jnd): target_size = feats_inp[0].shape[-2:] descriptor = None for level_module, rec_f, inp_f, jnd_f in zip(self.levels, feats_rec, feats_inp, feats_jnd): level_out = level_module(rec_f, inp_f, jnd_f) if level_out.shape[-2:] != target_size: level_out = F.interpolate(level_out, size=target_size, mode="bilinear", align_corners=False) descriptor = level_out if descriptor is None else descriptor + level_out return descriptor # Public alias so notebooks can `from models.model import DiffNeck` DiffNeck = DiffNeckMultiScale # --------------------------------------------------------------------------- # 5. Localization Sub-network (6 scales, dilated Basic+DCSA at scales 3–6, BN) # --------------------------------------------------------------------------- class LocalizationNet(nn.Module): def __init__(self, in_ch: int = 6, base_ch: int = 32, num_classes: int = 4, dilation_rates=(1, 1, 2, 4, 8, 16), norm_type: str = "batch", dcsa_factor: float = 0.1): super().__init__() chs = [base_ch, base_ch * 2, base_ch * 4, base_ch * 4, base_ch * 8, base_ch * 8] rates = list(dilation_rates) self.enc1 = BasicBlock(in_ch, chs[0], norm_type, dilation=rates[0]) self.dcsa1 = DCSABlock(chs[0], factor=dcsa_factor, norm_type=norm_type) self.enc2 = BasicBlock(chs[0], chs[1], norm_type, dilation=rates[1]) self.dcsa2 = DCSABlock(chs[1], factor=dcsa_factor, norm_type=norm_type) self.enc3 = BasicBlock(chs[1], chs[2], norm_type, dilation=rates[2]) self.dcsa3 = DCSABlock(chs[2], factor=dcsa_factor, norm_type=norm_type) self.enc4 = BasicBlock(chs[2], chs[3], norm_type, dilation=rates[3]) self.dcsa4 = DCSABlock(chs[3], factor=dcsa_factor, norm_type=norm_type) self.enc5 = BasicBlock(chs[3], chs[4], norm_type, dilation=rates[4]) self.dcsa5 = DCSABlock(chs[4], factor=dcsa_factor, norm_type=norm_type) self.enc6 = BasicBlock(chs[4], chs[5], norm_type, dilation=rates[5]) self.dcsa6 = DCSABlock(chs[5], factor=dcsa_factor, norm_type=norm_type) self.pool = nn.MaxPool2d(2) self.up5 = nn.ConvTranspose2d(chs[5], chs[4], 2, stride=2) self.dec5 = BasicBlock(chs[4] * 2, chs[4], norm_type) self.dcsa_d5 = DCSABlock(chs[4], factor=dcsa_factor, norm_type=norm_type) self.up4 = nn.ConvTranspose2d(chs[4], chs[3], 2, stride=2) self.dec4 = BasicBlock(chs[3] * 2, chs[3], norm_type) self.dcsa_d4 = DCSABlock(chs[3], factor=dcsa_factor, norm_type=norm_type) self.up3 = nn.ConvTranspose2d(chs[3], chs[2], 2, stride=2) self.dec3 = BasicBlock(chs[2] * 2, chs[2], norm_type) self.dcsa_d3 = DCSABlock(chs[2], factor=dcsa_factor, norm_type=norm_type) self.up2 = nn.ConvTranspose2d(chs[2], chs[1], 2, stride=2) self.dec2 = BasicBlock(chs[1] * 2, chs[1], norm_type) self.dcsa_d2 = DCSABlock(chs[1], factor=dcsa_factor, norm_type=norm_type) self.up1 = nn.ConvTranspose2d(chs[1], chs[0], 2, stride=2) self.dec1 = BasicBlock(chs[0] * 2, chs[0], norm_type) self.dcsa_d1 = DCSABlock(chs[0], factor=dcsa_factor, norm_type=norm_type) self.out_seg = nn.Conv2d(chs[0], num_classes, 1) self.enc_channels = chs def forward(self, x): e1 = self.dcsa1(self.enc1(x)) e2 = self.dcsa2(self.enc2(self.pool(e1))) e3 = self.dcsa3(self.enc3(self.pool(e2))) e4 = self.dcsa4(self.enc4(self.pool(e3))) e5 = self.dcsa5(self.enc5(self.pool(e4))) e6 = self.dcsa6(self.enc6(self.pool(e5))) d5 = self.dcsa_d5(self.dec5(torch.cat([self.up5(e6), e5], dim=1))) d4 = self.dcsa_d4(self.dec4(torch.cat([self.up4(d5), e4], dim=1))) d3 = self.dcsa_d3(self.dec3(torch.cat([self.up3(d4), e3], dim=1))) d2 = self.dcsa_d2(self.dec2(torch.cat([self.up2(d3), e2], dim=1))) d1 = self.dcsa_d1(self.dec1(torch.cat([self.up1(d2), e1], dim=1))) return self.out_seg(d1) # --------------------------------------------------------------------------- # 6. Full OmniAL model # --------------------------------------------------------------------------- class OmniAL(nn.Module): def __init__(self, cfg: dict): super().__init__() m = cfg["model"] base_ch = m["base_channels"] factor = m.get("diffneck_factor", 0.1) self.reconstruction = ReconstructionNet( in_ch=m["in_channels"], base_ch=base_ch, norm_type=m["reconstruction_norm"], dcsa_factor=factor, ) self.diffneck = DiffNeckMultiScale( channels_list=self.reconstruction.level_channels, out_ch=base_ch, factor=factor, norm_type=m["localization_norm"], ) # Loc input: RGB anomaly image + DiffNeck descriptor self.localization = LocalizationNet( in_ch=m["in_channels"] + base_ch, base_ch=base_ch, num_classes=4, dilation_rates=tuple(m["dilation_rates"]), norm_type=m["localization_norm"], dcsa_factor=factor, ) def forward(self, anomaly_img: torch.Tensor): rec_img, rec_jnd, feats_inp, feats_rec, feats_jnd = self.reconstruction(anomaly_img) descriptor = self.diffneck(feats_rec, feats_inp, feats_jnd) if descriptor.shape[-2:] != anomaly_img.shape[-2:]: descriptor = F.interpolate( descriptor, size=anomaly_img.shape[-2:], mode="bilinear", align_corners=False ) loc_input = torch.cat([anomaly_img, descriptor], dim=1) seg_logits = self.localization(loc_input) return { "rec_img": rec_img, "rec_jnd": rec_jnd, "seg_logits": seg_logits, }