app.py

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

# 使用专门的AI图像检测模型
model_name = "umm-maybe/AI-image-detector"
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

def detect_ai_image(image):
    # 处理图像
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    
    # 获取预测结果
    logits = outputs.logits
    predicted_class_idx = logits.argmax(-1).item()
    
    # 获取概率
    probabilities = torch.nn.functional.softmax(logits, dim=-1)
    
    # 获取AI生成概率
    # 通常索引1对应AI生成，索引0对应真实图像，但我们需要确认模型的标签
    ai_index = 1 if "ai" in model.config.id2label[1].lower() else 0
    ai_probability = probabilities[0][ai_index].item()
    
    # 分析图像特征
    features = analyze_image_features(image)
    
    return {
        "ai_probability": float(ai_probability),
        "features": features,
        "predicted_class": model.config.id2label[predicted_class_idx]
    }

def analyze_image_features(image):
    # 简单图像特征分析
    features = {}
    
    # 转换为numpy数组
    img_array = np.array(image)
    
    # 基本特征
    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]))
    
    return features

# 创建Gradio界面
iface = gr.Interface(
    fn=detect_ai_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.JSON(),
    title="AI图像检测API",
    description="检测图像是否由AI生成"
)

iface.launch()