Update app.py

app.py

@@ -5,13 +5,25 @@ 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
+        "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
     }
 }
 
@@ -26,7 +38,78 @@ 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
+        
+        # 修改后的标签解释逻辑 - 反转解释
+        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())
+    
+    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())
+    
+    return ai_probability
+
 def analyze_image_features(image):
+    """分析图像特征"""
     # 转换为OpenCV格式
     img_array = np.array(image)
     if len(img_array.shape) == 3 and img_array.shape[2] == 3:
@@ -46,6 +129,13 @@ def analyze_image_features(image):
         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]))
+        
+        # 颜色标准差 - 用于检测颜色分布是否自然
+        features["color_std"] = float(np.std([
+            features["avg_red"], 
+            features["avg_green"], 
+            features["avg_blue"]
+        ]))
     
     # 边缘一致性分析
     edges = cv2.Canny(img_cv, 100, 200)
@@ -64,16 +154,168 @@ def analyze_image_features(image):
         # 计算GLCM属性
         features["texture_contrast"] = float(np.mean(graycoprops(glcm, 'contrast')[0]))
         features["texture_homogeneity"] = float(np.mean(graycoprops(glcm, 'homogeneity')[0]))
+        features["texture_correlation"] = float(np.mean(graycoprops(glcm, 'correlation')[0]))
+        features["texture_energy"] = float(np.mean(graycoprops(glcm, 'energy')[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))
+        
+        # 噪声分布 - 用于检测噪声是否自然
+        features["noise_std"] = float(np.std(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
+    
+    # 频率域分析 - 检测不自然的频率分布
+    if len(img_array.shape) == 3:
+        gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
+        f_transform = np.fft.fft2(gray)
+        f_shift = np.fft.fftshift(f_transform)
+        magnitude = np.log(np.abs(f_shift) + 1)
+        
+        # 计算高频和低频成分的比例
+        h, w = magnitude.shape
+        center_h, center_w = h // 2, w // 2
+        
+        # 低频区域 (中心区域)
+        low_freq_region = magnitude[center_h-h//8:center_h+h//8, center_w-w//8:center_w+w//8]
+        low_freq_mean = np.mean(low_freq_region)
+        
+        # 高频区域 (边缘区域)
+        high_freq_mean = np.mean(magnitude) - low_freq_mean
+        
+        features["freq_ratio"] = float(high_freq_mean / max(low_freq_mean, 0.001))
     
     return features
 
+def detect_photoshop_signs(image_features):
+    """检测图像中的PS痕迹"""
+    ps_score = 0
+    ps_signs = []
+    
+    # 检查皮肤质感
+    if "texture_homogeneity" in image_features:
+        if image_features["texture_homogeneity"] > 0.4:
+            ps_score += 0.2
+            ps_signs.append("皮肤质感过于均匀")
+        elif image_features["texture_homogeneity"] > 0.3:
+            ps_score += 0.1
+            ps_signs.append("皮肤质感较为均匀")
+    
+    # 检查边缘不自然
+    if "edge_density" in image_features:
+        if image_features["edge_density"] < 0.01:
+            ps_score += 0.2
+            ps_signs.append("边缘过于平滑")
+        elif image_features["edge_density"] < 0.03:
+            ps_score += 0.1
+            ps_signs.append("边缘较为平滑")
+    
+    # 检查颜色不自然
+    if "color_std" in image_features:
+        if image_features["color_std"] > 50:
+            ps_score += 0.2
+            ps_signs.append("颜色分布极不自然")
+        elif image_features["color_std"] > 30:
+            ps_score += 0.1
+            ps_signs.append("颜色分布略不自然")
+    
+    # 检查噪点不一致
+    if "noise_level" in image_features and "noise_std" in image_features:
+        noise_ratio = image_features["noise_std"] / max(image_features["noise_level"], 0.001)
+        if noise_ratio < 0.5:
+            ps_score += 0.2
+            ps_signs.append("噪点分布不自然")
+        elif noise_ratio < 0.7:
+            ps_score += 0.1
+            ps_signs.append("噪点分布略不自然")
+    
+    # 检查频率分布不自然
+    if "freq_ratio" in image_features:
+        if image_features["freq_ratio"] < 0.2:
+            ps_score += 0.2
+            ps_signs.append("频率分布不自然，可能有过度模糊处理")
+        elif image_features["freq_ratio"] > 2.0:
+            ps_score += 0.2
+            ps_signs.append("频率分布不自然，可能有过度锐化处理")
+    
+    # 检查对称性不自然
+    if "horizontal_symmetry" in image_features and "vertical_symmetry" in image_features:
+        avg_symmetry = (image_features["horizontal_symmetry"] + image_features["vertical_symmetry"]) / 2
+        if avg_symmetry > 0.8:
+            ps_score += 0.2
+            ps_signs.append("图像对称性过高，可能经过对称处理")
+    
+    return min(ps_score, 1.0), ps_signs
+
+def get_detailed_analysis(ai_probability, ps_score, ps_signs, valid_models_count):
+    """提供更详细的分析结果"""
+    
+    # 根据有效模型数量调整置信度描述
+    confidence_prefix = ""
+    if valid_models_count >= 3:
+        confidence_prefix = "极高置信度："
+    elif valid_models_count == 2:
+        confidence_prefix = "高置信度："
+    elif valid_models_count == 1:
+        confidence_prefix = "中等置信度："
+    
+    if ai_probability > 0.8:
+        category = confidence_prefix + "高概率AI生成"
+        description = "图像几乎可以确定是由AI完全生成，几乎没有真人照片的特征。"
+    elif ai_probability > 0.6:
+        if ps_score > 0.5:
+            category = confidence_prefix + "中等概率AI生成，高概率PS修图"
+            description = "图像可能是真人照片经过大量后期处理，或是AI生成后经过修饰的图像。"
+        else:
+            category = confidence_prefix + "中等概率AI生成"
+            description = "图像有较多AI生成的特征，但也保留了一些真实照片的特点。"
+    elif ai_probability > 0.4:
+        if ps_score > 0.5:
+            category = confidence_prefix + "低概率AI生成，高概率PS修图"
+            description = "图像更可能是真人照片经过大量后期处理，PS痕迹明显。"
+        else:
+            category = confidence_prefix + "低概率AI生成"
+            description = "图像更可能是真人照片，但有一些AI生成或修饰的特征。"
+    else:
+        if ps_score > 0.6:
+            category = confidence_prefix + "真人照片，重度PS修图"
+            description = "图像基本是真人照片，但经过了大量后期处理，修饰痕迹明显。"
+        elif ps_score > 0.3:
+            category = confidence_prefix + "真人照片，中度PS修图"
+            description = "图像是真人照片，有明显的后期处理痕迹。"
+        elif ps_score > 0.1:
+            category = confidence_prefix + "真人照片，轻度PS修图"
+            description = "图像是真人照片，有少量后期处理。"
+        else:
+            category = confidence_prefix + "高概率真人照片，几乎无修图"
+            description = "图像几乎可以确定是未经大量处理的真人照片。"
+    
+    # 添加具体的PS痕迹描述
+    if ps_signs:
+        ps_details = "检测到的修图痕迹：" + "、".join(ps_signs)
+    else:
+        ps_details = "未检测到明显的修图痕迹。"
+    
+    return category, description, ps_details
+
 def detect_ai_image(image):
+    """主检测函数"""
     if image is None:
         return {"error": "未提供图像"}
     
@@ -90,36 +332,14 @@ 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(logits, dim=-1)
+                probabilities = torch.nn.functional.softmax(outputs.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())
+                # 使用适配器处理不同模型的输出
+                ai_probability = process_model_output(model_info, outputs, probabilities)
                 
                 # 添加到结果
+                predicted_class_idx = outputs.logits.argmax(-1).item()
                 results[key] = {
                     "model_name": model_info["name"],
                     "ai_probability": ai_probability,
@@ -150,32 +370,36 @@ def detect_ai_image(image):
     
     # 低边缘密度通常表示AI生成
     if image_features["edge_density"] < 0.01:
-        adjusted_probability += 0.2
+        adjusted_probability += 0.15
     
     # 高纹理均匀性通常表示AI生成
-    if image_features["texture_homogeneity"] > 0.5:
+    if "texture_homogeneity" in image_features and image_features["texture_homogeneity"] > 0.5:
         adjusted_probability += 0.1
     
     # 低噪声水平通常表示AI生成
     if image_features["noise_level"] < 0.5:
         adjusted_probability += 0.1
     
+    # 高对称性通常表示AI生成
+    if "horizontal_symmetry" in image_features and image_features["horizontal_symmetry"] > 0.8:
+        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 = "无法确定"
+    # 分析PS痕迹
+    ps_score, ps_signs = detect_photoshop_signs(image_features)
+    
+    # 获取详细分析
+    category, description, ps_details = get_detailed_analysis(adjusted_probability, ps_score, ps_signs, valid_models)
     
     # 构建最终结果
     final_result = {
         "ai_probability": adjusted_probability,
-        "original_ai_probability": final_ai_probability,
-        "confidence_level": confidence_level,
+        "ps_score": ps_score,
+        "category": category,
+        "description": description,
+        "ps_details": ps_details,
         "individual_model_results": results,
         "features": image_features
     }
@@ -188,7 +412,7 @@ iface = gr.Interface(
     inputs=gr.Image(type="pil"),
     outputs=gr.JSON(),
     title="增强型AI图像检测API",
-    description="多模型集成检测图像是否由AI生成",
+    description="多模型集成检测图像是否由AI生成，同时分析PS修图痕迹",
     examples=None,
     allow_flagging="never"
 )