Update app.py

app.py

@@ -1,13 +1,21 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoModelForImageClassification, AutoFeatureExtractor
+import torch
+from PIL import Image
 
-# 选择 Hugging Face 预训练模型
-classifier = pipeline("feature-extraction", model="facebook/deit-base-distilled-patch16-224")
+# 加载模型和特征提取器
+model_name = "microsoft/beit-base-patch16-224"
+model = AutoModelForImageClassification.from_pretrained(model_name)
+feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
 
 # 定义分类函数
 def classify_image(image):
-    predictions = classifier(image)
-    return {"feature_vector": predictions[0]}  # 返回特征向量
+    image = feature_extractor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**image)
+    logits = outputs.logits
+    predicted_class = logits.argmax(-1).item()
+    return f"Predicted class: {predicted_class}"
 
 # 创建 Gradio 界面
 demo = gr.Interface(fn=classify_image, inputs="image", outputs="text", title="Image Classification Demo")