Update app.py

app.py

@@ -5,25 +5,13 @@ 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.5
-    },
-    "model2": {
-        "name": "microsoft/resnet-50",  # 通用图像分类模型
-        "processor": None,
-        "model": None,
-        "weight": 0.25
-    },
-    "model3": {
-        "name": "google/vit-base-patch16-224",  # Vision Transformer模型
-        "processor": None,
-        "model": None,
-        "weight": 0.25
+        "weight": 1.0
     }
 }
 
@@ -38,74 +26,6 @@ for key in models:
         models[key]["processor"] = None
         models[key]["model"] = None
 
-def process_model_output(model_info, outputs, probabilities):
-    """处理不同模型的输出，统一返回AI生成概率"""
-    model_name = model_info["name"].lower()
-    
-    # 针对不同模型的特殊处理
-    if "ai-image-detector" in model_name:
-        # umm-maybe/AI-image-detector模型特殊处理
-        # 检查标签
-        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
-        
-        # 根据标签确定AI概率
-        if human_label_idx is not None:
-            # 如果有human标签，AI概率是1减去human概率
-            return 1 - float(probabilities[0][human_label_idx].item())
-        elif ai_label_idx is not None:
-            # 如果有AI标签
-            return float(probabilities[0][ai_label_idx].item())
-        else:
-            # 默认使用索引1作为AI标签
-            return float(probabilities[0][1].item())
-    
-    elif "resnet" in model_name:
-        # 通用图像分类模型，使用简单启发式方法
-        predicted_class_idx = outputs.logits.argmax(-1).item()
-        # 检查是否有与AI相关的类别
-        predicted_class = model_info["model"].config.id2label[predicted_class_idx].lower()
-        
-        # 简单启发式：检查类别名称是否包含与AI生成相关的关键词
-        ai_keywords = ["artificial", "generated", "synthetic", "fake", "computer"]
-        for keyword in ai_keywords:
-            if keyword in predicted_class:
-                return float(probabilities[0][predicted_class_idx].item())
-        
-        # 如果没有明确的AI类别，返回中等概率
-        return 0.5
-    
-    elif "vit" in model_name:
-        # Vision Transformer模型
-        predicted_class_idx = outputs.logits.argmax(-1).item()
-        # 同样检查类别名称
-        predicted_class = model_info["model"].config.id2label[predicted_class_idx].lower()
-        
-        # 简单启发式：检查类别名称是否包含与AI生成相关的关键词
-        ai_keywords = ["artificial", "generated", "synthetic", "fake", "computer"]
-        for keyword in ai_keywords:
-            if keyword in predicted_class:
-                return float(probabilities[0][predicted_class_idx].item())
-        
-        # 如果没有明确的AI类别，返回中等概率
-        return 0.5
-    
-    # 默认处理
-    predicted_class_idx = outputs.logits.argmax(-1).item()
-    predicted_class = model_info["model"].config.id2label[predicted_class_idx].lower()
-    
-    if "ai" in predicted_class or "generated" in predicted_class or "fake" in predicted_class:
-        return float(probabilities[0][predicted_class_idx].item())
-    else:
-        return 1 - float(probabilities[0][predicted_class_idx].item())
-
 def analyze_image_features(image):
     # 转换为OpenCV格式
     img_array = np.array(image)
@@ -151,54 +71,7 @@ def analyze_image_features(image):
         noise = cv2.absdiff(img_cv, blurred)
         features["noise_level"] = float(np.mean(noise))
     
-    # 分析对称性 - AI生成图像通常有更高的对称性
-    # 水平对称性
-    if img_cv.shape[1] % 2 == 0:  # 确保宽度是偶数
-        left_half = img_cv[:, :img_cv.shape[1]//2]
-        right_half = cv2.flip(img_cv[:, img_cv.shape[1]//2:], 1)
-        if left_half.shape == right_half.shape:
-            h_symmetry = 1 - float(np.mean(cv2.absdiff(left_half, right_half)) / 255)
-            features["horizontal_symmetry"] = h_symmetry
-    
-    # 垂直对称性
-    if img_cv.shape[0] % 2 == 0:  # 确保高度是偶数
-        top_half = img_cv[:img_cv.shape[0]//2, :]
-        bottom_half = cv2.flip(img_cv[img_cv.shape[0]//2:, :], 0)
-        if top_half.shape == bottom_half.shape:
-            v_symmetry = 1 - float(np.mean(cv2.absdiff(top_half, bottom_half)) / 255)
-            features["vertical_symmetry"] = v_symmetry
-    
-    # 分析颜色分布 - AI生成图像通常有更平滑的颜色过渡
-    if len(img_cv.shape) == 3:
-        hsv = cv2.cvtColor(img_cv, cv2.COLOR_BGR2HSV)
-        hue_std = float(np.std(hsv[:,:,0]))
-        sat_std = float(np.std(hsv[:,:,1]))
-        val_std = float(np.std(hsv[:,:,2])) = hue_std /一化
-        features["saturation_variation"] =_variation"] = val_std /_final_decision(ai_probability, image率和图像特征做出更准确的决策策
-    if ai_probability > 0.7:
-        base_decision = "高概率AI生成"
-    elif ai_probability < 0.3:
-        base_decision = "高概率人类创作"
-    
-    
-    # 特征调整
-    feature_score = 0
-    
-    # 检查对称性 - 高对称性通常表示AI生成
-    if "horizontal_symmetry" in image_features and image_features["horizontal_symmetry"] > 0.850.1
-    if "vertical_symmetry" in image_features and image_features["vertical_symmetry"] > 0.85:
-        # 检查边缘 - AI生成图像通常边缘密度较低_density"] < 0.01_score += 0.1 检查噪声 - AI生成图像通常噪声较低
-    "] < 0_score += 0.1 检查颜色变化 - AI生成图像通常颜色变化更ue_variation" in image_features and image_features["hue_variation"] < 0.05_score += uration_variation" in image_features and image_features["saturation_variation"] < _score += 0.1
-    率
-    adjusted_probability = min(1.0, max(0.0, ai_probability + feature_score))
-    
-    断
-    if adjusted_probability > 0.7:
-        return "高概率AI生成", adjusted_probability
-    elif adjusted_probability < 0.3:
-        return "高概率人类创作", adjusted_probability
-    else:
-        return "无法确定", adjusted_probability
+    return features
 
 def detect_ai_image(image):
     if image is None:
@@ -217,14 +90,36 @@ def detect_ai_image(image):
                 with torch.no_grad():
                     outputs = model_info["model"](**inputs)
                 
+                # 获取预测结果
+                logits = outputs.logits
+                predicted_class_idx = logits.argmax(-1).item()
+                
                 # 获取概率
-                probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
+                probabilities = torch.nn.functional.softmax(logits, dim=-1)
+                
+                # 确定AI生成概率
+                ai_label_idx = None
+                human_label_idx = None
                 
-                # 使用适配器处理不同模型的输出
-                ai_probability = process_model_output(model_info, outputs, probabilities)
+                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
+                
+                # 根据标签确定AI概率
+                if human_label_idx is not None:
+                    # 如果有human标签，AI概率是1减去human概率
+                    ai_probability = 1 - 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())
                 
                 # 添加到结果
-                predicted_class_idx = outputs.logits.argmax(-1).item()
                 results[key] = {
                     "model_name": model_info["name"],
                     "ai_probability": ai_probability,
@@ -250,13 +145,17 @@ def detect_ai_image(image):
     # 分析图像特征
     image_features = analyze_image_features(image)
     
-    # 做出最终决策
-    confidence_level, adjusted_probability = make_final_decision(final_ai_probability, image_features)
+    # 确定置信度级别
+    if final_ai_probability > 0.7:
+        confidence_level = "高概率AI生成"
+    elif final_ai_probability < 0.3:
+        confidence_level = "高概率人类创作"
+    else:
+        confidence_level = "无法确定"
     
     # 构建最终结果
     final_result = {
-        "ai_probability": adjusted_probability,
-        "original_ai_probability": final_ai_probability,
+        "ai_probability": final_ai_probability,
         "confidence_level": confidence_level,
         "individual_model_results": results,
         "features": image_features
@@ -276,3 +175,4 @@ iface = gr.Interface(
 )
 
 iface.launch()
+