app.py

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

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
        
        # 根据标签确定AI概率
        if human_label_idx is not None:
            # 如果有human标签，AI概率是1减去human概率
            return 1 - float(probabilities[0][human_label_idx].item())
        elif ai_label_idx is not None:
            # 如果有AI标签
            return float(probabilities[0][ai_label_idx].item())
        else:
            # 默认使用索引1作为AI标签
            return float(probabilities[0][1].item())
    
    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())

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]))
    
    # 边缘一致性分析
    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]))
    
    # 噪声分析
    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))
    
    # 分析对称性 - 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
    
    # 分析颜色分布 - AI生成图像通常有更平滑的颜色过渡
    if len(img_cv.shape) == 3:
        hsv = cv2.cvtColor(img_cv, cv2.COLOR_BGR2HSV)
        hue_std = float(np.std(hsv[:,:,0]))
        sat_std = float(np.std(hsv[:,:,1]))
        val_std = float(np.std(hsv[:,:,2]))"] = hue_std / 180  # 归一化
        features["saturation_variation"] =_variation"] = val_std /_final_decision(ai_probability, image率和图像特征做出更准确的决策础决策
    if ai_probability > 0.7:
        base_decision = "高概率AI生 < 0.3:
        base_decision = "高概率人类创 "无法确定"
    
    整
    feature_score = 0
    
    # 检查对称性 - 高对称性通常表示AI" in image_features and image_features["horizontal_symmetry"] > 0.0.1
    if "vertical_symmetry" in image_features and image_features["vertical_symmetry"] > 0. 0.1
    
    - AI生成图像通常边缘密度较低
    if image_features["edge_density"] < 0.01
    
    # 检查噪声 - AI生成图像通常噪声较低< 0.3:
        feature_score += 0.1
    
    # 检查颜色变化 - AI生成图像通常颜色变化ue_variation" in image_features and image_features["hue_variation"] < 0.1
    if "saturation_variation" in image_features and image_features["saturation_variation"] 0.05
    
    # 调整最终概率
    adjusted_probability = min(1.0, max(0.0, ai_probability + feature_score))新判断
    if adjusted_probability > 0.7:
        return "高概率AI生成", adjusted_probability
    elif adjusted_probability < 0.3:
        return "高概率人类创作", adjusted_probability
    else:
        return "无法确定", adjusted_probability

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)
    
    # 做出最终决策
    confidence_level, adjusted_probability = make_final_decision(final_ai_probability, image_features)
    
    # 构建最终结果
    final_result = {
        "ai_probability": adjusted_probability,
        "original_ai_probability": final_ai_probability,
        "confidence_level": confidence_level,
        "individual_model_results": results,
        "features": image_features
    }
    
    return final_result

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

iface.launch()