app.py

# -*- coding: utf-8 -*-
"""
Swin-Large   AI vs. Non-AI Detector   (基于注意力机制可视化)
"""
import os
import math
import torch
import torch.nn.functional as F
import torch.nn as nn
import timm
import numpy as np
from PIL import Image
import gradio as gr

from huggingface_hub import hf_hub_download
from torchvision import transforms
import matplotlib.pyplot as plt
import matplotlib.cm as cm

# --- Configuration ---------------------------------------------------------
REPO_ID          = "telecomadm1145/swin-ai-detection"
HF_FILENAME      = "swin_classifier_stage1_v2_epoch_3.pth"
LOCAL_CKPT_DIR   = "./checkpoints"
MODEL_NAME       = "swin_large_patch4_window12_384"     # ← 使用 large
NUM_CLASSES      = 2
SEED             = 4421
dropout_rate     = 0.1

class_names = ["Non-AI Generated", "AI Generated"]     # 0, 1

device = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(SEED);  np.random.seed(SEED)
print(f"Using device: {device}")

# ---------------------------------------------------------------------------
# 1. 修改模型结构以提取注意力
class SwinClassifierWithAttention(nn.Module):
    def __init__(self, model_name, num_classes, pretrained=True):
        super().__init__()
        self.backbone = timm.create_model(model_name, pretrained=pretrained,
                                          num_classes=0)
        self.data_config = timm.data.resolve_data_config({}, model=self.backbone)

        self.classifier = nn.Sequential(
            nn.Dropout(dropout_rate),
            nn.Linear(self.backbone.num_features, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Dropout(dropout_rate * 0.7),
            nn.Linear(512, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Dropout(dropout_rate * 0.5),
            nn.Linear(128, num_classes)
        )
        
        # 存储注意力权重的钩子
        self.attention_weights = {}
        self.register_attention_hooks()

    def register_attention_hooks(self):
        """注册钩子函数来提取注意力权重"""
        def hook_fn(name):
            def hook(module, input, output):
                # 对于Swin Transformer的窗口注意力机制
                # output通常是 (B, N, C) 格式
                if hasattr(module, 'attn'):
                    # 获取注意力权重
                    self.attention_weights[name] = module.attn.attention_weights
            return hook
        
        # 为每个stage的每个block注册钩子
        for stage_idx, stage in enumerate(self.backbone.layers):
            for block_idx, block in enumerate(stage.blocks):
                hook_name = f"stage_{stage_idx}_block_{block_idx}"
                if hasattr(block, 'attn'):
                    block.attn.register_forward_hook(hook_fn(hook_name))

    def forward(self, x):
        # 清空之前的注意力权重
        self.attention_weights = {}
        feats = self.backbone(x)
        return self.classifier(feats)

# ---------------------------------------------------------------------------
# 2. 注意力提取和可视化类
class AttentionExtractor:
    def __init__(self, model):
        self.model = model
        self.attention_maps = {}
        
    def extract_attention_weights(self, x):
        """提取所有层的注意力权重"""
        with torch.no_grad():
            _ = self.model(x)  # 前向传播以触发钩子
            return self.model.attention_weights.copy()
    
    def process_attention_for_visualization(self, attention_weights, input_size):
        """处理注意力权重用于可视化"""
        processed_maps = {}
        
        for layer_name, attn_weight in attention_weights.items():
            if attn_weight is None:
                continue
                
            # attn_weight shape: [batch_size, num_heads, seq_len, seq_len]
            if len(attn_weight.shape) == 4:
                # 取平均池化所有注意力头
                attn_map = attn_weight.mean(dim=1)  # [batch_size, seq_len, seq_len]
                
                # 取第一个样本
                attn_map = attn_map[0]  # [seq_len, seq_len]
                
                # 对于自注意力，我们通常关注CLS token对其他token的注意力
                # 或者计算所有token的平均注意力
                if attn_map.shape[0] > 1:
                    # 计算每个位置的平均注意力分数
                    avg_attention = attn_map.mean(dim=0)  # [seq_len]
                    
                    # 将注意力分数reshape为2D特征图
                    seq_len = avg_attention.shape[0]
                    grid_size = int(math.sqrt(seq_len))
                    
                    if grid_size * grid_size == seq_len:
                        attention_2d = avg_attention.reshape(grid_size, grid_size)
                        processed_maps[layer_name] = attention_2d
        
        return processed_maps

# ---------------------------------------------------------------------------
# 3. 下载 / 缓存 checkpoint
print("⏬ Download / cache checkpoint …")
ckpt_path = hf_hub_download(
    repo_id      = REPO_ID,
    filename     = HF_FILENAME,
    local_dir    = LOCAL_CKPT_DIR,
    force_download=False     # 已存在则直接用
)
print(f"Checkpoint path: {ckpt_path}")

# ---------------------------------------------------------------------------
# 4. 实例化 & 加载权重
model = SwinClassifierWithAttention(MODEL_NAME, NUM_CLASSES, pretrained=False).to(device)
state = torch.load(ckpt_path, map_location=device, weights_only=False)
model.load_state_dict(state.get("model_state_dict", state), strict=False)  # strict=False 因为添加了新的组件
model.eval()
print("✅  Model loaded.")

attention_extractor = AttentionExtractor(model)

# ---------------------------------------------------------------------------
# 5. 变换函数
def build_transform(is_training: bool, interpolation: str):
    """
    根据插值方式(双线性 / 三次等)构建 timm 默认变换
    """
    cfg = model.data_config.copy()
    cfg.update(dict(interpolation=interpolation))
    return timm.data.create_transform(**cfg, is_training=is_training)

# ---------------------------------------------------------------------------
# 6. 注意力可视化函数
def visualize_attention(attention_map, original_image, normalize=True):
    """将注意力图可视化到原始图像上"""
    if normalize:
        # 归一化注意力图到[0,1]
        attention_map = attention_map.cpu().numpy()
        attention_map = (attention_map - attention_map.min()) / (attention_map.max() - attention_map.min() + 1e-8)
    else:
        attention_map = attention_map.cpu().numpy()
    
    # 调整注意力图大小到原始图像大小
    attention_resized = Image.fromarray((attention_map * 255).astype(np.uint8)) \
                            .resize(original_image.size, Image.Resampling.BILINEAR)
    
    # 转换为热力图
    attention_array = np.array(attention_resized) / 255.0
    heatmap = cm.jet(attention_array)[:, :, :3]  # 去掉alpha通道
    
    # 叠加到原始图像
    original_array = np.array(original_image) / 255.0
    if len(original_array.shape) == 3:
        overlay = 0.6 * original_array + 0.4 * heatmap
    else:
        # 灰度图像处理
        original_array = np.stack([original_array] * 3, axis=-1)
        overlay = 0.6 * original_array + 0.4 * heatmap
    
    overlay = np.clip(overlay, 0, 1)
    return Image.fromarray((overlay * 255).astype(np.uint8))

# ---------------------------------------------------------------------------
# 7. 推理 + 注意力可视化
def infer_with_attention(image_pil: Image.Image,
                        interpolation: str = "bilinear",
                        attention_layer: str = "stage_3_block_1",
                        stage_average: bool = False,
                        normalize_attention: bool = True):
    if image_pil is None:
        return None, None

    transform = build_transform(is_training=False, interpolation=interpolation)
    input_tensor = transform(image_pil).unsqueeze(0).to(device)

    # (1) 分类预测
    with torch.no_grad():
        logits = model(input_tensor)
        probs  = F.softmax(logits, dim=1)[0]
        confidences = {class_names[i]: float(probs[i]) for i in range(NUM_CLASSES)}

    # (2) 提取注意力权重
    attention_weights = attention_extractor.extract_attention_weights(input_tensor)
    
    if not attention_weights:
        return confidences, None
    
    # (3) 处理注意力权重
    processed_attention = attention_extractor.process_attention_for_visualization(
        attention_weights, input_tensor.shape[-2:]
    )
    
    if not processed_attention:
        return confidences, None
    
    # (4) 选择要可视化的注意力层
    if stage_average:
        # 计算指定stage所有block的平均注意力
        stage_num = attention_layer.split('_')[1]
        stage_attentions = []
        
        for layer_name, attn_map in processed_attention.items():
            if f"stage_{stage_num}_" in layer_name:
                stage_attentions.append(attn_map)
        
        if stage_attentions:
            # 计算平均注意力
            avg_attention = torch.stack(stage_attentions).mean(dim=0)
            attention_vis = visualize_attention(avg_attention, image_pil, normalize_attention)
        else:
            return confidences, None
    else:
        # 使用指定层的注意力
        if attention_layer in processed_attention:
            attention_vis = visualize_attention(
                processed_attention[attention_layer], image_pil, normalize_attention
            )
        else:
            # 如果指定层不存在，使用第一个可用的层
            first_layer = list(processed_attention.keys())[0]
            attention_vis = visualize_attention(
                processed_attention[first_layer], image_pil, normalize_attention
            )
    
    return confidences, attention_vis

# ---------------------------------------------------------------------------
# 8. Gradio UI
def launch_app():
    with gr.Blocks() as demo:
        gr.Markdown("""
# 🖼️ AI vs. Non-AI Image Classifier  (Swin-Large + Attention Visualization)

🖼️ AI 鉴别器(基于 Swin-Large 视觉骨干，输出注意力热力图)

基于Swin Transformer的自注意力机制来可视化模型关注的区域。

Notice: 使用 bicubic 效果较好。请负责任地使用此工具。

此工具仅供研究和教育用途。
""")

        with gr.Row():
            in_img = gr.Image(type="pil", label="Upload an Image")
            out_attention = gr.Image(type="pil", label="Attention Heatmap")

        with gr.Row():
            out_lbl = gr.Label(num_top_classes=2, label="Predictions")

        with gr.Row():
            interp_choice = gr.Radio(
                ["bilinear", "bicubic", "nearest"], value="bicubic",
                label="Resize Interpolation (预处理插值)"
            )
            
        with gr.Row():
            attention_layer_choice = gr.Dropdown(
                choices=[
                    "stage_0_block_0", "stage_0_block_1",
                    "stage_1_block_0", "stage_1_block_1", 
                    "stage_2_block_0", "stage_2_block_1", "stage_2_block_2",
                    "stage_3_block_0", "stage_3_block_1"
                ],
                value="stage_3_block_1",
                label="选择注意力层 (Attention Layer)"
            )
            
        with gr.Row():
            stage_avg_toggle = gr.Checkbox(
                value=False, 
                label="计算整个Stage的平均注意力 (Average Stage Attention)"
            )
            normalize_toggle = gr.Checkbox(
                value=True, 
                label="归一化注意力 (Normalize Attention)"
            )

        run_btn = gr.Button("🚀 Run Analysis")

        def _run(img, inter, attn_layer, stage_avg, normalize):
            return infer_with_attention(
                img, 
                interpolation=inter, 
                attention_layer=attn_layer,
                stage_average=stage_avg,
                normalize_attention=normalize
            )

        run_btn.click(
            _run,
            inputs=[in_img, interp_choice, attention_layer_choice, stage_avg_toggle, normalize_toggle],
            outputs=[out_lbl, out_attention]
        )

        gr.Markdown("""
### 说明：
- **注意力层选择**: 可以选择不同的Swin Transformer层来查看注意力模式
- **Stage平均**: 勾选后会计算选中stage中所有block的平均注意力
- **归一化**: 将注意力值归一化到0-1范围内，便于可视化
- **热力图**: 红色区域表示模型更关注的区域，蓝色区域表示关注度较低的区域
        """)

    demo.launch()

# ---------------------------------------------------------------------------
if __name__ == "__main__":
    launch_app()