diff --git a/README.md b/README.md index 656a413..739c8db 100644 --- a/README.md +++ b/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。 diff --git a/app.py b/app.py index 79942c6..39a5001 100644 --- a/app.py +++ b/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 data.ndim > 4: + if len(img.shape) > 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: - data = data[::stride, ::stride, ::stride, :] + if len(img.shape) == 4: + data = np.asarray(img.dataobj[spatial_slices + (slice(None),)], dtype=np.float32) else: - data = data[::stride, ::stride, ::stride] - data = np.nan_to_num(data.astype(np.float32, copy=False), copy=False) + 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] @@ -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}") - c3.metric("MAE", f"{metrics['after_mae']:.5f}", delta=f"{-metrics['mae_improvement']:+.5f}") + 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"]] @@ -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"])) - ddf_path = str(st.session_state.get("last_result", {}).get("ddf_path", result_paths["ddf"])) + 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 "缺失") @@ -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) - fixed_xyz, fixed_spacing, fixed_stride = load_nifti_cached(fixed_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, 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")))