944 lines
35 KiB
Python
944 lines
35 KiB
Python
"""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
|
||
from scipy import ndimage
|
||
import streamlit as st
|
||
|
||
from config import (
|
||
CHECKPOINT_DIR,
|
||
DEFAULT_CHECKPOINT,
|
||
DEFAULT_FIXED_DICOM_DIR,
|
||
DEFAULT_FIXED_NIFTI,
|
||
DEFAULT_MOVING_DICOM_DIR,
|
||
DEFAULT_MOVING_NIFTI,
|
||
INFERENCE_DIR,
|
||
NIFTI_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="collapsed",
|
||
)
|
||
|
||
|
||
CUSTOM_CSS = """
|
||
<style>
|
||
:root {
|
||
--ink: #202124;
|
||
--muted: #5f6368;
|
||
--line: #d8dee4;
|
||
--panel: #f6f8fa;
|
||
--teal: #0f766e;
|
||
--amber: #b45309;
|
||
}
|
||
.main .block-container {
|
||
padding-top: 1.1rem;
|
||
padding-bottom: 2rem;
|
||
max-width: 1480px;
|
||
}
|
||
h1, h2, h3 {
|
||
letter-spacing: 0;
|
||
}
|
||
div[data-testid="stMetric"] {
|
||
border: 1px solid var(--line);
|
||
border-radius: 6px;
|
||
padding: 0.65rem 0.75rem;
|
||
background: #ffffff;
|
||
}
|
||
section[data-testid="stSidebar"] {
|
||
border-right: 1px solid var(--line);
|
||
}
|
||
.stTabs [data-baseweb="tab-list"] {
|
||
gap: 8px;
|
||
border-bottom: 1px solid var(--line);
|
||
}
|
||
.stTabs [data-baseweb="tab"] {
|
||
border-radius: 0;
|
||
padding-left: 4px;
|
||
padding-right: 4px;
|
||
}
|
||
</style>
|
||
"""
|
||
|
||
|
||
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, plane: str) -> np.ndarray:
|
||
"""按平面统一显示方向。
|
||
|
||
数组轴约定为 (X, Y, Z)，其中 Z 轴在当前患者1 DICOM 中从头侧到足侧
|
||
递增。冠状面和矢状面应让 Z 轴作为屏幕竖直方向，避免头脚倒置。
|
||
"""
|
||
|
||
if image.ndim == 3 and image.shape[-1] in (3, 4):
|
||
return np.swapaxes(image, 0, 1)
|
||
return image.T
|
||
|
||
|
||
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 plane_shift_from_xyz(plane: str, shift_xyz: Sequence[float]) -> Tuple[float, float]:
|
||
sx, sy, sz = (float(v) for v in shift_xyz)
|
||
if plane == "Axial":
|
||
return sx, sy
|
||
if plane == "Coronal":
|
||
return sx, sz
|
||
if plane == "Sagittal":
|
||
return sy, sz
|
||
raise ValueError(f"未知平面: {plane}")
|
||
|
||
|
||
def apply_manual_slice_transform(
|
||
image: np.ndarray,
|
||
plane: str,
|
||
shift_xyz: Sequence[float],
|
||
scale: float,
|
||
rotation_deg: float,
|
||
) -> np.ndarray:
|
||
"""对被叠加切片做人工 2D 微调，仅用于可视化。"""
|
||
|
||
scale = max(float(scale), 1e-3)
|
||
shift = np.asarray(plane_shift_from_xyz(plane, shift_xyz), dtype=np.float64)
|
||
theta = np.deg2rad(float(rotation_deg))
|
||
forward = scale * np.asarray(
|
||
[
|
||
[np.cos(theta), -np.sin(theta)],
|
||
[np.sin(theta), np.cos(theta)],
|
||
],
|
||
dtype=np.float64,
|
||
)
|
||
inverse = np.linalg.inv(forward)
|
||
center = (np.asarray(image.shape[:2], dtype=np.float64) - 1.0) / 2.0
|
||
offset = center - inverse @ center - shift
|
||
return ndimage.affine_transform(
|
||
image.astype(np.float32, copy=False),
|
||
matrix=inverse,
|
||
offset=offset,
|
||
output_shape=image.shape,
|
||
order=1,
|
||
mode="constant",
|
||
cval=0.0,
|
||
prefilter=False,
|
||
output=np.float32,
|
||
)
|
||
|
||
|
||
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 gray_fusion(fixed_slice: np.ndarray, warped_slice: np.ndarray, alpha: float = 0.5) -> np.ndarray:
|
||
fixed_norm = normalize_slice(fixed_slice)
|
||
warped_norm = normalize_slice(warped_slice)
|
||
return np.clip(fixed_norm * (1.0 - alpha) + warped_norm * alpha, 0.0, 1.0)
|
||
|
||
|
||
def color_difference_overlay(fixed_slice: np.ndarray, warped_slice: np.ndarray, alpha: float = 0.65) -> np.ndarray:
|
||
fixed_norm = normalize_slice(fixed_slice)
|
||
warped_norm = normalize_slice(warped_slice)
|
||
base = fixed_norm * (1.0 - alpha) + warped_norm * alpha
|
||
rgb = np.zeros((*fixed_norm.shape, 3), dtype=np.float32)
|
||
rgb[..., 0] = np.maximum(base, warped_norm * alpha)
|
||
rgb[..., 1] = base
|
||
rgb[..., 2] = np.maximum(base, fixed_norm * alpha)
|
||
return np.clip(rgb, 0.0, 1.0)
|
||
|
||
|
||
def signed_difference(fixed_slice: np.ndarray, warped_slice: np.ndarray) -> np.ndarray:
|
||
fixed_norm = normalize_slice(fixed_slice)
|
||
warped_norm = normalize_slice(warped_slice)
|
||
return warped_norm - fixed_norm
|
||
|
||
|
||
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", plane: str = "Axial") -> None:
|
||
fig, ax = plt.subplots(figsize=(4.8, 4.8), dpi=120)
|
||
ax.imshow(orient_for_display(image, plane), 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, plane: str = "Axial") -> None:
|
||
fig, ax = plt.subplots(figsize=(4.8, 4.8), dpi=120)
|
||
ax.imshow(orient_for_display(image, plane), interpolation="nearest")
|
||
ax.set_title(caption, fontsize=10)
|
||
ax.axis("off")
|
||
st.pyplot(fig, width="stretch")
|
||
plt.close(fig)
|
||
|
||
|
||
def render_signed_difference(image: np.ndarray, caption: str, plane: str = "Axial") -> None:
|
||
fig, ax = plt.subplots(figsize=(4.8, 4.8), dpi=120)
|
||
diff = orient_for_display(image, plane)
|
||
finite = diff[np.isfinite(diff)]
|
||
vmax = float(np.percentile(np.abs(finite), 99.0)) if finite.size else 1.0
|
||
vmax = max(vmax, 1e-6)
|
||
im = ax.imshow(diff, cmap="coolwarm", vmin=-vmax, vmax=vmax, interpolation="nearest")
|
||
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_heatmap_overlay(
|
||
base_slice: np.ndarray,
|
||
mag_slice: np.ndarray,
|
||
alpha: float,
|
||
caption: str,
|
||
plane: 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), plane), cmap="gray", interpolation="nearest")
|
||
heat = orient_for_display(mag_slice, plane)
|
||
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}", plane=plane)
|
||
|
||
|
||
def render_three_way_views(fixed_xyz: np.ndarray, moving_xyz: np.ndarray, warped_xyz: np.ndarray | None) -> None:
|
||
if warped_xyz is None:
|
||
fixed_xyz, moving_xyz = crop_to_common_shape(fixed_xyz, moving_xyz)
|
||
volumes = [("固定图像", fixed_xyz), ("移动图像", moving_xyz), ("配准后图像", None)]
|
||
st.warning("尚未生成配准后图像。")
|
||
else:
|
||
fixed_xyz, moving_xyz, warped_xyz = crop_to_common_shape(fixed_xyz, moving_xyz, warped_xyz)
|
||
volumes = [("固定图像", fixed_xyz), ("移动图像", moving_xyz), ("配准后图像", warped_xyz)]
|
||
|
||
st.caption(f"显示尺寸：{fixed_xyz.shape[0]} x {fixed_xyz.shape[1]} x {fixed_xyz.shape[2]}")
|
||
planes = [("Axial", "轴状面"), ("Coronal", "冠状面"), ("Sagittal", "矢状面")]
|
||
for plane, cn_name in planes:
|
||
axis = plane_axis(plane)
|
||
index = st.slider(
|
||
f"{cn_name}切片",
|
||
0,
|
||
fixed_xyz.shape[axis] - 1,
|
||
fixed_xyz.shape[axis] // 2,
|
||
key=f"three-way-{plane}",
|
||
)
|
||
st.markdown(f"#### {cn_name}")
|
||
columns = st.columns(3)
|
||
for col, (label, volume) in zip(columns, volumes):
|
||
with col:
|
||
if volume is None:
|
||
st.info("待生成")
|
||
else:
|
||
render_image(get_slice(volume, plane, index), f"{label} #{index}", plane=plane)
|
||
|
||
|
||
def render_overlay_views(fixed_xyz: np.ndarray, moving_xyz: np.ndarray, warped_xyz: np.ndarray) -> None:
|
||
fixed_xyz, moving_xyz, warped_xyz = crop_to_common_shape(fixed_xyz, moving_xyz, warped_xyz)
|
||
target_label = st.radio("叠加对象", ["配准后图像", "移动图像"], horizontal=True)
|
||
target_xyz = warped_xyz if target_label == "配准后图像" else moving_xyz
|
||
mode = st.radio("对比模式", ["灰度融合", "红蓝差异", "差异热图", "棋盘格"], horizontal=True)
|
||
if mode in {"灰度融合", "红蓝差异"}:
|
||
alpha = st.slider("被叠加图像权重", 0.0, 1.0, 0.5, 0.05)
|
||
else:
|
||
alpha = 0.5
|
||
if mode == "棋盘格":
|
||
tile = st.slider("棋盘格尺寸", 8, 64, 24, 4)
|
||
else:
|
||
tile = 24
|
||
|
||
if mode == "红蓝差异":
|
||
st.caption("红/金色偏向被叠加图像，蓝色偏向固定图像；白灰区域表示两者灰度接近。")
|
||
elif mode == "差异热图":
|
||
st.caption("差异热图显示被叠加图像减固定图像的灰度差，色条以当前切片的 99% 绝对差自动缩放。")
|
||
|
||
with st.expander("人工微调重叠", expanded=True):
|
||
enabled = st.toggle("启用人工微调", value=False)
|
||
c1, c2, c3, c4 = st.columns(4)
|
||
with c1:
|
||
shift_x = st.slider("X 平移", -80.0, 80.0, 0.0, 0.5)
|
||
with c2:
|
||
shift_y = st.slider("Y 平移", -80.0, 80.0, 0.0, 0.5)
|
||
with c3:
|
||
shift_z = st.slider("Z 平移", -80.0, 80.0, 0.0, 0.5)
|
||
with c4:
|
||
scale = st.slider("缩放", 0.85, 1.15, 1.0, 0.005)
|
||
r1, r2, r3 = st.columns(3)
|
||
with r1:
|
||
axial_rotation = st.slider("轴状面旋转", -15.0, 15.0, 0.0, 0.5)
|
||
with r2:
|
||
coronal_rotation = st.slider("冠状面旋转", -15.0, 15.0, 0.0, 0.5)
|
||
with r3:
|
||
sagittal_rotation = st.slider("矢状面旋转", -15.0, 15.0, 0.0, 0.5)
|
||
if enabled:
|
||
st.caption(
|
||
f"当前人工微调：X={shift_x:+.1f}, Y={shift_y:+.1f}, Z={shift_z:+.1f}, "
|
||
f"缩放={scale:.3f}。这些参数只影响当前重叠显示。"
|
||
)
|
||
|
||
shift_xyz = (shift_x, shift_y, shift_z) if enabled else (0.0, 0.0, 0.0)
|
||
rotation_by_plane = {
|
||
"Axial": axial_rotation if enabled else 0.0,
|
||
"Coronal": coronal_rotation if enabled else 0.0,
|
||
"Sagittal": sagittal_rotation if enabled else 0.0,
|
||
}
|
||
scale_value = scale if enabled else 1.0
|
||
|
||
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)
|
||
target_slice = get_slice(target_xyz, plane, index)
|
||
target_slice = apply_manual_slice_transform(
|
||
target_slice,
|
||
plane=plane,
|
||
shift_xyz=shift_xyz,
|
||
scale=scale_value,
|
||
rotation_deg=rotation_by_plane[plane],
|
||
)
|
||
if mode == "棋盘格":
|
||
render_image(checkerboard(fixed_slice, target_slice, tile=tile), f"{cn_name} 棋盘格 #{index}", plane=plane)
|
||
elif mode == "红蓝差异":
|
||
render_rgb(color_difference_overlay(fixed_slice, target_slice, alpha=alpha), f"{cn_name} 红蓝差异 #{index}", plane=plane)
|
||
elif mode == "差异热图":
|
||
render_signed_difference(signed_difference(fixed_slice, target_slice), f"{cn_name} 差异热图 #{index}", plane=plane)
|
||
else:
|
||
render_image(gray_fusion(fixed_slice, target_slice, alpha=alpha), f"{cn_name} 灰度融合 #{index}", plane=plane)
|
||
|
||
|
||
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,
|
||
plane=plane,
|
||
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 {
|
||
"prepared_moving": str(out / "moving_model_input.nii.gz"),
|
||
"prepared_fixed": str(out / "fixed_model_input.nii.gz"),
|
||
"warped": str(out / "warped_moving.nii.gz"),
|
||
"ddf": str(out / "ddf_mm.nii.gz"),
|
||
"metrics": str(out / "metrics.json"),
|
||
}
|
||
|
||
|
||
def read_json_dict(path: str | Path) -> Dict:
|
||
path = Path(path)
|
||
if not path.exists():
|
||
return {}
|
||
try:
|
||
data = json.loads(path.read_text(encoding="utf-8"))
|
||
except Exception:
|
||
return {}
|
||
return data if isinstance(data, dict) else {}
|
||
|
||
|
||
def same_path(left: str | Path | None, right: str | Path | None) -> bool:
|
||
if not left or not right:
|
||
return False
|
||
try:
|
||
return Path(left).expanduser().resolve(strict=False) == Path(right).expanduser().resolve(strict=False)
|
||
except Exception:
|
||
return str(left) == str(right)
|
||
|
||
|
||
def result_matches_inputs(result: Dict, moving_path: str, fixed_path: str) -> bool:
|
||
return same_path(result.get("moving_path"), moving_path) and same_path(result.get("fixed_path"), fixed_path)
|
||
|
||
|
||
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 prepare_patient1_neck_aligned_inputs() -> Dict:
|
||
from data_loader import convert_dicom_series_to_nifti
|
||
from prealign import prealign_pair
|
||
|
||
fixed_raw_path = NIFTI_DIR / "patient1_fixed.nii.gz"
|
||
moving_raw_path = NIFTI_DIR / "patient1_moving.nii.gz"
|
||
convert_dicom_series_to_nifti(
|
||
dicom_dir=DEFAULT_FIXED_DICOM_DIR,
|
||
output_path=fixed_raw_path,
|
||
max_memory_mb=8192,
|
||
)
|
||
convert_dicom_series_to_nifti(
|
||
dicom_dir=DEFAULT_MOVING_DICOM_DIR,
|
||
output_path=moving_raw_path,
|
||
max_memory_mb=8192,
|
||
)
|
||
meta = prealign_pair(
|
||
fixed_input_path=fixed_raw_path,
|
||
moving_input_path=moving_raw_path,
|
||
fixed_output_path=DEFAULT_FIXED_NIFTI,
|
||
moving_output_path=DEFAULT_MOVING_NIFTI,
|
||
max_memory_mb=8192,
|
||
)
|
||
return {
|
||
"moving_translation_mm": list(meta.moving_translation_mm),
|
||
"fixed_neck_center_world_mm": list(meta.fixed_neck_center_world_mm),
|
||
"moving_neck_center_world_mm": list(meta.moving_neck_center_world_mm),
|
||
"target_shape_xyz": list(meta.target_shape_xyz),
|
||
}
|
||
|
||
|
||
def run_training_from_ui(moving_path: str, fixed_path: str, checkpoint_path: str) -> None:
|
||
from model_and_train import train_pair
|
||
|
||
train_pair(
|
||
moving_path=moving_path,
|
||
fixed_path=fixed_path,
|
||
checkpoint_path=checkpoint_path,
|
||
epochs=80,
|
||
learning_rate=1e-4,
|
||
smooth_weight=0.01,
|
||
image_loss="mse",
|
||
nb_features=(8, 8, 8, 8, 8),
|
||
save_every=20,
|
||
)
|
||
|
||
|
||
def main() -> None:
|
||
st.markdown(CUSTOM_CSS, unsafe_allow_html=True)
|
||
st.title("VoxelMorph 颈部 CT 配准工作台")
|
||
|
||
moving_path = str(DEFAULT_MOVING_NIFTI)
|
||
fixed_path = str(DEFAULT_FIXED_NIFTI)
|
||
checkpoint_path = str(DEFAULT_CHECKPOINT)
|
||
out_dir = str(INFERENCE_DIR)
|
||
display_max_voxels = 30_000_000
|
||
|
||
st.caption(
|
||
"患者1专用：固定图像 = 平扫CT；移动图像 = 仰头CT。"
|
||
"先按颈部 foreground 做平移预配准，再进入 VoxelMorph 训练/推理。"
|
||
)
|
||
info_cols = st.columns(4)
|
||
info_cols[0].metric("Fixed", "患者1-平扫CT")
|
||
info_cols[1].metric("Moving", "患者1-仰头CT")
|
||
info_cols[2].metric("模型", Path(checkpoint_path).name)
|
||
info_cols[3].metric("输出", Path(out_dir).name)
|
||
|
||
action_cols = st.columns([1, 1, 4])
|
||
with action_cols[0]:
|
||
train_now = st.button("颈部预配准+重新训练", type="primary", width="stretch")
|
||
with action_cols[1]:
|
||
start = st.button("开始推理", width="stretch")
|
||
|
||
if train_now:
|
||
with st.spinner("患者1颈部预配准、模型训练和推理中"):
|
||
try:
|
||
load_nifti_cached.clear()
|
||
align_info = prepare_patient1_neck_aligned_inputs()
|
||
run_training_from_ui(moving_path, fixed_path, checkpoint_path)
|
||
result = run_inference_from_ui(moving_path, fixed_path, checkpoint_path, out_dir)
|
||
result["neck_alignment"] = align_info
|
||
load_nifti_cached.clear()
|
||
st.session_state["last_result"] = result
|
||
st.success("颈部预配准、训练和推理完成")
|
||
except Exception as exc:
|
||
st.error(str(exc))
|
||
|
||
if start:
|
||
with st.spinner("颈部预配准和推理运行中"):
|
||
try:
|
||
load_nifti_cached.clear()
|
||
align_info = prepare_patient1_neck_aligned_inputs()
|
||
result = run_inference_from_ui(moving_path, fixed_path, checkpoint_path, out_dir)
|
||
result["neck_alignment"] = align_info
|
||
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", {})
|
||
metrics_path = str(last_result.get("metrics_path", result_paths["metrics"]))
|
||
saved_result = read_json_dict(metrics_path)
|
||
result_info = {**saved_result, **last_result}
|
||
outputs_current = result_matches_inputs(result_info, moving_path, fixed_path)
|
||
|
||
prepared_moving_path = str(result_info.get("prepared_moving_path", result_paths["prepared_moving"]))
|
||
prepared_fixed_path = str(result_info.get("prepared_fixed_path", result_paths["prepared_fixed"]))
|
||
if outputs_current:
|
||
warped_path = str(result_info.get("warped_path", result_paths["warped"]))
|
||
ddf_path = str(result_info.get("ddf_path", result_paths["ddf"]))
|
||
display_moving_path = prepared_moving_path if Path(prepared_moving_path).exists() else moving_path
|
||
display_fixed_path = prepared_fixed_path if Path(prepared_fixed_path).exists() else fixed_path
|
||
else:
|
||
warped_path = ""
|
||
ddf_path = ""
|
||
display_moving_path = moving_path
|
||
display_fixed_path = fixed_path
|
||
|
||
status_cols = st.columns(4)
|
||
status_cols[0].metric("Moving", "存在" if Path(display_moving_path).exists() else "缺失")
|
||
status_cols[1].metric("Fixed", "存在" if Path(display_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 "缺失")
|
||
if result_info and not outputs_current:
|
||
st.warning("输出目录中的历史结果与当前输入不匹配，请重新开始推理。")
|
||
|
||
alignment_meta = read_json_dict(Path(DEFAULT_FIXED_NIFTI).parent / "patient1_neck_alignment.json")
|
||
translation = alignment_meta.get("moving_translation_mm")
|
||
if isinstance(translation, list) and len(translation) == 3:
|
||
st.caption(
|
||
"颈部预配准平移："
|
||
f"X={float(translation[0]):+.2f} mm，"
|
||
f"Y={float(translation[1]):+.2f} mm，"
|
||
f"Z={float(translation[2]):+.2f} mm"
|
||
)
|
||
|
||
try:
|
||
moving_xyz, moving_spacing, moving_stride = load_nifti_cached(display_moving_path, max_voxels=display_max_voxels)
|
||
fixed_xyz, fixed_spacing, fixed_stride = load_nifti_cached(display_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:
|
||
render_three_way_views(fixed_xyz, moving_xyz, warped_xyz)
|
||
|
||
with tab_overlay:
|
||
if warped_xyz is None:
|
||
st.warning("尚未生成配准后图像。")
|
||
else:
|
||
render_overlay_views(fixed_xyz, moving_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()
|