diff --git a/README.md b/README.md index 739c8db..4c86d83 100644 --- a/README.md +++ b/README.md @@ -115,10 +115,14 @@ python infer.py \ 推理会输出： +- `outputs/inference/moving_model_input.nii.gz` +- `outputs/inference/fixed_model_input.nii.gz` - `outputs/inference/warped_moving.nii.gz` - `outputs/inference/ddf_mm.nii.gz` - `outputs/inference/metrics.json` +`moving_model_input.nii.gz` 和 `fixed_model_input.nii.gz` 是进入 VoxelMorph 前的统一网格图像。即使两套 CT 原始层数不同，推理前也会按 checkpoint 的输入尺寸和目标 spacing 完成重采样、归一化、中心裁剪/填充。 + ## 5. Web 结果展示 ```bash @@ -130,7 +134,7 @@ streamlit run app.py - Moving/Fixed/模型权重/输出目录输入。 - 自动发现 `outputs/` 与项目目录下的 NIfTI 和 checkpoint，也支持手动输入路径。 - “开始推理”按钮。 -- Axial、Coronal、Sagittal 正交三视图。 +- Axial、Coronal、Sagittal 正交三视图；每个平面按行同时展示 Fixed、Moving、Warped。 - Fixed 与 Warped 的 Alpha 融合或棋盘格对比。 - DDF 位移强度热力图。 - NCC、MSE、MAE、逐切片误差曲线、DDF 位移分布等量化图。 diff --git a/app.py b/app.py index 588c626..4458f38 100644 --- a/app.py +++ b/app.py @@ -321,6 +321,36 @@ def render_three_views(volume_xyz: np.ndarray, title: str) -> None: render_image(get_slice(volume_xyz, plane, index), f"{cn_name} #{index}") +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}") + + def render_overlay_views(fixed_xyz: np.ndarray, warped_xyz: np.ndarray) -> None: fixed_xyz, warped_xyz = crop_to_common_shape(fixed_xyz, warped_xyz) mode = st.radio("对比模式", ["Alpha 融合", "棋盘格"], horizontal=True) @@ -535,12 +565,38 @@ def render_metric_charts(fixed_xyz: np.ndarray, moving_xyz: np.ndarray, warped_x 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 @@ -580,19 +636,35 @@ def main() -> None: result_paths = load_result_paths(out_dir) 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"])) + 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(moving_path).exists() else "缺失") - status_cols[1].metric("Fixed", "存在" if Path(fixed_path).exists() else "缺失") + 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("输出目录中的历史结果与当前输入不匹配，请重新开始推理。") try: - 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) + 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 @@ -615,15 +687,7 @@ def main() -> None: tab_views, tab_overlay, tab_ddf, tab_metrics = st.tabs(["正交三视图", "重叠对比", "形变场", "量化图"]) with tab_views: - view_target = st.radio("显示对象", ["固定图像", "移动图像", "配准后图像"], horizontal=True) - if view_target == "固定图像": - render_three_views(fixed_xyz, "固定图像") - elif view_target == "移动图像": - render_three_views(moving_xyz, "移动图像") - elif warped_xyz is not None: - render_three_views(warped_xyz, "配准后图像") - else: - st.warning("尚未生成配准后图像。") + render_three_way_views(fixed_xyz, moving_xyz, warped_xyz) with tab_overlay: if warped_xyz is None: diff --git a/infer.py b/infer.py index 73feb12..6b50167 100644 --- a/infer.py +++ b/infer.py @@ -1,6 +1,8 @@ """独立推理模块。 本模块加载官方 ``voxelmorph.nn.models.VxmPairwise`` 训练出的权重，输出： +- fixed_model_input.nii.gz：重采样、归一化、裁剪/填充后的 Fixed。 +- moving_model_input.nii.gz：重采样、归一化、裁剪/填充后的 Moving。 - warped_moving.nii.gz：形变后的 Moving。 - ddf_mm.nii.gz：Dense Displacement Field，shape = (X, Y, Z, 3)，单位 mm。 - metrics.json：配准前后 NCC/MSE/MAE 与 DDF 统计。 @@ -143,7 +145,7 @@ def run_inference( nb_features = checkpoint.get("nb_features", [16, 16, 16, 16, 16]) integration_steps = int(checkpoint.get("integration_steps", 0)) - moving_xyz, _ = prepare_nifti_for_model( + moving_xyz, moving_affine = prepare_nifti_for_model( moving_path, target_shape_xyz=target_shape_xyz, target_spacing=target_spacing, @@ -182,10 +184,14 @@ def run_inference( out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) + prepared_moving_path = out_dir / "moving_model_input.nii.gz" + prepared_fixed_path = out_dir / "fixed_model_input.nii.gz" warped_path = out_dir / "warped_moving.nii.gz" ddf_path = out_dir / "ddf_mm.nii.gz" metrics_path = out_dir / "metrics.json" + save_nifti(moving_xyz, moving_affine, prepared_moving_path) + save_nifti(fixed_xyz, fixed_affine, prepared_fixed_path) save_nifti(warped_xyz, fixed_affine, warped_path) save_ddf_nifti(ddf_xyz_mm, fixed_affine, ddf_path) @@ -195,6 +201,8 @@ def run_inference( { "moving_path": str(moving_path), "fixed_path": str(fixed_path), + "prepared_moving_path": str(prepared_moving_path), + "prepared_fixed_path": str(prepared_fixed_path), "checkpoint_path": str(checkpoint_path), "warped_path": str(warped_path), "ddf_path": str(ddf_path),