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

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