Improve Streamlit registration viewer robustness

README.md

@@ -128,6 +128,7 @@ streamlit run app.py
 网页提供：
 
 - Moving/Fixed/模型权重/输出目录输入。
+- 自动发现 `outputs/` 与项目目录下的 NIfTI 和 checkpoint，也支持手动输入路径。
 - “开始推理”按钮。
 - Axial、Coronal、Sagittal 正交三视图。
 - Fixed 与 Warped 的 Alpha 融合或棋盘格对比。
@@ -137,5 +138,5 @@ streamlit run app.py
 ## 内存注意事项
 
 - DICOM 转换和重采样都有 `--max-memory-mb` 防护。
-- Web 界面对超大 NIfTI 会自动 stride 下采样，只影响浏览器展示，不改变磁盘结果。
+- Web 界面对超大 NIfTI 会通过 nibabel proxy 按 stride 切片读取并下采样，只影响浏览器展示，不改变磁盘结果；侧栏可调整显示体素上限。
 - 训练阶段的主要瓶颈是 3D U-Net 显存；`160x192x224` 是较重的 3D 输入，建议优先使用 CUDA GPU。

app.py

@@ -8,13 +8,14 @@ from __future__ import annotations
 
 import json
 from pathlib import Path
-from typing import Dict, Iterable, List, Sequence, Tuple
+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,
@@ -91,10 +92,21 @@ def discover_nifti_files() -> List[str]:
     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 0
+    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)
@@ -113,19 +125,19 @@ def load_nifti_cached(path: str, max_voxels: int = 14_000_000) -> Tuple[np.ndarr
         raise ValueError("路径为空。")
     nib = _require_nibabel()
     img = nib.load(path, mmap=True)
-    data = np.asanyarray(img.dataobj, dtype=np.float32)
+    if len(img.shape) > 4:
-    if data.ndim > 4:
         raise ValueError(f"暂不支持超过 4D 的 NIfTI: {path}")
 
-    spatial_shape = data.shape[:3]
+    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 data.ndim == 4:
+    if len(img.shape) == 4:
-        data = data[::stride, ::stride, ::stride, :]
+        data = np.asarray(img.dataobj[spatial_slices + (slice(None),)], dtype=np.float32)
     else:
-        data = data[::stride, ::stride, ::stride]
+        data = np.asarray(img.dataobj[spatial_slices], dtype=np.float32)
-    data = np.nan_to_num(data.astype(np.float32, copy=False), copy=False)
+    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]
 
@@ -275,8 +287,18 @@ def render_metric_charts(fixed_xyz: np.ndarray, moving_xyz: np.ndarray, warped_x
 
     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['mse_improvement']:+.5f}")
+    c2.metric(
-    c3.metric("MAE", f"{metrics['after_mae']:.5f}", delta=f"{-metrics['mae_improvement']:+.5f}")
+        "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"]]
@@ -339,26 +361,32 @@ def main() -> None:
     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 = st.text_input("模型权重", value=str(DEFAULT_CHECKPOINT))
+        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)
-    warped_path = str(st.session_state.get("last_result", {}).get("warped_path", result_paths["warped"]))
+    last_result = st.session_state.get("last_result", {})
-    ddf_path = str(st.session_state.get("last_result", {}).get("ddf_path", result_paths["ddf"]))
+    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 "缺失")
@@ -367,8 +395,8 @@ def main() -> None:
     status_cols[3].metric("DDF", "存在" if Path(ddf_path).exists() else "缺失")
 
     try:
-        moving_xyz, moving_spacing, moving_stride = load_nifti_cached(moving_path)
+        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)
+        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
@@ -377,12 +405,12 @@ def main() -> None:
     ddf_xyz = None
     if Path(warped_path).exists():
         try:
-            warped_xyz, _, _ = load_nifti_cached(warped_path)
+            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)
+            ddf_xyz, _, _ = load_nifti_cached(ddf_path, max_voxels=display_max_voxels)
         except Exception as exc:
             st.warning(f"DDF 读取失败: {exc}")
 
@@ -418,7 +446,7 @@ def main() -> None:
             st.warning("尚未生成 Warped Image。")
         else:
             render_metric_charts(fixed_xyz, moving_xyz, warped_xyz, ddf_xyz)
-            metrics_file = Path(result_paths["metrics"])
+            metrics_file = Path(metrics_path)
             if metrics_file.exists():
                 try:
                     st.json(json.loads(metrics_file.read_text(encoding="utf-8")))