Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -11,19 +11,19 @@ models = {
|
|
11 |
"name": "umm-maybe/AI-image-detector",
|
12 |
"processor": None,
|
13 |
"model": None,
|
14 |
-
"weight": 0.
|
15 |
},
|
16 |
"model2": {
|
17 |
-
"name": "
|
18 |
"processor": None,
|
19 |
"model": None,
|
20 |
-
"weight": 0.
|
21 |
},
|
22 |
"model3": {
|
23 |
-
"name": "
|
24 |
"processor": None,
|
25 |
"model": None,
|
26 |
-
"weight": 0.
|
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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 |
-
|
120 |
-
|
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 |
-
|
|
|
|
|
|
|
143 |
|
144 |
# 分析图像特征
|
145 |
image_features = analyze_image_features(image)
|
146 |
|
147 |
-
#
|
148 |
-
|
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":
|
|
|
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()
|