133 lines
4.0 KiB
Python
133 lines
4.0 KiB
Python
"""配准质量评估工具。
|
||
|
||
这里刻意只依赖 NumPy，便于在训练、推理和 Streamlit 三处复用。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from typing import Dict, Iterable, Tuple
|
||
|
||
import numpy as np
|
||
|
||
|
||
def crop_to_common_shape(*arrays: np.ndarray) -> Tuple[np.ndarray, ...]:
|
||
"""把多个 3D/4D 数组中心裁剪到共同尺寸，避免形状不一致导致评估失败。"""
|
||
|
||
if not arrays:
|
||
return tuple()
|
||
|
||
spatial_shapes = [arr.shape[:3] for arr in arrays]
|
||
common = tuple(min(shape[axis] for shape in spatial_shapes) for axis in range(3))
|
||
cropped = []
|
||
|
||
for arr in arrays:
|
||
slices = []
|
||
for axis, target in enumerate(common):
|
||
start = max((arr.shape[axis] - target) // 2, 0)
|
||
slices.append(slice(start, start + target))
|
||
if arr.ndim > 3:
|
||
slices.append(slice(None))
|
||
cropped.append(arr[tuple(slices)])
|
||
|
||
return tuple(cropped)
|
||
|
||
|
||
def global_ncc(a: np.ndarray, b: np.ndarray, eps: float = 1e-8) -> float:
|
||
"""计算全局 NCC。训练时使用局部 NCC，这里用于快速量化展示。"""
|
||
|
||
a = np.asarray(a, dtype=np.float32)
|
||
b = np.asarray(b, dtype=np.float32)
|
||
a = a - float(np.mean(a))
|
||
b = b - float(np.mean(b))
|
||
denom = float(np.sqrt(np.sum(a * a) * np.sum(b * b)) + eps)
|
||
return float(np.sum(a * b) / denom)
|
||
|
||
|
||
def mse(a: np.ndarray, b: np.ndarray) -> float:
|
||
diff = np.asarray(a, dtype=np.float32) - np.asarray(b, dtype=np.float32)
|
||
return float(np.mean(diff * diff))
|
||
|
||
|
||
def mae(a: np.ndarray, b: np.ndarray) -> float:
|
||
diff = np.abs(np.asarray(a, dtype=np.float32) - np.asarray(b, dtype=np.float32))
|
||
return float(np.mean(diff))
|
||
|
||
|
||
def registration_metrics(
|
||
fixed: np.ndarray,
|
||
moving: np.ndarray,
|
||
warped: np.ndarray,
|
||
) -> Dict[str, float]:
|
||
"""输出配准前后可比较的常用指标。"""
|
||
|
||
fixed, moving, warped = crop_to_common_shape(fixed, moving, warped)
|
||
before_mse = mse(fixed, moving)
|
||
after_mse = mse(fixed, warped)
|
||
before_mae = mae(fixed, moving)
|
||
after_mae = mae(fixed, warped)
|
||
before_ncc = global_ncc(fixed, moving)
|
||
after_ncc = global_ncc(fixed, warped)
|
||
|
||
return {
|
||
"before_mse": before_mse,
|
||
"after_mse": after_mse,
|
||
"before_mae": before_mae,
|
||
"after_mae": after_mae,
|
||
"before_ncc": before_ncc,
|
||
"after_ncc": after_ncc,
|
||
"mse_improvement": before_mse - after_mse,
|
||
"mae_improvement": before_mae - after_mae,
|
||
"ncc_improvement": after_ncc - before_ncc,
|
||
}
|
||
|
||
|
||
def slice_metric_curve(
|
||
fixed: np.ndarray,
|
||
moving: np.ndarray,
|
||
warped: np.ndarray,
|
||
axis: int = 2,
|
||
metric: str = "mse",
|
||
) -> Dict[str, Iterable[float]]:
|
||
"""逐切片计算配准前后指标曲线。"""
|
||
|
||
fixed, moving, warped = crop_to_common_shape(fixed, moving, warped)
|
||
metric_funcs = {
|
||
"mse": mse,
|
||
"mae": mae,
|
||
"ncc": global_ncc,
|
||
}
|
||
if metric not in metric_funcs:
|
||
raise ValueError(f"不支持的逐切片指标: {metric}")
|
||
metric_fn = metric_funcs[metric]
|
||
|
||
before = []
|
||
after = []
|
||
|
||
for index in range(fixed.shape[axis]):
|
||
selector = [slice(None)] * 3
|
||
selector[axis] = index
|
||
selector = tuple(selector)
|
||
before.append(metric_fn(fixed[selector], moving[selector]))
|
||
after.append(metric_fn(fixed[selector], warped[selector]))
|
||
|
||
return {
|
||
"slice_index": list(range(fixed.shape[axis])),
|
||
f"before_{metric}": before,
|
||
f"after_{metric}": after,
|
||
}
|
||
|
||
|
||
def ddf_summary(ddf_xyz: np.ndarray) -> Dict[str, float]:
|
||
"""统计形变场向量大小。输入应为 X/Y/Z/3，单位可为 voxel 或 mm。"""
|
||
|
||
if ddf_xyz.ndim != 4 or ddf_xyz.shape[-1] != 3:
|
||
raise ValueError("DDF 必须是形状为 (X, Y, Z, 3) 的 4D 数组。")
|
||
|
||
mag = np.linalg.norm(ddf_xyz.astype(np.float32), axis=-1)
|
||
return {
|
||
"ddf_mean": float(np.mean(mag)),
|
||
"ddf_std": float(np.std(mag)),
|
||
"ddf_p95": float(np.percentile(mag, 95)),
|
||
"ddf_max": float(np.max(mag)),
|
||
}
|