459 lines
17 KiB
Python
459 lines
17 KiB
Python
"""Streamlit 交互式结果展示界面。
|
||
|
||
运行：
|
||
streamlit run app.py
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
from pathlib import Path
|
||
from typing import Dict, List, Sequence, Tuple
|
||
|
||
import matplotlib.pyplot as plt
|
||
import numpy as np
|
||
import streamlit as st
|
||
|
||
from config import (
|
||
CHECKPOINT_DIR,
|
||
DEFAULT_CHECKPOINT,
|
||
DEFAULT_FIXED_NIFTI,
|
||
DEFAULT_MOVING_NIFTI,
|
||
INFERENCE_DIR,
|
||
OUTPUT_ROOT,
|
||
PROJECT_ROOT,
|
||
)
|
||
from metrics import crop_to_common_shape, ddf_summary, registration_metrics, slice_metric_curve
|
||
|
||
|
||
st.set_page_config(
|
||
page_title="VoxelMorph 颈部 CT 配准工作台",
|
||
layout="wide",
|
||
initial_sidebar_state="expanded",
|
||
)
|
||
|
||
|
||
CUSTOM_CSS = """
|
||
<style>
|
||
:root {
|
||
--ink: #202124;
|
||
--muted: #5f6368;
|
||
--line: #d8dee4;
|
||
--panel: #f6f8fa;
|
||
--teal: #0f766e;
|
||
--amber: #b45309;
|
||
}
|
||
.main .block-container {
|
||
padding-top: 1.1rem;
|
||
padding-bottom: 2rem;
|
||
max-width: 1480px;
|
||
}
|
||
h1, h2, h3 {
|
||
letter-spacing: 0;
|
||
}
|
||
div[data-testid="stMetric"] {
|
||
border: 1px solid var(--line);
|
||
border-radius: 6px;
|
||
padding: 0.65rem 0.75rem;
|
||
background: #ffffff;
|
||
}
|
||
section[data-testid="stSidebar"] {
|
||
border-right: 1px solid var(--line);
|
||
}
|
||
.stTabs [data-baseweb="tab-list"] {
|
||
gap: 8px;
|
||
border-bottom: 1px solid var(--line);
|
||
}
|
||
.stTabs [data-baseweb="tab"] {
|
||
border-radius: 0;
|
||
padding-left: 4px;
|
||
padding-right: 4px;
|
||
}
|
||
</style>
|
||
"""
|
||
|
||
|
||
def _require_nibabel():
|
||
try:
|
||
import nibabel as nib # type: ignore
|
||
except Exception as exc:
|
||
raise RuntimeError("缺少 nibabel，请先运行: pip install -r requirements.txt") from exc
|
||
return nib
|
||
|
||
|
||
def discover_nifti_files() -> List[str]:
|
||
roots = [OUTPUT_ROOT, PROJECT_ROOT]
|
||
paths: List[Path] = []
|
||
for root in roots:
|
||
if root.exists():
|
||
paths.extend(root.rglob("*.nii"))
|
||
paths.extend(root.rglob("*.nii.gz"))
|
||
unique = sorted({str(path) for path in paths})
|
||
return unique
|
||
|
||
|
||
def discover_checkpoint_files() -> List[str]:
|
||
roots = [CHECKPOINT_DIR, OUTPUT_ROOT, PROJECT_ROOT]
|
||
paths: List[Path] = []
|
||
for root in roots:
|
||
if root.exists():
|
||
paths.extend(root.rglob("*.pt"))
|
||
paths.extend(root.rglob("*.pth"))
|
||
paths.extend(root.rglob("*.ckpt"))
|
||
return sorted({str(path) for path in paths})
|
||
|
||
|
||
def choose_path(label: str, default_path: Path, candidates: Sequence[str]) -> str:
|
||
options = ["手动输入"] + list(candidates)
|
||
default_str = str(default_path)
|
||
selected_index = options.index(default_str) if default_str in options else (1 if candidates else 0)
|
||
selected = st.selectbox(label, options=options, index=selected_index)
|
||
if selected == "手动输入":
|
||
return st.text_input(f"{label}路径", value=default_str)
|
||
return selected
|
||
|
||
|
||
@st.cache_data(max_entries=8, show_spinner=False)
|
||
def load_nifti_cached(path: str, max_voxels: int = 14_000_000) -> Tuple[np.ndarray, Tuple[float, float, float], int]:
|
||
"""为 Web 展示读取 NIfTI。
|
||
|
||
Streamlit/浏览器端不适合一次渲染超大体数据。如果体素数过大，按等比例
|
||
stride 下采样，只影响网页查看，不改变磁盘上的推理结果。
|
||
"""
|
||
|
||
if not path:
|
||
raise ValueError("路径为空。")
|
||
nib = _require_nibabel()
|
||
img = nib.load(path, mmap=True)
|
||
if len(img.shape) > 4:
|
||
raise ValueError(f"暂不支持超过 4D 的 NIfTI: {path}")
|
||
|
||
spatial_shape = img.shape[:3]
|
||
voxels = int(np.prod(spatial_shape))
|
||
stride = max(1, int(np.ceil((voxels / max_voxels) ** (1.0 / 3.0)))) if voxels > max_voxels else 1
|
||
spatial_slices = (slice(None, None, stride), slice(None, None, stride), slice(None, None, stride))
|
||
|
||
if len(img.shape) == 4:
|
||
data = np.asarray(img.dataobj[spatial_slices + (slice(None),)], dtype=np.float32)
|
||
else:
|
||
data = np.asarray(img.dataobj[spatial_slices], dtype=np.float32)
|
||
data = np.nan_to_num(data, copy=False)
|
||
spacing = tuple(float(v) * stride for v in img.header.get_zooms()[:3])
|
||
return data, spacing, stride # type: ignore[return-value]
|
||
|
||
|
||
def normalize_slice(image: np.ndarray, p_low: float = 1.0, p_high: float = 99.0) -> np.ndarray:
|
||
image = image.astype(np.float32, copy=False)
|
||
low, high = np.percentile(image, [p_low, p_high])
|
||
if high <= low:
|
||
return np.zeros_like(image, dtype=np.float32)
|
||
return np.clip((image - low) / (high - low), 0.0, 1.0)
|
||
|
||
|
||
def orient_for_display(image: np.ndarray) -> np.ndarray:
|
||
"""统一显示方向，让三视图在网页中更接近医学浏览器观感。"""
|
||
|
||
return np.rot90(image)
|
||
|
||
|
||
def get_slice(volume_xyz: np.ndarray, plane: str, index: int) -> np.ndarray:
|
||
if plane == "Axial":
|
||
return volume_xyz[:, :, index]
|
||
if plane == "Coronal":
|
||
return volume_xyz[:, index, :]
|
||
if plane == "Sagittal":
|
||
return volume_xyz[index, :, :]
|
||
raise ValueError(f"未知平面: {plane}")
|
||
|
||
|
||
def plane_axis(plane: str) -> int:
|
||
return {"Sagittal": 0, "Coronal": 1, "Axial": 2}[plane]
|
||
|
||
|
||
def checkerboard(fixed_slice: np.ndarray, warped_slice: np.ndarray, tile: int = 24) -> np.ndarray:
|
||
fixed_norm = normalize_slice(fixed_slice)
|
||
warped_norm = normalize_slice(warped_slice)
|
||
yy, xx = np.indices(fixed_norm.shape)
|
||
mask = ((yy // tile + xx // tile) % 2).astype(bool)
|
||
return np.where(mask, fixed_norm, warped_norm)
|
||
|
||
|
||
def alpha_overlay(fixed_slice: np.ndarray, warped_slice: np.ndarray, alpha: float = 0.45) -> np.ndarray:
|
||
fixed_norm = normalize_slice(fixed_slice)
|
||
warped_norm = normalize_slice(warped_slice)
|
||
rgb = np.zeros((*fixed_norm.shape, 3), dtype=np.float32)
|
||
rgb[..., 0] = warped_norm
|
||
rgb[..., 1] = fixed_norm * (1.0 - alpha) + warped_norm * alpha
|
||
rgb[..., 2] = fixed_norm
|
||
return np.clip(rgb, 0.0, 1.0)
|
||
|
||
|
||
def ddf_magnitude(ddf_xyz: np.ndarray) -> np.ndarray:
|
||
if ddf_xyz.ndim != 4 or ddf_xyz.shape[-1] != 3:
|
||
raise ValueError("DDF 应为 (X, Y, Z, 3)。")
|
||
return np.linalg.norm(ddf_xyz.astype(np.float32, copy=False), axis=-1)
|
||
|
||
|
||
def render_image(image: np.ndarray, caption: str, cmap: str = "gray") -> None:
|
||
fig, ax = plt.subplots(figsize=(4.8, 4.8), dpi=120)
|
||
ax.imshow(orient_for_display(image), cmap=cmap, interpolation="nearest")
|
||
ax.set_title(caption, fontsize=10)
|
||
ax.axis("off")
|
||
st.pyplot(fig, use_container_width=True)
|
||
plt.close(fig)
|
||
|
||
|
||
def render_rgb(image: np.ndarray, caption: str) -> None:
|
||
fig, ax = plt.subplots(figsize=(4.8, 4.8), dpi=120)
|
||
ax.imshow(orient_for_display(image), interpolation="nearest")
|
||
ax.set_title(caption, fontsize=10)
|
||
ax.axis("off")
|
||
st.pyplot(fig, use_container_width=True)
|
||
plt.close(fig)
|
||
|
||
|
||
def render_heatmap_overlay(base_slice: np.ndarray, mag_slice: np.ndarray, alpha: float, caption: str) -> None:
|
||
fig, ax = plt.subplots(figsize=(4.8, 4.8), dpi=120)
|
||
ax.imshow(orient_for_display(normalize_slice(base_slice)), cmap="gray", interpolation="nearest")
|
||
heat = orient_for_display(mag_slice)
|
||
im = ax.imshow(heat, cmap="inferno", alpha=alpha, interpolation="nearest")
|
||
ax.set_title(caption, fontsize=10)
|
||
ax.axis("off")
|
||
fig.colorbar(im, ax=ax, fraction=0.046, pad=0.02)
|
||
st.pyplot(fig, use_container_width=True)
|
||
plt.close(fig)
|
||
|
||
|
||
def render_three_views(volume_xyz: np.ndarray, title: str) -> None:
|
||
st.subheader(title)
|
||
columns = st.columns(3)
|
||
planes = [("Axial", "轴状面"), ("Coronal", "冠状面"), ("Sagittal", "矢状面")]
|
||
for col, (plane, cn_name) in zip(columns, planes):
|
||
axis = plane_axis(plane)
|
||
with col:
|
||
index = st.slider(f"{cn_name}", 0, volume_xyz.shape[axis] - 1, volume_xyz.shape[axis] // 2, key=f"{title}-{plane}")
|
||
render_image(get_slice(volume_xyz, plane, index), f"{cn_name} #{index}")
|
||
|
||
|
||
def render_overlay_views(fixed_xyz: np.ndarray, warped_xyz: np.ndarray) -> None:
|
||
fixed_xyz, warped_xyz = crop_to_common_shape(fixed_xyz, warped_xyz)
|
||
mode = st.radio("对比模式", ["Alpha 融合", "棋盘格"], horizontal=True)
|
||
alpha = st.slider("透明度", 0.0, 1.0, 0.45, 0.05)
|
||
tile = st.slider("棋盘格尺寸", 8, 64, 24, 4)
|
||
|
||
columns = st.columns(3)
|
||
planes = [("Axial", "轴状面"), ("Coronal", "冠状面"), ("Sagittal", "矢状面")]
|
||
for col, (plane, cn_name) in zip(columns, planes):
|
||
axis = plane_axis(plane)
|
||
with col:
|
||
index = st.slider(f"{cn_name}切片", 0, fixed_xyz.shape[axis] - 1, fixed_xyz.shape[axis] // 2, key=f"overlay-{plane}")
|
||
fixed_slice = get_slice(fixed_xyz, plane, index)
|
||
warped_slice = get_slice(warped_xyz, plane, index)
|
||
if mode == "棋盘格":
|
||
render_image(checkerboard(fixed_slice, warped_slice, tile=tile), f"{cn_name} 棋盘格 #{index}")
|
||
else:
|
||
render_rgb(alpha_overlay(fixed_slice, warped_slice, alpha=alpha), f"{cn_name} 融合 #{index}")
|
||
|
||
|
||
def render_ddf_views(fixed_xyz: np.ndarray, ddf_xyz: np.ndarray) -> None:
|
||
fixed_xyz, ddf_xyz = crop_to_common_shape(fixed_xyz, ddf_xyz)
|
||
mag = ddf_magnitude(ddf_xyz)
|
||
alpha = st.slider("热力图透明度", 0.0, 1.0, 0.55, 0.05)
|
||
|
||
stats = ddf_summary(ddf_xyz)
|
||
c1, c2, c3, c4 = st.columns(4)
|
||
c1.metric("平均位移 mm", f"{stats['ddf_mean']:.3f}")
|
||
c2.metric("P95 mm", f"{stats['ddf_p95']:.3f}")
|
||
c3.metric("最大 mm", f"{stats['ddf_max']:.3f}")
|
||
c4.metric("标准差", f"{stats['ddf_std']:.3f}")
|
||
|
||
columns = st.columns(3)
|
||
planes = [("Axial", "轴状面"), ("Coronal", "冠状面"), ("Sagittal", "矢状面")]
|
||
for col, (plane, cn_name) in zip(columns, planes):
|
||
axis = plane_axis(plane)
|
||
with col:
|
||
index = st.slider(f"{cn_name}DDF", 0, fixed_xyz.shape[axis] - 1, fixed_xyz.shape[axis] // 2, key=f"ddf-{plane}")
|
||
render_heatmap_overlay(
|
||
get_slice(fixed_xyz, plane, index),
|
||
get_slice(mag, plane, index),
|
||
alpha=alpha,
|
||
caption=f"{cn_name} 形变强度 #{index}",
|
||
)
|
||
|
||
|
||
def render_metric_charts(fixed_xyz: np.ndarray, moving_xyz: np.ndarray, warped_xyz: np.ndarray, ddf_xyz: np.ndarray | None) -> None:
|
||
fixed_xyz, moving_xyz, warped_xyz = crop_to_common_shape(fixed_xyz, moving_xyz, warped_xyz)
|
||
metrics = registration_metrics(fixed_xyz, moving_xyz, warped_xyz)
|
||
|
||
c1, c2, c3 = st.columns(3)
|
||
c1.metric("NCC", f"{metrics['after_ncc']:.4f}", delta=f"{metrics['ncc_improvement']:+.4f}")
|
||
c2.metric(
|
||
"MSE",
|
||
f"{metrics['after_mse']:.5f}",
|
||
delta=f"{metrics['after_mse'] - metrics['before_mse']:+.5f}",
|
||
delta_color="inverse",
|
||
)
|
||
c3.metric(
|
||
"MAE",
|
||
f"{metrics['after_mae']:.5f}",
|
||
delta=f"{metrics['after_mae'] - metrics['before_mae']:+.5f}",
|
||
delta_color="inverse",
|
||
)
|
||
|
||
labels = ["NCC", "MSE", "MAE"]
|
||
before = [metrics["before_ncc"], metrics["before_mse"], metrics["before_mae"]]
|
||
after = [metrics["after_ncc"], metrics["after_mse"], metrics["after_mae"]]
|
||
|
||
fig, axes = plt.subplots(1, 2, figsize=(10, 3.6), dpi=130)
|
||
x = np.arange(len(labels))
|
||
width = 0.36
|
||
axes[0].bar(x - width / 2, before, width, label="配准前", color="#64748b")
|
||
axes[0].bar(x + width / 2, after, width, label="配准后", color="#0f766e")
|
||
axes[0].set_xticks(x, labels)
|
||
axes[0].set_title("配准前后指标对比")
|
||
axes[0].legend(frameon=False)
|
||
axes[0].grid(axis="y", alpha=0.22)
|
||
|
||
curve = slice_metric_curve(fixed_xyz, moving_xyz, warped_xyz, axis=2)
|
||
axes[1].plot(curve["slice_index"], curve["before_mse"], label="配准前 MSE", color="#64748b", linewidth=1.8)
|
||
axes[1].plot(curve["slice_index"], curve["after_mse"], label="配准后 MSE", color="#b45309", linewidth=1.8)
|
||
axes[1].set_title("轴状面逐切片误差")
|
||
axes[1].set_xlabel("Slice")
|
||
axes[1].grid(alpha=0.22)
|
||
axes[1].legend(frameon=False)
|
||
st.pyplot(fig, use_container_width=True)
|
||
plt.close(fig)
|
||
|
||
if ddf_xyz is not None:
|
||
mag = ddf_magnitude(ddf_xyz)
|
||
fig2, ax = plt.subplots(figsize=(10, 3.2), dpi=130)
|
||
ax.hist(mag.ravel(), bins=80, color="#7c3f00", alpha=0.82)
|
||
ax.set_title("DDF 位移强度分布")
|
||
ax.set_xlabel("mm")
|
||
ax.set_ylabel("Voxel count")
|
||
ax.grid(axis="y", alpha=0.2)
|
||
st.pyplot(fig2, use_container_width=True)
|
||
plt.close(fig2)
|
||
|
||
|
||
def load_result_paths(out_dir: str) -> Dict[str, str]:
|
||
out = Path(out_dir)
|
||
return {
|
||
"warped": str(out / "warped_moving.nii.gz"),
|
||
"ddf": str(out / "ddf_mm.nii.gz"),
|
||
"metrics": str(out / "metrics.json"),
|
||
}
|
||
|
||
|
||
def run_inference_from_ui(moving_path: str, fixed_path: str, checkpoint_path: str, out_dir: str) -> Dict:
|
||
from infer import run_inference
|
||
|
||
return run_inference(
|
||
moving_path=moving_path,
|
||
fixed_path=fixed_path,
|
||
checkpoint_path=checkpoint_path,
|
||
out_dir=out_dir,
|
||
)
|
||
|
||
|
||
def main() -> None:
|
||
st.markdown(CUSTOM_CSS, unsafe_allow_html=True)
|
||
st.title("VoxelMorph 颈部 CT 配准工作台")
|
||
|
||
candidates = discover_nifti_files()
|
||
checkpoint_candidates = discover_checkpoint_files()
|
||
with st.sidebar:
|
||
st.header("输入")
|
||
moving_path = choose_path("Moving", DEFAULT_MOVING_NIFTI, candidates)
|
||
fixed_path = choose_path("Fixed", DEFAULT_FIXED_NIFTI, candidates)
|
||
checkpoint_path = choose_path("模型权重", DEFAULT_CHECKPOINT, checkpoint_candidates)
|
||
out_dir = st.text_input("输出目录", value=str(INFERENCE_DIR))
|
||
display_max_voxels = st.slider("Web显示体素上限", 2_000_000, 30_000_000, 14_000_000, 1_000_000)
|
||
|
||
start = st.button("开始推理", type="primary", use_container_width=True)
|
||
if start:
|
||
with st.spinner("推理运行中"):
|
||
try:
|
||
load_nifti_cached.clear()
|
||
result = run_inference_from_ui(moving_path, fixed_path, checkpoint_path, out_dir)
|
||
load_nifti_cached.clear()
|
||
st.session_state["last_result"] = result
|
||
st.success("推理完成")
|
||
except Exception as exc:
|
||
st.error(str(exc))
|
||
|
||
result_paths = load_result_paths(out_dir)
|
||
last_result = st.session_state.get("last_result", {})
|
||
warped_path = str(last_result.get("warped_path", result_paths["warped"]))
|
||
ddf_path = str(last_result.get("ddf_path", result_paths["ddf"]))
|
||
metrics_path = str(last_result.get("metrics_path", result_paths["metrics"]))
|
||
|
||
status_cols = st.columns(4)
|
||
status_cols[0].metric("Moving", "存在" if Path(moving_path).exists() else "缺失")
|
||
status_cols[1].metric("Fixed", "存在" if Path(fixed_path).exists() else "缺失")
|
||
status_cols[2].metric("Warped", "存在" if Path(warped_path).exists() else "缺失")
|
||
status_cols[3].metric("DDF", "存在" if Path(ddf_path).exists() else "缺失")
|
||
|
||
try:
|
||
moving_xyz, moving_spacing, moving_stride = load_nifti_cached(moving_path, max_voxels=display_max_voxels)
|
||
fixed_xyz, fixed_spacing, fixed_stride = load_nifti_cached(fixed_path, max_voxels=display_max_voxels)
|
||
except Exception as exc:
|
||
st.warning(f"待显示数据尚未就绪: {exc}")
|
||
return
|
||
|
||
warped_xyz = None
|
||
ddf_xyz = None
|
||
if Path(warped_path).exists():
|
||
try:
|
||
warped_xyz, _, _ = load_nifti_cached(warped_path, max_voxels=display_max_voxels)
|
||
except Exception as exc:
|
||
st.warning(f"Warped 读取失败: {exc}")
|
||
if Path(ddf_path).exists():
|
||
try:
|
||
ddf_xyz, _, _ = load_nifti_cached(ddf_path, max_voxels=display_max_voxels)
|
||
except Exception as exc:
|
||
st.warning(f"DDF 读取失败: {exc}")
|
||
|
||
if moving_stride > 1 or fixed_stride > 1:
|
||
st.info(f"网页显示已自动下采样：Moving stride={moving_stride}, Fixed stride={fixed_stride}")
|
||
|
||
tab_views, tab_overlay, tab_ddf, tab_metrics = st.tabs(["正交三视图", "重叠对比", "形变场", "量化图"])
|
||
with tab_views:
|
||
view_target = st.radio("显示对象", ["Fixed", "Moving", "Warped"], horizontal=True)
|
||
if view_target == "Fixed":
|
||
render_three_views(fixed_xyz, "Fixed Image")
|
||
elif view_target == "Moving":
|
||
render_three_views(moving_xyz, "Moving Image")
|
||
elif warped_xyz is not None:
|
||
render_three_views(warped_xyz, "Warped Image")
|
||
else:
|
||
st.warning("尚未生成 Warped Image。")
|
||
|
||
with tab_overlay:
|
||
if warped_xyz is None:
|
||
st.warning("尚未生成 Warped Image。")
|
||
else:
|
||
render_overlay_views(fixed_xyz, warped_xyz)
|
||
|
||
with tab_ddf:
|
||
if ddf_xyz is None:
|
||
st.warning("尚未生成 DDF。")
|
||
else:
|
||
render_ddf_views(fixed_xyz, ddf_xyz)
|
||
|
||
with tab_metrics:
|
||
if warped_xyz is None:
|
||
st.warning("尚未生成 Warped Image。")
|
||
else:
|
||
render_metric_charts(fixed_xyz, moving_xyz, warped_xyz, ddf_xyz)
|
||
metrics_file = Path(metrics_path)
|
||
if metrics_file.exists():
|
||
try:
|
||
st.json(json.loads(metrics_file.read_text(encoding="utf-8")))
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|