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

# 加载预训练的AI图像检测器
model_name = "SehwanHong/Stable-Diffusion-Detector"
feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

def detect_ai_image(image):
    # 处理图像
    inputs = feature_extractor(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_probability = probabilities[0][1].item()  # 假设索引1是AI生成类
    
    # 分析图像特征
    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()