admin/Voxelmorph_Head_CT
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Improve Streamlit registration viewer robustness

Browse Source
This commit is contained in:
admin
2026-06-03 00:36:24 +08:00
parent 2dba05ae4a
commit 68353b2148
2 changed files with 49 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
README.md
View File

@@ -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。

66
app.py
View File

@@ -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")))