MedSeg-LLM-Demo / app.py
Yunqy's picture
Create app.py
d4bbc03 verified
raw
history blame contribute delete
1.68 kB
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()