import gradio as gr
import torch
from PIL import Image
import numpy as np
import cv2
from transformers import AutoImageProcessor, AutoModelForImageClassification

# 加载多个检测模型
models = {
    "model1": {
        "name": "umm-maybe/AI-image-detector",
        "processor": None,
        "model": None,
        "weight": 0.5
    },
    "model2": {
        "name": "microsoft/resnet-50",  # 通用图像分类模型
        "processor": None,
        "model": None,
        "weight": 0.25
    },
    "model3": {
        "name": "google/vit-base-patch16-224",  # Vision Transformer模型
        "processor": None,
        "model": None,
        "weight": 0.25
    }
}

# 初始化模型
for key in models:
    try:
        models[key]["processor"] = AutoImageProcessor.from_pretrained(models[key]["name"])
        models[key]["model"] = AutoModelForImageClassification.from_pretrained(models[key]["name"])
        print(f"成功加载模型: {models[key]['name']}")
    except Exception as e:
        print(f"加载模型 {models[key]['name']} 失败: {str(e)}")
        models[key]["processor"] = None
        models[key]["model"] = None


## 2. 模型输出处理

python
def process_model_output(model_info, outputs, probabilities):
    """处理不同模型的输出，统一返回AI生成概率"""
    model_name = model_info["name"].lower()
    
    # 针对不同模型的特殊处理
    if "ai-image-detector" in model_name:
        # umm-maybe/AI-image-detector模型特殊处理
        # 检查标签
        ai_label_idx = None
        human_label_idx = None
        
        for idx, label in model_info["model"].config.id2label.items():
            label_lower = label.lower()
            if "ai" in label_lower or "generated" in label_lower or "fake" in label_lower:
                ai_label_idx = idx
            if "human" in label_lower or "real" in label_lower:
                human_label_idx = idx
        
        # 修正后的标签解释逻辑
        if human_label_idx is not None:
            # 如果预测为human，则AI概率应该低
            ai_probability = 1 - float(probabilities[0][human_label_idx].item())
        elif ai_label_idx is not None:
            # 如果预测为AI，则AI概率应该高
            ai_probability = float(probabilities[0][ai_label_idx].item())
        else:
            # 默认情况
            ai_probability = 0.5
    
    elif "resnet" in model_name:
        # 通用图像分类模型，使用简单启发式方法
        predicted_class_idx = outputs.logits.argmax(-1).item()
        # 检查是否有与AI相关的类别
        predicted_class = model_info["model"].config.id2label[predicted_class_idx].lower()
        
        # 简单启发式：检查类别名称是否包含与AI生成相关的关键词
        ai_keywords = ["artificial", "generated", "synthetic", "fake", "computer"]
        for keyword in ai_keywords:
            if keyword in predicted_class:
                return float(probabilities[0][predicted_class_idx].item())
        
        # 如果没有明确的AI类别，返回中等概率
        return 0.5
    
    elif "vit" in model_name:
        # Vision Transformer模型
        predicted_class_idx = outputs.logits.argmax(-1).item()
        # 同样检查类别名称
        predicted_class = model_info["model"].config.id2label[predicted_class_idx].lower()
        
        # 简单启发式：检查类别名称是否包含与AI生成相关的关键词
        ai_keywords = ["artificial", "generated", "synthetic", "fake", "computer"]
        for keyword in ai_keywords:
            if keyword in predicted_class:
                return float(probabilities[0][predicted_class_idx].item())
        
        # 如果没有明确的AI类别，返回中等概率
        return 0.5
    
    # 默认处理
    predicted_class_idx = outputs.logits.argmax(-1).item()
    predicted_class = model_info["model"].config.id2label[predicted_class_idx].lower()
    
    if "ai" in predicted_class or "generated" in predicted_class or "fake" in predicted_class:
        return float(probabilities[0][predicted_class_idx].item())
    else:
        return 1 - float(probabilities[0][predicted_class_idx].item())
    
    return ai_probability


## 3. 图像特征分析

python
def analyze_image_features(image):
    """分析图像特征"""
    # 转换为OpenCV格式
    img_array = np.array(image)
    if len(img_array.shape) == 3 and img_array.shape[2] == 3:
        img_cv = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
    else:
        img_cv = img_array
    
    features = {}
    
    # 基本特征
    features["width"] = image.width
    features["height"] = image.height
    features["aspect_ratio"] = image.width / max(1, image.height)
    
    # 颜色分析
    if len(img_array.shape) == 3:
        features["avg_red"] = float(np.mean(img_array[:,:,0]))
        features["avg_green"] = float(np.mean(img_array[:,:,1]))
        features["avg_blue"] = float(np.mean(img_array[:,:,2]))
        
        # 颜色标准差 - 用于检测颜色分布是否自然
        features["color_std"] = float(np.std([
            features["avg_red"], 
            features["avg_green"], 
            features["avg_blue"]
        ]))
    
    # 边缘一致性分析
    edges = cv2.Canny(img_cv, 100, 200)
    features["edge_density"] = float(np.sum(edges > 0) / (image.width * image.height))
    
    # 纹理分析 - 使用灰度共生矩阵
    if len(img_array.shape) == 3:
        gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
        from skimage.feature import graycomatrix, graycoprops
        
        # 计算GLCM
        distances = [5]
        angles = [0, np.pi/4, np.pi/2, 3*np.pi/4]
        glcm = graycomatrix(gray, distances=distances, angles=angles, symmetric=True, normed=True)
        
        # 计算GLCM属性
        features["texture_contrast"] = float(np.mean(graycoprops(glcm, 'contrast')[0]))
        features["texture_homogeneity"] = float(np.mean(graycoprops(glcm, 'homogeneity')[0]))
        features["texture_correlation"] = float(np.mean(graycoprops(glcm, 'correlation')[0]))
        features["texture_energy"] = float(np.mean(graycoprops(glcm, 'energy')[0]))
    
    # 噪声分析
    if len(img_array.shape) == 3:
        blurred = cv2.GaussianBlur(img_cv, (5, 5), 0)
        noise = cv2.absdiff(img_cv, blurred)
        features["noise_level"] = float(np.mean(noise))
        
        # 噪声分布 - 用于检测噪声是否自然
        features["noise_std"] = float(np.std(noise))
    
    # 对称性分析 - AI生成图像通常有更高的对称性
    if img_cv.shape[1] % 2 == 0:  # 确保宽度是偶数
        left_half = img_cv[:, :img_cv.shape[1]//2]
        right_half = cv2.flip(img_cv[:, img_cv.shape[1]//2:], 1)
        if left_half.shape == right_half.shape:
            h_symmetry = 1 - float(np.mean(cv2.absdiff(left_half, right_half)) / 255)
            features["horizontal_symmetry"] = h_symmetry
    
    if img_cv.shape[0] % 2 == 0:  # 确保高度是偶数
        top_half = img_cv[:img_cv.shape[0]//2, :]
        bottom_half = cv2.flip(img_cv[img_cv.shape[0]//2:, :], 0)
        if top_half.shape == bottom_half.shape:
            v_symmetry = 1 - float(np.mean(cv2.absdiff(top_half, bottom_half)) / 255)
            features["vertical_symmetry"] = v_symmetry
    
    # 频率域分析 - 检测不自然的频率分布
    if len(img_array.shape) == 3:
        gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
        f_transform = np.fft.fft2(gray)
        f_shift = np.fft.fftshift(f_transform)
        magnitude = np.log(np.abs(f_shift) + 1)
        
        # 计算高频和低频成分的比例
        h, w = magnitude.shape
        center_h, center_w = h // 2, w // 2
        
        # 低频区域 (中心区域)
        low_freq_region = magnitude[center_h-h//8:center_h+h//8, center_w-w//8:center_w+w//8]
        low_freq_mean = np.mean(low_freq_region)
        
        # 高频区域 (边缘区域)
        high_freq_mean = np.mean(magnitude) - low_freq_mean
        
        features["freq_ratio"] = float(high_freq_mean / max(low_freq_mean, 0.001))
    
    return features


## 4. AI特征检查

python
def check_ai_specific_features(image_features):
    """检查AI生成图像的典型特征"""
    ai_score = 0
    ai_signs = []
    
    # 检查对称性 - AI生成图像通常对称性高
    if "horizontal_symmetry" in image_features and "vertical_symmetry" in image_features:
        avg_symmetry = (image_features["horizontal_symmetry"] + image_features["vertical_symmetry"]) / 2
        if avg_symmetry > 0.7:
            ai_score += 0.3
            ai_signs.append("图像对称性异常高")
    
    # 检查纹理相关性 - AI生成图像通常纹理相关性高
    if "texture_correlation" in image_features and image_features["texture_correlation"] > 0.9:
        ai_score += 0.2
        ai_signs.append("纹理相关性异常高")
    
    # 检查边缘与噪声的关系 - AI生成图像通常边缘清晰但噪声不自然
    if "edge_density" in image_features and "noise_level" in image_features:
        edge_noise_ratio = image_features["edge_density"] / max(image_features["noise_level"], 0.001)
        if edge_noise_ratio < 0.01:
            ai_score += 0.2
            ai_signs.append("边缘与噪声分布不自然")
    
    # 检查颜色平滑度 - AI生成图像通常颜色过渡更平滑
    if "color_std" in image_features and image_features["color_std"] < 10:
        ai_score += 0.2
        ai_signs.append("颜色过渡异常平滑")
    
    # 检查纹理能量 - AI生成图像通常纹理能量分布不自然
    if "texture_energy" in image_features and image_features["texture_energy"] < 0.02:
        ai_score += 0.2
        ai_signs.append("纹理能量分布不自然")
    
    # 检查频率比例 - AI生成图像通常频率分布不自然
    if "freq_ratio" in image_features:
        if image_features["freq_ratio"] < 0.1 or image_features["freq_ratio"] > 2.0:
            ai_score += 0.2
            ai_signs.append("频率分布不自然")
    
    return min(ai_score, 1.0), ai_signs


## 5. PS痕迹检测

python
def detect_photoshop_signs(image_features):
    """检测图像中的PS痕迹"""
    ps_score = 0
    ps_signs = []
    
    # 检查皮肤质感
    if "texture_homogeneity" in image_features:
        if image_features["texture_homogeneity"] > 0.4:
            ps_score += 0.2
            ps_signs.append("皮肤质感过于均匀")
        elif image_features["texture_homogeneity"] > 0.3:
            ps_score += 0.1
            ps_signs.append("皮肤质感较为均匀")
    
    # 检查边缘不自然
    if "edge_density" in image_features:
        if image_features["edge_density"] < 0.01:
            ps_score += 0.2
            ps_signs.append("边缘过于平滑")
        elif image_features["edge_density"] < 0.03:
            ps_score += 0.1
            ps_signs.append("边缘较为平滑")
    
    # 检查颜色不自然
    if "color_std" in image_features:
        if image_features["color_std"] > 50:
            ps_score += 0.2
            ps_signs.append("颜色分布极不自然")
        elif image_features["color_std"] > 30:
            ps_score += 0.1
            ps_signs.append("颜色分布略不自然")
    
    # 检查噪点不一致
    if "noise_level" in image_features and "noise_std" in image_features:
        noise_ratio = image_features["noise_std"] / max(image_features["noise_level"], 0.001)
        if noise_ratio < 0.5:
            ps_score += 0.2
            ps_signs.append("噪点分布不自然")
        elif noise_ratio < 0.7:
            ps_score += 0.1
            ps_signs.append("噪点分布略不自然")
    
    # 检查频率分布不自然
    if "freq_ratio" in image_features:
        if image_features["freq_ratio"] < 0.2:
            ps_score += 0.2
            ps_signs.append("频率分布不自然，可能有过度模糊处理")
        elif image_features["freq_ratio"] > 2.0:
            ps_score += 0.2
            ps_signs.append("频率分布不自然，可能有过度锐化处理")
    
    return min(ps_score, 1.0), ps_signs


## 6. 结果分析与分类

python
def get_detailed_analysis(ai_probability, ps_score, ps_signs, ai_signs, valid_models_count):
    """提供更详细的分析结果"""
    
    # 根据有效模型数量调整置信度描述
    confidence_prefix = ""
    if valid_models_count >= 3:
        confidence_prefix = "极高置信度："
    elif valid_models_count == 2:
        confidence_prefix = "高置信度："
    elif valid_models_count == 1:
        confidence_prefix = "中等置信度："
    
    # 调整后的阈值判断
    if ai_probability > 0.6:  # 降低为0.6
        category = confidence_prefix + "高概率AI生成"
        description = "图像很可能是由AI完全生成，几乎没有真人照片的特征。"
    elif ai_probability > 0.4:  # 降低为0.4
        if ps_score > 0.5:
            category = confidence_prefix + "中等概率AI生成，高概率PS修图"
            description = "图像可能是真人照片经过大量后期处理，或是AI生成后经过修饰的图像。"
        else:
            category = confidence_prefix + "中等概率AI生成"
            description = "图像有较多AI生成的特征，但也保留了一些真实照片的特点。"
    elif ai_probability > 0.3:  # 降低为0.3
        if ps_score > 0.5:
            category = confidence_prefix + "低概率AI生成，高概率PS修图"
            description = "图像更可能是真人照片经过大量后期处理，PS痕迹明显。"
        else:
            category = confidence_prefix + "低概率AI生成"
            description = "图像更可能是真人照片，但有一些AI生成或修饰的特征。"
    else:
        if ps_score > 0.6:
            category = confidence_prefix + "真人照片，重度PS修图"
            description = "图像基本是真人照片，但经过了大量后期处理，修饰痕迹明显。"
        elif ps_score > 0.3:
            category = confidence_prefix + "真人照片，中度PS修图"
            description = "图像是真人照片，有明显的后期处理痕迹。"
        elif ps_score > 0.1:
            category = confidence_prefix + "真人照片，轻度PS修图"
            description = "图像是真人照片，有少量后期处理。"
        else:
            category = confidence_prefix + "高概率真人照片，几乎无修图"
            description = "图像几乎可以确定是未经大量处理的真人照片。"
    
    # 添加具体的PS痕迹描述
    if ps_signs:
        ps_details = "检测到的修图痕迹：" + "、".join(ps_signs)
    else:
        ps_details = "未检测到明显的修图痕迹。"
    
    # 添加AI特征描述
    if ai_signs:
        ai_details = "检测到的AI特征：" + "、".join(ai_signs)
    else:
        ai_details = "未检测到明显的AI生成特征。"
    
    return category, description, ps_details, ai_details


## 7. 主检测函数

python
def detect_ai_image(image):
    """主检测函数"""
    if image is None:
        return {"error": "未提供图像"}
    
    results = {}
    valid_models = 0
    weighted_ai_probability = 0
    
    # 使用每个模型进行预测
    for key, model_info in models.items():
        if model_info["processor"] is not None and model_info["model"] is not None:
            try:
                # 处理图像
                inputs = model_info["processor"](images=image, return_tensors="pt")
                with torch.no_grad():
                    outputs = model_info["model"](**inputs)
                
                # 获取概率
                probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
                
                # 使用适配器处理不同模型的输出
                ai_probability = process_model_output(model_info, outputs, probabilities)
                
                # 添加到结果
                predicted_class_idx = outputs.logits.argmax(-1).item()
                results[key] = {
                    "model_name": model_info["name"],
                    "ai_probability": ai_probability,
                    "predicted_class": model_info["model"].config.id2label[predicted_class_idx]
                }
                
                # 累加加权概率
                weighted_ai_probability += ai_probability * model_info["weight"]
                valid_models += 1
            
            except Exception as e:
                results[key] = {
                    "model_name": model_info["name"],
                    "error": str(e)
                }
    
    # 计算最终加权概率
    if valid_models > 0:
        final_ai_probability = weighted_ai_probability / sum(m["weight"] for k, m in models.items() if m["processor"] is not None and m["model"] is not None)
    else:
        return {"error": "所有模型加载失败"}
    
    # 分析图像特征
    image_features = analyze_image_features(image)
    
    # 检查AI特定特征
    ai_feature_score, ai_signs = check_ai_specific_features(image_features)
    
    # 分析PS痕迹
    ps_score, ps_signs = detect_photoshop_signs(image_features)
    
    # 应用特征权重调整AI概率
    adjusted_probability = final_ai_probability
    
    # 如果AI特征分数高，大幅提高AI概率
    if ai_feature_score > 0.5:
        adjusted_probability = max(adjusted_probability, 0.7)
    elif ai_feature_score > 0.3:
        adjusted_probability = max(adjusted_probability, 0.5)
    
    # 高对称性是AI生成的强烈指标
    if "horizontal_symmetry" in image_features and image_features["horizontal_symmetry"] > 0.7:
        adjusted_probability += 0.15
    if "vertical_symmetry" in image_features and image_features["vertical_symmetry"] > 0.7:
        adjusted_probability += 0.15
    
    # 高纹理相关性通常表示AI生成
    if "texture_correlation" in image_features and image_features["texture_correlation"] > 0.9:
        adjusted_probability += 0.1
    
    # 低边缘密度通常表示AI生成
    if image_features["edge_density"] < 0.01:
        adjusted_probability += 0.1
    
    # 确保概率在0-1范围内
    adjusted_probability = min(1.0, max(0.0, adjusted_probability))
    
    # 如果umm-maybe/AI-image-detector模型的预测与其他模型不一致，增加其权重
    if "model1" in results and "ai_probability" in results["model1"]:
        ai_detector_prob = results["model1"]["ai_probability"]
        # 如果专用AI检测器给出的概率与调整后概率差异大，增加其权重
        if abs(ai_detector_prob - adjusted_probability) > 0.3:
            adjusted_probability = (adjusted_probability + ai_detector_prob * 2) / 3
    
    # 获取详细分析
    category, description, ps_details, ai_details = get_detailed_analysis(
        adjusted_probability, ps_score, ps_signs, ai_signs, valid_models
    )
    
    # 构建最终结果
    final_result = {
        "ai_probability": adjusted_probability,
        "original_ai_probability": final_ai_probability,
        "ps_score": ps_score,
        "ai_feature_score": ai_feature_score,
        "category": category,
        "description": description,
        "ps_details": ps_details,
        "ai_details": ai_details,
        "individual_model_results": results,
        "features": image_features
    }
    
    return final_result


## 8. Gradio界面

python
# 创建Gradio界面
iface = gr.Interface(
    fn=detect_ai_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.JSON(),
    title="增强型AI图像检测API",
    description="多模型集成检测图像是否由AI生成，同时分析PS修图痕迹",
    examples=None,
    allow_flagging="never"
)

iface.launch()