import gradio as gr
from transformers import pipeline
import spaces

# 全局变量
classifier = None

@spaces.GPU
def predict_plant(image):
    global classifier
    if image is None:
        return None

    try:
        if classifier is None:
            print("正在加载模型 (Swin Tiny，绝对可用版)...")
            # 【终极修改】换成微软官方绝对有效的模型，识别叶子非常强
            classifier = pipeline(
                task="image-classification",
                model="microsoft/swin-tiny-patch4-window7-224",
                top_k=5
            )
            print("模型加载完成！")

        results = classifier(image)
        
        formatted_results = {}
        for res in results:
            if res["score"] > 0.001:
                formatted_results[res["label"]] = float(res["score"])
                
        if not formatted_results:
            return {"未能识别具体植物": 1.0}
            
        return formatted_results

    except Exception as e:
        return f"发生错误: {str(e)}"


demo = gr.Interface(
    fn=predict_plant,
    # 提示语加了特写建议
    inputs=gr.Image(type="pil", sources=["upload", "webcam"], label="🌿 请拍植物叶片特写 (避开泥土和花盆)"),
    outputs=gr.Label(num_top_classes=5, label="🔍 视觉识别结果 (Top 5)"),
    title="🌱 植物视觉识别 (微软 Swin Tiny)"
)

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