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": 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 } } # 初始化模型 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 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) 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)) # 分析对称性 - 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 / 180 # 归一化 features["saturation_variation"] =_variation"] = val_std /_final_decision(ai_probability, image率和图像特征做出更准确的决策础决策 if ai_probability > 0.7: base_decision = "高概率AI生 < 0.3: base_decision = "高概率人类创 "无法确定" 整 feature_score = 0 # 检查对称性 - 高对称性通常表示AI" in image_features and image_features["horizontal_symmetry"] > 0.0.1 if "vertical_symmetry" in image_features and image_features["vertical_symmetry"] > 0. 0.1 - AI生成图像通常边缘密度较低 if image_features["edge_density"] < 0.01 # 检查噪声 - AI生成图像通常噪声较低< 0.3: feature_score += 0.1 # 检查颜色变化 - AI生成图像通常颜色变化ue_variation" in image_features and image_features["hue_variation"] < 0.1 if "saturation_variation" in image_features and image_features["saturation_variation"] 0.05 # 调整最终概率 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 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) # 获取概率 probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1) # 使用适配器处理不同模型的输出 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, "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) # 做出最终决策 confidence_level, adjusted_probability = make_final_decision(final_ai_probability, image_features) # 构建最终结果 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()