Deploy production Patch-MFFT Tampering Forensics API

api/model_server.py

@@ -51,9 +51,9 @@ class TamperingModelServer:
             state = torch.load(path, map_location=self.device, weights_only=False)
 
         raw_state = state["model_state_dict"] if isinstance(state, dict) and "model_state_dict" in state else state
-        self.model.load_state_dict(raw_state, strict=False)
+        self.model.load_state_dict(raw_state, strict=True)
         self.is_loaded = True
-        print(f"[TamperingModelServer] Successfully loaded weights from {path}")
+        print(f"[TamperingModelServer] Successfully loaded weights with strict=True from {path}")
 
     def warmup(self, iters: int = 2):
         dummy = torch.randn(1, 3, self.img_size, self.img_size).to(self.device)
@@ -63,18 +63,11 @@ class TamperingModelServer:
         print(f"[TamperingModelServer] Warmup completed ({iters} iterations).")
 
     def _preprocess(self, image: Image.Image) -> torch.Tensor:
-        w, h = image.size
-        target = self.img_size + 16
-        if w < h:
-            new_w, new_h = target, max(target, int(h * target / w))
-        else:
-            new_h, new_w = target, max(target, int(w * target / h))
-        image = image.convert("RGB").resize((new_w, new_h), Image.Resampling.BILINEAR)
-        left = max(0, (new_w - self.img_size) // 2)
-        top = max(0, (new_h - self.img_size) // 2)
-        image = image.crop((left, top, left + self.img_size, top + self.img_size))
+        # Standardize directly to model input resolution via bilinear resampling,
+        # perfectly matching the training dataset pipeline (Cell 14 of kaggle_patch_mfft_pipeline)
+        img_res = image.convert("RGB").resize((self.img_size, self.img_size), Image.Resampling.BILINEAR)
 
-        img_array = np.array(image, dtype=np.float32) / 255.0
+        img_array = np.array(img_res, dtype=np.float32) / 255.0
         mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
         std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
         img_array = ((img_array - mean) / std).astype(np.float32)
@@ -102,6 +95,18 @@ class TamperingModelServer:
 
         # Mask to Base64 PNG (apply sigmoid to convert raw logits to probabilities)
         mask_np = torch.sigmoid(mask)[0, 0].cpu().numpy()
+
+        # Forensic mask gating:
+        # 1. Authentic Camera: Zero out tampering mask (untouched camera capture)
+        # 2. Full AI: Retain full frame generative anomaly heatmap
+        # 3. Spliced/Tampered: Suppress background noise (< 0.20) to sharply isolate forged regions
+        if pred_idx == 0:
+            mask_np = np.zeros_like(mask_np)
+        elif pred_idx == 1:
+            mask_np = np.clip(mask_np, 0.0, 1.0)
+        else:
+            mask_np = np.where(mask_np > 0.20, mask_np, 0.0)
+
         mask_uint8 = np.clip(mask_np * 255.0, 0, 255).astype(np.uint8)
         mask_pil = Image.fromarray(mask_uint8, mode="L")
         buf = io.BytesIO()

model/src/model.py

@@ -1,27 +1,24 @@
 """
-Patch-MFFT: Dual-Branch Frequency Attention Transformer for
+Patch-MFFT: Dual-Branch Radial Frequency Transformer for
 Localized Generative Inpainting, Splicing Extortion, and Style Tampering Forensics.
 
-Core Innovations:
-1. Radial 2D-FFT Annular Sub-band Decomposition with exact reconstruction.
-2. Dual-Branch Spatial-Max Gating (SMG):
-   Combines Global Average Pooling (macro scene context) with Spatial Max-Pooling
-   (localized tampering activation) so small spliced patches (~5-15% of frame)
-   are never washed out by the authentic camera background.
-3. Boundary Frequency-Discontinuity Attention (BFDA):
-   Isolates high-frequency phase and energy mismatch seams along splicing perimeters.
-4. Dual Task Output:
-   - 3-Class Categorical Head:
-     * 0: Authentic Photography (Untouched camera capture)
-     * 1: Fully AI-Generated (Synthesized end-to-end from random noise)
-     * 2: Locally Spliced / Inpainted / Tampered (Multi-person composites, generative fill)
-   - Pixel-Level Tampering Localization Mask Head (1 x H x W probability map).
+Core Architecture (as validated in production swarm training):
+1. Radial 2D-FFT Annular Sub-band Decomposition (Low, Mid, High concentric bands).
+2. Multi-Band Depthwise-Separable Feature Encoders (encoders.0, encoders.1, encoders.2).
+3. Cross-Attention Frequency Fusion (mha + LayerNorm blocks).
+4. Dual-Branch Spatial-Max Gating (SMG): Unifies Global Average Pooling and Spatial Max-Pooling.
+5. Dual Task Decoupled Output:
+   - 3-Class Triage Categorical Head:
+     * 0: Authentic Camera Photography
+     * 1: Holistic Generative AI Synthesis
+     * 2: Locally Spliced / Inpainted / Tampered
+   - Pixel-Level Tampering Localization Decoder:
+     * High-resolution tampering anomaly heatmap [0, 1]^{H x W}.
 """
 
 import math
 from typing import Dict, List, Optional, Tuple, Union
 
-import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -29,202 +26,211 @@ import torch.nn.functional as F
 
 class RadialFourierDecomposition(nn.Module):
     """
-    Decomposes an input image into K non-overlapping radial annular frequency bands
+    Decomposes an input image into K concentric radial annular frequency bands
     in the 2D Fast Fourier Transform (FFT) domain.
-    Guarantees the exact reconstruction property: sum_{k=1}^K B_k = X.
     """
 
-    def __init__(self, img_size: int = 384, num_bands: int = 3, band_configs: Optional[List[Tuple[float, float]]] = None):
+    def __init__(self, img_size: int = 384, num_bands: int = 3):
         super().__init__()
         self.img_size = img_size
         self.num_bands = num_bands
-
-        if band_configs is not None:
-            self.configs = band_configs
-            self.num_bands = len(band_configs)
-        elif num_bands == 2:
-            self.configs = [(0.0, 0.25), (0.25, 1.0)]
-        elif num_bands == 4:
-            self.configs = [(0.0, 0.125), (0.125, 0.25), (0.25, 0.5), (0.5, 1.0)]
-        else:
-            # Default 3 bands: Low, Mid, High
-            self.configs = [(0.0, 0.15), (0.15, 0.5), (0.5, 1.0)]
-
-        # Precompute radial coordinate grid
+        configs = (
+            [(0.0, 0.15), (0.15, 0.5), (0.5, 1.0)]
+            if num_bands == 3
+            else [(0.0, 0.125), (0.125, 0.25), (0.25, 0.5), (0.5, 1.0)]
+        )
         y = torch.linspace(-1, 1, img_size)
         x = torch.linspace(-1, 1, img_size)
         gy, gx = torch.meshgrid(y, x, indexing="ij")
-        r = torch.sqrt(gx**2 + gy**2)
-        r = r / r.max()  # normalize [0, 1]
-
+        r = torch.sqrt(gx**2 + gy**2) / math.sqrt(2)
         masks = []
-        for low_r, high_r in self.configs:
-            if high_r >= 1.0:
-                mask = (r >= low_r) & (r <= high_r)
-            else:
-                mask = (r >= low_r) & (r < high_r)
-            masks.append(mask.float())
-
-        # Ensure exact reconstruction without energy leakage
-        stacked = torch.stack(masks, dim=0)  # (K, H, W)
-        norm_factor = stacked.sum(dim=0, keepdim=True).clamp(min=1e-8)
-        stacked = stacked / norm_factor
-        self.register_buffer("masks", stacked)
+        for low_r, high_r in configs:
+            m = (r >= low_r) & (r <= high_r if high_r >= 1.0 else r < high_r)
+            masks.append(m.float())
+        stacked = torch.stack(masks, dim=0)
+        norm = stacked.sum(dim=0, keepdim=True).clamp(min=1e-8)
+        self.register_buffer("masks", stacked / norm)
 
     def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
-        """
-        Args:
-            x: Input image tensor (B, C, H, W)
-        Returns:
-            List of K decomposed spatial-frequency bands [B_1, B_2, ..., B_K]
-        """
-        B, C, H, W = x.shape
-        # Real-to-complex 2D FFT
         X_fft = torch.fft.fftshift(torch.fft.fft2(x, norm="ortho"), dim=(-2, -1))
-
         bands = []
         for k in range(self.num_bands):
-            mask = self.masks[k].unsqueeze(0).unsqueeze(0)  # (1, 1, H, W)
-            filtered_fft = X_fft * mask
-            unshifted = torch.fft.ifftshift(filtered_fft, dim=(-2, -1))
-            spatial_band = torch.fft.ifft2(unshifted, norm="ortho").real
-            bands.append(spatial_band)
-
+            m = self.masks[k].unsqueeze(0).unsqueeze(0)
+            bands.append(
+                torch.fft.ifft2(
+                    torch.fft.ifftshift(X_fft * m, dim=(-2, -1)), norm="ortho"
+                ).real
+            )
         return bands
 
 
-class DepthwiseSeparableConvBlock(nn.Module):
-    """Parameter-efficient depthwise separable convolutional residual block."""
-
-    def __init__(self, in_channels: int, out_channels: int, stride: int = 1):
+class DCT2DDecomposition(nn.Module):
+    def __init__(self, img_size: int = 384, num_bands: int = 3):
         super().__init__()
-        self.depthwise = nn.Conv2d(
-            in_channels, in_channels, kernel_size=3, stride=stride, padding=1,
-            groups=in_channels, bias=False
-        )
-        self.pointwise = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
-        self.bn1 = nn.BatchNorm2d(out_channels)
-        self.act1 = nn.GELU()
+        self.img_size = img_size
+        self.num_bands = num_bands
+        N = img_size
+        n = torch.arange(N).float()
+        k = torch.arange(N).float().unsqueeze(1)
+        D = torch.cos(math.pi * (2 * n + 1) * k / (2 * N)) * math.sqrt(2.0 / N)
+        D[0, :] = D[0, :] / math.sqrt(2.0)
+        self.register_buffer("D", D)
+
+        u = torch.linspace(0, 1, N)
+        v = torch.linspace(0, 1, N)
+        gu, gv = torch.meshgrid(u, v, indexing="ij")
+        d = (gu + gv) / 2.0
+        configs = [(0.0, 0.15), (0.15, 0.5), (0.5, 1.0)]
+        masks = []
+        for low_d, high_d in configs:
+            m = (d >= low_d) & (d <= high_d if high_d >= 1.0 else d < high_d)
+            masks.append(m.float())
+        stacked = torch.stack(masks, dim=0)
+        norm = stacked.sum(dim=0, keepdim=True).clamp(min=1e-8)
+        self.register_buffer("masks", stacked / norm)
 
-        self.depthwise2 = nn.Conv2d(
-            out_channels, out_channels, kernel_size=3, stride=1, padding=1,
-            groups=out_channels, bias=False
-        )
-        self.pointwise2 = nn.Conv2d(out_channels, out_channels, kernel_size=1, bias=False)
-        self.bn2 = nn.BatchNorm2d(out_channels)
+    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+        D = self.D
+        xD_T = torch.matmul(x, D.t())
+        Y = torch.matmul(D, xD_T)
+        bands = []
+        for k in range(self.num_bands):
+            m = self.masks[k].unsqueeze(0).unsqueeze(0)
+            Y_k = Y * m
+            D_T_Y = torch.matmul(D.t(), Y_k)
+            bands.append(torch.matmul(D_T_Y, D))
+        return bands
 
-        self.residual = (
-            nn.Sequential(
-                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False),
-                nn.BatchNorm2d(out_channels),
-            )
-            if (stride != 1 or in_channels != out_channels)
-            else nn.Identity()
-        )
-        self.act2 = nn.GELU()
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        res = self.residual(x)
-        out = self.act1(self.bn1(self.pointwise(self.depthwise(x))))
-        out = self.bn2(self.pointwise2(self.depthwise2(out)))
-        return self.act2(out + res)
+class HaarDWTDecomposition(nn.Module):
+    def __init__(self, in_c: int = 3):
+        super().__init__()
+        ll = torch.tensor([[1.0, 1.0], [1.0, 1.0]]) * 0.5
+        lh = torch.tensor([[-1.0, -1.0], [1.0, 1.0]]) * 0.5
+        hl = torch.tensor([[-1.0, 1.0], [-1.0, 1.0]]) * 0.5
+        hh = torch.tensor([[1.0, -1.0], [-1.0, 1.0]]) * 0.5
+        filters = torch.stack([ll, lh, hl, hh], dim=0).unsqueeze(1)
+        self.register_buffer("filters", filters)
+        self.in_c = in_c
 
+    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+        B, C, H, W = x.shape
+        bands = []
+        for i in range(4):
+            filt = self.filters[i : i + 1].repeat(C, 1, 1, 1)
+            y = F.conv2d(x, filt, stride=2, groups=C)
+            y_up = F.interpolate(y, size=(H, W), mode="bilinear", align_corners=False)
+            bands.append(y_up)
+        # 3 bands: Low (LL), Mid (LH + HL), High (HH)
+        return [bands[0], (bands[1] + bands[2]) * 0.5, bands[3]]
 
-class FrequencyGuidedAttention(nn.Module):
-    """
-    Physical Fourier Energy Prior Gating.
-    Dynamically modulates cross-attention weights based on logarithmic spectral energy.
-    """
 
-    def __init__(self, embed_dim: int, num_heads: int = 4):
+class SpatialOnlyDecomposition(nn.Module):
+    def __init__(self, num_bands: int = 3):
         super().__init__()
-        self.embed_dim = embed_dim
-        self.num_heads = num_heads
-        self.mha = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
-        self.norm = nn.LayerNorm(embed_dim)
-        self.freq_gate = nn.Sequential(
-            nn.Linear(1, embed_dim),
-            nn.Sigmoid(),
-        )
+        self.num_bands = num_bands
 
-    def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor,
-                energy_prior: Optional[torch.Tensor] = None) -> torch.Tensor:
-        # query, key, value: (B, N, D)
-        attn_out, _ = self.mha(query, key, value)
-        out = self.norm(query + attn_out)
+    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
+        return [x] * self.num_bands
 
-        if energy_prior is not None:
-            # energy_prior: (B, 1)
-            gate = self.freq_gate(energy_prior.unsqueeze(-1))  # (B, 1, D)
-            out = out * (1.0 + gate)
 
-        return out
+class DepthwiseSeparableConvBlock(nn.Module):
+    """Parameter-efficient depthwise separable convolutional residual block."""
 
+    def __init__(self, in_c: int, out_c: int, stride: int = 1):
+        super().__init__()
+        self.dw = nn.Conv2d(in_c, in_c, 3, stride=stride, padding=1, groups=in_c, bias=False)
+        self.pw = nn.Conv2d(in_c, out_c, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(out_c)
+        self.dw2 = nn.Conv2d(out_c, out_c, 3, 1, 1, groups=out_c, bias=False)
+        self.pw2 = nn.Conv2d(out_c, out_c, 1, bias=False)
+        self.bn2 = nn.BatchNorm2d(out_c)
+        self.act = nn.GELU()
+        self.res = (
+            nn.Sequential(
+                nn.Conv2d(in_c, out_c, 1, stride=stride, bias=False),
+                nn.BatchNorm2d(out_c),
+            )
+            if (stride != 1 or in_c != out_c)
+            else nn.Identity()
+        )
 
-class DualBranchSpatialMaxGating(nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.act(self.bn1(self.pw(self.dw(x))))
+        out = self.bn2(self.pw2(self.dw2(out)))
+        return self.act(out + self.res(x))
+
+
+class UniversalPoolingModule(nn.Module):
     """
-    Dual-Branch Spatial-Max Gating (SMG).
-    Crucial for localized tampering detection (splicing & inpainting).
-    Branch A: Global Average Pool -> captures holistic scene illumination & context.
-    Branch B: Patch Spatial Max Pool -> isolates localized high-frequency tampering spikes.
+    Universal feature aggregator supporting Dual-Branch Spatial-Max Gating (SMG),
+    Global Average Pooling (GAP), Global Max-Pooling (GMP), and GeM.
     """
 
-    def __init__(self, in_features: int, out_features: int):
+    def __init__(self, in_f: int, out_f: int, mode: str = "smg"):
         super().__init__()
-        # Concatenates average and max features (2 * in_features -> out_features)
-        self.proj = nn.Sequential(
-            nn.Linear(in_features * 2, out_features),
-            nn.BatchNorm1d(out_features),
-            nn.GELU(),
-            nn.Dropout(0.1),
-        )
+        self.mode = mode
+        if mode == "smg":
+            self.proj = nn.Sequential(
+                nn.Linear(in_f * 2, out_f),
+                nn.BatchNorm1d(out_f),
+                nn.GELU(),
+                nn.Dropout(0.1),
+            )
+        else:
+            self.proj = nn.Sequential(
+                nn.Linear(in_f, out_f),
+                nn.BatchNorm1d(out_f),
+                nn.GELU(),
+                nn.Dropout(0.1),
+            )
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Args:
-            x: Feature map (B, C, H, W)
-        Returns:
-            Aggregated feature vector (B, out_features)
-        """
-        feat_avg = x.mean(dim=[-2, -1])  # (B, C) - Macro scene distribution
-        feat_max = x.amax(dim=[-2, -1])  # (B, C) - Local tampering peak activation
-        combined = torch.cat([feat_avg, feat_max], dim=-1)  # (B, 2*C)
-        return self.proj(combined)
+    def forward(self, x: torch.Tensor, cls_token: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if self.mode == "smg":
+            avg_p = x.mean(dim=[-2, -1])
+            max_p = x.amax(dim=[-2, -1])
+            return self.proj(torch.cat([avg_p, max_p], dim=-1))
+        elif self.mode == "avg":
+            return self.proj(x.mean(dim=[-2, -1]))
+        elif self.mode == "max":
+            return self.proj(x.amax(dim=[-2, -1]))
+        elif self.mode == "gem":
+            p = 3.0
+            gem_p = (x.clamp(min=1e-6) ** p).mean(dim=[-2, -1]) ** (1.0 / p)
+            return self.proj(gem_p)
+        elif self.mode == "cls":
+            return self.proj(cls_token)
+        else:
+            return self.proj(x.mean(dim=[-2, -1]))
 
 
 class TamperingLocalizationDecoder(nn.Module):
     """
-    Pixel-level segmentation mask decoder.
-    Takes fused multi-scale frequency features and predicts a 1-channel probability heatmap
-    localizing the exact spliced/inpainted boundary and forged region.
+    Lightweight convolutional decoder for dense pixel-level tampering segmentation.
+    Produces uncompressed 1 x H x W probability maps.
     """
 
-    def __init__(self, in_channels: int, img_size: int = 384):
+    def __init__(self, in_c: int, img_size: int = 384):
         super().__init__()
-        self.img_size = img_size
-        self.decoder = nn.Sequential(
-            nn.Conv2d(in_channels, 64, kernel_size=3, padding=1),
+        self.dec = nn.Sequential(
+            nn.Conv2d(in_c, 64, 3, padding=1),
             nn.BatchNorm2d(64),
             nn.GELU(),
             nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
-            nn.Conv2d(64, 32, kernel_size=3, padding=1),
+            nn.Conv2d(64, 32, 3, padding=1),
             nn.BatchNorm2d(32),
             nn.GELU(),
             nn.Upsample(size=(img_size, img_size), mode="bilinear", align_corners=False),
-            nn.Conv2d(32, 1, kernel_size=1),
+            nn.Conv2d(32, 1, 1),
         )
 
-    def forward(self, feature_map: torch.Tensor) -> torch.Tensor:
-        # feature_map: (B, C, H', W')
-        return self.decoder(feature_map)
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.dec(x)
 
 
 class PatchMFFT(nn.Module):
     """
-    Patch-MFFT Neural Architecture.
-    Complete multi-frequency vision transformer for holistic synthesis detection,
-    localized inpainting detection, two-person composite splicing detection,
-    and pixel-wise tampering localization.
+    Patch-MFFT: Multi-Frequency Fusion Transformer for Image Tampering Forensics.
     """
 
     def __init__(
@@ -234,47 +240,58 @@ class PatchMFFT(nn.Module):
         depth: int = 3,
         num_bands: int = 3,
         img_size: int = 384,
-        num_classes: int = 3,  # 0: Authentic, 1: Fully AI-Generated, 2: Locally Spliced/Inpainted
-        use_fga: bool = True,
-        enable_localization_head: bool = True,
+        num_classes: int = 3,
+        ablation: Optional[dict] = None,
     ):
         super().__init__()
         self.embed_dim = embed_dim
         self.num_bands = num_bands
         self.img_size = img_size
-        self.num_classes = num_classes
-        self.use_fga = use_fga
-        self.enable_localization_head = enable_localization_head
-
-        # 1. 2D-FFT Radial Annular Spectral Decomposer
-        self.decomposer = RadialFourierDecomposition(img_size=img_size, num_bands=num_bands)
-
-        # 2. Multi-band convolutional feature extractors
-        channels_prog = [32, 64, 128, embed_dim]
-        self.band_encoders = nn.ModuleList()
-        for _ in range(num_bands):
-            blocks = []
-            curr_c = 3
-            for next_c in channels_prog:
-                blocks.append(DepthwiseSeparableConvBlock(curr_c, next_c, stride=2))
-                curr_c = next_c
-            self.band_encoders.append(nn.Sequential(*blocks))
-
-        # Spatial resolution after 4 stride-2 stages: 384 / 16 = 24x24
-        self.spatial_h = img_size // 16
-        self.spatial_w = img_size // 16
-
-        # 3. Cross-Attention Band Fusion (CAF) with FGA
-        self.fusion_blocks = nn.ModuleList([
-            FrequencyGuidedAttention(embed_dim, num_heads=num_heads)
-            for _ in range(depth)
+        self.ablation = ablation or {}
+
+        t_type = self.ablation.get("transform", "radial_fft")
+        if self.ablation.get("spatial_only", False) or t_type == "spatial":
+            self.decomposer = SpatialOnlyDecomposition(num_bands)
+        elif t_type == "dct":
+            self.decomposer = DCT2DDecomposition(img_size, num_bands)
+        elif t_type == "dwt":
+            self.decomposer = HaarDWTDecomposition(3)
+        else:
+            self.decomposer = RadialFourierDecomposition(img_size, num_bands)
+
+        # 1. Sub-band convolutional encoders (stride 2 x 4 stages = downscale 16x)
+        self.encoders = nn.ModuleList([
+            nn.Sequential(
+                DepthwiseSeparableConvBlock(3, 32, 2),
+                DepthwiseSeparableConvBlock(32, 64, 2),
+                DepthwiseSeparableConvBlock(64, 128, 2),
+                DepthwiseSeparableConvBlock(128, embed_dim, 2),
+            )
+            for _ in range(num_bands)
         ])
 
-        # 4. Dual-Branch Spatial-Max Gating (SMG)
-        self.smg_aggregator = DualBranchSpatialMaxGating(embed_dim, embed_dim)
+        # 2. Multi-head cross-attention & LayerNorm across bands
+        self.mha = nn.ModuleList([
+            nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
+            for _ in range(depth)
+        ])
+        self.norms = nn.ModuleList([nn.LayerNorm(embed_dim) for _ in range(depth)])
+
+        # 3. Dual-Branch SMG pooling module
+        pool_mode = self.ablation.get("pool_mode", "smg")
+        if self.ablation.get("no_smg", False):
+            pool_mode = "avg"
+        self.pool_mode = pool_mode
+        self.smg = UniversalPoolingModule(embed_dim, embed_dim, mode=pool_mode)
+
+        if pool_mode == "cls":
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+            nn.init.trunc_normal_(self.cls_token, std=0.02)
+        else:
+            self.cls_token = None
 
-        # 5. 3-Class Classification Head
-        self.classifier = nn.Sequential(
+        # 4. 3-Class Categorical Head (Authentic, Full AI, Spliced/Tampered)
+        self.cls = nn.Sequential(
             nn.Linear(embed_dim, embed_dim // 2),
             nn.BatchNorm1d(embed_dim // 2),
             nn.GELU(),
@@ -282,9 +299,12 @@ class PatchMFFT(nn.Module):
             nn.Linear(embed_dim // 2, num_classes),
         )
 
-        # 6. Pixel-Level Tampering Localization Head
-        if enable_localization_head:
-            self.localization_decoder = TamperingLocalizationDecoder(embed_dim, img_size=img_size)
+        # 5. Tampering Localization Decoder
+        self.loc = TamperingLocalizationDecoder(embed_dim, img_size)
+
+    @property
+    def pool(self) -> UniversalPoolingModule:
+        return self.smg
 
     def forward(
         self,
@@ -292,74 +312,41 @@ class PatchMFFT(nn.Module):
         return_mask: bool = False,
         return_heatmap: bool = False,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
-        """
-        Args:
-            x: Input image tensor (B, 3, H, W)
-            return_mask: Whether to return pixel-level localization mask
-            return_heatmap: Backward compatibility alias for return_mask
-        Returns:
-            logits: (B, num_classes)
-            mask (optional): (B, 1, H, W) probability heatmap of tampered regions
-        """
         B = x.shape[0]
-
-        # 1. Radial Fourier sub-band decomposition
         bands = self.decomposer(x)
-
-        # 2. Compute physical spectral energies for FGA gating
-        energies = [torch.log1p(band.abs().mean(dim=[-3, -2, -1], keepdim=True)) for band in bands]
-        total_energy = sum(energies) / len(energies)  # (B, 1, 1, 1)
-        energy_prior = total_energy.squeeze(-1).squeeze(-1)  # (B, 1)
-
-        # 3. Encode each frequency sub-band through depthwise separable networks
-        band_feats = []
-        for k, encoder in enumerate(self.band_encoders):
-            feat = encoder(bands[k])  # (B, embed_dim, 24, 24)
-            band_feats.append(feat)
-
-        # High-frequency band features carry dominant tampering cues
-        high_band_feat = band_feats[-1]
-
-        # 4. Convert spatial feature maps to token sequence for Cross-Attention
-        # Tokenize across bands: (B, num_bands * N_patches, embed_dim)
-        band_tokens = []
-        for feat in band_feats:
-            tok = feat.flatten(2).permute(0, 2, 1)  # (B, 576, embed_dim)
-            band_tokens.append(tok)
-
-        tokens = torch.cat(band_tokens, dim=1)  # (B, K * 576, embed_dim)
-
-        # Apply Cross-Attention Band Fusion modulated by FGA prior
-        for block in self.fusion_blocks:
-            tokens = block(
-                tokens, tokens, tokens,
-                energy_prior=energy_prior if self.use_fga else None,
+        feats = [enc(b) for enc, b in zip(self.encoders, bands)]
+        tokens = torch.cat([f.flatten(2).permute(0, 2, 1) for f in feats], dim=1)
+
+        if self.cls_token is not None:
+            cls_t = self.cls_token.expand(B, -1, -1)
+            tokens = torch.cat([cls_t, tokens], dim=1)
+            for mha, norm in zip(self.mha, self.norms):
+                attn, _ = mha(tokens, tokens, tokens, need_weights=False)
+                tokens = norm(tokens + attn)
+            extracted_cls = tokens[:, 0, :]
+            n_p = (self.img_size // 16) ** 2
+            fused = (
+                tokens[:, -n_p:, :]
+                .permute(0, 2, 1)
+                .reshape(B, self.embed_dim, self.img_size // 16, self.img_size // 16)
             )
+            pool = self.smg(fused, cls_token=extracted_cls)
+        else:
+            for mha, norm in zip(self.mha, self.norms):
+                attn, _ = mha(tokens, tokens, tokens, need_weights=False)
+                tokens = norm(tokens + attn)
+            n_p = (self.img_size // 16) ** 2
+            fused = (
+                tokens[:, -n_p:, :]
+                .permute(0, 2, 1)
+                .reshape(B, self.embed_dim, self.img_size // 16, self.img_size // 16)
+            )
+            pool = self.smg(fused)
 
-        # 5. Extract fused spatial feature map from token sequence
-        # Reshape the high-frequency slice back to (B, embed_dim, 24, 24)
-        n_patches = self.spatial_h * self.spatial_w
-        fused_high_spatial = tokens[:, -n_patches:, :].permute(0, 2, 1).reshape(
-            B, self.embed_dim, self.spatial_h, self.spatial_w
-        )
-
-        # 6. Dual-Branch Spatial-Max Gating (Average + Max pooling)
-        fused_vector = self.smg_aggregator(fused_high_spatial)  # (B, embed_dim)
-
-        # 7. Classification Logits (3-class)
-        logits = self.classifier(fused_vector)
+        logits = self.cls(pool)
 
-        # 8. Optional Pixel-Level Tampering Localization Mask
         if return_mask or return_heatmap:
-            if self.enable_localization_head:
-                mask = self.localization_decoder(fused_high_spatial)  # (B, 1, H, W)
-            else:
-                # Fallback: normalized spatial energy of high-frequency band
-                mask = high_band_feat.abs().mean(dim=1, keepdim=True)
-                mask = F.interpolate(mask, size=(self.img_size, self.img_size), mode="bilinear", align_corners=False)
-                mask = (mask - mask.amin(dim=[-2, -1], keepdim=True)) / (
-                    mask.amax(dim=[-2, -1], keepdim=True) - mask.amin(dim=[-2, -1], keepdim=True) + 1e-8
-                )
+            mask = self.loc(fused)
             return logits, mask
 
         return logits
@@ -372,35 +359,22 @@ def build_patch_mfft(
     ablation: Optional[dict] = None,
 ) -> PatchMFFT:
     """
-    Model factory for Patch-MFFT variants:
-    - tiny:   ~420K parameters  (embed_dim=128, heads=4, depth=2)
-    - base:   ~1.85M parameters (embed_dim=384, heads=6, depth=3)
-    - large:  ~7.10M parameters (embed_dim=768, heads=12, depth=4)
+    Factory builder for Patch-MFFT model variants:
+    - tiny:  embed_dim=128, heads=4,  depth=2 (563K params)
+    - base:  embed_dim=384, heads=6,  depth=3 (3.28M params)
+    - large: embed_dim=768, heads=12, depth=4 (13.94M params)
     """
-    v = (variant or "base").lower()
-    ablation = ablation or {}
-
-    configs = {
-        "tiny":  {"embed_dim": 128, "num_heads": 4,  "depth": 2, "num_bands": 3},
-        "base":  {"embed_dim": 384, "num_heads": 6,  "depth": 3, "num_bands": 3},
-        "large": {"embed_dim": 768, "num_heads": 12, "depth": 4, "num_bands": 3},
+    cfgs = {
+        "tiny": {"embed_dim": 128, "heads": 4, "depth": 2},
+        "base": {"embed_dim": 384, "heads": 6, "depth": 3},
+        "large": {"embed_dim": 768, "heads": 12, "depth": 4},
     }
-
-    if v not in configs:
-        raise ValueError(f"Unknown variant '{variant}'. Choose from: {list(configs.keys())}")
-
-    cfg = configs[v].copy()
-    if "num_bands" in ablation:
-        cfg["num_bands"] = ablation["num_bands"]
-
-    model = PatchMFFT(
-        embed_dim=cfg["embed_dim"],
-        num_heads=cfg["num_heads"],
-        depth=cfg["depth"],
-        num_bands=cfg["num_bands"],
+    c = cfgs.get(variant, cfgs["base"])
+    return PatchMFFT(
+        embed_dim=c["embed_dim"],
+        num_heads=c["heads"],
+        depth=c["depth"],
         img_size=img_size,
         num_classes=num_classes,
-        use_fga=ablation.get("use_fga", True),
-        enable_localization_head=True,
+        ablation=ablation,
     )
-    return model