Add neck-based CT prealignment before VoxelMorph

app.py

@@ -19,9 +19,12 @@ 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,
 )
@@ -608,6 +611,37 @@ def run_inference_from_ui(moving_path: str, fixed_path: str, checkpoint_path: st
     )
 
 
+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
 
@@ -636,7 +670,7 @@ def main() -> None:
 
     st.caption(
         "患者1专用：固定图像 = 平扫CT；移动图像 = 仰头CT。"
-        "推理前会重采样、归一化并裁剪/填充到同一模型网格。"
+        "先按颈部 foreground 做平移预配准，再进入 VoxelMorph 训练/推理。"
     )
     info_cols = st.columns(4)
     info_cols[0].metric("Fixed", "患者1-平扫CT")
@@ -646,27 +680,31 @@ def main() -> None:
 
     action_cols = st.columns([1, 1, 4])
     with action_cols[0]:
-        train_now = st.button("重新训练模型", type="primary", width="stretch")
+        train_now = st.button("颈部预配准+重新训练", type="primary", width="stretch")
     with action_cols[1]:
         start = st.button("开始推理", width="stretch")
 
     if train_now:
-        with st.spinner("患者1模型训练中"):
+        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("训练和推理完成")
+                st.success("颈部预配准、训练和推理完成")
             except Exception as exc:
                 st.error(str(exc))
 
     if start:
-        with st.spinner("推理运行中"):
+        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("推理完成")
@@ -701,6 +739,16 @@ def main() -> None:
     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)