import gradio as gr
from transformers import pipeline

# 1. 加载 Hugging Face Hub 上的图像分类模型
classifier = pipeline("image-classification", model="google/mobilenet_v2_1.0_224")

# 2. 推理函数
def predict(image):
    results = classifier(image)
    # 只取前 3 个结果更清晰
    return {r["label"]: float(r["score"]) for r in results[:3]}

# 3. Gradio 界面
demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=3),
    title="MobileNet Image Classification",
    description="上传一张图片，输出前 3 类及概率"
)

if __name__ == "__main__":
    demo.launch()