cursor review

config/config.yaml

@@ -42,19 +42,21 @@ synthesis:
   area_ratio_min: 0.5
   area_ratio_max: 0.8
   severity_levels: ["easy", "medium", "hard"]
-  # alpha_i, beta_i per severity level (i = 1..3), tune as needed
-  alpha: [0.3, 0.6, 0.9]
+  # α_i = Gaussian smooth kernel sizes (easy→smaller); β_i bias per level
+  alpha_kernels: [3, 7, 15]
   beta: [0.10, 0.05, 0.0]
   ssim_thresholds: [0.9, 0.7]   # boundaries splitting easy/medium/hard by SSIM
 
 model:
   in_channels: 3
   base_channels: 32
-  diffneck_factor: 0.1
+  diffneck_factor: 0.1           # separate-path balance (paper / Fig.8)
   localization_scales: 6
   dilation_rates: [1, 1, 2, 4, 8, 16]   # scale1..2 plain conv, scale3..6 dilated
   reconstruction_norm: "instance"        # instance norm in reconstruction sub-net
   localization_norm: "batch"             # batch norm in localization sub-net
+  anomaly_smooth_kernel: 7               # Softmax → Gaussian smooth → sum
+  anomaly_smooth_sigma: 1.5
 
 train:
   batch_size_mvtec: 20

data/data_loader.py

@@ -138,37 +138,73 @@ def get_panel_bbox(h: int, w: int, panel: str,
 
 
 def sample_texture_source(normal_img: np.ndarray, dtd_paths: List[str], size: Tuple[int, int]) -> np.ndarray:
+    """Anomaly texture T: DTD or augmented normal (resize, crop, flip, color-jitter)."""
+    h, w = size
     if len(dtd_paths) > 0 and random.random() < 0.5:
         path = random.choice(dtd_paths)
         tex = load_and_resize(path, size)
     else:
         tex = normal_img.copy()
-        if random.random() < 0.5:
-            tex = np.fliplr(tex).copy()
-        if random.random() < 0.5:
-            tex = np.flipud(tex).copy()
-        k = random.choice([0, 1, 2, 3])
-        tex = np.rot90(tex, k).copy()
-        jitter = random.uniform(0.8, 1.2)
-        tex = np.clip(tex * jitter, 0, 1)
-    return tex
+
+    # Random resize then crop back to (h, w)
+    scale = random.uniform(0.8, 1.4)
+    nh, nw = max(int(h * scale), h), max(int(w * scale), w)
+    tex = cv2.resize(tex, (nw, nh), interpolation=cv2.INTER_LINEAR)
+    y0 = random.randint(0, max(nh - h, 0))
+    x0 = random.randint(0, max(nw - w, 0))
+    tex = tex[y0:y0 + h, x0:x0 + w]
+    if tex.shape[0] != h or tex.shape[1] != w:
+        tex = cv2.resize(tex, (w, h), interpolation=cv2.INTER_LINEAR)
+
+    if random.random() < 0.5:
+        tex = np.fliplr(tex).copy()
+    if random.random() < 0.5:
+        tex = np.flipud(tex).copy()
+    # Color jitter
+    jitter = random.uniform(0.8, 1.2)
+    brightness = random.uniform(-0.1, 0.1)
+    tex = np.clip(tex * jitter + brightness, 0, 1)
+    return tex.astype(np.float32)
+
+
+def _smooth_weight(weight: np.ndarray, kernel_size: int) -> np.ndarray:
+    """α_i as a smooth kernel (paper): softens anomaly/normal boundary. Easy → smaller kernel."""
+    k = int(kernel_size)
+    if k < 1:
+        return weight
+    if k % 2 == 0:
+        k += 1
+    return cv2.GaussianBlur(weight, (k, k), 0)
 
 
 def synthesize_anomaly(normal_img: np.ndarray,
                         jnd_map: np.ndarray,
                         dtd_paths: List[str],
-                        alpha: List[float] = (0.3, 0.6, 0.9),
+                        alpha_kernels: List[int] = (3, 7, 15),
                         beta: List[float] = (0.10, 0.05, 0.0),
                         area_ratio_min: float = 0.5,
-                        area_ratio_max: float = 0.8) -> Dict[str, np.ndarray]:
+                        area_ratio_max: float = 0.8,
+                        ssim_thresholds: Tuple[float, float] = (0.9, 0.7),
+                        # backward-compat alias used by older notebook cells
+                        alpha: List = None) -> Dict[str, np.ndarray]:
+    """
+    Panel-guided synthesis (Eq. 1–2).
+    Wi = αi * (J + βi) where αi is a smoothing kernel (not a scalar).
+    A = mean_i (Wi*T + (1-Wi)*I)  — average keeps pixels in [0,1] (paper writes Σ).
+    Severity GT refined via SSIM(I, A) into easy/medium/hard.
+    """
+    if alpha is not None:
+        # If caller still passes scalar alphas, treat odds-ish values as kernel sizes
+        alpha_kernels = [max(3, int(round(a * 15)) | 1) if isinstance(a, float) else int(a)
+                         for a in alpha]
+
     h, w, _ = normal_img.shape
     panel = random.choice(PANELS)
     y0, y1, x0, x1 = get_panel_bbox(h, w, panel, area_ratio_min, area_ratio_max)
 
     texture = sample_texture_source(normal_img, dtd_paths, (h, w))
 
-    J = jnd_map[..., None] 
-
+    J = jnd_map.astype(np.float32)
     A = normal_img.copy()
     mask = np.zeros((h, w), dtype=np.float32)
 
@@ -177,30 +213,31 @@ def synthesize_anomaly(normal_img: np.ndarray,
     T_region = texture[region_slice]
     J_region = J[region_slice]
 
-
     A_region_accum = np.zeros_like(I_region)
-    W_last = None
     for i in range(3):
-        Wi = alpha[i] * (J_region + beta[i])
-        Wi = np.clip(Wi, 0.0, 1.0)
-        A_region_accum += Wi * T_region + (1 - Wi) * I_region
-        W_last = Wi
+        # Different random seeds / slight spatial jitter per level so regions may overlap
+        noise = np.random.rand(*J_region.shape).astype(np.float32) * 0.02
+        raw = np.clip(J_region + beta[i] + noise, 0.0, 1.0)
+        Wi = _smooth_weight(raw, alpha_kernels[i])
+        Wi = np.clip(Wi, 0.0, 1.0)[..., None]
+        A_region_accum += Wi * T_region + (1.0 - Wi) * I_region
     A_region = A_region_accum / 3.0
 
     A[region_slice] = A_region
     mask[y0:y1, x0:x1] = 1.0
 
-
+    # Refine & quantify into 3 defect levels via SSIM (paper Sec. 3.1)
     severity_mask = np.zeros((h, w), dtype=np.uint8)
     gray_orig = cv2.cvtColor((normal_img * 255).astype(np.uint8), cv2.COLOR_RGB2GRAY)
     gray_syn = cv2.cvtColor((A * 255).astype(np.uint8), cv2.COLOR_RGB2GRAY)
     _, ssim_map = ssim_fn(gray_orig, gray_syn, full=True)
     ssim_region = ssim_map[y0:y1, x0:x1]
 
+    thr_easy, thr_hard = ssim_thresholds  # e.g. 0.9, 0.7
     sev_region = np.zeros_like(ssim_region, dtype=np.uint8)
-    sev_region[ssim_region < 0.7] = 3   
-    sev_region[(ssim_region >= 0.7) & (ssim_region < 0.9)] = 2 
-    sev_region[(ssim_region >= 0.9) & (ssim_region < 0.995)] = 1  
+    sev_region[ssim_region < thr_hard] = 3                          # hard
+    sev_region[(ssim_region >= thr_hard) & (ssim_region < thr_easy)] = 2  # medium
+    sev_region[(ssim_region >= thr_easy) & (ssim_region < 0.995)] = 1     # easy
     severity_mask[y0:y1, x0:x1] = sev_region
 
     return {
@@ -217,16 +254,24 @@ def synthesize_anomaly(normal_img: np.ndarray,
 
 class OmniALDataset(Dataset):
     def __init__(self, root: str, categories: List[str], dtd_paths: List[str],
-                 size: Tuple[int, int], dataset_name: str = "mvtec"):
+                 size: Tuple[int, int], dataset_name: str = "mvtec",
+                 synthesis_cfg: dict = None):
         self.size = size
         self.dtd_paths = dtd_paths
         self.dataset_name = dataset_name
+        syn = synthesis_cfg or {}
+        self.alpha_kernels = syn.get("alpha_kernels", [3, 7, 15])
+        self.beta = syn.get("beta", [0.10, 0.05, 0.0])
+        self.area_ratio_min = syn.get("area_ratio_min", 0.5)
+        self.area_ratio_max = syn.get("area_ratio_max", 0.8)
+        thr = syn.get("ssim_thresholds", [0.9, 0.7])
+        self.ssim_thresholds = (float(thr[0]), float(thr[1]))
         self.samples = []  # list of (path, category)
         for cat in categories:
             paths = list_normal_images(root, cat, dataset_name)
             self.samples += [(p, cat) for p in paths]
         if len(self.samples) == 0:
-            print(f"[WARN] {dataset_name}  path {root} dosnt find ")
+            print(f"[WARN] {dataset_name}: no images found under {root}")
 
     def __len__(self):
         return len(self.samples)
@@ -236,9 +281,15 @@ class OmniALDataset(Dataset):
         img = load_and_resize(path, self.size)
         jnd = compute_jnd_map(img)
 
-        syn = synthesize_anomaly(img, jnd, self.dtd_paths)
+        syn = synthesize_anomaly(
+            img, jnd, self.dtd_paths,
+            alpha_kernels=self.alpha_kernels,
+            beta=self.beta,
+            area_ratio_min=self.area_ratio_min,
+            area_ratio_max=self.area_ratio_max,
+            ssim_thresholds=self.ssim_thresholds,
+        )
         anomaly_img = syn["anomaly_image"]
-        anomaly_jnd = compute_jnd_map(anomaly_img)
         severity = syn["severity_mask"]
 
         to_tensor = lambda x: torch.from_numpy(x.transpose(2, 0, 1)).float() if x.ndim == 3 \
@@ -248,7 +299,6 @@ class OmniALDataset(Dataset):
             "image": to_tensor(img),
             "jnd": to_tensor(jnd),
             "anomaly_img": to_tensor(anomaly_img),
-            "anomaly_jnd": to_tensor(anomaly_jnd),
             "severity": torch.from_numpy(severity).long(),
             "category": cat,
         }
@@ -272,7 +322,10 @@ def build_dataloader(cfg: dict, dataset_name: str = "mvtec", shuffle: bool = Tru
         size = tuple(cfg["data"]["visa_size"])
         batch_size = cfg["train"]["batch_size_visa"]
 
-    dataset = OmniALDataset(root, categories, dtd_paths, size, dataset_name)
+    dataset = OmniALDataset(
+        root, categories, dtd_paths, size, dataset_name,
+        synthesis_cfg=cfg.get("synthesis", {}),
+    )
     loader = DataLoader(
         dataset,
         batch_size=batch_size,

models/__init__.py

@@ -0,0 +1,19 @@
+from .model import (
+    OmniAL,
+    DCSABlock,
+    ReconstructionNet,
+    DiffNeck,
+    DiffNeckMultiScale,
+    LocalizationNet,
+    BasicBlock,
+)
+
+__all__ = [
+    "OmniAL",
+    "DCSABlock",
+    "ReconstructionNet",
+    "DiffNeck",
+    "DiffNeckMultiScale",
+    "LocalizationNet",
+    "BasicBlock",
+]

models/model.py

@@ -4,7 +4,7 @@ import torch.nn.functional as F
 
 
 # ---------------------------------------------------------------------------
-# 0. توابع کمکی نرمال‌سازی
+# 0. Normalization helpers
 # ---------------------------------------------------------------------------
 
 def get_norm_layer(norm_type: str, num_channels: int):
@@ -13,7 +13,7 @@ def get_norm_layer(norm_type: str, num_channels: int):
     elif norm_type == "batch":
         return nn.BatchNorm2d(num_channels)
     else:
-        raise ValueError(f"norm_type نامعتبر: {norm_type}")
+        raise ValueError(f"Invalid norm_type: {norm_type}")
 
 
 # ---------------------------------------------------------------------------
@@ -37,56 +37,68 @@ class BasicBlock(nn.Module):
 
 
 # ---------------------------------------------------------------------------
-# 2. DCSA Block: Dilated Channel and Spatial Attention (CBAM + dilation)
+# 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)
-        self.max_pool = nn.AdaptiveMaxPool2d(1)
         hidden = max(channels // reduction, 4)
         self.mlp = nn.Sequential(
             nn.Conv2d(channels, hidden, 1),
             nn.ReLU(inplace=True),
             nn.Conv2d(hidden, channels, 1),
         )
-        self.sigmoid = nn.Sigmoid()
 
     def forward(self, x):
-        avg_out = self.mlp(self.avg_pool(x))
-        max_out = self.mlp(self.max_pool(x))
-        attn = self.sigmoid(avg_out + max_out)
-        return x * attn
+        attn = self.mlp(self.avg_pool(x))
+        return x + attn
 
 
 class SpatialAttention(nn.Module):
-    def __init__(self, kernel_size: int = 3, dilation: int = 3):
+    """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__()
-        padding = dilation * (kernel_size - 1) // 2
-        self.conv = nn.Conv2d(2, 1, kernel_size, padding=padding, dilation=dilation)
-        self.sigmoid = nn.Sigmoid()
+        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.sigmoid(self.conv(concat))
-        return x * attn
+        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):
+    def __init__(self, channels: int, factor: float = 0.1, norm_type: str = "batch"):
         super().__init__()
         self.channel_attn = ChannelAttention(channels)
-        self.spatial_attn = SpatialAttention()
-        self.factor = factor
+        self.spatial_attn = SpatialAttention(channels, factor=factor, norm_type=norm_type)
 
     def forward(self, x):
-        residual = x
         x = self.channel_attn(x)
         x = self.spatial_attn(x)
-        return residual + self.factor * x
+        return x
 
 
 # ---------------------------------------------------------------------------
@@ -94,38 +106,42 @@ class DCSABlock(nn.Module):
 # ---------------------------------------------------------------------------
 
 class ReconstructionNet(nn.Module):
-    def __init__(self, in_ch: int = 3, base_ch: int = 32, norm_type: str = "instance"):
+    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]]  # کانال سطوح مشترک برای DiffNeck (d1/e1, d2/e2, d3/e3)
+        self.level_channels = [chs[0], chs[1], chs[2]]
 
-        # Encoder
         self.enc1 = BasicBlock(in_ch, chs[0], norm_type)
-        self.dcsa1 = DCSABlock(chs[0])
+        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])
+        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])
+        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])
+        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):
@@ -137,57 +153,50 @@ class ReconstructionNet(nn.Module):
         p3 = self.pool(e3)
         e4 = self.dcsa4(self.enc4(p3))
 
-        d3 = self.up3(e4)
-        d3 = self.dec3(torch.cat([d3, e3], dim=1))
-        d2 = self.up2(d3)
-        d2 = self.dec2(torch.cat([d2, e2], dim=1))
-        d1 = self.up1(d2)
-        d1 = self.dec1(torch.cat([d1, e1], dim=1))
+        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.up3_jnd(e4)
-        d3j = self.dec3_jnd(torch.cat([d3j, e3], dim=1))
-        d2j = self.up2_jnd(d3j)
-        d2j = self.dec2_jnd(torch.cat([d2j, e2], dim=1))
-        d1j = self.up1_jnd(d2j)
-        d1j = self.dec1_jnd(torch.cat([d1j, e1], dim=1))
+        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 Module 
+# 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, norm_type: str = "batch"):
+    def __init__(self, in_ch: int, out_ch: int, factor: float = 0.1, norm_type: str = "batch"):
         super().__init__()
-        self.block_rec_inp = BasicBlock(in_ch * 2, out_ch, norm_type)
-        self.block_jnd_inp = BasicBlock(in_ch * 2, out_ch, norm_type)
-        self.block_rec_jnd = BasicBlock(in_ch * 2, out_ch, norm_type)
+        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_rec_jnd = self.block_rec_jnd(torch.cat([rec_feat, jnd_feat], dim=1))
-        return d_rec_inp + d_jnd_inp + d_rec_jnd  # "all levels/pairs of differences are summed up"
+        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, norm_type) for ch in channels_list
+            DiffNeckLevel(ch, out_ch, factor=factor, norm_type=norm_type) for ch in channels_list
         ])
-        self.factor = factor
         self.out_ch = out_ch
 
     def forward(self, feats_rec, feats_inp, feats_jnd):
@@ -198,75 +207,78 @@ class DiffNeckMultiScale(nn.Module):
             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
 
-        return descriptor * self.factor
+
+# Public alias so notebooks can `from models.model import DiffNeck`
+DiffNeck = DiffNeckMultiScale
 
 
 # ---------------------------------------------------------------------------
-# 5. Localization Sub-network (۶ مقیاس، Dilated Conv در مقیاس ۳-۶، Batch Norm)
+# 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"):
+                 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]
-
-        self.enc1 = BasicBlock(in_ch, chs[0], norm_type, dilation=dilation_rates[0])
-        self.enc2 = BasicBlock(chs[0], chs[1], norm_type, dilation=dilation_rates[1])
-        self.enc3 = BasicBlock(chs[1], chs[2], norm_type, dilation=dilation_rates[2])
-        self.enc4 = BasicBlock(chs[2], chs[3], norm_type, dilation=dilation_rates[3])
-        self.enc5 = BasicBlock(chs[3], chs[4], norm_type, dilation=dilation_rates[4])
-        self.enc6 = BasicBlock(chs[4], chs[5], norm_type, dilation=dilation_rates[5])
+        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.enc1(x)
-        p1 = self.pool(e1)
-        e2 = self.enc2(p1)
-        p2 = self.pool(e2)
-        e3 = self.enc3(p2)
-        p3 = self.pool(e3)
-        e4 = self.enc4(p3)
-        p4 = self.pool(e4)
-        e5 = self.enc5(p4)
-        p5 = self.pool(e5)
-        e6 = self.enc6(p5)
-
-        d5 = self.up5(e6)
-        d5 = self.dec5(torch.cat([d5, e5], dim=1))
-        d4 = self.up4(d5)
-        d4 = self.dec4(torch.cat([d4, e4], dim=1))
-        d3 = self.up3(d4)
-        d3 = self.dec3(torch.cat([d3, e3], dim=1))
-        d2 = self.up2(d3)
-        d2 = self.dec2(torch.cat([d2, e2], dim=1))
-        d1 = self.up1(d2)
-        d1 = self.dec1(torch.cat([d1, e1], dim=1))
-
-        seg_logits = self.out_seg(d1)  # (B, num_classes, H, W)
-        return seg_logits
+        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. مدل کامل OmniAL
+# 6. Full OmniAL model
 # ---------------------------------------------------------------------------
 
 class OmniAL(nn.Module):
@@ -274,29 +286,39 @@ class OmniAL(nn.Module):
         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"]
+            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=m["diffneck_factor"],
+            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"]
+            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,

notebooks/exploratory.ipynb

@@ -1,306 +1,381 @@
 {
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# OmniAL — اجرای ماژولار در Jupyter\n",
-    "\n",
-    "این نوت‌بوک دقیقاً همان کارهایی را که `scripts/main.py` انجام می‌دهد، مرحله‌به‌مرحله و با فراخوانی توابع ماژولار از `data/`, `models/`, `utils/`, `scripts/` اجرا می‌کند.\n",
-    "\n",
-    "ترتیب اجرا:\n",
-    "1. بارگذاری config\n",
-    "2. تست فاز ۱: data loading و پیش‌پردازش\n",
-    "3. تست فاز ۲: JND و Panel-guided Synthesis (نمایش بصری)\n",
-    "4. ساخت DataLoader یکپارچه\n",
-    "5. ساخت مدل OmniAL و تست forward pass\n",
-    "6. آموزش (train loop)\n",
-    "7. ارزیابی (evaluate)\n",
-    "8. نمایش نتایج"
-   ]
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# OmniAL — اجرای ماژولار در Jupyter\n",
+        "\n",
+        "این نوت‌بوک دقیقاً همان کارهایی را که `scripts/main.py` انجام می‌دهد، مرحله‌به‌مرحله و با فراخوانی توابع ماژولار از `data/`, `models/`, `utils/`, `scripts/` اجرا می‌کند.\n",
+        "\n",
+        "ترتیب اجرا:\n",
+        "1. بارگذاری config\n",
+        "2. تست فاز ۱: data loading و پیش‌پردازش\n",
+        "3. تست فاز ۲: JND و Panel-guided Synthesis (نمایش بصری)\n",
+        "4. ساخت DataLoader یکپارچه\n",
+        "5. ساخت مدل OmniAL و تست forward pass\n",
+        "6. آموزش (train loop)\n",
+        "7. ارزیابی (evaluate)\n",
+        "8. نمایش نتایج"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 0. Setup: افزودن مسیر پروژه به sys.path و import ماژول‌ها"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "cc35b5f4",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import os, sys\n",
+        "sys.path.append(os.path.abspath(\"..\"))  # ریشه‌ی homework_code/\n",
+        "\n",
+        "import yaml\n",
+        "import torch\n",
+        "import numpy as np\n",
+        "\n",
+        "from data.data_loader import (\n",
+        "    list_normal_images, list_dtd_textures, load_and_resize,\n",
+        "    compute_jnd_map, synthesize_anomaly, OmniALDataset, build_dataloader\n",
+        ")\n",
+        "from models.model import OmniAL, DCSABlock, ReconstructionNet, DiffNeck, LocalizationNet\n",
+        "from utils.metrics import OmniALLoss, anomaly_score_from_seg, compute_image_auroc, compute_pixel_auroc, compute_pixel_ap\n",
+        "from utils.visualization import show_synthesis_result, show_model_output, plot_loss_curves\n",
+        "from scripts.train import build_model_and_optimizer, run_one_epoch, train_model\n",
+        "from scripts.evaluate import evaluate_category, evaluate_dataset\n",
+        "\n",
+        "with open(\"../config/config.yaml\", \"r\", encoding=\"utf-8\") as f:\n",
+        "    cfg = yaml.safe_load(f)\n",
+        "\n",
+        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+        "print(\"Device:\", device)\n",
+        "cfg[\"train\"][\"device\"] = device\n",
+        "print(\"reconstruction_norm:\", cfg[\"model\"][\"reconstruction_norm\"],\n",
+        "      \"| localization_norm:\", cfg[\"model\"][\"localization_norm\"])\n",
+        "print(\"alpha_kernels:\", cfg[\"synthesis\"][\"alpha_kernels\"], \"| beta:\", cfg[\"synthesis\"][\"beta\"])\n",
+        "cfg"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 1. فاز ۱ — تست بارگذاری داده\n",
+        "قبل از اجرا، مسیرهای `data/dataset/mvtec_ad`, `data/dataset/visa`, `data/dataset/dtd/images` را در `config/config.yaml` پر کن و دیتاست‌ها را دانلود/کپی کن."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# تست روی یک کلاس نمونه از MVTecAD\n",
+        "category = cfg[\"data\"][\"mvtec_categories\"][0]  # مثلا 'bottle'\n",
+        "root = cfg[\"data\"][\"mvtec_root\"]\n",
+        "size = tuple(cfg[\"data\"][\"mvtec_size\"])\n",
+        "\n",
+        "normal_paths = list_normal_images(root, category, \"mvtec\")\n",
+        "dtd_paths = list_dtd_textures(cfg[\"data\"][\"dtd_root\"])\n",
+        "\n",
+        "print(f\"تعداد تصاویر normal برای '{category}': {len(normal_paths)}\")\n",
+        "print(f\"تعداد تصاویر بافت DTD: {len(dtd_paths)}\")\n",
+        "\n",
+        "if len(normal_paths) > 0:\n",
+        "    img = load_and_resize(normal_paths[0], size)\n",
+        "    print(\"شکل تصویر لود‌شده:\", img.shape, \"| min/max:\", img.min(), img.max())\n",
+        "else:\n",
+        "    print(\"[!] هیچ تصویری پیدا نشد — مسیر data.mvtec_root را در config.yaml بررسی کن.\")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 2. فاز ۲ — تست JND Map و Panel-guided Synthesis (نمایش بصری)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "bf0c6c98",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "if len(normal_paths) > 0:\n",
+        "    img = load_and_resize(normal_paths[0], size)\n",
+        "    jnd = compute_jnd_map(img)\n",
+        "    syn = synthesize_anomaly(\n",
+        "        img, jnd, dtd_paths,\n",
+        "        alpha_kernels=cfg[\"synthesis\"][\"alpha_kernels\"],\n",
+        "        beta=cfg[\"synthesis\"][\"beta\"],\n",
+        "        area_ratio_min=cfg[\"synthesis\"][\"area_ratio_min\"],\n",
+        "        area_ratio_max=cfg[\"synthesis\"][\"area_ratio_max\"],\n",
+        "        ssim_thresholds=tuple(cfg[\"synthesis\"][\"ssim_thresholds\"]),\n",
+        "    )\n",
+        "    assert syn[\"anomaly_image\"].shape == img.shape\n",
+        "    assert syn[\"severity_mask\"].shape == img.shape[:2]\n",
+        "    assert syn[\"panel\"] in [\"left\", \"right\", \"top\", \"bottom\", \"all\"]\n",
+        "    print(\"panel:\", syn[\"panel\"], \"| severity unique:\", np.unique(syn[\"severity_mask\"]))\n",
+        "    show_synthesis_result(img, jnd, syn[\"anomaly_image\"], syn[\"severity_mask\"], panel_name=syn[\"panel\"])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# چند نمونه‌ی دیگر برای دیدن تنوع پنل‌ها\n",
+        "for i in range(3):\n",
+        "    if len(normal_paths) == 0:\n",
+        "        break\n",
+        "    img = load_and_resize(normal_paths[min(i, len(normal_paths)-1)], size)\n",
+        "    jnd = compute_jnd_map(img)\n",
+        "    syn = synthesize_anomaly(img, jnd, dtd_paths)\n",
+        "    show_synthesis_result(img, jnd, syn[\"anomaly_image\"], syn[\"severity_mask\"], panel_name=syn[\"panel\"])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 3. ساخت Dataset / DataLoader یکپارچه (Unified)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "train_loader, train_dataset = build_dataloader(cfg, dataset_name=\"mvtec\", shuffle=True)\n",
+        "print(\"تعداد کل نمونه‌های آموزشی (همه کلاس‌ها):\", len(train_dataset))\n",
+        "\n",
+        "batch = next(iter(train_loader))\n",
+        "for k, v in batch.items():\n",
+        "    if torch.is_tensor(v):\n",
+        "        print(k, v.shape, v.dtype)\n",
+        "    else:\n",
+        "        print(k, type(v), v[:5])"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 4. فاز ۳ — تست معماری شبکه (اجزا به‌صورت جدا)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "79a07eb2",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# تست DCSA Block به تنهایی (channel ADD + dual-path spatial)\n",
+        "dummy = torch.randn(2, 32, 64, 64)\n",
+        "dcsa = DCSABlock(32, factor=0.1, norm_type=\"instance\")\n",
+        "out = dcsa(dummy)\n",
+        "assert out.shape == dummy.shape\n",
+        "print(\"DCSA output shape:\", out.shape)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "1e2a7446",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# تست Reconstruction Sub-network به تنهایی\n",
+        "recon = ReconstructionNet(norm_type=cfg[\"model\"][\"reconstruction_norm\"])\n",
+        "x = torch.randn(2, 3, 256, 256)\n",
+        "rec_img, rec_jnd, feats_inp, feats_rec, feats_jnd = recon(x)\n",
+        "assert rec_img.shape == (2, 3, 256, 256)\n",
+        "assert rec_jnd.shape == (2, 1, 256, 256)\n",
+        "print(\"rec_img:\", rec_img.shape, \"| rec_jnd:\", rec_jnd.shape)\n",
+        "print(\"feats_inp:\", [f.shape for f in feats_inp])\n",
+        "print(\"feats_rec:\", [f.shape for f in feats_rec])\n",
+        "print(\"feats_jnd:\", [f.shape for f in feats_jnd])\n",
+        "\n",
+        "# تست DiffNeck: مسیرهای (Rec,Inp), (JND,Inp), (Rec,JND,Inp)\n",
+        "diffneck = DiffNeck(\n",
+        "    channels_list=recon.level_channels,\n",
+        "    out_ch=cfg[\"model\"][\"base_channels\"],\n",
+        "    factor=cfg[\"model\"][\"diffneck_factor\"],\n",
+        "    norm_type=cfg[\"model\"][\"localization_norm\"],\n",
+        ")\n",
+        "desc = diffneck(feats_rec, feats_inp, feats_jnd)\n",
+        "assert desc.shape[0] == 2 and desc.shape[1] == cfg[\"model\"][\"base_channels\"]\n",
+        "print(\"DiffNeck descriptor:\", desc.shape)\n",
+        "\n",
+        "# تست LocalizationNet به تنهایی\n",
+        "loc = LocalizationNet(\n",
+        "    in_ch=3 + cfg[\"model\"][\"base_channels\"],\n",
+        "    base_ch=cfg[\"model\"][\"base_channels\"],\n",
+        "    num_classes=4,\n",
+        "    dilation_rates=tuple(cfg[\"model\"][\"dilation_rates\"]),\n",
+        "    norm_type=cfg[\"model\"][\"localization_norm\"],\n",
+        ")\n",
+        "# pad descriptor to full res if needed\n",
+        "if desc.shape[-2:] != x.shape[-2:]:\n",
+        "    import torch.nn.functional as F\n",
+        "    desc_up = F.interpolate(desc, size=x.shape[-2:], mode=\"bilinear\", align_corners=False)\n",
+        "else:\n",
+        "    desc_up = desc\n",
+        "loc_in = torch.cat([x, desc_up], dim=1)\n",
+        "seg = loc(loc_in)\n",
+        "assert seg.shape == (2, 4, 256, 256)\n",
+        "print(\"Localization seg_logits:\", seg.shape)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# تست کامل مدل OmniAL (forward pass end-to-end)\n",
+        "model, optimizer, criterion, device = build_model_and_optimizer(cfg)\n",
+        "x = torch.randn(2, 3, *size).to(device)\n",
+        "out = model(x)\n",
+        "for k, v in out.items():\n",
+        "    print(k, v.shape)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 5. فاز ۴ — تست Loss روی یک batch واقعی"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "batch = next(iter(train_loader))\n",
+        "anomaly_img = batch[\"anomaly_img\"].to(device)\n",
+        "targets = {\n",
+        "    \"image\": batch[\"image\"].to(device),\n",
+        "    \"jnd\": batch[\"jnd\"].to(device),\n",
+        "    \"severity\": batch[\"severity\"].to(device),\n",
+        "}\n",
+        "\n",
+        "outputs = model(anomaly_img)\n",
+        "losses = criterion(outputs, targets)\n",
+        "losses"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 6. آموزش مدل (Training Loop)\n",
+        "برای تست سریع، ابتدا `epochs` را در `config.yaml` کم بگذار (مثلا 2) تا فقط pipeline را چک کنی."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# برای تست سریع فقط چند epoch:\n",
+        "cfg[\"train\"][\"epochs\"] = 2\n",
+        "history = train_model(model, train_loader, optimizer, criterion, device, cfg)\n",
+        "plot_loss_curves(history)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 7. نمایش خروجی مدل روی یک نمونه بعد از آموزش"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "83590624",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "model.eval()\n",
+        "with torch.no_grad():\n",
+        "    batch = next(iter(train_loader))\n",
+        "    anomaly_img = batch[\"anomaly_img\"].to(device)\n",
+        "    out = model(anomaly_img)\n",
+        "    seg_map = anomaly_score_from_seg(\n",
+        "        out[\"seg_logits\"],\n",
+        "        smooth_kernel=cfg[\"model\"].get(\"anomaly_smooth_kernel\", 7),\n",
+        "        smooth_sigma=cfg[\"model\"].get(\"anomaly_smooth_sigma\", 1.5),\n",
+        "    )\n",
+        "\n",
+        "idx = 0\n",
+        "show_model_output(\n",
+        "    anomaly_img[idx].cpu(), out[\"rec_img\"][idx].cpu(), out[\"rec_jnd\"][idx].cpu(), seg_map[idx].cpu()\n",
+        ")"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 8. فاز ۶ — ارزیابی روی مجموعه تست (Image-AUROC, Pixel-AUROC, Pixel-AP)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "results = evaluate_dataset(model, cfg, dataset_name=\"mvtec\", device=device)\n",
+        "results"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## 9. (اختیاری) ذخیره‌ی مدل نهایی"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "os.makedirs(cfg[\"logging\"][\"ckpt_dir\"], exist_ok=True)\n",
+        "save_path = os.path.join(cfg[\"logging\"][\"ckpt_dir\"], \"omnial_final.pt\")\n",
+        "torch.save(model.state_dict(), save_path)\n",
+        "print(\"مدل ذخیره شد در:\", save_path)"
+      ]
+    }
+  ],
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python",
+      "version": "3.10"
+    }
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": ["## 0. Setup: افزودن مسیر پروژه به sys.path و import ماژول‌ها"]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os, sys\n",
-    "sys.path.append(os.path.abspath(\"..\"))  # ریشه‌ی homework_code/\n",
-    "\n",
-    "import yaml\n",
-    "import torch\n",
-    "import numpy as np\n",
-    "\n",
-    "from data.data_loader import (\n",
-    "    list_normal_images, list_dtd_textures, load_and_resize,\n",
-    "    compute_jnd_map, synthesize_anomaly, OmniALDataset, build_dataloader\n",
-    ")\n",
-    "from models.model import OmniAL, DCSABlock, ReconstructionNet, DiffNeck, LocalizationNet\n",
-    "from utils.metrics import OmniALLoss, anomaly_score_from_seg, compute_image_auroc, compute_pixel_auroc, compute_pixel_ap\n",
-    "from utils.visualization import show_synthesis_result, show_model_output, plot_loss_curves\n",
-    "from scripts.train import build_model_and_optimizer, run_one_epoch, train_model\n",
-    "from scripts.evaluate import evaluate_category, evaluate_dataset\n",
-    "\n",
-    "with open(\"../config/config.yaml\", \"r\", encoding=\"utf-8\") as f:\n",
-    "    cfg = yaml.safe_load(f)\n",
-    "\n",
-    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
-    "print(\"Device:\", device)\n",
-    "cfg[\"train\"][\"device\"] = device\n",
-    "cfg"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": ["## 1. فاز ۱ — تست بارگذاری داده\n", "قبل از اجرا، مسیرهای `data/dataset/mvtec_ad`, `data/dataset/visa`, `data/dataset/dtd/images` را در `config/config.yaml` پر کن و دیتاست‌ها را دانلود/کپی کن."]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# تست روی یک کلاس نمونه از MVTecAD\n",
-    "category = cfg[\"data\"][\"mvtec_categories\"][0]  # مثلا 'bottle'\n",
-    "root = cfg[\"data\"][\"mvtec_root\"]\n",
-    "size = tuple(cfg[\"data\"][\"mvtec_size\"])\n",
-    "\n",
-    "normal_paths = list_normal_images(root, category, \"mvtec\")\n",
-    "dtd_paths = list_dtd_textures(cfg[\"data\"][\"dtd_root\"])\n",
-    "\n",
-    "print(f\"تعداد تصاویر normal برای '{category}': {len(normal_paths)}\")\n",
-    "print(f\"تعداد تصاویر بافت DTD: {len(dtd_paths)}\")\n",
-    "\n",
-    "if len(normal_paths) > 0:\n",
-    "    img = load_and_resize(normal_paths[0], size)\n",
-    "    print(\"شکل تصویر لود‌شده:\", img.shape, \"| min/max:\", img.min(), img.max())\n",
-    "else:\n",
-    "    print(\"[!] هیچ تصویری پیدا نشد — مسیر data.mvtec_root را در config.yaml بررسی کن.\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": ["## 2. فاز ۲ — تست JND Map و Panel-guided Synthesis (نمایش بصری)"]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "if len(normal_paths) > 0:\n",
-    "    img = load_and_resize(normal_paths[0], size)\n",
-    "    jnd = compute_jnd_map(img)\n",
-    "    syn = synthesize_anomaly(img, jnd, dtd_paths,\n",
-    "                              alpha=cfg[\"synthesis\"][\"alpha\"],\n",
-    "                              beta=cfg[\"synthesis\"][\"beta\"],\n",
-    "                              area_ratio_min=cfg[\"synthesis\"][\"area_ratio_min\"],\n",
-    "                              area_ratio_max=cfg[\"synthesis\"][\"area_ratio_max\"])\n",
-    "\n",
-    "    show_synthesis_result(img, jnd, syn[\"anomaly_image\"], syn[\"severity_mask\"], panel_name=syn[\"panel\"])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# چند نمونه‌ی دیگر برای دیدن تنوع پنل‌ها\n",
-    "for i in range(3):\n",
-    "    if len(normal_paths) == 0:\n",
-    "        break\n",
-    "    img = load_and_resize(normal_paths[min(i, len(normal_paths)-1)], size)\n",
-    "    jnd = compute_jnd_map(img)\n",
-    "    syn = synthesize_anomaly(img, jnd, dtd_paths)\n",
-    "    show_synthesis_result(img, jnd, syn[\"anomaly_image\"], syn[\"severity_mask\"], panel_name=syn[\"panel\"])"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": ["## 3. ساخت Dataset / DataLoader یکپارچه (Unified)"]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "train_loader, train_dataset = build_dataloader(cfg, dataset_name=\"mvtec\", shuffle=True)\n",
-    "print(\"تعداد کل نمونه‌های آموزشی (همه کلاس‌ها):\", len(train_dataset))\n",
-    "\n",
-    "batch = next(iter(train_loader))\n",
-    "for k, v in batch.items():\n",
-    "    if torch.is_tensor(v):\n",
-    "        print(k, v.shape, v.dtype)\n",
-    "    else:\n",
-    "        print(k, type(v), v[:5])"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": ["## 4. فاز ۳ — تست معماری شبکه (اجزا به‌صورت جدا)"]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# تست DCSA Block به تنهایی\n",
-    "dummy = torch.randn(2, 32, 64, 64)\n",
-    "dcsa = DCSABlock(32)\n",
-    "out = dcsa(dummy)\n",
-    "print(\"DCSA output shape:\", out.shape)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# تست Reconstruction Sub-network به تنهایی\n",
-    "recon = ReconstructionNet(norm_type=cfg[\"model\"][\"reconstruction_norm\"])\n",
-    "x = torch.randn(2, 3, 256, 256)\n",
-    "rec_img, rec_jnd, feats = recon(x)\n",
-    "print(\"rec_img:\", rec_img.shape, \"| rec_jnd:\", rec_jnd.shape)\n",
-    "print(\"feature shapes:\", [f.shape for f in feats])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# تست کامل مدل OmniAL (forward pass end-to-end)\n",
-    "model, optimizer, criterion, device = build_model_and_optimizer(cfg)\n",
-    "x = torch.randn(2, 3, *size).to(device)\n",
-    "out = model(x)\n",
-    "for k, v in out.items():\n",
-    "    print(k, v.shape)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": ["## 5. فاز ۴ — تست Loss روی یک batch واقعی"]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "batch = next(iter(train_loader))\n",
-    "anomaly_img = batch[\"anomaly_img\"].to(device)\n",
-    "targets = {\n",
-    "    \"image\": batch[\"image\"].to(device),\n",
-    "    \"jnd\": batch[\"jnd\"].to(device),\n",
-    "    \"severity\": batch[\"severity\"].to(device),\n",
-    "}\n",
-    "\n",
-    "outputs = model(anomaly_img)\n",
-    "losses = criterion(outputs, targets)\n",
-    "losses"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": ["## 6. آموزش مدل (Training Loop)\n", "برای تست سریع، ابتدا `epochs` را در `config.yaml` کم بگذار (مثلا 2) تا فقط pipeline را چک کنی."]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# برای تست سریع فقط چند epoch:\n",
-    "cfg[\"train\"][\"epochs\"] = 2\n",
-    "history = train_model(model, train_loader, optimizer, criterion, device, cfg)\n",
-    "plot_loss_curves(history)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": ["## 7. نمایش خروجی مدل روی یک نمونه بعد از آموزش"]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model.eval()\n",
-    "with torch.no_grad():\n",
-    "    batch = next(iter(train_loader))\n",
-    "    anomaly_img = batch[\"anomaly_img\"].to(device)\n",
-    "    out = model(anomaly_img)\n",
-    "    seg_map = anomaly_score_from_seg(out[\"seg_logits\"])\n",
-    "\n",
-    "idx = 0\n",
-    "show_model_output(\n",
-    "    anomaly_img[idx].cpu(), out[\"rec_img\"][idx].cpu(), out[\"rec_jnd\"][idx].cpu(), seg_map[idx].cpu()\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": ["## 8. فاز ۶ — ارزیابی روی مجموعه تست (Image-AUROC, Pixel-AUROC, Pixel-AP)"]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "results = evaluate_dataset(model, cfg, dataset_name=\"mvtec\", device=device)\n",
-    "results"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": ["## 9. (اختیاری) ذخیره‌ی مدل نهایی"]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "os.makedirs(cfg[\"logging\"][\"ckpt_dir\"], exist_ok=True)\n",
-    "save_path = os.path.join(cfg[\"logging\"][\"ckpt_dir\"], \"omnial_final.pt\")\n",
-    "torch.save(model.state_dict(), save_path)\n",
-    "print(\"مدل ذخیره شد در:\", save_path)"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "name": "python",
-   "version": "3.10"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
+  "nbformat": 4,
+  "nbformat_minor": 5
 }

scripts/evaluate.py

@@ -16,7 +16,8 @@ from utils.metrics import anomaly_score_from_seg, compute_image_auroc, compute_p
 
 
 @torch.no_grad()
-def evaluate_category(model, root: str, category: str, size, device: str, dataset_name: str = "mvtec"):
+def evaluate_category(model, root: str, category: str, size, device: str, dataset_name: str = "mvtec",
+                      smooth_kernel: int = 7, smooth_sigma: float = 1.5):
     """
     مدل را روی تصاویر تست یک کلاس اجرا می‌کند و امتیازهای سطح تصویر/پیکسل را برمی‌گرداند.
     خروجی: dict شامل image_scores, image_labels, pixel_scores(list), pixel_labels(list)
@@ -32,7 +33,9 @@ def evaluate_category(model, root: str, category: str, size, device: str, datase
         img_t = torch.from_numpy(img.transpose(2, 0, 1)).float().unsqueeze(0).to(device)
 
         out = model(img_t)
-        seg_map = anomaly_score_from_seg(out["seg_logits"])[0].cpu().numpy()  # (H,W)
+        seg_map = anomaly_score_from_seg(
+            out["seg_logits"], smooth_kernel=smooth_kernel, smooth_sigma=smooth_sigma
+        )[0].cpu().numpy()  # (H,W)
 
         label = 0 if defect_type == "good" else 1
         image_scores.append(seg_map.max())
@@ -67,9 +70,15 @@ def evaluate_dataset(model, cfg: dict, dataset_name: str = "mvtec", device: str
         categories = cfg["data"]["visa_categories"]
         size = tuple(cfg["data"]["visa_size"])
 
+    smooth_k = cfg["model"].get("anomaly_smooth_kernel", 7)
+    smooth_s = cfg["model"].get("anomaly_smooth_sigma", 1.5)
+
     results = {}
     for cat in categories:
-        res = evaluate_category(model, root, cat, size, device, dataset_name)
+        res = evaluate_category(
+            model, root, cat, size, device, dataset_name,
+            smooth_kernel=smooth_k, smooth_sigma=smooth_s,
+        )
         if len(np.unique(res["image_labels"])) < 2:
             print(f"[WARN] کلاس {cat}: هر دو لیبل normal/anomaly برای محاسبه AUROC موجود نیست.")
             continue

utils/metrics.py

@@ -113,14 +113,30 @@ class OmniALLoss(nn.Module):
 # 2. معیارهای ارزیابی
 # ---------------------------------------------------------------------------
 
-def anomaly_score_from_seg(seg_logits: torch.Tensor) -> torch.Tensor:
+def _gaussian_kernel_2d(kernel_size: int = 7, sigma: float = 1.5, device="cpu", dtype=torch.float32):
+    coords = torch.arange(kernel_size, dtype=dtype, device=device) - kernel_size // 2
+    g = torch.exp(-(coords ** 2) / (2 * sigma ** 2))
+    g = g / g.sum()
+    return g.unsqueeze(1) @ g.unsqueeze(0)
+
+
+def anomaly_score_from_seg(seg_logits: torch.Tensor,
+                            smooth_kernel: int = 7,
+                            smooth_sigma: float = 1.5) -> torch.Tensor:
     """
-    نقشه‌ی امتیاز آنومالی پیوسته را از خروجی سگمنتیشن می‌سازد:
-    مجموع احتمال کلاس‌های easy+medium+hard (یعنی 1 - prob(background)).
+    Paper Sec. 3.3: Softmax → smoothness → sum of anomaly-level maps.
+    Classes: 0=background, 1=easy, 2=medium, 3=hard.
     """
     probs = F.softmax(seg_logits, dim=1)
-    anomaly_prob = 1 - probs[:, 0, :, :]  # (B,H,W)
-    return anomaly_prob
+    anomaly = probs[:, 1:, :, :].sum(dim=1, keepdim=True)  # (B,1,H,W)
+
+    if smooth_kernel and smooth_kernel > 1:
+        k = smooth_kernel if smooth_kernel % 2 == 1 else smooth_kernel + 1
+        kernel = _gaussian_kernel_2d(k, smooth_sigma, device=anomaly.device, dtype=anomaly.dtype)
+        kernel = kernel.view(1, 1, k, k)
+        anomaly = F.conv2d(anomaly, kernel, padding=k // 2)
+
+    return anomaly.squeeze(1)  # (B,H,W)
 
 
 def compute_image_auroc(image_scores: np.ndarray, image_labels: np.ndarray) -> float: