cursor review

8ca6224 2 months ago

14.9 kB

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,
        }

	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,
	}