# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ # app.py # # 说明：这个脚本实现了 # 1. 用户上传病灶图像（JPEG/PNG） # 2. 调用 Hugging Face Inference API 的预训练“medical-sam”分割模型，返回 mask # 3. 使用 LangChain 将 “knowledge” 目录下所有 PDF 文本做向量化，构建 FAISS 索引 # 4. 将分割结果特征（面积、位置）与用户图像基本信息一起，拼在一起生成检索+生成提示（RAG+Agent） # 5. 用 OpenAI GPT-3.5/GPT-4 生成 “几句话” 分割细节描述 # 6. Gradio 前端展示：并排展示原图 + 分割 mask + 文字描述 # ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ import os import io import tempfile import numpy as np from PIL import Image import torch from transformers import SegformerForSemanticSegmentation, SegformerFeatureExtractor import gradio as gr from langchain.document_loaders import PDFPlumberLoader from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain.embeddings import SentenceTransformerEmbeddings from langchain.vectorstores import FAISS from langchain.chat_models import ChatOpenAI from langchain.chains import RetrievalQA # ─── 一、加载环境变量 ────────────────────────────────────────────────────── HF_API_TOKEN = os.getenv("HF_API_TOKEN") # 从 Space Settings → Secrets 填入 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") # 从 Space Settings → Secrets 填入 MODEL_CHECKPOINT = "salesforce/segformer-b0-finetuned-ade-512-512" # 示例通用分割，可根据需要换成医学专用模型 EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2" os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY # ─── 二、初始化分割模型 ──────────────────────────────────────────────────── feature_extractor = SegformerFeatureExtractor.from_pretrained(MODEL_CHECKPOINT, use_auth_token=HF_API_TOKEN) seg_model = SegformerForSemanticSegmentation.from_pretrained(MODEL_CHECKPOINT, use_auth_token=HF_API_TOKEN) device = "cuda" if torch.cuda.is_available() else "cpu" seg_model.to(device) # ─── 三、构建 RAG 向量索引 ───────────────────────────────────────────────── def build_knowledge_vectorstore(pdf_folder="knowledge"): """把 knowledge/ 下所有 PDF 加载、拆分页、做 Embedding、用 FAISS 建索引""" docs = [] for fn in os.listdir(pdf_folder): if fn.lower().endswith(".pdf"): loader = PDFPlumberLoader(os.path.join(pdf_folder, fn)) for page in loader.load(): docs.append(page) # 按页拆分，保证每个 chunk 不超过 1000 字 text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100) splits = text_splitter.split_documents(docs) # 用 Sentence-Transformers 做 Embedding embeddings = SentenceTransformerEmbeddings(model_name=EMBEDDING_MODEL) vs = FAISS.from_documents(splits, embeddings) return vs # 当 Space 启动时，先构建一次索引 vectorstore = build_knowledge_vectorstore(pdf_folder="knowledge") qa_chain = RetrievalQA.from_chain_type( llm=ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0), chain_type="stuff", retriever=vectorstore.as_retriever(search_kwargs={"k": 3}), # 检索 Top-3 相关段落 ) # ─── 四、辅助函数：分割 + 特征提取 ──────────────────────────────────────────── def run_segmentation(image_pil: Image.Image): """调用 SegFormer 模型做语义分割，返回 mask numpy array""" image_rgb = image_pil.convert("RGB") inputs = feature_extractor(images=image_rgb, return_tensors="pt").to(device) with torch.no_grad(): outputs = seg_model(**inputs) logits = outputs.logits # shape: (1, num_classes, H, W) # 取每像素最大概率的类别 preds = torch.argmax(logits, dim=1)[0].cpu().numpy() # shape: (H, W) return preds def extract_mask_stats(mask_np: np.ndarray): """ 简单示例：假设分割模型只分两类（背景=0，病灶=1） 1. 统计病灶像素个数占比 2. 计算病灶包围盒（xmin,ymin,xmax,ymax） """ lesion_pixels = (mask_np == 1) total_pixels = mask_np.size lesion_count = lesion_pixels.sum() area_ratio = float(lesion_count) / total_pixels ys, xs = np.where(lesion_pixels) if len(xs) > 0 and len(ys) > 0: xmin, xmax = int(xs.min()), int(xs.max()) ymin, ymax = int(ys.min()), int(ys.max()) else: xmin = ymin = xmax = ymax = 0 return { "area_ratio": round(area_ratio, 4), "bbox": (xmin, ymin, xmax, ymax), "lesion_pixels": int(lesion_count), } # ─── 五、Gradio 回调函数：上传→分割→RAG+Agent 生成描述 → 返回前端 ────────────────── def segment_and_describe(image_file): # 1. 把上传文件转成 PIL Image image_pil = Image.open(image_file).convert("RGB") # 2. 运行分割模型 mask_np = run_segmentation(image_pil) # shape: (H, W) # 3. 提取分割特征 stats = extract_mask_stats(mask_np) area_pct = stats["area_ratio"] * 100 # 换成百分比 bbox = stats["bbox"] # 4. 准备给 LLM 的提示语（Prompt）；先做一次知识检索 # 检索用户关心的关键上下文，比如“病灶形态”“病灶大小”等 query_text = ( f"请根据以下信息撰写医学分割解析：\n" f"- 病灶像素占比约 {area_pct:.2f}%。\n" f"- 病灶包围盒坐标 (xmin, ymin, xmax, ymax) = {bbox}。\n" f"- 用户上传的图像类型为病灶图像。\n" f"请结合医学知识，对此病灶分割结果做 2–3 句详细专业描述。" ) # 使用 RAG 检索相关片段 rag_answer = qa_chain.run(query_text) # 5. 将原始 Prompt + RAG 检索内容 + Query 一起传给 ChatOpenAI，让它“整合”成最终描述 llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0.2) full_prompt = ( f"下面是检索到的医学文档片段：\n{rag_answer}\n\n" f"请结合上述片段和统计信息，对病灶分割结果撰写 2–3 句专业描述：\n" f"- 病灶像素占比：{area_pct:.2f}%\n" f"- 病灶包围盒：{bbox}\n" ) description = llm(full_prompt).content.strip() # 6. 把 mask 变成 RGBA 图，叠加到原图上 mask_rgba = Image.fromarray((mask_np * 255).astype(np.uint8)).convert("L") mask_rgba = mask_rgba.resize(image_pil.size) # 生成红色蒙版 red_mask = Image.new("RGBA", image_pil.size, color=(255, 0, 0, 0)) red_mask.putalpha(mask_rgba) overlay = Image.alpha_composite(image_pil.convert("RGBA"), red_mask) return overlay, description # ─── 六、Gradio 前端界面构建 ────────────────────────────────────────────────── with gr.Blocks() as demo: gr.Markdown("## 🏥 医学病灶分割 + RAG 专业解读 Demo\n\n" "1. 上传你的病灶图像（JPEG/PNG）\n" "2. 点击下方按钮，自动返回 **分割结果** + **几句话详细描述**\n") with gr.Row(): img_in = gr.Image(type="file", label="上传病灶图像") with gr.Column(): btn = gr.Button("开始分割并生成描述") text_out = gr.Textbox(label="分割+解读结果", lines=4) img_out = gr.Image(label="叠加分割结果") btn.click(fn=segment_and_describe, inputs=img_in, outputs=[img_out, text_out]) if __name__ == "__main__": demo.launch()