Update app.py

app.py

@@ -1,64 +1,52 @@
 import gradio as gr
-import torch
-from PIL import Image
-import numpy as np
-from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+from transformers import pipeline
 
-# 加载预训练的AI图像检测器
-model_name = "SehwanHong/Stable-Diffusion-Detector"
-feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
-model = AutoModelForImageClassification.from_pretrained(model_name)
+# 使用公开可用的AI文本检测模型
+# 这个模型专门用于检测AI生成文本
+detector = pipeline("text-classification", model="Xenova/distilbert-base-ai-generated-text-detection")
 
-def detect_ai_image(image):
-    # 处理图像
-    inputs = feature_extractor(images=image, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model(**inputs)
+def detect_ai_text(text):
+    if not text or len(text.strip()) < 50:
+        return {"error": "文本太短，无法可靠检测"}
     
-    # 获取预测结果
-    logits = outputs.logits
-    predicted_class_idx = logits.argmax(-1).item()
+    result = detector(text)
     
-    # 获取概率
-    probabilities = torch.nn.functional.softmax(logits, dim=-1)
-    ai_probability = probabilities[0][1].item()  # 假设索引1是AI生成类
+    # 提取结果
+    label = result[0]["label"]
+    score = result[0]["score"]
     
-    # 分析图像特征
-    features = analyze_image_features(image)
+    # 格式化为人类可读结果
+    if "ai" in label.lower():  # AI生成
+        ai_probability = score
+    else:  # 人类撰写
+        ai_probability = 1 - score
+    
+    # 分析特征
+    features = analyze_text_features(text)
     
     return {
         "ai_probability": float(ai_probability),
         "features": features,
-        "predicted_class": model.config.id2label[predicted_class_idx]
+        "confidence": float(score),
+        "label": label
     }
 
-def analyze_image_features(image):
-    # 简单图像特征分析
+def analyze_text_features(text):
+    # 简单文本特征分析
     features = {}
-    
-    # 转换为numpy数组
-    img_array = np.array(image)
-    
-    # 基本特征
-    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]))
+    features["length"] = len(text)
+    features["avg_word_length"] = sum(len(word) for word in text.split()) / max(1, len(text.split()))
+    features["unique_words_ratio"] = len(set(text.lower().split())) / max(1, len(text.split()))
     
     return features
 
 # 创建Gradio界面
 iface = gr.Interface(
-    fn=detect_ai_image,
-    inputs=gr.Image(type="pil"),
+    fn=detect_ai_text,
+    inputs=gr.Textbox(lines=10, placeholder="粘贴要检测的文本..."),
     outputs=gr.JSON(),
-    title="AI图像检测API",
-    description="检测图像是否由AI生成"
+    title="AI文本检测API",
+    description="检测文本是否由AI生成"
 )
 
 iface.launch()