import gradio as gr
import numpy as np
import torch
from transformers import pipeline, AutoModel, AutoTokenizer
from monai.transforms import Compose, LoadImage, ScaleIntensity, EnsureChannelFirst
import SimpleITK as sitk

# 初始化组件
llm = pipeline("text-generation", model="microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract")
seg_model = AutoModel.from_pretrained("Project-MONAI/model-zoo/hnfnet_brats21").eval() 

# 医学图像预处理
preprocess = Compose([
    LoadImage(image_only=True),
    EnsureChannelFirst(channel_dim='no_channel'),
    ScaleIntensity(minv=0.0, maxv=1.0)
])

def analyze_image(image_path, clinical_note):
    # 生成分割提示
    prompt = f"根据临床报告生成分割提示：{clinical_note}"
    guidance = llm(prompt, max_length=200)[0]['generated_text']
    
    # 图像处理
    img = preprocess(image_path)
    
    # 分割推理
    with torch.no_grad():
        seg = seg_model(img.unsqueeze(0))[0]
    
    # 后处理
    result = sitk.GetArrayFromImage(seg.squeeze().numpy())
    return (result > 0.5).astype(np.uint8), guidance

# 创建交互界面
demo = gr.Interface(
    fn=analyze_image,
    inputs=[
        gr.File(label="上传DICOM/NIfTI文件"),
        gr.Textbox(label="临床描述", placeholder="输入影像学检查报告...")
    ],
    outputs=[
        gr.Image(label="分割结果", colormap="viridis"),
        gr.Textbox(label="生成的分割提示")
    ],
    examples=[
        ["assets/sample1.nii.gz", "左侧基底节区可见直径2cm占位，T1低信号，T2高信号"],
        ["assets/sample2.dcm", "右肺上叶结节，边缘毛刺，考虑恶性肿瘤可能"]
    ]
)

demo.launch()