aaappp7878 commited on
Commit
071497b
·
verified ·
1 Parent(s): 0d119a8

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +137 -32
app.py CHANGED
@@ -11,19 +11,19 @@ models = {
11
  "name": "umm-maybe/AI-image-detector",
12
  "processor": None,
13
  "model": None,
14
- "weight": 0.4
15
  },
16
  "model2": {
17
- "name": "sayakpaul/convnext-base-finetuned-ai-generated-detection",
18
  "processor": None,
19
  "model": None,
20
- "weight": 0.3
21
  },
22
  "model3": {
23
- "name": "Xenova/clip-image-classification-ai-generated",
24
  "processor": None,
25
  "model": None,
26
- "weight": 0.3
27
  }
28
  }
29
 
@@ -38,6 +38,74 @@ for key in models:
38
  models[key]["processor"] = None
39
  models[key]["model"] = None
40
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
41
  def analyze_image_features(image):
42
  # 转换为OpenCV格式
43
  img_array = np.array(image)
@@ -83,7 +151,57 @@ def analyze_image_features(image):
83
  noise = cv2.absdiff(img_cv, blurred)
84
  features["noise_level"] = float(np.mean(noise))
85
 
86
- return features
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
87
 
88
  def detect_ai_image(image):
89
  if image is None:
@@ -102,26 +220,14 @@ def detect_ai_image(image):
102
  with torch.no_grad():
103
  outputs = model_info["model"](**inputs)
104
 
105
- # 获取预测结果
106
- logits = outputs.logits
107
- predicted_class_idx = logits.argmax(-1).item()
108
-
109
  # 获取概率
110
- probabilities = torch.nn.functional.softmax(logits, dim=-1)
111
-
112
- # 确定AI生成概率
113
- ai_label_idx = None
114
- for idx, label in model_info["model"].config.id2label.items():
115
- if "ai" in label.lower() or "generated" in label.lower() or "fake" in label.lower():
116
- ai_label_idx = idx
117
- break
118
 
119
- if ai_label_idx is None:
120
- ai_label_idx = 1 # 默认索引1为AI生成
121
-
122
- ai_probability = float(probabilities[0][ai_label_idx].item())
123
 
124
  # 添加到结果
 
125
  results[key] = {
126
  "model_name": model_info["name"],
127
  "ai_probability": ai_probability,
@@ -139,22 +245,21 @@ def detect_ai_image(image):
139
  }
140
 
141
  # 计算最终加权概率
142
- final_ai_probability = weighted_ai_probability / max(sum(m["weight"] for k, m in models.items() if m["processor"] is not None), 1)
 
 
 
143
 
144
  # 分析图像特征
145
  image_features = analyze_image_features(image)
146
 
147
- # 确定置信度级别
148
- if final_ai_probability > 0.7:
149
- confidence_level = "高概率AI生成"
150
- elif final_ai_probability < 0.3:
151
- confidence_level = "高概率人类创作"
152
- else:
153
- confidence_level = "无法确定"
154
 
155
  # 构建最终结果
156
  final_result = {
157
- "ai_probability": final_ai_probability,
 
158
  "confidence_level": confidence_level,
159
  "individual_model_results": results,
160
  "features": image_features
@@ -173,4 +278,4 @@ iface = gr.Interface(
173
  allow_flagging="never"
174
  )
175
 
176
- iface.launch()
 
11
  "name": "umm-maybe/AI-image-detector",
12
  "processor": None,
13
  "model": None,
14
+ "weight": 0.5
15
  },
16
  "model2": {
17
+ "name": "microsoft/resnet-50", # 通用图像分类模型
18
  "processor": None,
19
  "model": None,
20
+ "weight": 0.25
21
  },
22
  "model3": {
23
+ "name": "google/vit-base-patch16-224", # Vision Transformer模型
24
  "processor": None,
25
  "model": None,
26
+ "weight": 0.25
27
  }
28
  }
29
 
 
38
  models[key]["processor"] = None
39
  models[key]["model"] = None
40
 
41
+ def process_model_output(model_info, outputs, probabilities):
42
+ """处理不同模型的输出,统一返回AI生成概率"""
43
+ model_name = model_info["name"].lower()
44
+
45
+ # 针对不同模型的特殊处理
46
+ if "ai-image-detector" in model_name:
47
+ # umm-maybe/AI-image-detector模型特殊处理
48
+ # 检查标签
49
+ ai_label_idx = None
50
+ human_label_idx = None
51
+
52
+ for idx, label in model_info["model"].config.id2label.items():
53
+ label_lower = label.lower()
54
+ if "ai" in label_lower or "generated" in label_lower or "fake" in label_lower:
55
+ ai_label_idx = idx
56
+ if "human" in label_lower or "real" in label_lower:
57
+ human_label_idx = idx
58
+
59
+ # 根据标签确定AI概率
60
+ if human_label_idx is not None:
61
+ # 如果有human标签,AI概率是1减去human概率
62
+ return 1 - float(probabilities[0][human_label_idx].item())
63
+ elif ai_label_idx is not None:
64
+ # 如果有AI标签
65
+ return float(probabilities[0][ai_label_idx].item())
66
+ else:
67
+ # 默认使用索引1作为AI标签
68
+ return float(probabilities[0][1].item())
69
+
70
+ elif "resnet" in model_name:
71
+ # 通用图像分类模型,使用简单启发式方法
72
+ predicted_class_idx = outputs.logits.argmax(-1).item()
73
+ # 检查是否有与AI相关的类别
74
+ predicted_class = model_info["model"].config.id2label[predicted_class_idx].lower()
75
+
76
+ # 简单启发式:检查类别名称是否包含与AI生成相关的关键词
77
+ ai_keywords = ["artificial", "generated", "synthetic", "fake", "computer"]
78
+ for keyword in ai_keywords:
79
+ if keyword in predicted_class:
80
+ return float(probabilities[0][predicted_class_idx].item())
81
+
82
+ # 如果没有明确的AI类别,返回中等概率
83
+ return 0.5
84
+
85
+ elif "vit" in model_name:
86
+ # Vision Transformer模型
87
+ predicted_class_idx = outputs.logits.argmax(-1).item()
88
+ # 同样检查类别名称
89
+ predicted_class = model_info["model"].config.id2label[predicted_class_idx].lower()
90
+
91
+ # 简单启发式:检查类别名称是否包含与AI生成相关的关键词
92
+ ai_keywords = ["artificial", "generated", "synthetic", "fake", "computer"]
93
+ for keyword in ai_keywords:
94
+ if keyword in predicted_class:
95
+ return float(probabilities[0][predicted_class_idx].item())
96
+
97
+ # 如果没有明确的AI类别,返回中等概率
98
+ return 0.5
99
+
100
+ # 默认处理
101
+ predicted_class_idx = outputs.logits.argmax(-1).item()
102
+ predicted_class = model_info["model"].config.id2label[predicted_class_idx].lower()
103
+
104
+ if "ai" in predicted_class or "generated" in predicted_class or "fake" in predicted_class:
105
+ return float(probabilities[0][predicted_class_idx].item())
106
+ else:
107
+ return 1 - float(probabilities[0][predicted_class_idx].item())
108
+
109
  def analyze_image_features(image):
110
  # 转换为OpenCV格式
111
  img_array = np.array(image)
 
151
  noise = cv2.absdiff(img_cv, blurred)
152
  features["noise_level"] = float(np.mean(noise))
153
 
154
+ # 分析对称性 - AI生成图像通常有更高的对称性
155
+ # 水平对称性
156
+ if img_cv.shape[1] % 2 == 0: # 确保宽度是偶数
157
+ left_half = img_cv[:, :img_cv.shape[1]//2]
158
+ right_half = cv2.flip(img_cv[:, img_cv.shape[1]//2:], 1)
159
+ if left_half.shape == right_half.shape:
160
+ h_symmetry = 1 - float(np.mean(cv2.absdiff(left_half, right_half)) / 255)
161
+ features["horizontal_symmetry"] = h_symmetry
162
+
163
+ # 垂直对称性
164
+ if img_cv.shape[0] % 2 == 0: # 确保高度是偶数
165
+ top_half = img_cv[:img_cv.shape[0]//2, :]
166
+ bottom_half = cv2.flip(img_cv[img_cv.shape[0]//2:, :], 0)
167
+ if top_half.shape == bottom_half.shape:
168
+ v_symmetry = 1 - float(np.mean(cv2.absdiff(top_half, bottom_half)) / 255)
169
+ features["vertical_symmetry"] = v_symmetry
170
+
171
+ # 分析��色分布 - AI生成图像通常有更平滑的颜色过渡
172
+ if len(img_cv.shape) == 3:
173
+ hsv = cv2.cvtColor(img_cv, cv2.COLOR_BGR2HSV)
174
+ hue_std = float(np.std(hsv[:,:,0]))
175
+ sat_std = float(np.std(hsv[:,:,1]))
176
+ val_std = float(np.std(hsv[:,:,2]))"] = hue_std / 180 # 归一化
177
+ features["saturation_variation"] =_variation"] = val_std /_final_decision(ai_probability, image率和图像特征做出更准确的决策础决策
178
+ if ai_probability > 0.7:
179
+ base_decision = "高概率AI生 < 0.3:
180
+ base_decision = "高概率人类创 "无法确定"
181
+
182
+
183
+ feature_score = 0
184
+
185
+ # 检查对称性 - 高对称性通常表示AI" in image_features and image_features["horizontal_symmetry"] > 0.0.1
186
+ if "vertical_symmetry" in image_features and image_features["vertical_symmetry"] > 0. 0.1
187
+
188
+ - AI生成图像通常边缘密度较低
189
+ if image_features["edge_density"] < 0.01
190
+
191
+ # 检查噪声 - AI生成图像通常噪声较低< 0.3:
192
+ feature_score += 0.1
193
+
194
+ # 检查颜色变化 - AI生成图像通常颜色变化ue_variation" in image_features and image_features["hue_variation"] < 0.1
195
+ if "saturation_variation" in image_features and image_features["saturation_variation"] 0.05
196
+
197
+ # 调整最终概率
198
+ adjusted_probability = min(1.0, max(0.0, ai_probability + feature_score))新判断
199
+ if adjusted_probability > 0.7:
200
+ return "高概率AI生成", adjusted_probability
201
+ elif adjusted_probability < 0.3:
202
+ return "高概率人类创作", adjusted_probability
203
+ else:
204
+ return "无法确定", adjusted_probability
205
 
206
  def detect_ai_image(image):
207
  if image is None:
 
220
  with torch.no_grad():
221
  outputs = model_info["model"](**inputs)
222
 
 
 
 
 
223
  # 获取概率
224
+ probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
 
 
 
 
 
 
 
225
 
226
+ # 使用适配器处理不同模型的输出
227
+ ai_probability = process_model_output(model_info, outputs, probabilities)
 
 
228
 
229
  # 添加到结果
230
+ predicted_class_idx = outputs.logits.argmax(-1).item()
231
  results[key] = {
232
  "model_name": model_info["name"],
233
  "ai_probability": ai_probability,
 
245
  }
246
 
247
  # 计算最终加权概率
248
+ if valid_models > 0:
249
+ 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)
250
+ else:
251
+ return {"error": "所有模型加载失败"}
252
 
253
  # 分析图像特征
254
  image_features = analyze_image_features(image)
255
 
256
+ # 做出最终决策
257
+ confidence_level, adjusted_probability = make_final_decision(final_ai_probability, image_features)
 
 
 
 
 
258
 
259
  # 构建最终结果
260
  final_result = {
261
+ "ai_probability": adjusted_probability,
262
+ "original_ai_probability": final_ai_probability,
263
  "confidence_level": confidence_level,
264
  "individual_model_results": results,
265
  "features": image_features
 
278
  allow_flagging="never"
279
  )
280
 
281
+ iface.launch()