# app.py
import gradio as gr
from PIL import Image
from transformers import CLIPProcessor, CLIPModel
import time

# 模型加载（约需2分钟）
def load_model():
    start = time.time()
    model = CLIPModel.from_pretrained("vinid/plip")
    processor = CLIPProcessor.from_pretrained("vinid/plip") 
    print(f"模型加载完成，耗时：{time.time()-start:.1f}秒")
    return model, processor

model, processor = load_model()

# 预测函数
def classify_image(image, text_options):
    try:
        # 图像预处理
        processed_img = image.convert("RGB").resize((224,224))
        
        # 文本处理（自动分割逗号分隔的标签）
        labels = [t.strip() for t in text_options.split(',')]
        
        # 模型推理
        inputs = processor(
            text=labels,
            images=processed_img,
            return_tensors="pt",
            padding=True
        )
        outputs = model(**inputs)
        probs = outputs.logits_per_image.softmax(dim=1).tolist()[0]
        
        # 结果格式化
        return {label: round(prob,3) for label, prob in zip(labels, probs)}
    
    except Exception as e:
        return f"处理错误：{str(e)}"

# 界面布局
with gr.Blocks(theme=gr.themes.Soft()) as app:
    gr.Markdown("# 🩺 医学影像智能诊断系统")
    
    with gr.Row():
        image_input = gr.Image(type="pil", label="上传病理切片")
        text_input = gr.Textbox(
            label="诊断标签（逗号分隔）",
            value="恶性肿瘤, 良性病变, 炎症反应"
        )
    
    submit_btn = gr.Button("开始分析", variant="primary")
    
    output = gr.Label(label="诊断概率分布", num_top_classes=3)
    
    # 示例数据
    gr.Examples(
        examples=[
            ["sample1.jpg", "肺癌, 肺结核, 正常组织"],
            ["sample2.png", "胃癌, 胃炎, 胃溃疡"]
        ],
        inputs=[image_input, text_input]
    )

    submit_btn.click(
        fn=classify_image,
        inputs=[image_input, text_input],
        outputs=output
    )

app.launch(debug=True)