File size: 793 Bytes
0e4149c 3e61f53 0e4149c 3e61f53 0e4149c 3e61f53 0e4149c 93ca441 0e4149c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
import gradio as gr
from transformers import AutoModelForImageClassification, AutoFeatureExtractor
import torch
from PIL import Image
# 加载模型和特征提取器
model_name = "microsoft/beit-base-patch16-224"
model = AutoModelForImageClassification.from_pretrained(model_name)
feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
# 定义分类函数
def classify_image(image):
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")
demo.launch()
|