Release original LTF v3.3 epoch 20 inference model

.gitattributes

@@ -1,35 +1,2 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text

CLIP_LICENSE

@@ -0,0 +1,22 @@
+MIT License
+
+Copyright (c) 2021 OpenAI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+

backbones/ViT-L-14.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836
+size 932768134

configs/eval.yaml

@@ -0,0 +1,58 @@
+version: v3_3
+backbone: backbones/ViT-L-14.pt
+seed: 42
+epochs: 20
+batch_size: 8
+test_batch_size: 1
+learning_rate: 0.0004
+weight_decay: 0.05
+dropout: 0.1
+jpeg_augment:
+- 60
+- 100
+jpeg_p: 0.1
+blur_p: 0.1
+blur_augment:
+- 0.1
+- 2.5
+resize_augment:
+- 0.25
+- 2
+resize_p: 0.1
+replace_p: 0.1
+replace_ratio:
+- 0.2
+- 0.98
+replace_grid: 16
+num_workers: 8
+train_set: genimage_dr3_rec
+test_set: genimage
+train_set_size: 50000
+test_set_size: null
+use_patch_transformer: false
+linear_classifier: true
+transformer_layers: 0
+transformer_heads: 2
+tile_layers: 1
+tile_heads: 2
+default_train_tiles: 16
+use_global_image: true
+mil_heads: 8
+fusion_hidden: 512
+v2_l1_cross_shuffle_p: 0.9
+v2_l1_replace_ratio_range:
+- 0.5
+- 0.9
+v2_l1_with_replacement_p: 0.5
+v2_l2_cross_shuffle_p: 0.9
+v2_l2_replace_ratio_range:
+- 0.1
+- 0.8
+v2_l2_with_replacement_p: 0.5
+v2_bag_token_shuffle_p: 0.5
+v2_cross_shuffle_min_keep_original: 1
+v2_use_clip_intermediate_for_mil: false
+v3_3_use_pool_gate_fusion: true
+v3_3_pool_gate_hidden: 16
+weights: weights/epoch_20.pth
+eval_amp: false

inference.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+"""Predict real/fake for one image or a directory, without torchrun."""
+import argparse
+import json
+import os
+from pathlib import Path
+import sys
+
+from ltf_final import LTFDetector
+from ltf_final.detector import ROOT, BRANCHES
+from utils.tile_dataset_srm import IMG_EXTENSIONS
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--input', type=Path, required=True, help='Image file or directory, scanned recursively')
+    parser.add_argument('--output', type=Path, help='JSONL predictions; stdout is used when omitted')
+    parser.add_argument('--config', type=Path, default=ROOT / 'configs/eval.yaml')
+    parser.add_argument('--weights', type=Path, help='Override the bundled epoch_20.pth')
+    parser.add_argument('--device', default='auto', help='auto, cpu, cuda:0, etc.')
+    parser.add_argument('--branch', choices=BRANCHES, default='fusion')
+    parser.add_argument('--threshold', type=float, default=.5)
+    parser.add_argument('--chunk-size', type=int, default=256, help='CLIP views per forward chunk; tile coverage is preserved')
+    parser.add_argument('--threads', type=int, default=4, help='PyTorch CPU thread count')
+    args = parser.parse_args()
+    if args.threads < 1:
+        parser.error('--threads must be positive')
+    if not 0 <= args.threshold <= 1:
+        parser.error('--threshold must lie in [0,1]')
+    source = args.input.expanduser().resolve()
+    if source.is_file():
+        paths = [source]
+    elif source.is_dir():
+        paths = sorted(p for p in source.rglob('*') if p.is_file() and p.suffix.lower() in IMG_EXTENSIONS)
+    else:
+        parser.error(f'Input does not exist: {source}')
+    if not paths:
+        parser.error('No images found')
+    import torch
+    torch.set_num_threads(args.threads)
+    detector = LTFDetector(args.config, args.weights.resolve() if args.weights else None,
+                           args.device, args.chunk_size)
+    stream, temporary = sys.stdout, None
+    if args.output:
+        output = args.output.expanduser().resolve()
+        output.parent.mkdir(parents=True, exist_ok=True)
+        temporary = output.with_name(output.name + f'.{os.getpid()}.tmp')
+        stream = temporary.open('w', encoding='utf-8')
+    try:
+        for path in paths:
+            result = detector.predict(path, args.branch, args.threshold)
+            stream.write(json.dumps(result, ensure_ascii=False, allow_nan=False) + '\n')
+        stream.flush()
+        if temporary:
+            stream.close()
+            os.replace(temporary, output)
+    except BaseException:
+        if temporary:
+            stream.close()
+            temporary.unlink(missing_ok=True)
+        raise
+    print(f'Predicted {len(paths)} image(s); branch={args.branch}; threshold={args.threshold}', file=sys.stderr)
+
+
+if __name__ == '__main__':
+    main()

ltf_final/__init__.py

@@ -0,0 +1,5 @@
+"""Inference API for the released LTF v3.3 model."""
+
+from .detector import LTFDetector, load_config, prepare_image
+
+__all__ = ['LTFDetector', 'load_config', 'prepare_image']

ltf_final/detector.py

@@ -0,0 +1,125 @@
+"""Offline model loading and the original experiment's image preprocessing."""
+from __future__ import annotations
+
+import contextlib
+import sys
+from pathlib import Path
+
+from PIL import Image
+import torch
+import yaml
+
+from model.LTF_v3_3 import LTFNetV3_3
+from utils.tile_dataset_srm import clip_transform, crop_image_to_tiles
+
+ROOT = Path(__file__).resolve().parents[1]
+BRANCHES = ('level1', 'level2', 'fusion')
+
+
+def load_config(path=None):
+    path = Path(path) if path else ROOT / 'configs/eval.yaml'
+    with path.expanduser().resolve().open(encoding='utf-8') as stream:
+        config = yaml.safe_load(stream)
+    if not isinstance(config, dict):
+        raise ValueError('Configuration must be a YAML mapping')
+    return config
+
+
+def package_path(path):
+    path = Path(path).expanduser()
+    return path.resolve() if path.is_absolute() else (ROOT / path).resolve()
+
+
+def model_kwargs(config):
+    return dict(
+        num_heads=config.get('transformer_heads', 2), num_classes=2,
+        num_layers=config.get('transformer_layers', 0), dropout=config.get('dropout', .1),
+        tile_heads=config.get('tile_heads', 2), tile_layers=config.get('tile_layers', 1),
+        use_global_image=config.get('use_global_image', True),
+        use_patch_transformer=config.get('use_patch_transformer', False),
+        linear_classifier=config.get('linear_classifier', True),
+        mil_heads=config.get('mil_heads', 8), fusion_hidden=config.get('fusion_hidden', 512),
+        l1_cross_shuffle_p=float(config.get('v2_l1_cross_shuffle_p', .9)),
+        l1_replace_ratio_range=tuple(config.get('v2_l1_replace_ratio_range', [.5, .9])),
+        l1_with_replacement_p=float(config.get('v2_l1_with_replacement_p', .5)),
+        l2_cross_shuffle_p=float(config.get('v2_l2_cross_shuffle_p', .9)),
+        l2_replace_ratio_range=tuple(config.get('v2_l2_replace_ratio_range', [.1, .8])),
+        l2_with_replacement_p=float(config.get('v2_l2_with_replacement_p', .5)),
+        bag_token_shuffle_p=float(config.get('v2_bag_token_shuffle_p', .5)),
+        cross_shuffle_min_keep_original=int(config.get('v2_cross_shuffle_min_keep_original', 1)),
+        use_clip_intermediate_for_mil=config.get('v2_use_clip_intermediate_for_mil', False),
+        clip_mil_stop_after_block_1based=int(config.get('v2_clip_mil_stop_after_block_1based', 12)),
+        pool_gate_hidden=int(config.get('v3_3_pool_gate_hidden', 16)),
+        use_pool_gate_fusion=config.get('v3_3_use_pool_gate_fusion', True),
+        use_mil_cross_shuffle=config.get('mil_use_cross_shuffle', True),
+    )
+
+
+def read_core_weights(path):
+    weights = torch.load(path, map_location='cpu', weights_only=True)
+    if not isinstance(weights, dict) or not weights or not all(torch.is_tensor(v) for v in weights.values()):
+        raise ValueError('Expected an LTF v3.3 core state_dict, such as epoch_20.pth')
+    return weights
+
+
+def create_model(config, weights=None, device='cpu'):
+    import clip
+    backbone = package_path(config['backbone'])
+    if not backbone.is_file():
+        raise FileNotFoundError(f'Bundled CLIP backbone not found: {backbone}')
+    # Local checkpoint loading is explicit: this path never downloads a model.
+    with contextlib.redirect_stdout(sys.stderr):
+        clip_model, _ = clip.load(str(backbone), device='cpu', jit=False)
+        model = LTFNetV3_3(clip_model.visual, **model_kwargs(config))
+    del clip_model
+    if weights is not None:
+        model.core.load_state_dict(read_core_weights(package_path(weights)), strict=True)
+    return model.to(device)
+
+
+def prepare_image(path, tile_size=224):
+    """Same RGB, tiling, Resize/CenterCrop and CLIP normalization as original eval."""
+    path = Path(path).expanduser().resolve()
+    with Image.open(path) as source:
+        image = source.convert('RGB')
+    transform = clip_transform(tile_size, blur_radius=None)
+    local_tiles = crop_image_to_tiles(image, tile_size)
+    return torch.stack([transform(image)] + [transform(tile) for tile in local_tiles])
+
+
+class LTFDetector:
+    def __init__(self, config=None, weights=None, device='auto', chunk_size=256):
+        self.config = dict(config) if isinstance(config, dict) else load_config(config)
+        if device == 'auto':
+            device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+        self.device = torch.device(device)
+        if self.device.type == 'cuda' and not torch.cuda.is_available():
+            raise RuntimeError('CUDA is unavailable; use device="cpu" for CPU inference')
+        if int(chunk_size) < 1:
+            raise ValueError('chunk_size must be positive')
+        torch.manual_seed(int(self.config.get('seed', 42)))
+        self.weights = package_path(weights or self.config.get('weights', 'weights/epoch_20.pth'))
+        self.model = create_model(self.config, self.weights, self.device).eval()
+        self.model.eval_chunk_size = int(chunk_size)
+
+    @torch.inference_mode()
+    def predict(self, path, branch='fusion', threshold=.5):
+        if branch not in BRANCHES:
+            raise ValueError(f'branch must be one of {BRANCHES}')
+        if not 0 <= threshold <= 1:
+            raise ValueError('threshold must be between zero and one')
+        path = Path(path).expanduser().resolve()
+        images = prepare_image(path, int(self.model.input_resolution))
+        self.model.eval()
+        output = self.model(images.unsqueeze(0).to(self.device))
+        probabilities = {}
+        for name in BRANCHES:
+            probability = output[f'logits_{name}'].float().softmax(-1)[0].cpu()
+            if not torch.isfinite(probability).all():
+                raise ValueError(f'Nonfinite {name} prediction for {path}')
+            probabilities[name] = {'real': float(probability[0]), 'fake': float(probability[1])}
+        score = probabilities[branch]['fake']
+        label = int(score >= threshold)
+        return dict(path=str(path), prediction='fake' if label else 'real', label=label,
+                    branch=branch, threshold=float(threshold), score_fake=score,
+                    views=int(images.shape[0]), probabilities=probabilities)

mil_module/__init__.py

@@ -0,0 +1,31 @@
+"""
+Standalone MIL 组件（Cross Shuffle、MIL Head、MILCore 串联），便于拷贝到其它工程。
+
+- 训练主仓 Level3 仍使用 ``mil.cross_shuffle`` + ``mil.mil_head``；本目录为 **导出副本**。
+- ``mil_module_export.MILHead`` 与 ``MILCore`` 默认使用 **Beta 头**（``mil_head_beta``），与 ``mil/mil_head.py`` 行为略有差异（见 README）。
+- 需与旧 checkpoint 键完全一致时，请使用 ``MILHeadLegacy``（``mil_head.py``）。
+
+Note: ``mil.fusion_module.HighResAggregate`` 使用 ``mil.cross_shuffle_v2``；迁移该路径请单独拷贝 v2。
+"""
+
+from .cross_shuffle import CrossTokenShuffle
+from .mil_head import MILHead as MILHeadLegacy
+from .mil_head_beta import MILHead
+from .mil_core import (
+    CrossShuffleStage,
+    CrossShuffleTransMIL,
+    MILCore,
+    MILHeadStage,
+    TransStage,
+)
+
+__all__ = [
+    "CrossTokenShuffle",
+    "MILHead",
+    "MILHeadLegacy",
+    "MILCore",
+    "CrossShuffleStage",
+    "TransStage",
+    "MILHeadStage",
+    "CrossShuffleTransMIL",
+]

mil_module/cross_shuffle.py

@@ -0,0 +1,145 @@
+import torch
+import torch.nn as nn
+from typing import Optional
+
+class CrossTokenShuffle(nn.Module):
+    """
+    CrossTokenShuffle: 跨图像 token 级 shuffle/replace 的增广模块（零 Python for-loop）
+
+    Donor 规则（与实现一致）：
+      - 候选 token **仅来自标签为 real（y==0）的图像** 的 token 池。
+      - **real 样本**：只能从「其他 real 图像」的 token 中抽取（排除本图）。
+      - **fake 样本**：只能从 **real 图像** 的 token 中抽取（fake 图自身无 real token 入池；
+        等价于不允许用 fake token 作 donor）。
+
+    替换数量：
+      - 先取 K = floor(N * replace_ratio)；若 ``min_keep_original > 0``，再令
+        ``K = min(K, N - min_keep_original)``，保证至少 ``min_keep_original`` 个位置仍为本图原 token。
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    @torch.no_grad()
+    def forward(
+        self,
+        x: torch.Tensor,              # (B, N, D)
+        y: torch.Tensor,              # (B,) 0/1
+        replace_ratio: float = 0.5,
+        with_replacement: bool = True,
+        detach_inserted: bool = True,
+        generator: Optional[torch.Generator] = None,
+        min_keep_original: int = 0,
+    ) -> torch.Tensor:
+        """
+        返回与 x 同形状的增广张量，严格遵循候选源与比例规则。
+        备注：无任何 Python 层的 for 循环；内部只用张量广播/索引/聚合算子。
+        """
+        assert x.dim() == 3, "x must be (B, N, D)"
+        B, N, D = x.shape
+        if B == 0 or N == 0:
+            return x
+
+        device = x.device
+        dtype = x.dtype
+
+        # 计算每图替换数 K
+        K = int(N * replace_ratio)
+        mk = int(min_keep_original)
+        if mk > 0:
+            K = min(K, max(0, N - mk))
+        if K <= 0:
+            return x  # 无需替换
+
+        # 展平成 token 表 (B*N, D)，并构造每个 token 的标签
+        x_flat = x.reshape(B * N, D)
+        y = y.to(device).view(B)  # (B,)
+        token_labels = y.repeat_interleave(N)  # (B*N,)
+
+        # 仅「真实图」的 token 可作为 donor（与 allowed_real 语义一致）。
+        # 用 M 列紧凑矩阵 (B, M) 代替 (B, B*N)：M = N * num_real <= B*N，fake 多时显著省显存/算力。
+        real_flat = torch.nonzero(token_labels == 0, as_tuple=False).view(-1)
+        M = int(real_flat.numel())
+        if M == 0:
+            return x.clone()
+
+        owner = real_flat // N  # (M,) donor token 所属 batch 下标（仅 y==0 的图）
+        b_ar = torch.arange(B, device=device, dtype=torch.long).unsqueeze(1)  # (B, 1)
+        rows_fake = (y == 1).unsqueeze(1)  # (B, 1)
+        # fake 行：donor 仅来自 real 池（任意 real 图的 token）；real 行：排除本图 real token
+        allowed_m = rows_fake | (owner.unsqueeze(0) != b_ar)  # (B, M)
+
+        has_cand = allowed_m.any(dim=1)  # (B,)
+
+        # —— 选择用于替换的“位置索引”（每图从 [0..N-1] 中采样 K 个，不放回）——
+        # 使用随机排序后取前 K（零循环、逐行独立）
+        r_pos = torch.rand(B, N, device=device, generator=generator)
+        replace_pos_idx = r_pos.argsort(dim=1)[:, :K].long()  # [FIX] 显式为 long  # (B, K), 每行无放回 K 个位置
+
+        # —— 选择用于插入的“候选 token 索引”（在扁平 token 表上的行号）——
+        # 两种策略：
+        #   - with_replacement=True：用 torch.multinomial（对 rows_with 有候选的子矩阵）；
+        #   - with_replacement=False：用 Gumbel-TopK（全矩阵），自然无放回。
+
+        # 预分配容器
+        cand_flat_idx = torch.zeros(B, K, device=device, dtype=torch.long)
+
+        if with_replacement:
+            rows_with = has_cand  # 只对有候选的行采样
+            if rows_with.any():
+                weights = allowed_m[rows_with].float()  # (Rw, M)
+                idx_m = torch.multinomial(
+                    weights, num_samples=K, replacement=True, generator=generator
+                )  # (Rw, K) in [0..M-1]
+                cand_flat_idx[rows_with] = real_flat[idx_m]
+            # rows_without: 保持 0（占位），后续不会写入（通过行掩码屏蔽）
+        else:
+            scores = torch.where(
+                allowed_m,
+                torch.zeros(B, M, device=device, dtype=torch.float32),
+                torch.full((B, M), float("-inf"), device=device, dtype=torch.float32),
+            )
+            U = torch.rand(B, M, device=device, generator=generator).clamp_(1e-6, 1 - 1e-6)
+            g = -torch.log(-torch.log(U))
+            cand_m = (scores + g).topk(k=K, dim=1).indices  # (B, K) in [0..M-1]
+            cand_flat_idx = real_flat[cand_m]
+
+        # 根据 cand_flat_idx 取回 token 内容 (B, K, D)
+        inserted_tokens = x_flat.index_select(0, cand_flat_idx.reshape(-1)).reshape(B, K, D)
+        if detach_inserted:
+            inserted_tokens = inserted_tokens.detach()
+
+        # 构造输出并写入被替换的位置（仅对有候选的行）
+        x_out = x.clone()
+        if has_cand.any():
+            # -------------------- 修  复  开  始 --------------------
+            rows = torch.where(has_cand)[0]                                   # [FIX] 用行号而不是布尔掩码
+            pos_idx = replace_pos_idx.index_select(0, rows)                   # (Rw, K)
+            batch_idx = rows.unsqueeze(1).expand(-1, pos_idx.size(1))         # (Rw, K) [FIX] 批次索引与位置索引同形状
+            ins = inserted_tokens.index_select(0, rows)                       # (Rw, K, D)
+
+            # 稳定的高级索引：两组索引 (batch_idx, pos_idx) 形状一致
+            x_out[batch_idx, pos_idx, :] = ins                                # [FIX] 正确回写，不再广播冲突
+            # -------------------- 修  复  结  束 --------------------
+
+        # 没有候选的行（如 B==1）保持原样
+        return x_out.to(dtype=dtype, device=device)
+
+
+# -------------------------
+# Quick self-check (toy)
+# -------------------------
+if __name__ == "__main__":
+    torch.manual_seed(0)
+    B, N, D = 256, 256, 1024
+    x = torch.arange(B*N*D, dtype=torch.float32).reshape(B, N, D)
+    y = torch.tensor([0, 1, 0, 1]*64)  # real, fake, real, fake
+
+    aug = CrossTokenShuffle()
+    out = aug(x, y, replace_ratio=0.5, with_replacement=False, detach_inserted=True)
+
+    # 形状不变
+    assert out.shape == x.shape
+    # 对 real 行（0/2），只应来自 real 池（行 0 和 2 的 token）且不能用自己：可在调试中检查
+    # 对 fake 行（1/3），可来自全池（排除自身）
+    print("OK, shape:", out.shape)

mil_module/mil_core.py

@@ -0,0 +1,219 @@
+"""
+Shuffle -> Trans -> MILHead，由 MILCore 串联；各段为独立 nn.Module 子类。
+
+MIL 头默认使用 ``mil_head_beta.MILHead``（attention 模式下支持 ``embed_cls``：
+先 bag embedding 再分类）。旧版行为见 ``mil_head.py`` 中的 ``MILHead``。
+"""
+from __future__ import annotations
+
+import inspect
+from typing import Any, Dict, Optional, Set, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from .cross_shuffle import CrossTokenShuffle
+from .mil_head_beta import MILHead
+
+
+def _forward_param_names(module: nn.Module) -> Set[str]:
+    try:
+        return set(inspect.signature(module.forward).parameters.keys())
+    except (TypeError, ValueError):
+        return set()
+
+
+class CrossShuffleStage(nn.Module):
+    """跨 batch token 替换增广；``use=False`` 时不建内部模块，forward 恒等。"""
+
+    def __init__(self, use: bool = True) -> None:
+        super().__init__()
+        self._use = use
+        self.op: Optional[CrossTokenShuffle]
+        if use:
+            self.op = CrossTokenShuffle()
+        else:
+            self.op = None
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        y: Optional[torch.Tensor],
+        *,
+        apply: bool,
+        replace_ratio: float = 0.5,
+        with_replacement: bool = False,
+        detach_inserted: bool = True,
+        min_keep_original: int = 0,
+    ) -> torch.Tensor:
+        if not self._use or not apply or self.op is None or y is None:
+            return x
+        if y.dim() != 1 or y.shape[0] != x.shape[0]:
+            raise ValueError(
+                f"y must be 1D [B], got {tuple(y.shape)} for B={x.shape[0]}"
+            )
+        return self.op(
+            x, y,
+            replace_ratio=replace_ratio,
+            with_replacement=with_replacement,
+            detach_inserted=detach_inserted,
+            min_keep_original=min_keep_original,
+        )
+
+
+class TransStage(nn.Module):
+    """可选中间变换；``core=None`` 时恒等。自动适配 ArtifactSetBranchB（p_shuffle / p_drop / compress_dropped_tokens）或 Encoder（src_key_padding_mask）。
+
+    固定 ``[B,N,D]`` 长度、仅用掩码做随机性时：传 ``key_padding_mask``（与 ``core`` 约定一致，一般为 True=padding 忽略），
+    勿依赖 ``p_drop`` 改长度；MIL 侧再用 ``MILHead`` 的 ``mask``（True=有效 token）。
+
+    若 ``core`` 返回 ``(x, mil_valid)``（例如压紧 token 后的分支），本阶段原样返回该二元组；
+    ``MILCore`` 会将 ``mil_valid`` 作为 MIL 的 ``mask``（与 ``mask`` 参数二选一语义见 ``MILCore.forward``）。
+    """
+
+    def __init__(self, core: Optional[nn.Module] = None) -> None:
+        super().__init__()
+        self.core = core
+        self._params = _forward_param_names(core) if core is not None else set()
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        *,
+        p_shuffle: float = 0.0,
+        p_drop: float = 0.0,
+        key_padding_mask: Optional[torch.Tensor] = None,
+        **kwargs: Any,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        if self.core is None:
+            return x
+        if "p_shuffle" in self._params:
+            call_kw: Dict[str, Any] = {}
+            for name, val in (
+                ("p_shuffle", p_shuffle),
+                ("p_drop", p_drop),
+                ("key_padding_mask", key_padding_mask),
+            ):
+                if name in self._params:
+                    call_kw[name] = val
+            for k, v in kwargs.items():
+                if k in self._params:
+                    call_kw[k] = v
+            return self.core(x, **call_kw)
+        kw2: Dict[str, Any] = {}
+        if key_padding_mask is not None and "src_key_padding_mask" in self._params:
+            kw2["src_key_padding_mask"] = key_padding_mask
+        for k, v in kwargs.items():
+            if k in self._params:
+                kw2[k] = v
+        return self.core(x, **kw2) if kw2 else self.core(x)
+
+
+class MILHeadStage(nn.Module):
+    """对 ``[B,N,D]`` 做实例分类 + bag 池化；内部为 ``mil_head_beta.MILHead``，forward 返回相同结构的 dict。"""
+
+    def __init__(self, **mil_kw: Any) -> None:
+        super().__init__()
+        self.head = MILHead(**mil_kw)
+
+    def forward(self, x: torch.Tensor, mask: Optional[torch.Tensor] = None) -> Dict[str, Any]:
+        return self.head(x, mask=mask)
+
+
+class MILCore(nn.Module):
+    """
+    ``shuffle`` → ``trans`` → ``head`` 三段组合，等价于原 CrossShuffleTransMIL。
+
+    **固定 token 数 N（常见导出用法）**：输入始终 ``[B, N, D]``，不需要在 trans 里做 random token drop。
+    随机性可用 **定长掩码**：``trans_key_padding_mask`` 控制 Transformer 忽略的位置；``mask`` 传给 MILHead（True=参与池化的有效 token）。
+    默认 ``allow_trans_token_drop=False``，forward 里传给 trans 的 ``p_drop`` 会被置 0；若要与训练期 L1 的 ``ArtifactSetBranchB`` 完全一致再打开。
+
+    ��� ``trans`` 返回 ``(x, mil_valid)``（压紧序列场景），``forward`` 会 **忽略** 入参 ``mask`` 并改用 ``mil_valid``；此时请勿再传 ``mask``。
+    其它传给中间模块的参数可通过 ``trans_extra`` 传入（如 ``compress_dropped_tokens=True``）。
+    """
+
+    def __init__(
+        self,
+        d_in: int,
+        num_classes: int = 2,
+        mil_head_pooling: str = "attn",
+        mil_head_kwargs: Optional[Dict[str, Any]] = None,
+        *,
+        use_cross_shuffle: bool = True,
+        trans: Optional[nn.Module] = None,
+        allow_trans_token_drop: bool = False,
+    ) -> None:
+        super().__init__()
+        self.allow_trans_token_drop = bool(allow_trans_token_drop)
+        self.shuffle = CrossShuffleStage(use=use_cross_shuffle)
+        self.trans = TransStage(core=trans)
+        mil_kw: Dict[str, Any] = {"d_in": d_in, "num_classes": num_classes, "pooling": mil_head_pooling}
+        if mil_head_kwargs:
+            mil_kw.update(mil_head_kwargs)
+        self.head = MILHeadStage(**mil_kw)
+
+    def forward(
+        self,
+        tokens: torch.Tensor,
+        y: Optional[torch.Tensor] = None,
+        mask: Optional[torch.Tensor] = None,
+        *,
+        do_cross_shuffle: bool = False,
+        replace_ratio: float = 0.5,
+        with_replacement: bool = False,
+        detach_inserted: bool = True,
+        cross_shuffle_min_keep: int = 0,
+        p_shuffle: float = 0.0,
+        p_drop: float = 0.0,
+        trans_key_padding_mask: Optional[torch.Tensor] = None,
+        trans_extra: Optional[Dict[str, Any]] = None,
+        return_dict: bool = False,
+    ):
+        x = self.shuffle(
+            tokens, y,
+            apply=do_cross_shuffle,
+            replace_ratio=replace_ratio,
+            with_replacement=with_replacement,
+            detach_inserted=detach_inserted,
+            min_keep_original=cross_shuffle_min_keep,
+        )
+        trans_p_drop = float(p_drop) if self.allow_trans_token_drop else 0.0
+        tex = dict(trans_extra or {})
+        trans_out = self.trans(
+            x,
+            p_shuffle=p_shuffle,
+            p_drop=trans_p_drop,
+            key_padding_mask=trans_key_padding_mask,
+            **tex,
+        )
+        if isinstance(trans_out, tuple) and len(trans_out) == 2:
+            x, trans_mil_mask = trans_out
+            if mask is not None:
+                raise ValueError(
+                    "MILCore: `trans` returned (x, mil_valid); omit `mask` and use the trans-stage mask."
+                )
+            mask = trans_mil_mask
+        else:
+            x = trans_out  # type: ignore[assignment]
+        if mask is not None and mask.shape != x.shape[:2]:
+            raise ValueError(
+                f"mask must be [B,N], got {tuple(mask.shape)} vs x {tuple(x.shape[:2])}"
+            )
+        mil_out = self.head(x, mask=mask)
+        if not return_dict:
+            return mil_out["bag_logits"]
+        return {"bag_logits": mil_out["bag_logits"], "mil_head": mil_out, "features": x}
+
+    def forward_logits(
+        self,
+        tokens: torch.Tensor,
+        y: Optional[torch.Tensor] = None,
+        mask: Optional[torch.Tensor] = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        extra = dict(kwargs)
+        extra.pop("return_dict", None)
+        return self.forward(tokens, y, mask, return_dict=False, **extra)
+
+
+CrossShuffleTransMIL = MILCore

mil_module/mil_head.py

@@ -0,0 +1,211 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class MILHead(nn.Module):
+    """
+    MIL Head supporting max / LogSumExp / top-k / attention-gated pooling.
+
+    Args:
+        d_in (int): token feature dim (ViT最后一层的D)
+        num_classes (int): 类别数；二分类用1更自然（配合 BCEWithLogits）
+        pooling (str): 'max' | 'lsep' | 'topk' | 'attn'
+        lsep_temp (float): LogSumExp 的温度 (softmax 温度的倒数); 有效时 bag = logsumexp(logits/T)
+        lsep_norm (bool): 是否减去 log(N) 做长度归一
+        topk (int | None): 固定 k；若为 None 则使用 ratio
+        topk_ratio (float): 使用比例选 top-k，例如 0.05 表示前 5%
+        attn_hidden (int): gated-attention 的隐层维度 L
+        attn_dropout (float): gated-attention 的 dropout
+    """
+    def __init__(
+        self,
+        d_in: int = 1024,
+        num_classes: int = 2,
+        pooling: str = "lsep",
+        lsep_temp: float = 1.0,
+        lsep_norm: bool = True,
+        topk: int = None,
+        topk_ratio: float = 0.1,
+        attn_hidden: int = 128,
+        attn_dropout: float = 0.0,
+    ):
+        super().__init__()
+        self.num_classes = num_classes
+        self.pooling = pooling.lower()
+        self.lsep_temp = float(lsep_temp)
+        self.lsep_norm = bool(lsep_norm)
+        self.topk = topk
+        self.topk_ratio = float(topk_ratio)
+
+        # per-token classifier (shared across tokens)
+        self.classifier = nn.Linear(d_in, num_classes)
+
+        # components for gated-attention MIL (Ilse et al., 2018)
+        if self.pooling == "attn":
+            print(f'Using gated-attention pooling with hidden size {attn_hidden}')
+            self.attn_V = nn.Linear(d_in, attn_hidden, bias=True)
+            self.attn_U = nn.Linear(d_in, attn_hidden, bias=True)
+            self.attn_w = nn.Linear(attn_hidden, 1, bias=False)
+            self.attn_drop = nn.Dropout(attn_dropout)
+
+    @staticmethod
+    def _masked_fill_min(x, mask, fill_value=None):
+        """Utility: set masked positions to a very small number before max/topk."""
+        if mask is None:
+            return x
+        if fill_value is None:
+            fill_value = torch.finfo(x.dtype).min
+        return x.masked_fill(~mask, fill_value)
+
+    def _pool_max(self, logits, mask=None):
+        # logits: [B, N, C] -> [B, C]
+        logits = self._masked_fill_min(logits, mask.unsqueeze(-1) if mask is not None else None)
+        bag = logits.max(dim=1).values
+        return bag
+
+    def _pool_lsep(self, logits, mask=None):
+        # logits: [B, N, C] -> [B, C]
+        # temperature T: divide logits by T, smaller T -> sharper
+        z = logits / self.lsep_temp
+        if mask is not None:
+            # 将被mask掉的设为 -inf，避免影响 logsumexp
+            neg_inf = torch.finfo(z.dtype).min
+            z = z.masked_fill(~mask.unsqueeze(-1), neg_inf)
+        bag = torch.logsumexp(z, dim=1)
+        if self.lsep_norm:
+            # 归一：减去 log(有效实例数)
+            if mask is None:
+                n = logits.size(1)
+                bag = bag - math.log(max(n, 1))
+            else:
+                n_eff = mask.sum(dim=1).clamp_min(1).unsqueeze(-1)  # [B,1]
+                bag = bag - n_eff.log()
+        return bag
+
+    def _pool_topk(self, logits, mask=None):
+        # logits: [B, N, C] -> [B, C]
+        B, N, C = logits.shape
+        if mask is not None:
+            logits = self._masked_fill_min(logits, mask.unsqueeze(-1))
+            n_eff = mask.sum(dim=1)  # [B]
+        else:
+            n_eff = logits.new_full((B,), N)
+
+        k = self.topk
+        if k is None:
+            # ratio -> at least 1
+            k = torch.clamp((n_eff.float() * self.topk_ratio).long(), min=1)
+        else:
+            k = torch.clamp(torch.full_like(n_eff, k), min=1, max=N)
+
+        # 为了支持不同 batch 内不同 k，这里用排序 + gather
+        # scores: [B, N, C] -> 先在 N 维排序，取 top-k 再平均
+        # 我们对每个类独立做 topk：展开为 [B, C, N]
+        scores = logits.transpose(1, 2)  # [B, C, N]
+        # 对每个 (B,C) 维度上选择前 k 个
+        topk_vals = []
+        for b in range(B):
+            # 如果某些样本 k > 有效实例数，用有效实例数
+            kb = int(min(int(k[b].item()), int(n_eff[b].item())))
+            vals, _ = scores[b].topk(kb, dim=-1)
+            topk_vals.append(vals.mean(dim=-1, keepdim=True))  # [C,1]
+        bag = torch.cat(topk_vals, dim=1).transpose(0,1).contiguous().squeeze(1)  # [B, C]
+        return bag
+    # v1
+    # def _pool_attn(self, feats, logits, mask=None):
+    #     """
+    #     Gated-attention pooling:
+    #         a_i ∝ w^T ( tanh(V h_i) ⊙ sigmoid(U h_i) )
+    #         bag = Σ a_i * logits_i
+    #     feats:  [B, N, D]
+    #     logits: [B, N, C]
+    #     mask:   [B, N] (True for valid)
+    #     Returns:
+    #         bag_logits: [B, C]
+    #         attn:       [B, N, 1]
+    #     """
+    #     Vh = self.attn_V(feats)           # [B, N, L]
+    #     Uh = self.attn_U(feats)           # [B, N, L]
+    #     gate = torch.tanh(Vh) * torch.sigmoid(Uh)
+    #     gate = self.attn_drop(gate)
+    #     attn = self.attn_w(gate)          # [B, N, 1]
+
+    #     if mask is not None:
+    #         neg_inf = torch.finfo(attn.dtype).min
+    #         attn = attn.masked_fill(~mask.unsqueeze(-1), neg_inf)
+
+    #     attn = F.softmax(attn, dim=1)     # [B, N, 1]
+    #     bag = (attn * logits).sum(dim=1)  # [B, C]
+    #     return bag, attn
+    # v2 新增了返回加权之后的embedding
+    def _pool_attn(self, feats, logits, mask=None):
+        """
+        Gated-attention pooling:
+            a_i ∝ w^T ( tanh(V h_i) ⊙ sigmoid(U h_i) )
+            bag_logits = Σ a_i * logits_i
+            bag_embed  = Σ a_i * h_i
+        feats:  [B, N, D]
+        logits: [B, N, C]
+        mask:   [B, N] (True for valid)
+        Returns:
+            bag_logits: [B, C]
+            attn:       [B, N, 1]
+            bag_embed:  [B, D]   (mix之后的embedding)
+        """
+        Vh = self.attn_V(feats)           # [B, N, L]
+        Uh = self.attn_U(feats)           # [B, N, L]
+        gate = torch.tanh(Vh) * torch.sigmoid(Uh)
+        gate = self.attn_drop(gate)
+        attn = self.attn_w(gate)          # [B, N, 1]
+
+        if mask is not None:
+            neg_inf = torch.finfo(attn.dtype).min
+            attn = attn.masked_fill(~mask.unsqueeze(-1), neg_inf)
+
+        attn = F.softmax(attn, dim=1)        # [B, N, 1]
+
+        # 原来的 bag-level logits
+        bag_logits = (attn * logits).sum(dim=1)  # [B, C]
+
+        # 新增：mix 之后的 bag-level embedding
+        bag_embed = (attn * feats).sum(dim=1)    # [B, D]
+
+        return bag_logits, attn, bag_embed
+    def forward(self, feats, mask=None):
+        """
+        Args:
+            feats: [B, N, D]  — ViT最后一层每patch的token特征
+            mask: [B, N]     — 可选，True表示有效token（padding用）
+        Returns:
+            dict{
+              'inst_logits': [B, N, C],
+              'bag_logits':  [B, C],
+              'attn':        [B, N, 1] or None
+            }
+        """
+        inst_logits = self.classifier(feats)  # [B, N, C]
+
+        if self.pooling == "max":
+            bag_logits = self._pool_max(inst_logits, mask)
+            attn = None
+        elif self.pooling == "lsep":
+            bag_logits = self._pool_lsep(inst_logits, mask)
+            attn = None
+        elif self.pooling == "topk":
+            bag_logits = self._pool_topk(inst_logits, mask)
+            attn = None
+        elif self.pooling == "attn":
+            # v1
+            # bag_logits, attn = self._pool_attn(feats, inst_logits, mask)
+            # v2
+            bag_logits, attn, bag_embed = self._pool_attn(feats, inst_logits, mask) # 返回bag_embed是为了后续的聚合
+        else:
+            raise ValueError(f"Unknown pooling mode: {self.pooling}")
+
+        return {
+            "inst_logits": inst_logits,
+            "bag_logits": bag_logits,
+            "attn": attn,
+            "bag_embed": bag_embed if self.pooling == "attn" else None,
+        }

mil_module/mil_head_beta.py

@@ -0,0 +1,212 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class MILHead(nn.Module):
+    """
+    MIL Head supporting max / LogSumExp / top-k / attention-gated pooling.
+
+    Args:
+        d_in (int): token feature dim
+        num_classes (int): number of classes
+        pooling (str): 'max' | 'lsep' | 'topk' | 'attn'
+        lsep_temp (float): temperature for LogSumExp pooling
+        lsep_norm (bool): whether to subtract log(N_eff)
+        topk (int | None): fixed k for top-k pooling
+        topk_ratio (float): ratio for top-k if topk is None
+        attn_hidden (int): hidden dim for gated-attention
+        attn_dropout (float): dropout for gated-attention
+        attn_bag_mode (str): only used when pooling='attn'
+            - 'logit_pool': bag_logits = sum(attn * inst_logits)
+            - 'embed_cls' : bag_embed = sum(attn * feats), then bag_logits = bag_classifier(bag_embed)
+    """
+
+    def __init__(
+        self,
+        d_in: int = 1024,
+        num_classes: int = 2,
+        pooling: str = "lsep",
+        lsep_temp: float = 1.0,
+        lsep_norm: bool = True,
+        topk: int = None,
+        topk_ratio: float = 0.1,
+        attn_hidden: int = 128,
+        attn_dropout: float = 0.0,
+        attn_bag_mode: str = "embed_cls",
+    ):
+        super().__init__()
+        self.num_classes = num_classes
+        self.pooling = pooling.lower()
+        self.lsep_temp = float(lsep_temp)
+        self.lsep_norm = bool(lsep_norm)
+        self.topk = topk
+        self.topk_ratio = float(topk_ratio)
+        self.attn_bag_mode = attn_bag_mode.lower()
+
+        if self.pooling not in {"max", "lsep", "topk", "attn"}:
+            raise ValueError(f"Unknown pooling mode: {self.pooling}")
+
+        if self.attn_bag_mode not in {"logit_pool", "embed_cls"}:
+            raise ValueError(f"Unknown attn_bag_mode: {self.attn_bag_mode}")
+
+        # Instance-level classifier: shared across tokens
+        self.classifier = nn.Linear(d_in, num_classes)
+
+        # Attention pooling modules
+        if self.pooling == "attn":
+            print(
+                f"Using gated-attention pooling with hidden size {attn_hidden}, "
+                f"bag mode = {self.attn_bag_mode}"
+            )
+            self.attn_V = nn.Linear(d_in, attn_hidden, bias=True)
+            self.attn_U = nn.Linear(d_in, attn_hidden, bias=True)
+            self.attn_w = nn.Linear(attn_hidden, 1, bias=False)
+            self.attn_drop = nn.Dropout(attn_dropout)
+
+            # Decoupled bag classifier
+            if self.attn_bag_mode == "embed_cls":
+                self.bag_classifier = nn.Linear(d_in, num_classes)
+            else:
+                self.bag_classifier = None
+        else:
+            self.bag_classifier = None
+
+    @staticmethod
+    def _masked_fill_min(x, mask, fill_value=None):
+        """Set masked positions to a very small number before max/topk."""
+        if mask is None:
+            return x
+        if fill_value is None:
+            fill_value = torch.finfo(x.dtype).min
+        return x.masked_fill(~mask, fill_value)
+
+    def _pool_max(self, logits, mask=None):
+        logits = self._masked_fill_min(
+            logits, mask.unsqueeze(-1) if mask is not None else None
+        )
+        bag = logits.max(dim=1).values
+        return bag
+
+    def _pool_lsep(self, logits, mask=None):
+        z = logits / self.lsep_temp
+        if mask is not None:
+            neg_inf = torch.finfo(z.dtype).min
+            z = z.masked_fill(~mask.unsqueeze(-1), neg_inf)
+
+        bag = torch.logsumexp(z, dim=1)
+
+        if self.lsep_norm:
+            if mask is None:
+                n = logits.size(1)
+                bag = bag - math.log(max(n, 1))
+            else:
+                n_eff = mask.sum(dim=1).clamp_min(1).unsqueeze(-1)
+                bag = bag - n_eff.log()
+
+        return bag
+
+    def _pool_topk(self, logits, mask=None):
+        B, N, C = logits.shape
+
+        if mask is not None:
+            logits = self._masked_fill_min(logits, mask.unsqueeze(-1))
+            n_eff = mask.sum(dim=1)
+        else:
+            n_eff = logits.new_full((B,), N, dtype=torch.long)
+
+        k = self.topk
+        if k is None:
+            k = torch.clamp((n_eff.float() * self.topk_ratio).long(), min=1)
+        else:
+            k = torch.clamp(torch.full_like(n_eff, k), min=1, max=N)
+
+        scores = logits.transpose(1, 2)  # [B, C, N]
+        topk_vals = []
+        for b in range(B):
+            kb = int(min(int(k[b].item()), int(n_eff[b].item())))
+            vals, _ = scores[b].topk(kb, dim=-1)      # [C, kb]
+            topk_vals.append(vals.mean(dim=-1))       # [C]
+
+        bag = torch.stack(topk_vals, dim=0)           # [B, C]
+        return bag
+
+    def _compute_attn(self, feats, mask=None):
+        """
+        feats: [B, N, D]
+        mask : [B, N]
+        returns:
+            attn: [B, N, 1]
+        """
+        Vh = self.attn_V(feats)                       # [B, N, L]
+        Uh = self.attn_U(feats)                       # [B, N, L]
+        gate = torch.tanh(Vh) * torch.sigmoid(Uh)
+        gate = self.attn_drop(gate)
+        attn = self.attn_w(gate)                      # [B, N, 1]
+
+        if mask is not None:
+            neg_inf = torch.finfo(attn.dtype).min
+            attn = attn.masked_fill(~mask.unsqueeze(-1), neg_inf)
+
+        attn = F.softmax(attn, dim=1)
+        return attn
+
+    def _pool_attn(self, feats, logits, mask=None):
+        """
+        Returns:
+            bag_logits: [B, C]
+            attn:       [B, N, 1]
+            bag_embed:  [B, D]
+        """
+        attn = self._compute_attn(feats, mask)         # [B, N, 1]
+        bag_embed = (attn * feats).sum(dim=1)          # [B, D]
+
+        if self.attn_bag_mode == "logit_pool":
+            bag_logits = (attn * logits).sum(dim=1)    # [B, C]
+        elif self.attn_bag_mode == "embed_cls":
+            bag_logits = self.bag_classifier(bag_embed)  # [B, C]
+        else:
+            raise ValueError(f"Unknown attn_bag_mode: {self.attn_bag_mode}")
+
+        return bag_logits, attn, bag_embed
+
+    def forward(self, feats, mask=None):
+        """
+        Args:
+            feats: [B, N, D]
+            mask:  [B, N], True means valid token
+        Returns:
+            dict{
+                'inst_logits': [B, N, C],
+                'bag_logits':  [B, C],
+                'attn':        [B, N, 1] or None,
+                'bag_embed':   [B, D] or None
+            }
+        """
+        inst_logits = self.classifier(feats)  # [B, N, C]
+
+        attn = None
+        bag_embed = None
+
+        if self.pooling == "max":
+            bag_logits = self._pool_max(inst_logits, mask)
+
+        elif self.pooling == "lsep":
+            bag_logits = self._pool_lsep(inst_logits, mask)
+
+        elif self.pooling == "topk":
+            bag_logits = self._pool_topk(inst_logits, mask)
+
+        elif self.pooling == "attn":
+            bag_logits, attn, bag_embed = self._pool_attn(feats, inst_logits, mask)
+
+        else:
+            raise ValueError(f"Unknown pooling mode: {self.pooling}")
+
+        return {
+            "inst_logits": inst_logits,
+            "bag_logits": bag_logits,
+            "attn": attn,
+            "bag_embed": bag_embed,
+        }

model/LTF.py

@@ -0,0 +1,206 @@
+import torch
+import torch.nn as nn
+
+class LTFCore(nn.Module):
+    """
+    Core module.
+    Input:  x (B, N, T, D)
+    Output: logits (B, num_classes), token_scores (B, N, T-1, 1), jpeg_score (B, 1)
+    """
+    def __init__(
+        self,
+        embed_dim=1024,
+        num_heads=8,
+        num_classes=2,
+        num_layers=1,
+        dropout=0.1,
+        tile_heads=2,
+        tile_layers=2,
+        use_global_image=True,
+        use_patch_transformer=True,
+        linear_classifier=False,
+    ):
+        super().__init__()
+        self.use_global_image = use_global_image
+        self.use_patch_transformer = use_patch_transformer
+        # Learnable CLS for tile-level aggregation
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        # Tile-level encoder (across tiles)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=embed_dim,
+            dropout=dropout,
+            nhead=tile_heads,
+            batch_first=True  
+        )
+        self.tile_transformer = nn.TransformerEncoder(encoder_layer, num_layers=tile_layers)
+        # Token-level encoder (within each tile); v1_7 等可关闭以降低容量
+        if use_patch_transformer:
+            token_layer = nn.TransformerEncoderLayer(
+                d_model=embed_dim,
+                dropout=dropout,
+                nhead=num_heads,
+                batch_first=True 
+            )
+            self.patch_transformer = nn.TransformerEncoder(token_layer, num_layers=num_layers)
+        else:
+            self.patch_transformer = None
+        # Image-level classifier: with global (main tile + aggregated tiles) or only aggregated tiles (ablation)
+        classifier_in_dim = embed_dim * 2 if use_global_image else embed_dim
+        if linear_classifier:
+            self.classifier = nn.Linear(classifier_in_dim, num_classes)
+        else:
+            self.classifier = nn.Sequential(
+                nn.Linear(classifier_in_dim, 512),
+                nn.ReLU(),
+                nn.Dropout(p=dropout),
+                nn.Linear(512, num_classes)
+            )
+        # Per-token head (for non-CLS tokens)
+        self.patch_classifier = nn.Sequential(
+            nn.Linear(embed_dim, 512),
+            nn.ReLU(),
+            nn.Dropout(p=dropout),
+            nn.Linear(512, 1)
+        )
+        # Scalar head (e.g., JPEG-related score)
+        self.jpeg_estimate = nn.Sequential(
+            nn.Linear(embed_dim, 256),
+            nn.ReLU(),
+            nn.Linear(256, 128),
+            nn.ReLU(),
+            nn.Linear(128, 128),
+            nn.ReLU(),
+            nn.Linear(128, 1),
+        )
+
+        nn.init.trunc_normal_(self.cls_token, std=0.02)
+    
+    def forward(self, x):
+        """
+        x: (B, N, T, D)
+        returns: (B, num_classes), (B, N, T-1, 1), (B, 1)
+        """
+        B_2, N_2, T_2, D_2 = x.shape
+
+        # Token-level encoding per tile (optional)
+        x = x.reshape(B_2*N_2, T_2, D_2)       # [B*N, T, D]
+        if self.patch_transformer is not None:
+            x = self.patch_transformer(x)
+        x = x.reshape(B_2, N_2, T_2, D_2)
+        
+        # Split CLS vs non-CLS tokens
+        new_cls = x[:, :, 0, :]
+        new_x = x[:, :, 1:, :]
+        # Per-token head on non-CLS tokens
+        x_flat = new_x.reshape(-1, D_2)       # [B*N*T, D]
+        x_result = self.patch_classifier(x_flat)       # [B*N*(T-1), 1]
+        x_result = x_result.reshape(B_2, N_2, T_2 - 1, 1)   # [B, N, T-1, 1]
+
+        main_feat = new_cls[:, 0]              # (B, D)
+        tile_feats = new_cls[:, 1:]            # (B, N-1, D)
+        # print(main_feat.shape, tile_feats.shape)
+
+        # Tile-level aggregation with a learnable CLS
+        cls_token = self.cls_token.expand(B_2, 1, -1)              # (B, 1, D)
+        feature_tile = torch.cat([cls_token, tile_feats], dim=1)       # (B, N, D)
+
+        feature_encoded = self.tile_transformer(feature_tile)                     # (B, N, D)
+        cls_out = feature_encoded[:, 0]                                # (B, D)
+
+        # Image-level classification: with or without global (main tile) feature (ablation v1_1)
+        if self.use_global_image:
+            concat_feat = torch.cat([main_feat, cls_out], dim=1)     # (B, 2D)
+        else:
+            concat_feat = cls_out                                      # (B, D)
+        return self.classifier(concat_feat), x_result, torch.sigmoid(self.jpeg_estimate(cls_out))  # (B, num_classes) (B)
+
+class LTFNet(nn.Module):
+    """
+    Wrapper with a frozen visual backbone (e.g., CLIP visual) + LTFCore.
+    Input:  images x (B, N, 3, H, W)
+    Output: same as LTFCore
+    """
+    def __init__(
+        self,
+        clip_visual,
+        num_heads=8,
+        num_classes=2,
+        num_layers=1,
+        dropout=0.2,
+        tile_heads=2,
+        tile_layers=2,
+        use_global_image=True,
+        use_patch_transformer=True,
+        linear_classifier=False,
+    ):
+        super().__init__()
+        self.visual = clip_visual 
+        self.dtype = self.visual.conv1.weight.dtype
+
+        for p in self.visual.parameters():
+            p.requires_grad = False
+        self.visual.eval()
+
+        self.input_resolution = self.visual.input_resolution
+        # Token feature dim = transformer hidden size (ViT-L -> 1024). Use conv1 out_channels;
+        # do NOT use output_dim which is the projection-head dim (often 768) and would cause
+        # "expecting 768 but got 1024" in patch_transformer.
+        self.embed_dim = self.visual.conv1.weight.shape[0]
+
+        self.core = LTFCore(
+            embed_dim=self.embed_dim,
+            num_heads=num_heads,
+            num_classes=num_classes,
+            num_layers=num_layers,
+            dropout=dropout,
+            tile_heads=tile_heads,
+            tile_layers=tile_layers,
+            use_global_image=use_global_image,
+            use_patch_transformer=use_patch_transformer,
+            linear_classifier=linear_classifier,
+        )
+
+    def get_tokens(self, x) -> torch.Tensor:
+        """
+        Extract token sequence (with CLS) from the visual backbone.
+        Assumes CLIP ViT structure: conv1 -> patch embed -> class_embedding + pos_embed -> ln_pre -> transformer.
+        x: (B, 3, H, W) -> tokens: (B, T, D)
+        """
+        x = x.type(self.dtype)
+        x = self.visual.conv1(x)  # [B, C, H, W]
+        x = x.reshape(x.shape[0], x.shape[1], -1)  # [B, C, HW]
+        x = x.permute(0, 2, 1)  # [B, HW, C]
+        cls = self.visual.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[2], dtype=x.dtype, device=x.device)
+        x = torch.cat([cls, x], dim=1)  # [B, 1+HW, C]
+        x = x + self.visual.positional_embedding.to(x.dtype)
+        x = self.visual.ln_pre(x)
+        x = x.permute(1, 0, 2)  # [sequence_len, B, C]
+        x = self.visual.transformer(x)
+        x = x.permute(1, 0, 2)  # [B, sequence_len, C]
+        return x 
+
+    def encode_image(self, image) -> torch.Tensor:
+        # image: (B, 3, H, W) -> feature: (B, D)
+        return self.visual(image.type(self.dtype))
+
+    def forward(self, x: torch.Tensor, chunk_size=256):
+        # Tensor[B, 5, 3, 224, 224]
+        """
+        x: Tensor of shape (B, N, 3, H, W)
+        chunk_size: number of tiles to pass through the visual backbone per batch. Same tiles are
+                    computed and concatenated; only speed/memory change, no effect on outputs.
+                    Eval can set model.eval_chunk_size (e.g. 512) in config for faster run.
+        """
+        chunk_size = getattr(self, "eval_chunk_size", chunk_size)
+        B, N, C, H, W = x.shape
+        imgs = x.reshape(B * N, C, H, W)
+        feats_chunks = []
+        with torch.no_grad():
+            for start in range(0, imgs.size(0), chunk_size):
+                end = min(start + chunk_size, imgs.size(0))
+                feats_chunk = self.get_tokens(imgs[start:end])   # (chunk, T, D)
+                feats_chunks.append(feats_chunk)
+        feats_token = torch.cat(feats_chunks, dim=0)
+        feats = feats_token.reshape(B, N, feats_token.shape[-2], -1)  # (B, N, T, D)
+        feats = feats.float()
+        return self.core(feats)          # logits: (B, num_classes)  (B)

model/LTF_v2.py

@@ -0,0 +1,427 @@
+"""
+LTF V2: multi-branch (baseline + Level-1 MIL on patch tokens + Level-2 MIL on CLS + fusion).
+V1 code in LTF.py is unchanged; this module is a separate copy for version V2.
+"""
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from model.LTF import LTFCore
+from mil_module.mil_core import MILCore
+
+
+def _clip_vit_resblock_modules(transformer) -> list:
+    """OpenAI CLIP ViT: ``transformer.resblocks`` 为 Sequential 或 ModuleList。"""
+    rb = transformer.resblocks
+    if isinstance(rb, nn.Sequential):
+        return list(rb.children())
+    return list(rb)
+
+
+def _count_clip_vit_depth(visual) -> int:
+    return len(_clip_vit_resblock_modules(visual.transformer))
+
+
+class BagTokenShuffle(nn.Module):
+    """袋内 token 随机重排；训练时以概率 p_shuffle 触发（由 MILCore 传入）。"""
+
+    def forward(self, x, p_shuffle: float = 0.0, p_drop: float = 0.0, key_padding_mask=None, **kwargs):
+        del p_drop, key_padding_mask, kwargs
+        if (not self.training) or p_shuffle <= 0:
+            return x
+        if float(torch.rand((), device=x.device)) >= float(p_shuffle):
+            return x
+        B, N, D = x.shape
+        if N <= 1:
+            return x
+        perm = torch.rand(B, N, device=x.device).argsort(dim=1)
+        return x.gather(1, perm.unsqueeze(-1).expand(B, N, D))
+
+
+class LTFCoreV2(nn.Module):
+    """
+    Baseline: identical logic to LTFCore (shared weights via inner LTFCore modules).
+    Level 1: per-tile patch tokens (no CLS, no global tile) -> MILCore -> Nt bag embeddings
+             -> learnable query token + MultiheadAttention -> level1 embedding（要求 N>=2）。
+    Level 2: all CLS tokens (global + every tile) -> MILCore -> bag embedding.
+    Fusion: MLP on concat(level1, level2, baseline_concat_feat).
+    """
+
+    def __init__(
+        self,
+        embed_dim=1024,
+        num_heads=8,
+        num_classes=2,
+        num_layers=1,
+        dropout=0.1,
+        tile_heads=2,
+        tile_layers=2,
+        use_global_image=True,
+        use_patch_transformer=True,
+        linear_classifier=False,
+        mil_heads=8,
+        fusion_hidden=512,
+        # Level-1 CrossTokenShuffle (仅 MIL 支路；baseline 不经过 shuffle / replace)
+        l1_cross_shuffle_p: float = 0.9,
+        l1_replace_ratio_range: Tuple[float, float] = (0.5, 0.9),
+        l1_with_replacement_p: float = 0.5,
+        # Level-2 CrossTokenShuffle
+        l2_cross_shuffle_p: float = 0.5,
+        l2_replace_ratio_range: Tuple[float, float] = (0.1, 0.5),
+        l2_with_replacement_p: float = 0.5,
+        bag_token_shuffle_p: float = 0.5,
+        cross_shuffle_min_keep_original: int = 0,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.use_global_image = use_global_image
+        self.num_classes = num_classes
+
+        self.l1_cross_shuffle_p = float(l1_cross_shuffle_p)
+        self.l1_replace_ratio_range = (float(l1_replace_ratio_range[0]), float(l1_replace_ratio_range[1]))
+        self.l1_with_replacement_p = float(l1_with_replacement_p)
+        self.l2_cross_shuffle_p = float(l2_cross_shuffle_p)
+        self.l2_replace_ratio_range = (float(l2_replace_ratio_range[0]), float(l2_replace_ratio_range[1]))
+        self.l2_with_replacement_p = float(l2_with_replacement_p)
+        self.bag_token_shuffle_p = float(bag_token_shuffle_p)
+        self.cross_shuffle_min_keep_original = int(cross_shuffle_min_keep_original)
+
+        self.baseline_core = LTFCore(
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            num_classes=num_classes,
+            num_layers=num_layers,
+            dropout=dropout,
+            tile_heads=tile_heads,
+            tile_layers=tile_layers,
+            use_global_image=use_global_image,
+            use_patch_transformer=use_patch_transformer,
+            linear_classifier=linear_classifier,
+        )
+
+        mil_kw = dict(
+            attn_hidden=min(256, embed_dim),
+            attn_dropout=dropout,
+            attn_bag_mode="embed_cls",
+        )
+        self.mil_level1 = MILCore(
+            d_in=embed_dim,
+            num_classes=num_classes,
+            mil_head_pooling="attn",
+            mil_head_kwargs=mil_kw,
+            use_cross_shuffle=True,
+            trans=BagTokenShuffle(),
+            allow_trans_token_drop=False,
+        )
+        self.mil_level2 = MILCore(
+            d_in=embed_dim,
+            num_classes=num_classes,
+            mil_head_pooling="attn",
+            mil_head_kwargs=mil_kw,
+            use_cross_shuffle=True,
+            trans=BagTokenShuffle(),
+            allow_trans_token_drop=False,
+        )
+
+        self.level1_query = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        mh = mil_heads if embed_dim % mil_heads == 0 else 8
+        while embed_dim % mh != 0:
+            mh -= 1
+        if mh < 1:
+            mh = 1
+        self.level1_aggregate_attn = nn.MultiheadAttention(
+            embed_dim, mh, dropout=dropout, batch_first=True
+        )
+
+        self.head_level1 = nn.Linear(embed_dim, num_classes)
+        self.head_level2 = nn.Linear(embed_dim, num_classes)
+
+        classifier_in_dim = embed_dim * 2 if use_global_image else embed_dim
+        fusion_in = embed_dim + embed_dim + classifier_in_dim
+        self.fusion_head = nn.Sequential(
+            nn.Linear(fusion_in, fusion_hidden),
+            nn.ReLU(),
+            nn.Dropout(p=dropout),
+            nn.Linear(fusion_hidden, num_classes),
+        )
+
+        nn.init.trunc_normal_(self.level1_query, std=0.02)
+
+    @staticmethod
+    def _sample_replace_ratio(low: float, high: float, device: torch.device) -> float:
+        lo, hi = min(low, high), max(low, high)
+        lo = max(0.0, lo)
+        hi = min(1.0, hi)
+        if hi <= lo:
+            return float(lo)
+        u = torch.rand((), device=device)
+        return float(lo + (hi - lo) * u)
+
+    def _apply_patch_transformer(self, x):
+        """x: (B, N, T, D) -> same shape after optional token-level encoder."""
+        B_2, N_2, T_2, D_2 = x.shape
+        xf = x.reshape(B_2 * N_2, T_2, D_2)
+        if self.baseline_core.patch_transformer is not None:
+            xf = self.baseline_core.patch_transformer(xf)
+        return xf.reshape(B_2, N_2, T_2, D_2)
+
+    def _forward_baseline_from_xfeat(self, x_feat):
+        """Run LTF classifier heads on processed tokens (same as LTFCore tail)."""
+        B_2, N_2, T_2, D_2 = x_feat.shape
+        bc = self.baseline_core
+        new_cls = x_feat[:, :, 0, :]
+        new_x = x_feat[:, :, 1:, :]
+        x_flat = new_x.reshape(-1, D_2)
+        x_result = bc.patch_classifier(x_flat)
+        x_result = x_result.reshape(B_2, N_2, T_2 - 1, 1)
+
+        main_feat = new_cls[:, 0]
+        tile_feats = new_cls[:, 1:]
+        cls_token = bc.cls_token.expand(B_2, 1, -1)
+        feature_tile = torch.cat([cls_token, tile_feats], dim=1)
+        feature_encoded = bc.tile_transformer(feature_tile)
+        cls_out = feature_encoded[:, 0]
+
+        if bc.use_global_image:
+            concat_feat = torch.cat([main_feat, cls_out], dim=1)
+        else:
+            concat_feat = cls_out
+        logits_b = bc.classifier(concat_feat)
+        jpeg_b = torch.sigmoid(bc.jpeg_estimate(cls_out))
+        return logits_b, x_result, jpeg_b, concat_feat
+
+    def forward(
+        self,
+        x,
+        labels: Optional[torch.Tensor] = None,
+        x_mil: Optional[torch.Tensor] = None,
+    ):
+        """
+        x: (B, N, T, D)，要求 N>=2（global + 至少一个 tile），否则 Level1 无 patch 袋。
+        x_mil: 可选，与 x 同形状；若提供则 Level1/Level2 用其 token（如 CLIP 中间层），baseline 仍用 x。
+        labels: (B,) 0/1，训练时用于 Level1/2 CrossTokenShuffle；baseline 路径不使用 replace/shuffle。
+        Returns dict with logits per branch + baseline aux outputs.
+        """
+        x_feat = self._apply_patch_transformer(x)
+        if x_mil is not None:
+            if x_mil.shape != x.shape:
+                raise ValueError(f"x_mil shape {tuple(x_mil.shape)} != x shape {tuple(x.shape)}")
+            x_mil_feat = self._apply_patch_transformer(x_mil)
+        else:
+            x_mil_feat = x_feat
+        logits_baseline, patch_logits, jpeg_pred, concat_feat = self._forward_baseline_from_xfeat(
+            x_feat
+        )
+
+        B, N, T, D = x_mil_feat.shape
+        if N < 2:
+            raise ValueError(f"LTFCoreV2 expects N>=2 (global + tiles), got N={N}")
+        Nt = N - 1
+        train_shuffle = self.training and labels is not None
+        p_bag_shuffle = self.bag_token_shuffle_p if train_shuffle else 0.0
+
+        patch_only = x_mil_feat[:, 1:, 1:, :].reshape(B * Nt, T - 1, D)
+        do1 = (
+            train_shuffle
+            and float(torch.rand((), device=x.device)) < self.l1_cross_shuffle_p
+        )
+        rr1 = self._sample_replace_ratio(*self.l1_replace_ratio_range, x.device)
+        wr1 = bool(float(torch.rand((), device=x.device)) < self.l1_with_replacement_p)
+        y1 = labels.unsqueeze(1).expand(-1, Nt).reshape(-1) if labels is not None else None
+        mk = self.cross_shuffle_min_keep_original
+        mil1 = self.mil_level1(
+            patch_only,
+            y=y1,
+            return_dict=True,
+            do_cross_shuffle=do1,
+            replace_ratio=rr1,
+            with_replacement=wr1,
+            detach_inserted=True,
+            cross_shuffle_min_keep=mk,
+            p_shuffle=p_bag_shuffle,
+        )
+        be = mil1["mil_head"]["bag_embed"]
+        if be is None:
+            raise RuntimeError("MIL level1 expected bag_embed (use attn + embed_cls).")
+        bag_seq = be.reshape(B, Nt, D)
+        q = self.level1_query.expand(B, -1, -1)
+        agg_out, _ = self.level1_aggregate_attn(q, bag_seq, bag_seq)
+        level1_emb = agg_out.squeeze(1)
+
+        cls_seq = x_mil_feat[:, :, 0, :]
+        do2 = (
+            train_shuffle
+            and float(torch.rand((), device=x.device)) < self.l2_cross_shuffle_p
+        )
+        rr2 = self._sample_replace_ratio(*self.l2_replace_ratio_range, x.device)
+        wr2 = bool(float(torch.rand((), device=x.device)) < self.l2_with_replacement_p)
+        mil2 = self.mil_level2(
+            cls_seq,
+            y=labels,
+            return_dict=True,
+            do_cross_shuffle=do2,
+            replace_ratio=rr2,
+            with_replacement=wr2,
+            detach_inserted=True,
+            cross_shuffle_min_keep=mk,
+            p_shuffle=p_bag_shuffle,
+        )
+        be2 = mil2["mil_head"]["bag_embed"]
+        if be2 is None:
+            raise RuntimeError("MIL level2 expected bag_embed (use attn + embed_cls).")
+        level2_emb = be2
+
+        logits_level1 = self.head_level1(level1_emb)
+        logits_level2 = self.head_level2(level2_emb)
+
+        fusion_in = torch.cat([level1_emb, level2_emb, concat_feat], dim=-1)
+        logits_fusion = self.fusion_head(fusion_in)
+
+        return {
+            "logits_baseline": logits_baseline,
+            "logits_level1": logits_level1,
+            "logits_level2": logits_level2,
+            "logits_fusion": logits_fusion,
+            "patch_logits": patch_logits,
+            "jpeg_pred": jpeg_pred,
+        }
+
+
+class LTFNetV2(nn.Module):
+    """Frozen CLIP visual + LTFCoreV2. Input: (B, N, 3, H, W)."""
+
+    def __init__(
+        self,
+        clip_visual,
+        num_heads=8,
+        num_classes=2,
+        num_layers=1,
+        dropout=0.2,
+        tile_heads=2,
+        tile_layers=2,
+        use_global_image=True,
+        use_patch_transformer=True,
+        linear_classifier=False,
+        mil_heads=8,
+        fusion_hidden=512,
+        l1_cross_shuffle_p: float = 0.9,
+        l1_replace_ratio_range: Tuple[float, float] = (0.5, 0.9),
+        l1_with_replacement_p: float = 0.5,
+        l2_cross_shuffle_p: float = 0.5,
+        l2_replace_ratio_range: Tuple[float, float] = (0.1, 0.5),
+        l2_with_replacement_p: float = 0.5,
+        bag_token_shuffle_p: float = 0.5,
+        cross_shuffle_min_keep_original: int = 0,
+        use_clip_intermediate_for_mil: bool = False,
+        clip_mil_stop_after_block_1based: int = 12,
+    ):
+        super().__init__()
+        self.visual = clip_visual
+        self.dtype = self.visual.conv1.weight.dtype
+
+        for p in self.visual.parameters():
+            p.requires_grad = False
+        self.visual.eval()
+
+        self.input_resolution = self.visual.input_resolution
+        self.embed_dim = self.visual.conv1.weight.shape[0]
+
+        self.use_clip_intermediate_for_mil = bool(use_clip_intermediate_for_mil)
+        self.clip_mil_stop_after_block_1based = int(clip_mil_stop_after_block_1based)
+        if self.use_clip_intermediate_for_mil:
+            n_bl = _count_clip_vit_depth(self.visual)
+            stop = self.clip_mil_stop_after_block_1based
+            if stop < 1 or stop > n_bl:
+                raise ValueError(
+                    f"use_clip_intermediate_for_mil requires 1 <= v2_clip_mil_stop_after_block_1based "
+                    f"<= {n_bl} (ViT depth), got {stop}"
+                )
+
+        self.core = LTFCoreV2(
+            embed_dim=self.embed_dim,
+            num_heads=num_heads,
+            num_classes=num_classes,
+            num_layers=num_layers,
+            dropout=dropout,
+            tile_heads=tile_heads,
+            tile_layers=tile_layers,
+            use_global_image=use_global_image,
+            use_patch_transformer=use_patch_transformer,
+            linear_classifier=linear_classifier,
+            mil_heads=mil_heads,
+            fusion_hidden=fusion_hidden,
+            l1_cross_shuffle_p=l1_cross_shuffle_p,
+            l1_replace_ratio_range=l1_replace_ratio_range,
+            l1_with_replacement_p=l1_with_replacement_p,
+            l2_cross_shuffle_p=l2_cross_shuffle_p,
+            l2_replace_ratio_range=l2_replace_ratio_range,
+            l2_with_replacement_p=l2_with_replacement_p,
+            bag_token_shuffle_p=bag_token_shuffle_p,
+            cross_shuffle_min_keep_original=cross_shuffle_min_keep_original,
+        )
+
+    def _stem_to_lnd(self, x_bchw: torch.Tensor) -> torch.Tensor:
+        """(B,3,H,W) -> (L,N,D) CLIP transformer 输入。"""
+        x = x_bchw.type(self.dtype)
+        x = self.visual.conv1(x)
+        x = x.reshape(x.shape[0], x.shape[1], -1)
+        x = x.permute(0, 2, 1)
+        cls = self.visual.class_embedding.to(x.dtype) + torch.zeros(
+            x.shape[0], 1, x.shape[2], dtype=x.dtype, device=x.device
+        )
+        x = torch.cat([cls, x], dim=1)
+        x = x + self.visual.positional_embedding.to(x.dtype)
+        x = self.visual.ln_pre(x)
+        return x.permute(1, 0, 2)
+
+    def _visual_tokens_from_imgs(
+        self, imgs_bchw: torch.Tensor
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """
+        返回 (tokens_final_NLD, tokens_mil_NLD | None)。
+        若未启用中间层 MIL，第二项为 None（与旧版单路前向一致）。
+        """
+        x = self._stem_to_lnd(imgs_bchw)
+        if not self.use_clip_intermediate_for_mil:
+            x = self.visual.transformer(x)
+            return x.permute(1, 0, 2), None
+        blocks = _clip_vit_resblock_modules(self.visual.transformer)
+        stop = self.clip_mil_stop_after_block_1based
+        for i in range(stop):
+            x = blocks[i](x)
+        x_mil = x.permute(1, 0, 2).clone()
+        for i in range(stop, len(blocks)):
+            x = blocks[i](x)
+        return x.permute(1, 0, 2), x_mil
+
+    def get_tokens(self, x) -> torch.Tensor:
+        """与旧版一致：完整 ViT 输出 (B, T, D)，不拆中间层。"""
+        x = self._stem_to_lnd(x)
+        x = self.visual.transformer(x)
+        return x.permute(1, 0, 2)
+
+    def encode_image(self, image) -> torch.Tensor:
+        return self.visual(image.type(self.dtype))
+
+    def forward(self, x: torch.Tensor, labels: Optional[torch.Tensor] = None, chunk_size=256):
+        chunk_size = getattr(self, "eval_chunk_size", chunk_size)
+        B, N, C, H, W = x.shape
+        imgs = x.reshape(B * N, C, H, W)
+        feats_chunks = []
+        mil_chunks = []
+        with torch.no_grad():
+            for start in range(0, imgs.size(0), chunk_size):
+                end = min(start + chunk_size, imgs.size(0))
+                f_tok, m_tok = self._visual_tokens_from_imgs(imgs[start:end])
+                feats_chunks.append(f_tok)
+                if m_tok is not None:
+                    mil_chunks.append(m_tok)
+        feats_token = torch.cat(feats_chunks, dim=0)
+        T = feats_token.shape[-2]
+        D = feats_token.shape[-1]
+        feats = feats_token.reshape(B, N, T, D).float()
+        if not mil_chunks:
+            return self.core(feats, labels=labels, x_mil=None)
+        mil_token = torch.cat(mil_chunks, dim=0).reshape(B, N, T, D).float()
+        return self.core(feats, labels=labels, x_mil=mil_token)

model/LTF_v3_3.py

@@ -0,0 +1,461 @@
+"""
+LTF V3.3: V3 + pool-logit gated fusion.
+
+Motivation:
+- Keep V3 outputs (logits_level1 / logits_level2 / logits_fusion) unchanged for training tools.
+- Use "logit_pool" signals (attention-weighted instance logits) as a reliability cue to gate fusion
+  between level1_emb and level2_emb, while keeping fusion_head input dim = 2*embed_dim.
+"""
+
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from model.LTF import LTFCore
+from model.LTF_v2 import (
+    BagTokenShuffle,
+    _clip_vit_resblock_modules,
+    _count_clip_vit_depth,
+)
+from mil_module.mil_core import MILCore
+
+
+class LTFCoreV3_3(nn.Module):
+    """
+    V3 core + a tiny gate that uses logit_pool(level1_tiles) and logit_pool(level2_cls-bag).
+
+    - Level1: patch tokens (tiles only) -> MIL per tile -> tile bag embeddings -> query attention -> level1_emb
+      Additionally: compute per-tile instance logits and attention weights -> logit_pool across tiles.
+    - Level2: CLS tokens (global + tiles) -> MIL (attn pooling) -> bag_embed = level2_emb
+      Additionally: compute logit_pool from MILHead's attn * inst_logits.
+    - Fusion: g = sigmoid(MLP([l1_pool_logits, l2_pool_logits])) in (0,1),
+      fusion_in = [g*level1_emb, (1-g)*level2_emb] -> fusion_head (same as V3 input dim).
+    """
+
+    def __init__(
+        self,
+        embed_dim=1024,
+        num_heads=8,
+        num_classes=2,
+        num_layers=1,
+        dropout=0.1,
+        tile_heads=2,
+        tile_layers=2,
+        use_global_image=True,
+        use_patch_transformer=True,
+        linear_classifier=False,
+        mil_heads=8,
+        fusion_hidden=512,
+        l1_cross_shuffle_p: float = 0.9,
+        l1_replace_ratio_range: Tuple[float, float] = (0.5, 0.9),
+        l1_with_replacement_p: float = 0.5,
+        l2_cross_shuffle_p: float = 0.5,
+        l2_replace_ratio_range: Tuple[float, float] = (0.1, 0.5),
+        l2_with_replacement_p: float = 0.5,
+        bag_token_shuffle_p: float = 0.5,
+        cross_shuffle_min_keep_original: int = 0,
+        use_mil_cross_shuffle: bool = True,
+        pool_gate_hidden: int = 16,
+        use_pool_gate_fusion: bool = True,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.use_global_image = use_global_image
+        self.num_classes = num_classes
+        self.use_mil_cross_shuffle = bool(use_mil_cross_shuffle)
+
+        self.l1_cross_shuffle_p = float(l1_cross_shuffle_p)
+        self.l1_replace_ratio_range = (
+            float(l1_replace_ratio_range[0]),
+            float(l1_replace_ratio_range[1]),
+        )
+        self.l1_with_replacement_p = float(l1_with_replacement_p)
+        self.l2_cross_shuffle_p = float(l2_cross_shuffle_p)
+        self.l2_replace_ratio_range = (
+            float(l2_replace_ratio_range[0]),
+            float(l2_replace_ratio_range[1]),
+        )
+        self.l2_with_replacement_p = float(l2_with_replacement_p)
+        self.bag_token_shuffle_p = float(bag_token_shuffle_p)
+        self.cross_shuffle_min_keep_original = int(cross_shuffle_min_keep_original)
+
+        self.use_pool_gate_fusion = bool(use_pool_gate_fusion)
+
+        # Keep the same core structure and frozen baseline classifier as V3 for compatibility.
+        self.baseline_core = LTFCore(
+            embed_dim=embed_dim,
+            num_heads=num_heads,
+            num_classes=num_classes,
+            num_layers=num_layers,
+            dropout=dropout,
+            tile_heads=tile_heads,
+            tile_layers=tile_layers,
+            use_global_image=use_global_image,
+            use_patch_transformer=use_patch_transformer,
+            linear_classifier=linear_classifier,
+        )
+        for p in self.baseline_core.classifier.parameters():
+            p.requires_grad = False
+
+        mil_kw = dict(
+            attn_hidden=min(256, embed_dim),
+            attn_dropout=dropout,
+            attn_bag_mode="embed_cls",
+        )
+        self.mil_level1 = MILCore(
+            d_in=embed_dim,
+            num_classes=num_classes,
+            mil_head_pooling="attn",
+            mil_head_kwargs=mil_kw,
+            use_cross_shuffle=self.use_mil_cross_shuffle,
+            trans=BagTokenShuffle(),
+            allow_trans_token_drop=False,
+        )
+        self.mil_level2 = MILCore(
+            d_in=embed_dim,
+            num_classes=num_classes,
+            mil_head_pooling="attn",
+            mil_head_kwargs=mil_kw,
+            use_cross_shuffle=self.use_mil_cross_shuffle,
+            trans=BagTokenShuffle(),
+            allow_trans_token_drop=False,
+        )
+
+        # Level-1 aggregation across tiles (query attention) + tile-level instance logits for logit_pool.
+        self.level1_query = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        mh = mil_heads if embed_dim % mil_heads == 0 else 8
+        while embed_dim % mh != 0:
+            mh -= 1
+        if mh < 1:
+            mh = 1
+        self.level1_aggregate_attn = nn.MultiheadAttention(
+            embed_dim, mh, dropout=dropout, batch_first=True
+        )
+
+        # Heads
+        self.head_level1 = nn.Linear(embed_dim, num_classes)
+        self.head_level2 = nn.Linear(embed_dim, num_classes)
+        self.tile_inst_head = nn.Linear(embed_dim, num_classes)
+
+        fusion_in = embed_dim + embed_dim
+        self.fusion_head = nn.Sequential(
+            nn.Linear(fusion_in, fusion_hidden),
+            nn.ReLU(),
+            nn.Dropout(p=dropout),
+            nn.Linear(fusion_hidden, num_classes),
+        )
+
+        # Gate uses pooled logits from L1/L2: concat => [B, 2*C]
+        gate_in = 2 * num_classes
+        self.pool_gate = nn.Sequential(
+            nn.Linear(gate_in, int(pool_gate_hidden)),
+            nn.ReLU(),
+            nn.Linear(int(pool_gate_hidden), 1),
+        )
+
+        nn.init.trunc_normal_(self.level1_query, std=0.02)
+
+    @staticmethod
+    def _sample_replace_ratio(low: float, high: float, device: torch.device) -> float:
+        lo, hi = min(low, high), max(low, high)
+        lo = max(0.0, lo)
+        hi = min(1.0, hi)
+        if hi <= lo:
+            return float(lo)
+        u = torch.rand((), device=device)
+        return float(lo + (hi - lo) * u)
+
+    def _apply_patch_transformer(self, x):
+        b2, n2, t2, d2 = x.shape
+        xf = x.reshape(b2 * n2, t2, d2)
+        if self.baseline_core.patch_transformer is not None:
+            xf = self.baseline_core.patch_transformer(xf)
+        return xf.reshape(b2, n2, t2, d2)
+
+    @staticmethod
+    def _logit_pool_from_mil_head(mil_head_out: dict) -> torch.Tensor:
+        """
+        Compute logit_pool = sum(attn * inst_logits) over instances.
+        Requires pooling='attn' so that 'attn' is not None.
+        """
+        inst = mil_head_out.get("inst_logits")
+        attn = mil_head_out.get("attn")
+        if inst is None or attn is None:
+            raise RuntimeError("logit_pool requires MILHead outputs: inst_logits and attn.")
+        return (attn * inst).sum(dim=1)
+
+    def forward(
+        self,
+        x,
+        labels: Optional[torch.Tensor] = None,
+        x_mil: Optional[torch.Tensor] = None,
+        return_viz: bool = False,
+    ):
+        x_feat = self._apply_patch_transformer(x)
+        if x_mil is not None:
+            if x_mil.shape != x.shape:
+                raise ValueError(f"x_mil shape {tuple(x_mil.shape)} != x shape {tuple(x.shape)}")
+            x_mil_feat = self._apply_patch_transformer(x_mil)
+        else:
+            x_mil_feat = x_feat
+
+        b, n, t, d = x_mil_feat.shape
+        if n < 2:
+            raise ValueError(f"LTFCoreV3_3 expects N>=2 (global + tiles), got N={n}")
+        nt = n - 1
+        train_shuffle = self.training and labels is not None
+        p_bag_shuffle = self.bag_token_shuffle_p if train_shuffle else 0.0
+
+        # --- Level 1: patch tokens (tiles only, no CLS) ---
+        patch_only = x_mil_feat[:, 1:, 1:, :].reshape(b * nt, t - 1, d)
+        do1 = train_shuffle and float(torch.rand((), device=x.device)) < self.l1_cross_shuffle_p
+        rr1 = self._sample_replace_ratio(*self.l1_replace_ratio_range, x.device)
+        wr1 = bool(float(torch.rand((), device=x.device)) < self.l1_with_replacement_p)
+        y1 = labels.unsqueeze(1).expand(-1, nt).reshape(-1) if labels is not None else None
+        mk = self.cross_shuffle_min_keep_original
+        mil1 = self.mil_level1(
+            patch_only,
+            y=y1,
+            return_dict=True,
+            do_cross_shuffle=do1,
+            replace_ratio=rr1,
+            with_replacement=wr1,
+            detach_inserted=True,
+            cross_shuffle_min_keep=mk,
+            p_shuffle=p_bag_shuffle,
+        )
+        be = mil1["mil_head"]["bag_embed"]
+        if be is None:
+            raise RuntimeError("MIL level1 expected bag_embed (use attn + embed_cls).")
+        bag_seq = be.reshape(b, nt, d)  # per-tile embeddings
+
+        # Query-attn aggregate -> level1_emb
+        q = self.level1_query.expand(b, -1, -1)
+        agg_out, attn_w = self.level1_aggregate_attn(
+            q, bag_seq, bag_seq, need_weights=True, average_attn_weights=False
+        )
+        level1_emb = agg_out.squeeze(1)  # [B, D]
+
+        # logit_pool across tiles using the same attention weights.
+        # attn_w: [B, heads, 1, Nt] when average_attn_weights=False
+        if attn_w is None:
+            raise RuntimeError("Expected attention weights from level1_aggregate_attn.")
+        tile_attn = attn_w.mean(dim=1).squeeze(1)  # [B, Nt]
+        tile_inst_logits = self.tile_inst_head(bag_seq)  # [B, Nt, C]
+        level1_logit_pool = (tile_attn.unsqueeze(-1) * tile_inst_logits).sum(dim=1)  # [B, C]
+
+        # --- Level 2: CLS tokens (global + tiles) ---
+        cls_seq = x_mil_feat[:, :, 0, :]  # [B, N, D]
+        do2 = train_shuffle and float(torch.rand((), device=x.device)) < self.l2_cross_shuffle_p
+        rr2 = self._sample_replace_ratio(*self.l2_replace_ratio_range, x.device)
+        wr2 = bool(float(torch.rand((), device=x.device)) < self.l2_with_replacement_p)
+        mil2 = self.mil_level2(
+            cls_seq,
+            y=labels,
+            return_dict=True,
+            do_cross_shuffle=do2,
+            replace_ratio=rr2,
+            with_replacement=wr2,
+            detach_inserted=True,
+            cross_shuffle_min_keep=mk,
+            p_shuffle=p_bag_shuffle,
+        )
+        mil2_head = mil2["mil_head"]
+        be2 = mil2_head["bag_embed"]
+        if be2 is None:
+            raise RuntimeError("MIL level2 expected bag_embed (use attn + embed_cls).")
+        level2_emb = be2  # [B, D]
+        level2_logit_pool = self._logit_pool_from_mil_head(mil2_head)  # [B, C]
+
+        logits_level1 = self.head_level1(level1_emb)
+        logits_level2 = self.head_level2(level2_emb)
+
+        # --- Fusion: pool-logit gate -> keep fusion input dim unchanged ---
+        if self.use_pool_gate_fusion:
+            gate_in = torch.cat([level1_logit_pool, level2_logit_pool], dim=-1)  # [B, 2C]
+            g = torch.sigmoid(self.pool_gate(gate_in))  # [B, 1]
+            fusion_in = torch.cat([g * level1_emb, (1.0 - g) * level2_emb], dim=-1)  # [B, 2D]
+        else:
+            fusion_in = torch.cat([level1_emb, level2_emb], dim=-1)
+        logits_fusion = self.fusion_head(fusion_in)
+
+        out = {
+            "logits_level1": logits_level1,
+            "logits_level2": logits_level2,
+            "logits_fusion": logits_fusion,
+        }
+        if return_viz:
+            mil1_head = mil1.get("mil_head", {})
+            out["viz"] = {
+                # Branch-level raw outputs
+                "logits_level1": logits_level1,
+                "logits_level2": logits_level2,
+                "logits_fusion": logits_fusion,
+                # MIL pooled logits and gate inputs
+                "level1_logit_pool": level1_logit_pool,
+                "level2_logit_pool": level2_logit_pool,
+                # Embeddings before branch heads / fusion
+                "level1_emb": level1_emb,
+                "level2_emb": level2_emb,
+                "fusion_in": fusion_in,
+                # Attention maps and instance logits
+                "tile_attn": tile_attn,
+                "tile_inst_logits": tile_inst_logits,
+                "level1_patch_attn": mil1_head.get("attn"),
+                "level1_patch_inst_logits": mil1_head.get("inst_logits"),
+                "level2_attn": mil2_head.get("attn"),
+                "level2_inst_logits": mil2_head.get("inst_logits"),
+                "num_tiles": torch.full(
+                    (b, 1), nt, device=logits_fusion.device, dtype=torch.int32
+                ),
+                "num_patches_per_tile": torch.full(
+                    (b, 1), (t - 1), device=logits_fusion.device, dtype=torch.int32
+                ),
+                # Scalar gate
+                "gate": g if self.use_pool_gate_fusion else torch.full(
+                    (b, 1), 0.5, device=logits_fusion.device, dtype=logits_fusion.dtype
+                ),
+                "use_pool_gate_fusion": torch.tensor(
+                    [1 if self.use_pool_gate_fusion else 0],
+                    device=logits_fusion.device,
+                    dtype=logits_fusion.dtype,
+                ),
+            }
+        return out
+
+
+class LTFNetV3_3(nn.Module):
+    """Frozen CLIP visual + LTFCoreV3_3. Input: (B, N, 3, H, W)."""
+
+    def __init__(
+        self,
+        clip_visual,
+        num_heads=8,
+        num_classes=2,
+        num_layers=1,
+        dropout=0.2,
+        tile_heads=2,
+        tile_layers=2,
+        use_global_image=True,
+        use_patch_transformer=True,
+        linear_classifier=False,
+        mil_heads=8,
+        fusion_hidden=512,
+        l1_cross_shuffle_p: float = 0.9,
+        l1_replace_ratio_range: Tuple[float, float] = (0.5, 0.9),
+        l1_with_replacement_p: float = 0.5,
+        l2_cross_shuffle_p: float = 0.5,
+        l2_replace_ratio_range: Tuple[float, float] = (0.1, 0.5),
+        l2_with_replacement_p: float = 0.5,
+        bag_token_shuffle_p: float = 0.5,
+        cross_shuffle_min_keep_original: int = 0,
+        use_mil_cross_shuffle: bool = True,
+        use_clip_intermediate_for_mil: bool = False,
+        clip_mil_stop_after_block_1based: int = 12,
+        pool_gate_hidden: int = 16,
+        use_pool_gate_fusion: bool = True,
+    ):
+        super().__init__()
+        self.visual = clip_visual
+        self.dtype = self.visual.conv1.weight.dtype
+
+        for p in self.visual.parameters():
+            p.requires_grad = False
+        self.visual.eval()
+
+        self.input_resolution = self.visual.input_resolution
+        self.embed_dim = self.visual.conv1.weight.shape[0]
+
+        self.use_clip_intermediate_for_mil = bool(use_clip_intermediate_for_mil)
+        self.clip_mil_stop_after_block_1based = int(clip_mil_stop_after_block_1based)
+        if self.use_clip_intermediate_for_mil:
+            n_bl = _count_clip_vit_depth(self.visual)
+            stop = self.clip_mil_stop_after_block_1based
+            if stop < 1 or stop > n_bl:
+                raise ValueError(
+                    f"use_clip_intermediate_for_mil requires 1 <= clip_mil_stop_after_block_1based "
+                    f"<= {n_bl} (ViT depth), got {stop}"
+                )
+
+        self.core = LTFCoreV3_3(
+            embed_dim=self.embed_dim,
+            num_heads=num_heads,
+            num_classes=num_classes,
+            num_layers=num_layers,
+            dropout=dropout,
+            tile_heads=tile_heads,
+            tile_layers=tile_layers,
+            use_global_image=use_global_image,
+            use_patch_transformer=use_patch_transformer,
+            linear_classifier=linear_classifier,
+            mil_heads=mil_heads,
+            fusion_hidden=fusion_hidden,
+            l1_cross_shuffle_p=l1_cross_shuffle_p,
+            l1_replace_ratio_range=l1_replace_ratio_range,
+            l1_with_replacement_p=l1_with_replacement_p,
+            l2_cross_shuffle_p=l2_cross_shuffle_p,
+            l2_replace_ratio_range=l2_replace_ratio_range,
+            l2_with_replacement_p=l2_with_replacement_p,
+            bag_token_shuffle_p=bag_token_shuffle_p,
+            cross_shuffle_min_keep_original=cross_shuffle_min_keep_original,
+            use_mil_cross_shuffle=use_mil_cross_shuffle,
+            pool_gate_hidden=pool_gate_hidden,
+            use_pool_gate_fusion=use_pool_gate_fusion,
+        )
+
+    def _stem_to_lnd(self, x_bchw: torch.Tensor) -> torch.Tensor:
+        x = x_bchw.type(self.dtype)
+        x = self.visual.conv1(x)
+        x = x.reshape(x.shape[0], x.shape[1], -1)
+        x = x.permute(0, 2, 1)
+        cls = self.visual.class_embedding.to(x.dtype) + torch.zeros(
+            x.shape[0], 1, x.shape[2], dtype=x.dtype, device=x.device
+        )
+        x = torch.cat([cls, x], dim=1)
+        x = x + self.visual.positional_embedding.to(x.dtype)
+        x = self.visual.ln_pre(x)
+        return x.permute(1, 0, 2)
+
+    def _visual_tokens_from_imgs(self, imgs_bchw: torch.Tensor):
+        x = self._stem_to_lnd(imgs_bchw)
+        if not self.use_clip_intermediate_for_mil:
+            x = self.visual.transformer(x)
+            return x.permute(1, 0, 2), None
+        blocks = _clip_vit_resblock_modules(self.visual.transformer)
+        stop = self.clip_mil_stop_after_block_1based
+        for i in range(stop):
+            x = blocks[i](x)
+        x_mil = x.permute(1, 0, 2).clone()
+        for i in range(stop, len(blocks)):
+            x = blocks[i](x)
+        return x.permute(1, 0, 2), x_mil
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        labels: Optional[torch.Tensor] = None,
+        chunk_size=256,
+        return_viz: bool = False,
+    ):
+        chunk_size = getattr(self, "eval_chunk_size", chunk_size)
+        b, n, c, h, w = x.shape
+        imgs = x.reshape(b * n, c, h, w)
+        feats_chunks = []
+        mil_chunks = []
+        with torch.no_grad():
+            for start in range(0, imgs.size(0), chunk_size):
+                end = min(start + chunk_size, imgs.size(0))
+                f_tok, m_tok = self._visual_tokens_from_imgs(imgs[start:end])
+                feats_chunks.append(f_tok)
+                if m_tok is not None:
+                    mil_chunks.append(m_tok)
+        feats_token = torch.cat(feats_chunks, dim=0)
+        t = feats_token.shape[-2]
+        d = feats_token.shape[-1]
+        feats = feats_token.reshape(b, n, t, d).float()
+        if not mil_chunks:
+            return self.core(feats, labels=labels, x_mil=None, return_viz=return_viz)
+        mil_token = torch.cat(mil_chunks, dim=0).reshape(b, n, t, d).float()
+        return self.core(feats, labels=labels, x_mil=mil_token, return_viz=return_viz)
+

model/__init__.py

@@ -0,0 +1 @@
+"""LTF v3.3 historical model package."""

release_manifest.json

@@ -0,0 +1,80 @@
+{
+  "model": "LTF v3.3",
+  "checkpoint": "epoch_20.pth",
+  "historical_run": "20260326_012209_ViT-L-14_genimage_v3_v3_3",
+  "training_run": "20260325_092132_ViT-L-14_genimage_sd14_pure_test_v3_3_float",
+  "historical_fusion_macro8_accuracy": 0.9839,
+  "historical_images": 100000,
+  "threshold": 0.5,
+  "source_note": "Historical release core reconstructed from git 65c18ee5868ced87a27d048261e237632ad5d704; original training-day source snapshot unavailable.",
+  "adaptations": [
+    "Tensor-only checkpoint loading",
+    "Portable evaluation configuration"
+  ],
+  "files": {
+    "ltf_final/__init__.py": {
+      "sha256": "027e133874edcab1a9fa1f7786618e91e5bad141d87d5f592cdc526c379b6fbb",
+      "bytes": 175
+    },
+    "ltf_final/detector.py": {
+      "sha256": "22b6b8ce000d0f02896f1d10df65fbfaa2107bf2ed83d218df306241279f0e91",
+      "bytes": 6073
+    },
+    "model/LTF.py": {
+      "sha256": "4e76ed3dcf03cf56ca5c843220e15cddc84d1db2b95c823d8fe8bdab0cc0e279",
+      "bytes": 8147
+    },
+    "model/LTF_v2.py": {
+      "sha256": "aea134d55d942c606b3876a7c98eb39092a709c1d2a863e3e3074523ab388de9",
+      "bytes": 16763
+    },
+    "model/LTF_v3_3.py": {
+      "sha256": "f771d49441a2702f241fd9debc10ee49da20414db8759d2b8a551b18674a4202",
+      "bytes": 18523
+    },
+    "model/__init__.py": {
+      "sha256": "e49d18603c3f0259d1c37da6ba7fabeb986f8e44e2b853d7b86cca1c5f94a433",
+      "bytes": 41
+    },
+    "mil_module/__init__.py": {
+      "sha256": "a64b516cada39fe04a117297a47f6bec29d34e7e6e7462ee80d4f2a1e180628b",
+      "bytes": 1025
+    },
+    "mil_module/cross_shuffle.py": {
+      "sha256": "4a995e4408708b1bc784bfa252e02e00ed5c91b27629543cf5eaca59352279c7",
+      "bytes": 6697
+    },
+    "mil_module/mil_core.py": {
+      "sha256": "87c85d888540858d45f869ba1e24efcfb553ea525d35518790b54e33277e784e",
+      "bytes": 8350
+    },
+    "mil_module/mil_head.py": {
+      "sha256": "bf383e4bf91fec6fc6788af71cc2b6e81191cf5e2137b0b443f6a2d9dada1e4f",
+      "bytes": 8167
+    },
+    "mil_module/mil_head_beta.py": {
+      "sha256": "46e3da092b20de6f2bc8400590d2aa3a3156584d1c79254b7a11abb2fd5ceced",
+      "bytes": 7189
+    },
+    "utils/__init__.py": {
+      "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+      "bytes": 0
+    },
+    "utils/tile_dataset_srm.py": {
+      "sha256": "49106f01b4fda4d10e7db8364967988a68a00c39daa2dde8fe89f52506531642",
+      "bytes": 14623
+    },
+    "configs/eval.yaml": {
+      "sha256": "4ce443731e70a4ac4eb71fb56103f33373e179fefe766c3e6a02635efdb94efc",
+      "bytes": 1039
+    },
+    "weights/epoch_20.pth": {
+      "sha256": "32e9fa2e38e04d6d0221c844d650f0d45f7235bbf898b5ad665543d8be580031",
+      "bytes": 62249626
+    },
+    "backbones/ViT-L-14.pt": {
+      "sha256": "b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836",
+      "bytes": 932768134
+    }
+  }
+}

requirements.txt

@@ -0,0 +1,9 @@
+torch==2.1.2
+torchvision==0.16.2
+numpy>=1.24,<2
+Pillow>=10.2
+PyYAML>=6.0
+scikit-learn>=1.2
+tqdm>=4.66
+tensorboard>=2.15
+openai-clip==1.0.1

results/genimage8_logged_metrics.csv

@@ -0,0 +1,9 @@
+subset,n_real,n_fake,acc_level1,acc_real_level1,acc_fake_level1,auc_level1,ap_level1,best_thr_level1,best_acc_level1,acc_level2,acc_real_level2,acc_fake_level2,auc_level2,ap_level2,best_thr_level2,best_acc_level2,acc_fusion,acc_real_fusion,acc_fake_fusion,auc_fusion,ap_fusion,best_thr_fusion,best_acc_fusion
+imagenet_midjourney,6000,6000,0.9194,0.9288,0.91,0.9755,0.9755,0.46542627,0.921,0.9692,0.956,0.9825,0.9953,0.9938,0.64407796,0.9722,0.9734,0.9812,0.9657,0.9962,0.995,0.37580317,0.9758
+imagenet_ai_0419_sdv4,6000,6000,0.962,0.924,1.0,0.9999,0.9999,0.99286216,0.9978,0.9763,0.9527,1.0,1.0,1.0,0.974082,0.9979,0.989,0.978,1.0,1.0,1.0,0.99918264,0.9999
+imagenet_ai_0424_sdv5,8000,8000,0.96,0.92,1.0,0.9998,0.9998,0.9958075,0.9961,0.9748,0.9507,0.9989,0.9998,0.9998,0.9699074,0.9963,0.9862,0.9735,0.9989,0.9999,0.9999,0.9994531,0.9988
+imagenet_ai_0508_adm,6000,6000,0.9501,0.9177,0.9825,0.9929,0.9928,0.7772057,0.9609,0.9727,0.9543,0.991,0.9976,0.9975,0.7358004,0.9783,0.9841,0.976,0.9922,0.9988,0.9986,0.57335776,0.9858
+imagenet_glide,6000,6000,0.956,0.9275,0.9845,0.9941,0.9944,0.79664856,0.9657,0.9754,0.9525,0.9983,0.9992,0.9992,0.8827087,0.9882,0.9816,0.9768,0.9863,0.9981,0.9978,0.70165557,0.9822
+imagenet_ai_0424_wukong,6000,6000,0.9627,0.9257,0.9998,0.9998,0.9998,0.9886967,0.9959,0.9741,0.9483,0.9998,0.9999,0.9999,0.952306,0.9965,0.989,0.978,1.0,1.0,1.0,0.98438674,0.9988
+imagenet_ai_0419_vqdm,6000,6000,0.9597,0.9223,0.9972,0.9983,0.9983,0.92264795,0.9829,0.9736,0.9538,0.9933,0.9985,0.9986,0.812101,0.9818,0.9868,0.9755,0.998,0.9997,0.9997,0.8302674,0.9928
+imagenet_ai_0419_biggan,6000,6000,0.9512,0.9265,0.9758,0.9888,0.9869,0.5910971,0.9535,0.9752,0.9507,0.9998,0.9998,0.9998,0.9196974,0.9932,0.9809,0.9753,0.9865,0.9981,0.998,0.54691106,0.9816

results/summary.json

@@ -0,0 +1,13 @@
+{
+  "source_log": "/raid5/chr/AIGCD/realguard_new/evaluate_logs/20260326_012209_ViT-L-14_genimage_v3_v3_3/test.log",
+  "source": "historical logged metrics (rounded to 4 decimals), not newly re-evaluated",
+  "images": 100000,
+  "subsets": 8,
+  "threshold": 0.5,
+  "average": "unweighted mean over the eight subsets",
+  "acc_level1": 0.9526,
+  "acc_level2": 0.9739,
+  "acc_fusion": 0.9839,
+  "weights": "weights/epoch_20.pth",
+  "weights_sha256": "32e9fa2e38e04d6d0221c844d650f0d45f7235bbf898b5ad665543d8be580031"
+}

utils/tile_dataset_srm.py

@@ -0,0 +1,362 @@
+from PIL import Image, ImageOps, ImageFilter
+import random
+import io, os
+import math
+import torch
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from torchvision.transforms import InterpolationMode
+BICUBIC = InterpolationMode.BICUBIC
+
+def resize_with_random_scale(image: Image.Image, p: float = 0.5, scale_range=(0.25, 2.0)) -> Image.Image:
+    if random.random() > p:
+        return image 
+
+    w, h = image.size
+    scale = random.uniform(*scale_range)
+    new_w = max(1, int(w * scale))
+    new_h = max(1, int(h * scale))
+    return image.resize((new_w, new_h), resample=Image.BICUBIC)
+
+def pad_crop_image(image, target_width, target_height):
+    src_width, src_height = image.size
+    # 0
+    pad_left = pad_right = pad_top = pad_bottom = 0
+    crop_left = crop_top = 0
+
+    if src_width < target_width:
+        pad_left = random.randint(0, target_width - src_width)
+        pad_right = target_width - src_width - pad_left
+    else:
+        crop_left = random.randint(0, src_width - target_width)
+
+    if src_height < target_height:
+        pad_top = random.randint(0, target_height - src_height)
+        pad_bottom = target_height - src_height - pad_top
+    else:
+        crop_top = random.randint(0, src_height - target_height)
+
+    if src_width > target_width or src_height > target_height:
+        left = crop_left if src_width > target_width else 0
+        top = crop_top if src_height > target_height else 0
+        right = left + target_width
+        bottom = top + target_height
+        image = image.crop((left, top, right, bottom))
+    # Padding
+    if src_width < target_width or src_height < target_height:
+        image = ImageOps.expand(
+            image,
+            border=(pad_left, pad_top, pad_right, pad_bottom),
+            fill=(0, 0, 0)
+        )
+    return image
+
+def get_start_positions(image_length, block_size, bias=0):
+    if image_length <= block_size:
+        return [0+bias]
+    num_blocks = image_length // block_size + (image_length % block_size > 0)
+    if num_blocks == 1:
+        return [0+bias]
+    total_stride = (image_length - block_size) / (num_blocks - 1)
+    positions = [int(round(i * total_stride) + bias) for i in range(num_blocks - 1)]
+    positions.append(image_length - block_size + bias)
+    return positions
+
+def crop_image_to_tiles(image, block_size):
+    """
+    crop to block_size x block_size
+    """
+    width, height = image.size
+    x_positions = get_start_positions(width, block_size)
+    y_positions = get_start_positions(height, block_size)
+    tiles = []
+    for y in y_positions:
+        for x in x_positions:
+            tiles.append(image.crop((x, y, x + block_size, y + block_size)))
+    return tiles
+
+def crop_image_to_tiles_new(image, block_size):
+    width, height = image.size
+    x_positions = get_start_positions(width, block_size)
+    y_positions = get_start_positions(height, block_size)
+    tiles = []
+    for y in y_positions:
+        for x in x_positions:
+            tiles.append(image.crop((x, y, x + block_size, y + block_size)))
+    if width > block_size and height > block_size:
+        leave_space = min(block_size, width-block_size)
+        x_positions = get_start_positions(width-leave_space, block_size, leave_space//2)
+        leave_space = min(block_size, height-block_size)
+        y_positions = get_start_positions(height-leave_space, block_size, leave_space//2)
+        for y in y_positions:
+            for x in x_positions:
+                tiles.append(image.crop((x, y, x + block_size, y + block_size)))
+    return tiles
+
+def random_crop(image: Image.Image, size: int = 224, is_pad_black: bool = False) -> Image.Image:
+    w, h = image.size
+    crop_size = min(w, h, size)
+    if w == crop_size:
+        left = 0
+    else:
+        left = random.randint(0, w - crop_size)
+    if h == crop_size:
+        top = 0
+    else:
+        top = random.randint(0, h - crop_size)
+
+    cropped = image.crop((left, top, left + crop_size, top + crop_size))
+
+    if is_pad_black and crop_size < size:
+        padded_img = Image.new("RGB", (size, size), (0, 0, 0))
+        offset = ((size - crop_size) // 2, (size - crop_size) // 2)
+        padded_img.paste(cropped, offset)
+        return padded_img
+    else:
+        return cropped
+
+def _convert_image_to_rgb(image):
+    return image.convert("RGB")
+
+
+def clip_transform_allresize(n_px, blur_radius=None, blur_prob=0.5):
+    if blur_radius is None:
+        return transforms.Compose([
+            transforms.Resize((n_px, n_px), interpolation=BICUBIC),
+            _convert_image_to_rgb,
+            transforms.ToTensor(),
+            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+        ])
+    else:
+        return transforms.Compose([
+            transforms.Resize((n_px, n_px), interpolation=BICUBIC),
+            transforms.RandomApply(
+                [transforms.GaussianBlur(kernel_size=5, sigma=blur_radius)], p=blur_prob  # (0.1, 3)
+            ),
+            _convert_image_to_rgb,
+            transforms.ToTensor(),
+            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+        ])
+
+def clip_transform(n_px, blur_radius=None, blur_prob=0.5):
+    if blur_radius is None:
+        return transforms.Compose([
+            transforms.Resize(n_px, interpolation=BICUBIC),
+            transforms.CenterCrop(n_px),
+            _convert_image_to_rgb,
+            transforms.ToTensor(),
+            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+        ])
+    else:
+        return transforms.Compose([
+            transforms.Resize(n_px, interpolation=BICUBIC),
+            transforms.CenterCrop(n_px),
+            transforms.RandomApply(
+                [transforms.GaussianBlur(kernel_size=5, sigma=blur_radius)], p=blur_prob  # (0.1, 3)
+            ),
+            _convert_image_to_rgb,
+            transforms.ToTensor(),
+            transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+        ])
+
+class JpegAugment():
+    def __init__(self, quality=None, p=1, seed=42):
+        self.local_random = random.Random(seed)
+        self.quality = quality
+        self.p = p
+    
+    def __call__(self, img):
+        if self.quality is None:
+            return img
+        if self.local_random.random() < self.p:
+            buffer = io.BytesIO()
+            if self.quality == -1:
+                img.save(buffer, format='JPEG', quality=self.local_random.randint(30, 100))
+            else:
+                if isinstance(self.quality, tuple):
+                    img.save(buffer, format='JPEG', quality=self.local_random.randint(*self.quality))
+                else:
+                    img.save(buffer, format='JPEG', quality=self.quality)
+            buffer.seek(0)
+            return Image.open(buffer).convert('RGB')
+        return img
+
+def webp_compress(img, quality=80, method=4):
+    buffer = io.BytesIO() 
+    img.save(buffer, format='WEBP', quality=quality, method=method)
+    buffer.seek(0)
+    return Image.open(buffer).convert('RGB')
+
+# Image extensions supported when scanning real/fake dirs (case-insensitive).
+# Missing .JPG / .webp etc. can cause real folder to show 0 images if only those exist.
+IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.webp', '.bmp', '.tif', '.tiff', '.ppm')
+
+
+def find_files_with_extensions(directory, extensions=None, max_num=None):
+    """Find image paths under directory. If extensions is None, use IMG_EXTENSIONS (case-insensitive)."""
+    if extensions is None:
+        extensions = IMG_EXTENSIONS
+    ext_lower = tuple(e.lower() for e in extensions)
+    matching_files = []
+    check_num = 0
+    for root, dirs, files in os.walk(directory):
+        for file in files:
+            if file.lower().endswith(ext_lower):
+                matching_files.append(os.path.join(root, file))
+                check_num += 1
+                if max_num is not None and check_num >= max_num:
+                    return matching_files
+    return matching_files
+
+class TiledRFDataset(Dataset):
+    """ Dataset -> (img, label)"""
+    def __init__(self, real_dirs:list, fake_dirs:list, tile_size=224, set_size=None, shuffle_seed=42, balance=True, 
+                 jpeg_augment=None, jpeg_p=0.1,
+                 webp_augment=None, webp_p=0.1,
+                 resize_augment=None, resize_p=0.5,
+                 blur_augment=None, blur_p=0.5,
+                 default_train_tiles=4, is_train=False):
+        self.set_size = set_size
+        self.shuffle_seed = shuffle_seed
+        self.samples = []
+        self.real_samples = []
+        self.fake_samples = []
+        self.tile_size=tile_size
+        self.transform = clip_transform(tile_size, blur_augment, blur_p)
+        self.default_train_tiles = default_train_tiles
+        self.is_train = is_train
+        
+        self.jpeg_augment=jpeg_augment
+        self.jpeg_p = jpeg_p
+        self.JpegEnhance = JpegAugment(quality=jpeg_augment, p=jpeg_p, seed=shuffle_seed)
+        self.webp_augment = webp_augment
+        self.webp_p = webp_p
+        self.resize_augment = resize_augment
+        self.resize_p = resize_p
+
+        for p in real_dirs:
+            find_paths_real = find_files_with_extensions(p)
+            self.real_samples += [(f, 0) for f in find_paths_real]
+        for p in fake_dirs:
+            find_paths_fake = find_files_with_extensions(p)
+            self.fake_samples += [(f, 1) for f in find_paths_fake]
+        if balance:
+            target_len = min(len(self.real_samples), len(self.fake_samples))
+            self.real_samples = self.real_samples[:target_len]
+            self.fake_samples = self.fake_samples[:target_len]
+        self.samples = self.real_samples + self.fake_samples
+        # if self.set_size is not None and len(self.samples) < self.set_size:
+        #     n_repeats = math.ceil(self.set_size / len(self.samples))
+        #     self.samples = (self.samples * n_repeats)[:self.set_size]
+        if self.set_size is not None and len(self.samples) > self.set_size:
+            local_random = random.Random(self.shuffle_seed)
+            self.samples = local_random.sample(self.samples, self.set_size)
+    def __len__(self):
+        return len(self.samples)
+    
+    def _get_train_item(self, idx):
+        path, label = self.samples[idx]
+        try:
+            img = Image.open(path).convert('RGB')
+        except Exception as e:
+            print(f"Skipping corrupted image: {path}")
+            return self.__getitem__((idx + 1) % len(self.samples))
+        if self.resize_augment is not None:
+            if self.resize_augment == -1:
+                img = resize_with_random_scale(img, p=self.resize_p)
+            else:
+                if isinstance(self.resize_augment, tuple):
+                    img = resize_with_random_scale(img, p=self.resize_p, scale_range=self.resize_augment)
+                else:
+                    img = resize_with_random_scale(img, p=self.resize_p, scale_range=(self.resize_augment, self.resize_augment))
+        if self.jpeg_augment is not None:
+            img = self.JpegEnhance(img)
+        
+        # crops_one_side = 2
+        # img = pad_crop_image(img, crops_one_side*self.tile_size, crops_one_side*self.tile_size)
+        # tiles = crop_image_to_tiles(img, self.tile_size)[:crops_one_side**2]
+        tiles = []
+        for i in range(self.default_train_tiles):  # 10
+            tiles.append(random_crop(img, self.tile_size, is_pad_black=False))
+        img_resize = self.transform(img)
+        img_list = [img_resize]
+        for tile in tiles:
+            img_list.append(self.transform(tile))
+        return img_list, img, label
+
+    def _get_test_item(self, idx):
+        path, label = self.samples[idx]
+        try:
+            img = Image.open(path).convert('RGB')
+        except Exception as e:
+            print(f"Skipping corrupted image: {path}")
+            return self.__getitem__((idx + 1) % len(self.samples))
+        if self.resize_augment is not None:
+            if self.resize_augment == -1:
+                img = resize_with_random_scale(img, p=self.resize_p)
+            else:
+                if isinstance(self.resize_augment, tuple):
+                    img = resize_with_random_scale(img, p=self.resize_p, scale_range=self.resize_augment)
+                else:
+                    img = resize_with_random_scale(img, p=self.resize_p, scale_range=(self.resize_augment, self.resize_augment))
+        if self.jpeg_augment is not None:
+            img = self.JpegEnhance(img)
+        if self.webp_augment is not None:
+            if random.random() < self.webp_p:
+                if isinstance(self.webp_augment, tuple):
+                    img = webp_compress(img, quality=random.randint(*self.webp_augment))
+                else:
+                    img = webp_compress(img, quality=self.webp_augment)
+        # testtime batchsize=1
+        tiles = crop_image_to_tiles(img, self.tile_size)
+        img_resize = self.transform(img)
+        img_list = [img_resize]
+        for tile in tiles:
+            img_list.append(self.transform(tile))
+        return img_list, img, label, path
+
+    def __getitem__(self, idx):
+        if self.is_train:
+            return self._get_train_item(idx)
+        else:
+            return self._get_test_item(idx)
+        
+
+def collate_img_list_to_tensor_train(batch):
+    """
+    batch = list of (img_list, label)
+    - img_list: List of N Tensor[B, 3, H, W]
+    return：
+    - images: Tensor[B, N, 3, H, W]
+    - labels: Tensor[B]
+    """
+    img_lists, raw_imgs, labels = zip(*batch)  # len = B
+    # img_lists: list of length B, each is List[N x (3,H,W)]
+    stacked = torch.stack([torch.stack(imgs) for imgs in img_lists])  # (B, N, 3, H, W)
+    raw_imgs = list(raw_imgs)
+    labels = torch.tensor(labels)
+    return stacked, raw_imgs, labels
+
+def collate_img_list_to_tensor_test(batch):
+    """
+    batch = list of (img_list, label)
+    - img_list: List of N Tensor[B, 3, H, W]
+    return：
+    - images: Tensor[B, N, 3, H, W]
+    - labels: Tensor[B]
+    """
+    img_lists, raw_imgs, labels, paths = zip(*batch)  # len = B
+    # img_lists: list of length B, each is List[N x (3,H,W)]
+    stacked = torch.stack([torch.stack(imgs) for imgs in img_lists])  # (B, N, 3, H, W)
+    raw_imgs = list(raw_imgs)
+    labels = torch.tensor(labels)
+    return stacked, raw_imgs, labels, paths
+
+
+if __name__ == '__main__':
+    block_size = 224
+    width = 512
+    leave_space = min(block_size, width-block_size)
+    x_positions = get_start_positions(width-leave_space, block_size, leave_space//2)
+    print(leave_space, x_positions)
+

weights/epoch_20.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:32e9fa2e38e04d6d0221c844d650f0d45f7235bbf898b5ad665543d8be580031
+size 62249626