"""Streamlit 交互式结果展示界面。 运行: streamlit run app.py """ from __future__ import annotations import json from pathlib import Path from typing import Dict, List, Sequence, Tuple import matplotlib as mpl import matplotlib.pyplot as plt from matplotlib import font_manager 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 = """ """ def configure_matplotlib_cjk_font() -> None: """配置 Matplotlib 中文字体。 Streamlit 自身用浏览器字体渲染中文,通常没有问题;Matplotlib 图像会走 后端字体管理,如果默认 DejaVu Sans 不含中文字形,就会显示成方块。 """ candidate_paths = [ "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc", "/usr/share/fonts/opentype/noto/NotoSansCJK-Bold.ttc", "/usr/share/fonts/truetype/droid/DroidSansFallbackFull.ttf", ] candidate_names = [ "Noto Sans CJK SC", "Noto Sans CJK JP", "Droid Sans Fallback", "WenQuanYi Micro Hei", "Source Han Sans SC", "SimHei", ] discovered_names: List[str] = [] for path_str in candidate_paths: path = Path(path_str) if not path.exists(): continue try: font_manager.fontManager.addfont(str(path)) discovered_names.append(font_manager.FontProperties(fname=str(path)).get_name()) except Exception: continue mpl.rcParams["font.family"] = "sans-serif" mpl.rcParams["font.sans-serif"] = [*discovered_names, *candidate_names, "DejaVu Sans"] mpl.rcParams["axes.unicode_minus"] = False configure_matplotlib_cjk_font() 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.array(img.dataobj[spatial_slices + (slice(None),)], dtype=np.float32, copy=True) else: data = np.array(img.dataobj[spatial_slices], dtype=np.float32, copy=True) 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 ddf_color_limits(mag_xyz: np.ndarray, scale_mode: str) -> Tuple[float, float]: values = mag_xyz[np.isfinite(mag_xyz)] if values.size == 0: return 0.0, 1.0 if scale_mode == "高位移增强": vmin, vmax = np.percentile(values, [50.0, 99.5]) elif scale_mode == "绝对范围": vmin, vmax = np.min(values), np.max(values) else: vmin, vmax = np.percentile(values, [1.0, 99.0]) vmin = float(vmin) vmax = float(vmax) if not np.isfinite(vmin) or not np.isfinite(vmax) or vmax <= vmin: vmax = vmin + 1e-6 return vmin, vmax def format_mm(value: float) -> str: value = float(value) if abs(value) < 0.01: return f"{value:.4f}" if abs(value) < 1.0: return f"{value:.3f}" return f"{value:.2f}" 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, width="stretch") 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, width="stretch") plt.close(fig) def render_heatmap_overlay( base_slice: np.ndarray, mag_slice: np.ndarray, alpha: float, caption: str, vmin: float, vmax: float, show_base: bool, ) -> None: fig, ax = plt.subplots(figsize=(4.8, 4.8), dpi=120) if show_base: ax.imshow(orient_for_display(normalize_slice(base_slice)), cmap="gray", interpolation="nearest") heat = orient_for_display(mag_slice) im = ax.imshow( heat, cmap="turbo", alpha=alpha if show_base else 1.0, interpolation="nearest", vmin=vmin, vmax=vmax, ) ax.set_title(caption, fontsize=10) ax.axis("off") fig.colorbar(im, ax=ax, fraction=0.046, pad=0.02) st.pyplot(fig, width="stretch") 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) stats = ddf_summary(ddf_xyz) c1, c2, c3, c4 = st.columns(4) c1.metric("平均位移 mm", format_mm(stats["ddf_mean"])) c2.metric("P95 mm", format_mm(stats["ddf_p95"])) c3.metric("最大 mm", format_mm(stats["ddf_max"])) c4.metric("标准差", format_mm(stats["ddf_std"])) if stats["ddf_max"] < 0.1: st.info(f"当前最大位移只有 {format_mm(stats['ddf_max'])} mm,通常说明这是 smoke 权重或模型还没有充分训练。") controls = st.columns(3) with controls[0]: display_mode = st.radio("显示方式", ["仅位移热力图", "CT叠加热力图"], horizontal=True) with controls[1]: scale_mode = st.selectbox("色阶范围", ["相对增强", "高位移增强", "绝对范围"]) with controls[2]: alpha = st.slider("叠加透明度", 0.0, 1.0, 0.55, 0.05) vmin, vmax = ddf_color_limits(mag, scale_mode) show_base = display_mode == "CT叠加热力图" st.caption(f"当前色阶范围:{format_mm(vmin)} - {format_mm(vmax)} mm") 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, vmin=vmin, vmax=vmax, show_base=show_base, caption=f"{cn_name} 形变强度 #{index}", ) METRIC_SPECS = [ ("NCC", "before_ncc", "after_ncc", "越高越好", True), ("MSE", "before_mse", "after_mse", "越低越好", False), ("MAE", "before_mae", "after_mae", "越低越好", False), ] def metric_improvement(before: float, after: float, higher_is_better: bool) -> float: return after - before if higher_is_better else before - after def format_metric(value: float) -> str: return f"{float(value):.6f}" def build_metric_table(metrics: Dict[str, float]) -> List[Dict[str, str]]: rows = [] for label, before_key, after_key, direction, higher_is_better in METRIC_SPECS: before = float(metrics[before_key]) after = float(metrics[after_key]) improvement = metric_improvement(before, after, higher_is_better) improvement_rate = improvement / (abs(before) + 1e-8) * 100.0 rows.append( { "指标": label, "方向": direction, "配准前": format_metric(before), "配准后": format_metric(after), "改善量": f"{improvement:+.6f}", "改善率": f"{improvement_rate:+.4f}%", } ) return rows def render_metric_value_panels(metrics: Dict[str, float]) -> None: fig, axes = plt.subplots(1, 3, figsize=(11.5, 3.1), dpi=130) for ax, (label, before_key, after_key, direction, higher_is_better) in zip(axes, METRIC_SPECS): before = float(metrics[before_key]) after = float(metrics[after_key]) improvement = metric_improvement(before, after, higher_is_better) values = [before, after] x = np.arange(2) ax.bar(x, values, width=0.5, color=["#64748b", "#0f766e"]) ax.plot(x, values, color="#111827", linewidth=1.2, marker="o", markersize=4) ax.set_xticks(x, ["配准前", "配准后"]) ax.set_title(f"{label}({direction})") ymin = min(0.0, min(values) * 1.18) ymax = max(0.0, max(values) * 1.18, 1e-5) ax.set_ylim(ymin, ymax) ax.grid(axis="y", alpha=0.22) for index, value in enumerate(values): ax.text(index, value, format_metric(value), ha="center", va="bottom", fontsize=8) ax.text(0.5, max(values) * 1.06, f"改善 {improvement:+.6f}", ha="center", fontsize=8, color="#374151") fig.tight_layout() st.pyplot(fig, width="stretch") plt.close(fig) 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) metric_rows = build_metric_table(metrics) c1, c2, c3 = st.columns(3) c1.metric("NCC", format_metric(metrics["after_ncc"]), delta=f"{metrics['ncc_improvement']:+.6f}") c2.metric( "MSE", format_metric(metrics["after_mse"]), delta=f"{metrics['after_mse'] - metrics['before_mse']:+.6f}", delta_color="inverse", ) c3.metric( "MAE", format_metric(metrics["after_mae"]), delta=f"{metrics['after_mae'] - metrics['before_mae']:+.6f}", delta_color="inverse", ) max_improvement_rate = max(abs(float(row["改善率"].rstrip("%"))) for row in metric_rows) if max_improvement_rate < 0.1: st.info("当前配准前后指标变化低于 0.1%,曲线和柱子高度接近是正常的;这更像是 smoke 权重的验证结果,不是充分训练后的配准效果。") st.caption("量化结果由当前加载的 Fixed、Moving、Warped 体数据现场计算;逐切片曲线未平滑、未随机采样。") st.table(metric_rows) render_metric_value_panels(metrics) controls = st.columns(2) with controls[0]: plane_label = st.selectbox("逐切片平面", ["轴状面", "冠状面", "矢状面"]) with controls[1]: metric_label = st.selectbox("逐切片指标", ["MSE", "MAE", "NCC"]) axis = {"矢状面": 0, "冠状面": 1, "轴状面": 2}[plane_label] metric_key = metric_label.lower() higher_is_better = metric_key == "ncc" curve = slice_metric_curve(fixed_xyz, moving_xyz, warped_xyz, axis=axis, metric=metric_key) slice_index = np.asarray(curve["slice_index"], dtype=np.int32) before_curve = np.asarray(curve[f"before_{metric_key}"], dtype=np.float32) after_curve = np.asarray(curve[f"after_{metric_key}"], dtype=np.float32) improvement_curve = after_curve - before_curve if higher_is_better else before_curve - after_curve fig, axes = plt.subplots( 2, 1, figsize=(10.5, 5.4), dpi=130, sharex=True, gridspec_kw={"height_ratios": [2.0, 1.0]}, ) axes[0].plot(slice_index, before_curve, label=f"配准前 {metric_label}", color="#64748b", linewidth=1.5) axes[0].plot(slice_index, after_curve, label=f"配准后 {metric_label}", color="#0f766e", linewidth=1.5) axes[0].set_title(f"{plane_label}逐切片{metric_label}") axes[0].set_ylabel(metric_label) axes[0].grid(alpha=0.22) axes[0].legend(frameon=False) axes[1].axhline(0.0, color="#9ca3af", linewidth=0.9) axes[1].plot(slice_index, improvement_curve, color="#b45309", linewidth=1.4) axes[1].set_title(f"{metric_label}逐切片改善量") axes[1].set_xlabel("切片序号") axes[1].set_ylabel("改善量") axes[1].grid(alpha=0.22) st.pyplot(fig, width="stretch") plt.close(fig) if ddf_xyz is not None: mag = ddf_magnitude(ddf_xyz) values = mag[np.isfinite(mag)] if values.size == 0: return fig2, ax = plt.subplots(figsize=(10, 3.2), dpi=130) ax.hist(values, bins=80, color="#7c3f00", alpha=0.82) for percentile, color in [(50, "#111827"), (95, "#0f766e"), (99, "#dc2626")]: marker = float(np.percentile(values, percentile)) ax.axvline(marker, color=color, linewidth=1.2, label=f"P{percentile}={format_mm(marker)} mm") ax.set_title("DDF 位移强度分布") ax.set_xlabel("mm") ax.set_ylabel("体素数") ax.grid(axis="y", alpha=0.2) ax.legend(frameon=False) st.pyplot(fig2, width="stretch") 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", width="stretch") 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("显示对象", ["固定图像", "移动图像", "配准后图像"], horizontal=True) if view_target == "固定图像": render_three_views(fixed_xyz, "固定图像") elif view_target == "移动图像": render_three_views(moving_xyz, "移动图像") elif warped_xyz is not None: render_three_views(warped_xyz, "配准后图像") else: st.warning("尚未生成配准后图像。") with tab_overlay: if warped_xyz is None: st.warning("尚未生成配准后图像。") 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("尚未生成配准后图像。") 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()