Upload app.py

app.py

@@ -0,0 +1,211 @@
+import gradio as gr
+import torch
+import numpy as np
+from pathlib import Path
+import re
+from Model import OmniPathWithInterTaskAttention
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import tempfile
+import os
+
+# 设备设置
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+print(f"使用设备: {device}")
+
+# 预加载模型（避免重复加载）
+@torch.no_grad()
+def load_models():
+    """预加载必要的模型"""
+    # 1. 加载分类模型
+    ckpt_path = "best_model.pth"
+    if not Path(ckpt_path).exists():
+        raise FileNotFoundError(f"找不到模型文件: {ckpt_path}")
+    
+    ckpt = torch.load(ckpt_path, map_location=device)
+    label_mappings = ckpt.get('label_mappings', None)
+    if not label_mappings:
+        raise ValueError("checkpoint 中缺少 label_mappings")
+    
+    ck_cfg = ckpt.get('config', {})
+    feature_dim = 512  # 根据你的实际特征维度调整
+    hidden_dim = int(ck_cfg.get('hidden_dim', 256))
+    dropout = float(ck_cfg.get('dropout', 0.3))
+    use_inter_task_attention = bool(ck_cfg.get('use_inter_task_attention', True))
+    inter_task_heads = int(ck_cfg.get('inter_task_heads', 4))
+
+    classification_model = OmniPathWithInterTaskAttention(
+        label_mappings=label_mappings,
+        feature_dim=feature_dim,
+        hidden_dim=hidden_dim,
+        dropout=dropout,
+        use_inter_task_attention=use_inter_task_attention,
+        inter_task_heads=inter_task_heads
+    ).to(device)
+    classification_model.load_state_dict(ckpt['model_state_dict'], strict=False)
+    classification_model.eval()
+    
+    # 2. 加载文本生成模型
+    llm_model_name = "Qwen/Qwen3-0.6B"
+    tokenizer = AutoTokenizer.from_pretrained(llm_model_name)
+    llm_model = AutoModelForCausalLM.from_pretrained(
+        llm_model_name,
+        torch_dtype="auto",
+        device_map="auto"
+    )
+    
+    return classification_model, llm_model, tokenizer, label_mappings
+
+# 预加载模型
+classification_model, llm_model, tokenizer, label_mappings = load_models()
+
+def build_prompt(pred_names, pred_scores):
+    """构建提示词"""
+    def get_pred(task_name):
+        name = pred_names.get(task_name, "N/A")
+        score = pred_scores.get(task_name, 0.0)
+        return f"{name} (confidence: {score:.1%})"
+
+    cancer_type = get_pred('cancer_type')
+    pathologic_stage = get_pred('pathologic_stage')
+    clinical_stage = get_pred('clinical_stage')
+    histological_type = get_pred('histological_type')
+
+    prompt = (
+        "You are a professional medical report generator. "
+        "Based on the patient's pathological classification and diagnostic model results, "
+        f"the cancer_type is {cancer_type}, "
+        f"pathologic_stage is {pathologic_stage}, "
+        f"clinical_stage is {clinical_stage}, "
+        f"and histological_type is {histological_type}. "
+        "Please write a concise English summary describing the diagnosis, staging interpretation, and general clinical implications as a short paragraph. "
+        "Avoid placeholders and avoid repeating words."
+    )
+    return prompt
+
+def generate_description(predictions, confidences):
+    """生成描述"""
+    prompt = build_prompt(predictions, confidences)
+    
+    messages = [{"role": "user", "content": prompt}]
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=False
+    )
+    
+    model_inputs = tokenizer([text], return_tensors="pt").to(llm_model.device)
+    generated_ids = llm_model.generate(
+        **model_inputs,
+        max_new_tokens=500,  # 减少token数量以加快生成速度
+        do_sample=True,
+        temperature=0.7,
+    )
+    
+    output_ids = generated_ids[0][len(model_inputs.input_ids[0]):].tolist()
+    try:
+        index = len(output_ids) - output_ids[::-1].index(151668)
+    except ValueError:
+        index = 0
+        
+    content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
+    return {"prompt": prompt, "description": content}
+
+def process_npy_file(npy_file):
+    """处理上传的NPY文件"""
+    if npy_file is None:
+        return "请先上传NPY文件", "", ""
+    
+    try:
+        # 读取NPY文件
+        arr = np.load(npy_file.name, allow_pickle=False)
+        if not isinstance(arr, np.ndarray) or arr.ndim != 2:
+            return "错误: NPY文件必须是二维特征矩阵", "", ""
+        
+        features = torch.from_numpy(arr).float()
+        
+        # 提取短ID
+        p = Path(npy_file.name)
+        m = re.search(r'(TCGA-[A-Z0-9]{2}-[A-Z0-9]{4})', p.name.upper())
+        short_id = m.group(1) if m else p.stem[:12]
+        
+        # 推理
+        feat_batch = features.unsqueeze(0).to(device)
+        outputs = classification_model(feat_batch)
+        
+        # 解码结果
+        pred_names, pred_scores = {}, {}
+        for task_name, logits in outputs.items():
+            probs = torch.softmax(logits[0], dim=-1)
+            idx = int(torch.argmax(probs).item())
+            classes = label_mappings[task_name]['classes']
+            class_name = classes[idx] if 0 <= idx < len(classes) else str(idx)
+            pred_names[task_name] = class_name
+            pred_scores[task_name] = float(probs[idx].item())
+        
+        # 生成描述
+        desc_result = generate_description(pred_names, pred_scores)
+        
+        # 格式化结果显示
+        results_text = f"患者ID: {short_id}\n\n预测结果:\n"
+        for task, name in pred_names.items():
+            results_text += f"- {task}: {name} (置信度: {pred_scores.get(task, 0.0):.3f})\n"
+        
+        return results_text, desc_result['prompt'], desc_result['description']
+        
+    except Exception as e:
+        return f"处理过程中出现错误: {str(e)}", "", ""
+
+# 创建Gradio界面
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 🏥 医学病理诊断分析系统
+    
+    上传病理特征NPY文件，获取AI辅助诊断结果和医学描述。
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            file_input = gr.File(
+                label="上传NPY特征文件",
+                file_types=[".npy"],
+                type="filepath"
+            )
+            analyze_btn = gr.Button("开始分析", variant="primary")
+        
+        with gr.Column():
+            results_output = gr.Textbox(
+                label="分析结果",
+                lines=10,
+                max_lines=20
+            )
+    
+    with gr.Row():
+        prompt_output = gr.Textbox(
+            label="提示词 (Prompt)",
+            lines=4,
+            max_lines=6
+        )
+    
+    with gr.Row():
+        description_output = gr.Textbox(
+            label="AI生成描述",
+            lines=6,
+            max_lines=10
+        )
+    
+    # 示例文件
+    gr.Examples(
+        examples=[["example.npy"]],  # 你需要提供一个示例NPY文件
+        inputs=file_input,
+        label="点击使用示例文件"
+    )
+    
+    analyze_btn.click(
+        fn=process_npy_file,
+        inputs=file_input,
+        outputs=[results_output, prompt_output, description_output]
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True)