app.py

# ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
# 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()