app.py

import os
import gc
import cv2
import torch
import timm
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
import gradio as gr
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

# --- 引入新的可视化库 ---
# pip install grad-cam
from pytorch_grad_cam import GradCAM, HiResCAM, ScoreCAM, GradCAMPlusPlus, AblationCAM, XGradCAM, EigenCAM, FullGrad, LayerCAM
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
from pytorch_grad_cam.utils.image import show_cam_on_image

# --- 配置 ---
SEED = 4421
DROP_RATE = 0.1
LOCAL_CKPT_DIR = "./checkpoints"
DEFAULT_CKPT = "convformer_b36.ai_cls"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]

# 设置随机种子
torch.manual_seed(SEED)
np.random.seed(SEED)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(SEED)

CKPT_META = {
    "caformer_b36.old": {
        "num_classes": 4,
        "head": "v7",
        "backbone": "caformer_b36.sail_in22k_ft_in1k_384",
        "repo_id": "telecomadm1145/swin-ai-detection",
        "filename": "caformer_b36_4class_96.safetensors",
        "labels": ["non_ai", "ai", "ani_non_ai", "ani_ai"],
        "input_size": 384,
    },
    "caformer_b36.v2": {
        "num_classes": 2,
        "head": "timm_cross_entropy",
        "backbone_timm_name": "hf-hub:animetimm/caformer_b36.dbv4-full",
        "repo_id": "telecomadm1145/danbooru-real-vs-ai-caformer-b36-v2",
        "filename": "pytorch_model.bin",
        "labels": ["AI", "Non-AI"],
        "input_size": 384,
    },
    "convformer_b36.ai_cls": {
        "num_classes": 2,
        "head": "timm_cross_entropy",
        "backbone_timm_name": "convformer_b36.sail_in1k",
        "repo_id": "telecomadm1145/convformer_b36.ai_cls",
        "filename": "best_checkpoint.pth",
        "labels": ["AI", "Non-AI"],
        "input_size": 224,
    },
    "deepghs/cls-ai-check-1m.caformer_s36.r512": {
        "num_classes": 2,
        "head": "timm_cross_entropy",
        "backbone_timm_name": "caformer_s36.sail_in22k_ft_in1k_384",
        "repo_id": "deepghs/cls-ai-check-1m.caformer_s36.r512",
        "filename": "model.safetensors",
        "labels": ["AI", "Non-AI"],
        "input_size": 512,
    },
}

# --- 模型定义 ---

class TimmClassifierWithHead(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.classifier = nn.Sequential(
            nn.Dropout(DROP_RATE),
            nn.Linear(self.backbone.num_features, 64),
            nn.BatchNorm1d(64),
            nn.GELU(),
            nn.Dropout(DROP_RATE * 0.8),
            nn.Linear(64, num_classes),
        )

    def forward(self, x):
        features = self.backbone(x)
        return self.classifier(features)

# --- 关键工具函数：处理 Transformer/Hybrid 输出维度 ---

def reshape_transform(tensor):
    """
    处理 Timm 新版模型 (CaFormer/ConvFormer/Swin) 的输出维度。
    如果 tensor 是 [B, H, W, C] (Channels Last)，转换为 [B, C, H, W] 给 CAM 库使用。
    """
    # 如果是 4D 张量
    if tensor.ndim == 4:
        # 启发式判断：通常 C 通道数在几百上千，而 H, W 较小。
        # 如果最后一个维度特别大，或者第1维度(索引)比第3维度小很多，那可能是 BHWC
        h, w, c = tensor.shape[1], tensor.shape[2], tensor.shape[3]
        
        # 很多 timm 模型内部 block 输出是 BHWC (Permuted)
        # 例如: tensor.shape = [1, 56, 56, 384] -> 需要变为 [1, 384, 56, 56]
        if c > h and c > w: 
            return tensor.permute(0, 3, 1, 2)
    
    # 如果是 3D 张量 (ViT 输出: [B, Tokens, C])，需要还原成 2D 图片
    if tensor.ndim == 3:
        # 这里需要知道原始 feature map 的 h, w，稍显复杂
        # 简化处理：假设是方形
        b, num_tokens, c = tensor.shape
        # 除去 cls_token (如果有)
        if num_tokens % 2 != 0: 
             # 简单假设去掉第一个
            tensor = tensor[:, 1:, :]
            num_tokens -= 1
        side = int(np.sqrt(num_tokens))
        result = tensor.reshape(b, side, side, c).permute(0, 3, 1, 2)
        return result

    return tensor

# --- 核心逻辑类 ---

class AIImageDetector:
    def __init__(self):
        self.model = None
        self.current_ckpt_name = None
        self.current_meta = None
        self.cam_engine = None # 替代原来的 grad_cam

    def _get_target_layers(self, model):
        """准确找到适合可视化的层"""
        backbone = getattr(model, 'backbone', model)

        # 优先定位到最后一个 Stage 的最后一个 Block
        if hasattr(backbone, 'stages'):
            # CaFormer / ConvFormer / ResNet 结构
            # 这是一个 ModuleList，取最后一个 stage
            last_stage = backbone.stages[-1]
            # 取最后一个 block
            if hasattr(last_stage, 'blocks'):
                 return [last_stage.blocks[-1]]
            return [last_stage]
        
        # 针对 Swin 等
        if hasattr(backbone, 'layers'):
             return [backbone.layers[-1].blocks[-1]]
        
        # 兜底
        return [list(backbone.children())[-1]]

    def load_model(self, ckpt_name):
        if ckpt_name == self.current_ckpt_name and self.model is not None:
            return

        print(f"Loading model: {ckpt_name}...")
        meta = CKPT_META[ckpt_name]
        
        # 显存清理
        if self.model is not None:
            del self.model
            del self.cam_engine
            torch.cuda.empty_cache()
            gc.collect()

        ckpt_file = hf_hub_download(repo_id=meta["repo_id"], filename=meta["filename"], local_dir=LOCAL_CKPT_DIR, force_download=False)

        if meta.get("head") == "timm_cross_entropy":
            model = timm.create_model(meta["backbone_timm_name"], pretrained=False, num_classes=meta["num_classes"])
        else:
            model = TimmClassifierWithHead(meta["backbone"], num_classes=meta["num_classes"], pretrained=False)

        model = model.to(DEVICE)

        # 权重加载逻辑
        if meta["filename"].endswith(".safetensors"):
            state_dict = load_file(ckpt_file, device=DEVICE)
        else:
            state_dict = torch.load(ckpt_file, map_location=DEVICE)

        if "model_state_dict" in state_dict:
            state_dict = state_dict["model_state_dict"]
        elif "state_dict" in state_dict:
            state_dict = state_dict["state_dict"]
        
        state_dict = {k.replace("module.", "").replace("model.", ""): v for k, v in state_dict.items()}

        model.load_state_dict(state_dict, strict=False)
        model.eval()
        
        self.model = model
        self.current_ckpt_name = ckpt_name
        self.current_meta = meta
        self.cam_engine = None # 懒加载

    def preprocess(self, image: Image.Image):
        if image.mode != "RGB":
            image = image.convert("RGB") if image.mode != "RGBA" else Image.new("RGBA", image.size, (255, 255, 255)).alpha_composite(image).convert("RGB")

        meta = self.current_meta
        size = int(meta.get("input_size", 384))
        
        image_resized = image.resize((size, size), Image.BICUBIC)
        
        # 归一化并转 Tensor
        img_np = np.array(image_resized).astype(np.float32) / 255.0
        
        # 供可视化使用 (0-1 float, RGB)
        img_viz_np = img_np.copy()

        mean = np.array(meta.get("mean", IMAGENET_MEAN), dtype=np.float32).reshape(1, 1, 3)
        std = np.array(meta.get("std", IMAGENET_STD), dtype=np.float32).reshape(1, 1, 3)
        
        img_norm = (img_np - mean) / std
        img_tensor = torch.from_numpy(img_norm).permute(2, 0, 1).unsqueeze(0).to(DEVICE)
        
        return img_tensor, img_viz_np

    def predict(self, image, ckpt_name, enable_viz, method_name):
        if image is None:
            return None, None

        self.load_model(ckpt_name)
        input_tensor, img_viz_np = self.preprocess(image)
        labels = self.current_meta["labels"]

        if enable_viz:
            target_layers = self._get_target_layers(self.model)
            
            # 动态选择 CAM 方法
            # LayerCAM 在检测 AI 伪影（纹理异常）方面通常比 GradCAM 更好
            cam_cls = LayerCAM if method_name == "LayerCAM" else GradCAMPlusPlus
            if method_name == "GradCAM": cam_cls = GradCAM
            
            # 初始化 CAM
            # 注意：reshape_transform 对于 CaFormer/ConvFormer 至关重要
            cam = cam_cls(model=self.model, target_layers=target_layers, reshape_transform=reshape_transform)
            
            # 获取预测类别作为目标
            with torch.no_grad():
                logits = self.model(input_tensor)
                probs = F.softmax(logits, dim=1)
                predicted_class = torch.argmax(probs, dim=1).item()
            
            # 生成热力图
            targets = [ClassifierOutputTarget(predicted_class)]
            grayscale_cam = cam(input_tensor=input_tensor, targets=targets)
            grayscale_cam = grayscale_cam[0, :]
            
            # 叠加
            visualization = show_cam_on_image(img_viz_np, grayscale_cam, use_rgb=True)
            viz_image = Image.fromarray(visualization)
            
            result_dict = {labels[i]: float(probs.detach().cpu().numpy()[0][i]) for i in range(len(labels))}
            
            # 清理 Hook (防止显存泄露)
            # cam.del_hooks() # 旧版本可能需要手动删，新版通常会自动管理或复用
            
            return result_dict, viz_image

        else:
            with torch.no_grad():
                logits = self.model(input_tensor)
                probs = F.softmax(logits, dim=1).cpu().numpy()[0]
            
            result_dict = {labels[i]: float(probs[i]) for i in range(len(labels))}
            return result_dict, None

detector = AIImageDetector()

# --- Gradio 界面 ---

def gr_predict(image, ckpt_name, enable_viz, method):
    return detector.predict(image, ckpt_name, enable_viz, method)

def launch():
    detector.load_model(DEFAULT_CKPT)

    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown("## 🕵️ AI Generated Image Detection & Explainability")
        
        with gr.Row():
            with gr.Column(scale=1):
                in_img = gr.Image(type="pil", label="Input Image")
                sel_ckpt = gr.Dropdown(list(CKPT_META.keys()), value=DEFAULT_CKPT, label="Model Checkpoint")
                
                with gr.Group():
                    enable_viz = gr.Checkbox(label="Visualize Heatmap", value=True)
                    # 增加了方法选择
                    viz_method = gr.Dropdown(["LayerCAM", "GradCAM++", "GradCAM"], value="LayerCAM", label="Method (LayerCAM recommended for texture)")
                    run_btn = gr.Button("Analyze", variant="primary")

            with gr.Column(scale=1):
                out_lbl = gr.Label(num_top_classes=4, label="Prediction Confidence")
                out_viz = gr.Image(type="pil", label="Attention / Heatmap")

        run_btn.click(
            fn=gr_predict, 
            inputs=[in_img, sel_ckpt, enable_viz, viz_method], 
            outputs=[out_lbl, out_viz]
        )

    demo.launch()

if __name__ == "__main__":
    launch()