Create app.py

app.py

@@ -0,0 +1,65 @@
+python
+import gradio as gr
+import torch
+from PIL import Image
+import numpy as np
+from transformers import AutoFeatureExtractor, AutoModelForImageClassification
+
+# 加载预训练的AI图像检测器
+model_name = "SehwanHong/Stable-Diffusion-Detector"
+feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
+model = AutoModelForImageClassification.from_pretrained(model_name)
+
+def detect_ai_image(image):
+    # 处理图像
+    inputs = feature_extractor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    
+    # 获取预测结果
+    logits = outputs.logits
+    predicted_class_idx = logits.argmax(-1).item()
+    
+    # 获取概率
+    probabilities = torch.nn.functional.softmax(logits, dim=-1)
+    ai_probability = probabilities[0][1].item()  # 假设索引1是AI生成类
+    
+    # 分析图像特征
+    features = analyze_image_features(image)
+    
+    return {
+        "ai_probability": float(ai_probability),
+        "features": features,
+        "predicted_class": model.config.id2label[predicted_class_idx]
+    }
+
+def analyze_image_features(image):
+    # 简单图像特征分析
+    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]))
+    
+    return features
+
+# 创建Gradio界面
+iface = gr.Interface(
+    fn=detect_ai_image,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.JSON(),
+    title="AI图像检测API",
+    description="检测图像是否由AI生成"
+)
+
+iface.launch()