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.4
    },
    "model2": {
        "name": "sayakpaul/convnext-base-finetuned-ai-generated-detection",
        "processor": None,
        "model": None,
        "weight": 0.3
    },
    "model3": {
        "name": "Xenova/clip-image-classification-ai-generated",
        "processor": None,
        "model": None,
        "weight": 0.3
    }
}

# 初始化模型
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 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)
        glcm = cv2.createGLCM(gray, 5)  # 创建灰度共生矩阵
        features["texture_contrast"] = cv2.GLCMContrast(glcm)
        features["texture_homogeneity"] = cv2.GLCMHomogeneity(glcm)
    
    # 噪声分析
    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))
    
    return features

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)
                
                # 获取预测结果
                logits = outputs.logits
                predicted_class_idx = logits.argmax(-1).item()
                
                # 获取概率
                probabilities = torch.nn.functional.softmax(logits, dim=-1)
                
                # 确定AI生成概率
                ai_label_idx = None
                for idx, label in model_info["model"].config.id2label.items():
                    if "ai" in label.lower() or "generated" in label.lower() or "fake" in label.lower():
                        ai_label_idx = idx
                        break
                
                if ai_label_idx is None:
                    ai_label_idx = 1  # 默认索引1为AI生成
                
                ai_probability = float(probabilities[0][ai_label_idx].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)
                }
    
    # 计算最终加权概率
    final_ai_probability = weighted_ai_probability / max(sum(m["weight"] for k, m in models.items() if m["processor"] is not None), 1)
    
    # 分析图像特征
    image_features = analyze_image_features(image)
    
    # 确定置信度级别
    if final_ai_probability > 0.7:
        confidence_level = "高概率AI生成"
    elif final_ai_probability < 0.3:
        confidence_level = "高概率人类创作"
    else:
        confidence_level = "无法确定"
    
    # 构建最终结果
    final_result = {
        "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=[
        ["example1.jpg"],
        ["example2.jpg"]
    ],
    allow_flagging="never"
)

iface.launch()