Update app.py

app.py

@@ -2,46 +2,51 @@ import gradio as gr
 import torch
 from PIL import Image
 import numpy as np
+import cv2
 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]
+# 加载多个检测模型
+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):
-    # 简单图像特征分析
-    features = {}
-    
-    # 转换为numpy数组
+    # 转换为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
@@ -49,20 +54,119 @@ def analyze_image_features(image):
     features["aspect_ratio"] = image.width / max(1, image.height)
     
     # 颜色分析
-    if len(img_array.shape) == 3:  # 彩色图像
+    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生成"
+    title="增强型AI图像检测API",
+    description="多模型集成检测图像是否由AI生成",
+    examples=[
+        ["example1.jpg"],
+        ["example2.jpg"]
+    ],
+    allow_flagging="never"
 )
 
 iface.launch()