Update app.py

app.py

@@ -2,16 +2,16 @@ import gradio as gr
 import torch
 from PIL import Image
 import numpy as np
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+from transformers import AutoImageProcessor, AutoModelForImageClassification
 
-# 加载预训练的AI图像检测器
-model_name = "SehwanHong/Stable-Diffusion-Detector"
-feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+# 使用专门的AI图像检测模型
+model_name = "umm-maybe/AI-image-detector"
+processor = AutoImageProcessor.from_pretrained(model_name)
 model = AutoModelForImageClassification.from_pretrained(model_name)
 
 def detect_ai_image(image):
     # 处理图像
-    inputs = feature_extractor(images=image, return_tensors="pt")
+    inputs = processor(images=image, return_tensors="pt")
     with torch.no_grad():
         outputs = model(**inputs)
     
@@ -21,7 +21,11 @@ def detect_ai_image(image):
     
     # 获取概率
     probabilities = torch.nn.functional.softmax(logits, dim=-1)
-    ai_probability = probabilities[0][1].item()  # 假设索引1是AI生成类
+    
+    # 获取AI生成概率
+    # 通常索引1对应AI生成，索引0对应真实图像，但我们需要确认模型的标签
+    ai_index = 1 if "ai" in model.config.id2label[1].lower() else 0
+    ai_probability = probabilities[0][ai_index].item()
     
     # 分析图像特征
     features = analyze_image_features(image)
@@ -62,4 +66,3 @@ iface = gr.Interface(
 )
 
 iface.launch()
-