src/streamlit_app.py

import streamlit as st
import torch
import numpy as np
import cv2
from PIL import Image
from transformers import AutoModelForImageClassification, AutoImageProcessor
from torchvision import transforms
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
import torch.nn.functional as F

# ==========================================
# 1. 页面配置与标题
# ==========================================
st.set_page_config(
    page_title="RSNA Intracranial Hemorrhage AI",
    page_icon="🧠",
    layout="wide"
)

st.title("🧠 AI-Assisted Intracranial Hemorrhage Detection")
st.markdown("""
**Workflow:** `Input CT Scan` $\\rightarrow$ `ResNet50 Inference` $\\rightarrow$ `Grad-CAM XAI` $\\rightarrow$ `Radiologist Review`
""")

# ==========================================
# 2. 模型加载 (使用 @st.cache_resource 加速)
# ==========================================
MODEL_ID = "dongqinggeng/rsna"  # 你的 Hugging Face Model ID

@st.cache_resource
def load_model_and_processor():
    try:
        # 加载模型
        model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
        model.eval()
        
        # 尝试加载配套的 processor，如果失败则使用默认逻辑
        try:
            processor = AutoImageProcessor.from_pretrained(MODEL_ID)
        except:
            processor = None
            
        return model, processor
    except Exception as e:
        st.error(f"Error loading model from Hugging Face: {e}")
        return None, None

model, processor = load_model_and_processor()

# 定义标签映射 (确保顺序与训练时一致)
LABELS = ['epidural', 'intraparenchymal', 'intraventricular', 'subarachnoid', 'subdural', 'any']
id2label = {i: label for i, label in enumerate(LABELS)}

# ==========================================
# 3. 辅助类与函数
# ==========================================

# Grad-CAM 需要的 Wrapper (适配 HF 模型输出)
class HuggingFaceModelWrapper(torch.nn.Module):
    def __init__(self, model):
        super(HuggingFaceModelWrapper, self).__init__()
        self.model = model
    def forward(self, x):
        return self.model(x).logits

# 图像预处理 (必须与训练时保持一致!)
def process_image(image):
    # 定义变换
    transform = transforms.Compose([
        transforms.Resize((224, 224)), # 如果你之前是 128，这里用 128，ResNet 默认通常适配 224 效果更好
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    # 确保是 RGB
    image = image.convert("RGB")
    input_tensor = transform(image).unsqueeze(0) # Add batch dimension
    return input_tensor, image

# 生成 Grad-CAM
def generate_gradcam(model, input_tensor, target_layer):
    cam = GradCAM(model=HuggingFaceModelWrapper(model), target_layers=[target_layer])
    # target=None 表示解释概率最高的那个类
    grayscale_cam = cam(input_tensor=input_tensor, targets=None)
    return grayscale_cam[0, :]

# ==========================================
# 4. 主界面逻辑
# ==========================================

# --- 侧边栏：上传图片 ---
st.sidebar.header("1. Input Data")
uploaded_file = st.sidebar.file_uploader("Upload a CT Slice (PNG/JPG/DICOM)", type=["png", "jpg", "jpeg"])

# 提供示例图片下载链接或按钮 (可选)
if st.sidebar.button("Load Demo Image (Simulated)"):
    # 这里只是为了演示，实际应用中你可以放一个示例文件的 URL
    st.sidebar.info("Please upload a real file to test.")

# --- 主区域 ---
col1, col2 = st.columns([1, 1.5])

if uploaded_file is not None and model is not None:
    # A. 显示原图
    image_pil = Image.open(uploaded_file)
    
    with col1:
        st.subheader("Original CT Scan")
        st.image(image_pil, use_container_width=True, caption="Input Slice")

    # B. 推理 (Inference)
    with st.spinner("Running AI Model & Generating Explanations..."):
        # 预处理
        input_tensor, image_rgb_pil = process_image(image_pil)
        
        # 预测
        with torch.no_grad():
            outputs = model(input_tensor)
            probs = torch.sigmoid(outputs.logits).cpu().numpy()[0]
        
        # 找到概率最高的特定类型 (排除 'any')
        # labels 顺序: epidural, intraparenchymal, intraventricular, subarachnoid, subdural, any
        # any 是 index 5
        specific_probs = probs[:5]
        top_idx = np.argmax(specific_probs)
        top_label = LABELS[top_idx]
        top_prob = specific_probs[top_idx]
        any_prob = probs[5]

        # C. XAI: Grad-CAM
        # 获取 ResNet 最后一层
        target_layer = model.resnet.encoder.stages[-1].layers[-1]
        cam_mask = generate_gradcam(model, input_tensor, target_layer)
        
        # 图像叠加处理
        # 1. 归一化原图用于显示
        img_np = np.array(image_rgb_pil)
        img_np = img_np.astype(np.float32) / 255.0
        # 2. 叠加
        visualization = show_cam_on_image(img_np, cam_mask, use_rgb=True)

    # D. 显示结果
    with col2:
        st.subheader("XAI Output (Grad-CAM)")
        st.image(visualization, use_container_width=True, caption=f"Model Focus Area (Red = High Attention)")
        
    # E. 详细报告与指标
    st.divider()
    st.header("📝 AI Analysis Report")
    
    # 动态生成的“放射科医生评论”
    if any_prob > 0.5:
        status_color = "red"
        status_text = "Hemorrhage Detected"
        confidence_text = f"High Confidence ({any_prob:.2%})"
        
        review_text = (
            f"**AI Findings:** The model detected specific features consistent with **{top_label} hemorrhage** "
            f"(Probability: {top_prob:.2%}).\n\n"
            f"**XAI Localization:** The Grad-CAM heatmap highlights a region of interest. "
            f"Please verify if this corresponds to a hyperdense area in the brain parenchyma or extra-axial space."
        )
    else:
        status_color = "green"
        status_text = "No Hemorrhage Detected"
        confidence_text = f"({1-any_prob:.2%} sure)"
        review_text = "**AI Findings:** No significant signs of intracranial hemorrhage were detected. Heatmap shows diffuse or non-specific activation."

    # 使用 Metrics 展示
    m1, m2, m3 = st.columns(3)
    m1.metric("Overall Prediction", status_text, delta=confidence_text, delta_color="inverse" if any_prob > 0.5 else "normal")
    m2.metric("Primary Subtype", top_label.capitalize() if any_prob > 0.3 else "N/A", f"{top_prob:.2%}")
    
    st.markdown(f"""
    > **Radiologist Review Note:** > {review_text}
    """)
    
    # F. 详细概率条形图
    st.subheader("Detailed Class Probabilities")
    st.bar_chart({label: prob for label, prob in zip(LABELS, probs)})

else:
    # 初始欢迎界面
    st.info("👈 Please upload a CT image from the sidebar to start the analysis.")
    st.markdown("### How to interpret the heatmap?")
    st.markdown("""
    * **Red Areas**: Regions that contributed most to the AI's decision (High Importance).
    * **Blue Areas**: Regions the AI ignored.
    * *Note: If the heatmap highlights the skull or background, the prediction might be an artifact.*
    """)