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": 1.0
    }
}

# 初始化模型
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)
        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))
    
    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
                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 = float(probabilities[0][human_label_idx].item())
                elif ai_label_idx is not None:
                    # 如果有AI标签
                    ai_probability = float(probabilities[0][ai_label_idx].item())
                else:
                    # 默认使用索引1作为AI标签
                    ai_probability = float(probabilities[0][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概率
    adjusted_probability = final_ai_probability
    
    # 低边缘密度通常表示AI生成
    if image_features["edge_density"] < 0.01:
        adjusted_probability += 0.2
    
    # 高纹理均匀性通常表示AI生成
    if image_features["texture_homogeneity"] > 0.5:
        adjusted_probability += 0.1
    
    # 低噪声水平通常表示AI生成
    if image_features["noise_level"] < 0.5:
        adjusted_probability += 0.1
    
    # 确保概率在0-1范围内
    adjusted_probability = min(1.0, max(0.0, adjusted_probability))
    
    # 调整后的阈值判断
    if adjusted_probability > 0.5:  # 降低AI判定阈值
        confidence_level = "高概率AI生成"
    elif adjusted_probability < 0.2:  # 提高人类判定要求
        confidence_level = "高概率人类创作"
    else:
        confidence_level = "无法确定"
    
    # 构建最终结果
    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()