diff --git a/README.md b/README.md index 0320ebc..ecff070 100644 --- a/README.md +++ b/README.md @@ -25,7 +25,8 @@ Voxelmorph_Head_CT/ ├── infer.py # 独立推理，输出 warped image 与 DDF ├── metrics.py # NCC/MSE/MAE/DDF 等量化指标 ├── model_and_train.py # 官方 VoxelMorph 训练适配器 -├── preprocess.py # 重采样、窗宽窗位、裁剪/填充 +├── prealign.py # 颈部基准 CT 预配准 +├── preprocess.py # 单体 CT 重采样、窗宽窗位、裁剪/填充 ├── requirements.txt └── outputs/ ├── nifti/ # DICOM 转换结果 @@ -69,23 +70,20 @@ python data_loader.py \ `data_loader.py` 会优先按 `InstanceNumber` 排序，其次按 `SliceLocation` 排序，并保存 spacing、层厚、排序依据等元数据 JSON。 -## 2. 预处理 +## 2. 颈部预配准 ```bash -python preprocess.py \ - --input "outputs/nifti/patient1_fixed.nii.gz" \ - --output "outputs/preprocessed/patient1_fixed_preprocessed.nii.gz" \ +python prealign.py \ + --fixed "outputs/nifti/patient1_fixed.nii.gz" \ + --moving "outputs/nifti/patient1_moving.nii.gz" \ + --fixed-output "outputs/preprocessed/patient1_fixed_preprocessed.nii.gz" \ + --moving-output "outputs/preprocessed/patient1_moving_preprocessed.nii.gz" \ --target-spacing 1 1 1 \ - --target-shape 256 256 352 - -python preprocess.py \ - --input "outputs/nifti/patient1_moving.nii.gz" \ - --output "outputs/preprocessed/patient1_moving_preprocessed.nii.gz" \ - --target-spacing 1 1 1 \ - --target-shape 256 256 352 + --target-shape 320 320 352 \ + --metadata-json "outputs/preprocessed/patient1_neck_alignment.json" ``` -默认窗口为 `W=400, L=40`，适合观察颈部软组织和气道。 +该步骤会先把 Fixed/Moving 重采样到 1mm，再基于颈部 foreground mask 估计中心，并将 Moving 平移到 Fixed 的颈部中心，最后在共同颈部空间中裁剪/填充和归一化。默认窗口为 `W=400, L=40`，适合观察颈部软组织和气道。 ## 3. 训练 @@ -121,7 +119,7 @@ python infer.py \ - `outputs/inference/ddf_mm.nii.gz` - `outputs/inference/metrics.json` -`moving_model_input.nii.gz` 和 `fixed_model_input.nii.gz` 是进入 VoxelMorph 前的统一网格图像。即使两套 CT 原始层数不同，推理前也会按 checkpoint 的输入尺寸和目标 spacing 完成重采样、归一化、中心裁剪/填充。 +`moving_model_input.nii.gz` 和 `fixed_model_input.nii.gz` 是进入 VoxelMorph 前的统一网格图像。患者1默认先由 `prealign.py` 按颈部中心完成预配准，再进入训练/推理。 ## 5. Web 结果展示 @@ -132,7 +130,7 @@ streamlit run app.py 网页提供： - 患者1专用固定路径：Fixed 为 `患者1-平扫CT`，Moving 为 `患者1-仰头CT`。 -- “重新训练模型”和“开始推理”按钮。 +- “颈部预配准+重新训练”和“开始推理”按钮；推理前也会刷新颈部预配准输入。 - Axial、Coronal、Sagittal 正交三视图；每个平面按行同时展示 Fixed、Moving、Warped。 - Fixed 与 Warped 的 Alpha 融合或棋盘格对比。 - DDF 位移强度热力图。 @@ -142,4 +140,4 @@ streamlit run app.py - DICOM 转换和重采样都有 `--max-memory-mb` 防护。 - Web 界面对超大 NIfTI 会通过 nibabel proxy 按 stride 切片读取并下采样，只影响浏览器展示，不改变磁盘结果；侧栏可调整显示体素上限。 -- 训练阶段的主要瓶颈是 3D U-Net 显存；`256x256x352` 是较重的 3D 输入，建议优先使用 CUDA GPU。当前患者1默认使用较轻的 `8 8 8 8 8` 特征配置。 +- 训练阶段的主要瓶颈是 3D U-Net 显存；`320x320x352` 是较重的 3D 输入，建议优先使用 CUDA GPU。当前患者1默认使用较轻的 `8 8 8 8 8` 特征配置。 diff --git a/app.py b/app.py index c811c4a..d42ffa3 100644 --- a/app.py +++ b/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) diff --git a/config.py b/config.py index 8ee094f..3796539 100644 --- a/config.py +++ b/config.py @@ -23,8 +23,8 @@ DEFAULT_MOVING_NIFTI = PREPROCESSED_DIR / "patient1_moving_preprocessed.nii.gz" DEFAULT_FIXED_NIFTI = PREPROCESSED_DIR / "patient1_fixed_preprocessed.nii.gz" DEFAULT_CHECKPOINT = CHECKPOINT_DIR / "vxm_head_ct_patient1.pt" -# VoxelMorph 的 3D U-Net 多次下采样，三维尺寸建议均为 16 的倍数。 -DEFAULT_TARGET_SHAPE = (256, 256, 352) # NIfTI 轴顺序: X, Y, Z +# VoxelMorph 的 3D U-Net 多次下采样，三维尺寸建议均为 32 的倍数。 +DEFAULT_TARGET_SHAPE = (320, 320, 352) # NIfTI 轴顺序: X, Y, Z DEFAULT_TARGET_SPACING = (1.0, 1.0, 1.0) # mm, X/Y/Z # 颈部软组织/气道观察常用窗口：W=400, L=40。 diff --git a/outputs/checkpoints/vxm_head_ct_patient1.history.json b/outputs/checkpoints/vxm_head_ct_patient1.history.json index b604be2..98019d3 100644 --- a/outputs/checkpoints/vxm_head_ct_patient1.history.json +++ b/outputs/checkpoints/vxm_head_ct_patient1.history.json @@ -1,482 +1,482 @@ [ { "epoch": 1.0, - "loss": 0.058744531124830246, - "image_loss": 0.058744531124830246, - "smooth_loss": 1.5868747288427798e-12 + "loss": 0.04345163702964783, + "image_loss": 0.04345163702964783, + "smooth_loss": 5.719839749410149e-13 }, { "epoch": 2.0, - "loss": 0.058717112988233566, - "image_loss": 0.058717112988233566, - "smooth_loss": 2.6513908846226286e-09 + "loss": 0.043440207839012146, + "image_loss": 0.043440207839012146, + "smooth_loss": 1.4629444322622476e-09 }, { "epoch": 3.0, - "loss": 0.058689698576927185, - "image_loss": 0.058689698576927185, - "smooth_loss": 1.0584165544003099e-08 + "loss": 0.04342950880527496, + "image_loss": 0.04342950880527496, + "smooth_loss": 5.5231410556189076e-09 }, { "epoch": 4.0, - "loss": 0.058662042021751404, - "image_loss": 0.058662042021751404, - "smooth_loss": 2.401168686105848e-08 + "loss": 0.04341849684715271, + "image_loss": 0.04341849684715271, + "smooth_loss": 1.247603442777745e-08 }, { "epoch": 5.0, - "loss": 0.05863436311483383, - "image_loss": 0.05863436311483383, - "smooth_loss": 4.28750439596115e-08 + "loss": 0.043407391756772995, + "image_loss": 0.043407391756772995, + "smooth_loss": 2.2369590624293778e-08 }, { "epoch": 6.0, - "loss": 0.05860644578933716, - "image_loss": 0.05860644578933716, - "smooth_loss": 6.765733928659756e-08 + "loss": 0.04339618235826492, + "image_loss": 0.04339618235826492, + "smooth_loss": 3.538238146916228e-08 }, { "epoch": 7.0, - "loss": 0.058578286319971085, - "image_loss": 0.058578286319971085, - "smooth_loss": 9.85552759402708e-08 + "loss": 0.04338495060801506, + "image_loss": 0.04338495060801506, + "smooth_loss": 5.16029103891924e-08 }, { "epoch": 8.0, - "loss": 0.05855008214712143, - "image_loss": 0.05855008214712143, - "smooth_loss": 1.3519149888452375e-07 + "loss": 0.0433737076818943, + "image_loss": 0.0433737076818943, + "smooth_loss": 7.094314469213714e-08 }, { "epoch": 9.0, - "loss": 0.05852168798446655, - "image_loss": 0.05852168798446655, - "smooth_loss": 1.7822365805386653e-07 + "loss": 0.04336243495345116, + "image_loss": 0.04336243495345116, + "smooth_loss": 9.347723306518674e-08 }, { "epoch": 10.0, - "loss": 0.05849294736981392, - "image_loss": 0.05849294364452362, - "smooth_loss": 2.282957041188638e-07 + "loss": 0.04335121437907219, + "image_loss": 0.04335121437907219, + "smooth_loss": 1.1908265662441408e-07 }, { "epoch": 11.0, - "loss": 0.058464165776968, - "image_loss": 0.058464162051677704, - "smooth_loss": 2.843589470558072e-07 + "loss": 0.043339941650629044, + "image_loss": 0.043339941650629044, + "smooth_loss": 1.482868015045824e-07 }, { "epoch": 12.0, - "loss": 0.05843484401702881, - "image_loss": 0.05843484029173851, - "smooth_loss": 3.4935141002279124e-07 + "loss": 0.043328724801540375, + "image_loss": 0.043328724801540375, + "smooth_loss": 1.804429388130302e-07 }, { "epoch": 13.0, - "loss": 0.05840571969747543, - "image_loss": 0.058405715972185135, - "smooth_loss": 4.197248131276865e-07 + "loss": 0.04331750050187111, + "image_loss": 0.04331749677658081, + "smooth_loss": 2.166258354918682e-07 }, { "epoch": 14.0, - "loss": 0.05837589502334595, - "image_loss": 0.05837589129805565, - "smooth_loss": 5.000206897420867e-07 + "loss": 0.043306246399879456, + "image_loss": 0.04330624267458916, + "smooth_loss": 2.5594988528609974e-07 }, { "epoch": 15.0, - "loss": 0.058346040546894073, - "image_loss": 0.05834603309631348, - "smooth_loss": 5.874838393538084e-07 + "loss": 0.04329502210021019, + "image_loss": 0.04329501837491989, + "smooth_loss": 2.9920690280960116e-07 }, { "epoch": 16.0, - "loss": 0.058315880596637726, - "image_loss": 0.05831587314605713, - "smooth_loss": 6.829379231021449e-07 + "loss": 0.043283771723508835, + "image_loss": 0.043283767998218536, + "smooth_loss": 3.4602453524712473e-07 }, { "epoch": 17.0, - "loss": 0.05828540772199631, - "image_loss": 0.05828540027141571, - "smooth_loss": 7.876523113736766e-07 + "loss": 0.043272584676742554, + "image_loss": 0.043272580951452255, + "smooth_loss": 3.9618711866751255e-07 }, { "epoch": 18.0, - "loss": 0.05825469642877579, - "image_loss": 0.05825468897819519, - "smooth_loss": 9.008151096168149e-07 + "loss": 0.043261270970106125, + "image_loss": 0.043261267244815826, + "smooth_loss": 4.4989539560447156e-07 }, { "epoch": 19.0, - "loss": 0.05822316184639931, - "image_loss": 0.05822315067052841, - "smooth_loss": 1.0293301784258801e-06 + "loss": 0.04325007647275925, + "image_loss": 0.04325007274746895, + "smooth_loss": 5.079763809590077e-07 }, { "epoch": 20.0, - "loss": 0.05819175764918327, - "image_loss": 0.05819174647331238, - "smooth_loss": 1.1623228601820301e-06 + "loss": 0.04323875159025192, + "image_loss": 0.043238744139671326, + "smooth_loss": 5.703672059098608e-07 }, { "epoch": 21.0, - "loss": 0.058159757405519485, - "image_loss": 0.05815974250435829, - "smooth_loss": 1.3093066399960662e-06 + "loss": 0.04322751984000206, + "image_loss": 0.04322751238942146, + "smooth_loss": 6.375527732416231e-07 }, { "epoch": 22.0, - "loss": 0.058127511292696, - "image_loss": 0.058127496391534805, - "smooth_loss": 1.465943910261558e-06 + "loss": 0.04321633651852608, + "image_loss": 0.04321632906794548, + "smooth_loss": 7.064697911118856e-07 }, { "epoch": 23.0, - "loss": 0.058094657957553864, - "image_loss": 0.05809464305639267, - "smooth_loss": 1.63818526743853e-06 + "loss": 0.043204840272665024, + "image_loss": 0.04320483282208443, + "smooth_loss": 7.851771215428016e-07 }, { "epoch": 24.0, - "loss": 0.0580613799393177, - "image_loss": 0.05806136131286621, - "smooth_loss": 1.824913169912179e-06 + "loss": 0.0431935153901577, + "image_loss": 0.0431935079395771, + "smooth_loss": 8.643701221444644e-07 }, { "epoch": 25.0, - "loss": 0.05802749842405319, - "image_loss": 0.0580274797976017, - "smooth_loss": 2.0261988993297564e-06 + "loss": 0.043182097375392914, + "image_loss": 0.04318208619952202, + "smooth_loss": 9.485592045166413e-07 }, { "epoch": 26.0, - "loss": 0.05799323692917824, - "image_loss": 0.05799321457743645, - "smooth_loss": 2.2423303107643733e-06 + "loss": 0.043170712888240814, + "image_loss": 0.04317070171236992, + "smooth_loss": 1.0357232440583175e-06 }, { "epoch": 27.0, - "loss": 0.057958219200372696, - "image_loss": 0.05795819312334061, - "smooth_loss": 2.477491534591536e-06 + "loss": 0.043159086257219315, + "image_loss": 0.04315907508134842, + "smooth_loss": 1.1295738886474282e-06 }, { "epoch": 28.0, - "loss": 0.05792278051376343, - "image_loss": 0.05792275443673134, - "smooth_loss": 2.7321416382619645e-06 + "loss": 0.043147530406713486, + "image_loss": 0.04314751923084259, + "smooth_loss": 1.2269233593542594e-06 }, { "epoch": 29.0, - "loss": 0.05788649618625641, - "image_loss": 0.05788646638393402, - "smooth_loss": 3.0071300898271147e-06 + "loss": 0.04313572868704796, + "image_loss": 0.043135713785886765, + "smooth_loss": 1.3320702691999031e-06 }, { "epoch": 30.0, - "loss": 0.057849571108818054, - "image_loss": 0.05784953758120537, - "smooth_loss": 3.304860001662746e-06 + "loss": 0.04312397539615631, + "image_loss": 0.04312396049499512, + "smooth_loss": 1.4407385151571361e-06 }, { "epoch": 31.0, - "loss": 0.05781206861138344, - "image_loss": 0.057812031358480453, - "smooth_loss": 3.6249500681151403e-06 + "loss": 0.04311199113726616, + "image_loss": 0.043111976236104965, + "smooth_loss": 1.5546060012638918e-06 }, { "epoch": 32.0, - "loss": 0.05777355656027794, - "image_loss": 0.057773515582084656, - "smooth_loss": 3.9738624764140695e-06 + "loss": 0.04309987276792526, + "image_loss": 0.04309985786676407, + "smooth_loss": 1.6753501768107526e-06 }, { "epoch": 33.0, - "loss": 0.057734414935112, - "image_loss": 0.05773437023162842, - "smooth_loss": 4.350997187430039e-06 + "loss": 0.04308755695819855, + "image_loss": 0.043087538331747055, + "smooth_loss": 1.8031345234703622e-06 }, { "epoch": 34.0, - "loss": 0.057694148272275925, - "image_loss": 0.057694099843502045, - "smooth_loss": 4.76802779303398e-06 + "loss": 0.04307503253221512, + "image_loss": 0.043075013905763626, + "smooth_loss": 1.9371573216631077e-06 }, { "epoch": 35.0, - "loss": 0.05765282362699509, - "image_loss": 0.05765277147293091, - "smooth_loss": 5.221423634793609e-06 + "loss": 0.0430622436106205, + "image_loss": 0.04306222125887871, + "smooth_loss": 2.0795719137822744e-06 }, { "epoch": 36.0, - "loss": 0.05761110037565231, - "image_loss": 0.057611044496297836, - "smooth_loss": 5.697434062312823e-06 + "loss": 0.04304904118180275, + "image_loss": 0.04304901883006096, + "smooth_loss": 2.2306476239464246e-06 }, { "epoch": 37.0, - "loss": 0.05756793171167374, - "image_loss": 0.05756786838173866, - "smooth_loss": 6.224857315828558e-06 + "loss": 0.04303569346666336, + "image_loss": 0.04303567111492157, + "smooth_loss": 2.3862683065090096e-06 }, { "epoch": 38.0, - "loss": 0.05752401426434517, - "image_loss": 0.0575239472091198, - "smooth_loss": 6.778077477065381e-06 + "loss": 0.043021950870752335, + "image_loss": 0.043021924793720245, + "smooth_loss": 2.5546942197252065e-06 }, { "epoch": 39.0, - "loss": 0.05747854709625244, - "image_loss": 0.05747847259044647, - "smooth_loss": 7.391299732262269e-06 + "loss": 0.04300765320658684, + "image_loss": 0.04300762712955475, + "smooth_loss": 2.7329397198627703e-06 }, { "epoch": 40.0, - "loss": 0.057432375848293304, - "image_loss": 0.05743229389190674, - "smooth_loss": 8.039884960453492e-06 + "loss": 0.04299306869506836, + "image_loss": 0.04299303889274597, + "smooth_loss": 2.9173079383326694e-06 }, { "epoch": 41.0, - "loss": 0.0573841817677021, - "image_loss": 0.05738409236073494, - "smooth_loss": 8.764058293309063e-06 + "loss": 0.04297754168510437, + "image_loss": 0.04297751188278198, + "smooth_loss": 3.1240251701092348e-06 }, { "epoch": 42.0, - "loss": 0.057334933429956436, - "image_loss": 0.057334836572408676, - "smooth_loss": 9.537392543279566e-06 + "loss": 0.042961329221725464, + "image_loss": 0.04296129569411278, + "smooth_loss": 3.3366113711963408e-06 }, { "epoch": 43.0, - "loss": 0.057283785194158554, - "image_loss": 0.0572836808860302, - "smooth_loss": 1.0389785529696383e-05 + "loss": 0.04294413700699806, + "image_loss": 0.04294409975409508, + "smooth_loss": 3.5652053611556767e-06 }, { "epoch": 44.0, - "loss": 0.05723174661397934, - "image_loss": 0.057231634855270386, - "smooth_loss": 1.1303329301881604e-05 + "loss": 0.04292583465576172, + "image_loss": 0.042925797402858734, + "smooth_loss": 3.8064451928221388e-06 }, { "epoch": 45.0, - "loss": 0.057176779955625534, - "image_loss": 0.057176657021045685, - "smooth_loss": 1.2341002729954198e-05 + "loss": 0.04290642961859703, + "image_loss": 0.04290638864040375, + "smooth_loss": 4.066736892127665e-06 }, { "epoch": 46.0, - "loss": 0.057120226323604584, - "image_loss": 0.05712009221315384, - "smooth_loss": 1.3473189028445631e-05 + "loss": 0.04288554936647415, + "image_loss": 0.04288550466299057, + "smooth_loss": 4.346161858848063e-06 }, { "epoch": 47.0, - "loss": 0.057062555104494095, - "image_loss": 0.057062409818172455, - "smooth_loss": 1.4687300790683366e-05 + "loss": 0.04286346957087517, + "image_loss": 0.042863424867391586, + "smooth_loss": 4.647703008231474e-06 }, { "epoch": 48.0, - "loss": 0.05700160190463066, - "image_loss": 0.05700144171714783, - "smooth_loss": 1.6045180018409155e-05 + "loss": 0.042839810252189636, + "image_loss": 0.042839761823415756, + "smooth_loss": 4.978734978067223e-06 }, { "epoch": 49.0, - "loss": 0.056939497590065, - "image_loss": 0.056939322501420975, - "smooth_loss": 1.7514281353214756e-05 + "loss": 0.04281448572874069, + "image_loss": 0.042814433574676514, + "smooth_loss": 5.339582457963843e-06 }, { "epoch": 50.0, - "loss": 0.05687557905912399, - "image_loss": 0.05687538906931877, - "smooth_loss": 1.911029903567396e-05 + "loss": 0.04278765991330147, + "image_loss": 0.04278760403394699, + "smooth_loss": 5.742196663049981e-06 }, { "epoch": 51.0, - "loss": 0.056809235364198685, - "image_loss": 0.05680902674794197, - "smooth_loss": 2.0875198970315978e-05 + "loss": 0.04275878146290779, + "image_loss": 0.04275871813297272, + "smooth_loss": 6.183304776641307e-06 }, { "epoch": 52.0, - "loss": 0.05674029514193535, - "image_loss": 0.05674006789922714, - "smooth_loss": 2.2824997358839028e-05 + "loss": 0.04272802174091339, + "image_loss": 0.04272795468568802, + "smooth_loss": 6.67552012600936e-06 }, { "epoch": 53.0, - "loss": 0.056669920682907104, - "image_loss": 0.05666967108845711, - "smooth_loss": 2.4954817490652204e-05 + "loss": 0.042694948613643646, + "image_loss": 0.042694877833127975, + "smooth_loss": 7.233465112221893e-06 }, { "epoch": 54.0, - "loss": 0.0565965436398983, - "image_loss": 0.05659627169370651, - "smooth_loss": 2.7323461836203933e-05 + "loss": 0.04265972599387169, + "image_loss": 0.04265964776277542, + "smooth_loss": 7.846486369089689e-06 }, { "epoch": 55.0, - "loss": 0.056521233171224594, - "image_loss": 0.05652093514800072, - "smooth_loss": 2.9880573492846452e-05 + "loss": 0.04262172058224678, + "image_loss": 0.042621634900569916, + "smooth_loss": 8.557326509617269e-06 }, { "epoch": 56.0, - "loss": 0.05644303560256958, - "image_loss": 0.056442707777023315, - "smooth_loss": 3.275117705925368e-05 + "loss": 0.04258149117231369, + "image_loss": 0.04258139804005623, + "smooth_loss": 9.357276212540455e-06 }, { "epoch": 57.0, - "loss": 0.05636239051818848, - "image_loss": 0.056362032890319824, - "smooth_loss": 3.594645386328921e-05 + "loss": 0.042538225650787354, + "image_loss": 0.042538121342659, + "smooth_loss": 1.0288686098647304e-05 }, { "epoch": 58.0, - "loss": 0.056279003620147705, - "image_loss": 0.05627860873937607, - "smooth_loss": 3.945654316339642e-05 + "loss": 0.042492616921663284, + "image_loss": 0.04249250143766403, + "smooth_loss": 1.136589708039537e-05 }, { "epoch": 59.0, - "loss": 0.05619405582547188, - "image_loss": 0.056193623691797256, - "smooth_loss": 4.33080131188035e-05 + "loss": 0.042444635182619095, + "image_loss": 0.04244450852274895, + "smooth_loss": 1.261604847968556e-05 }, { "epoch": 60.0, - "loss": 0.056105729192495346, - "image_loss": 0.05610525235533714, - "smooth_loss": 4.769213410327211e-05 + "loss": 0.042394258081912994, + "image_loss": 0.04239411652088165, + "smooth_loss": 1.4080704204388894e-05 }, { "epoch": 61.0, - "loss": 0.0560150220990181, - "image_loss": 0.056014496833086014, - "smooth_loss": 5.2607003453886136e-05 + "loss": 0.04234214872121811, + "image_loss": 0.042341988533735275, + "smooth_loss": 1.583471203048248e-05 }, { "epoch": 62.0, - "loss": 0.05592316761612892, - "image_loss": 0.05592258647084236, - "smooth_loss": 5.80552987230476e-05 + "loss": 0.042289428412914276, + "image_loss": 0.04228924959897995, + "smooth_loss": 1.7939397366717458e-05 }, { "epoch": 63.0, - "loss": 0.05583030730485916, - "image_loss": 0.055829666554927826, - "smooth_loss": 6.423084414564073e-05 + "loss": 0.04223756119608879, + "image_loss": 0.042237356305122375, + "smooth_loss": 2.041880725300871e-05 }, { "epoch": 64.0, - "loss": 0.055736664682626724, - "image_loss": 0.05573595315217972, - "smooth_loss": 7.123746036086231e-05 + "loss": 0.04218873754143715, + "image_loss": 0.042188502848148346, + "smooth_loss": 2.3403612431138754e-05 }, { "epoch": 65.0, - "loss": 0.0556456595659256, - "image_loss": 0.055644869804382324, - "smooth_loss": 7.909350097179413e-05 + "loss": 0.04214523360133171, + "image_loss": 0.04214496165513992, + "smooth_loss": 2.7044574380852282e-05 }, { "epoch": 66.0, - "loss": 0.05555867403745651, - "image_loss": 0.055557794868946075, - "smooth_loss": 8.807046106085181e-05 + "loss": 0.042111899703741074, + "image_loss": 0.042111586779356, + "smooth_loss": 3.138223473797552e-05 }, { "epoch": 67.0, - "loss": 0.055479418486356735, - "image_loss": 0.05547843500971794, - "smooth_loss": 9.828618203755468e-05 + "loss": 0.04209335148334503, + "image_loss": 0.04209298640489578, + "smooth_loss": 3.650638245744631e-05 }, { "epoch": 68.0, - "loss": 0.05541320890188217, - "image_loss": 0.055412109941244125, - "smooth_loss": 0.00010981244122376665 + "loss": 0.042094651609659195, + "image_loss": 0.04209422692656517, + "smooth_loss": 4.229835394653492e-05 }, { "epoch": 69.0, - "loss": 0.05536733567714691, - "image_loss": 0.05536610633134842, - "smooth_loss": 0.00012293846521060914 + "loss": 0.0421144962310791, + "image_loss": 0.0421140156686306, + "smooth_loss": 4.809771417058073e-05 }, { "epoch": 70.0, - "loss": 0.05535025894641876, - "image_loss": 0.05534888431429863, - "smooth_loss": 0.00013746831973548979 + "loss": 0.04213815554976463, + "image_loss": 0.04213763028383255, + "smooth_loss": 5.256465738057159e-05 }, { "epoch": 71.0, - "loss": 0.0553678423166275, - "image_loss": 0.05536631867289543, - "smooth_loss": 0.00015228186384774745 + "loss": 0.04214761406183243, + "image_loss": 0.042147066444158554, + "smooth_loss": 5.478101593325846e-05 }, { "epoch": 72.0, - "loss": 0.05540918558835983, - "image_loss": 0.055407531559467316, - "smooth_loss": 0.00016536397743038833 + "loss": 0.04213856905698776, + "image_loss": 0.04213802143931389, + "smooth_loss": 5.4818134231027216e-05 }, { "epoch": 73.0, - "loss": 0.05544228106737137, - "image_loss": 0.05544053763151169, - "smooth_loss": 0.00017419658252038062 + "loss": 0.04211747646331787, + "image_loss": 0.04211694374680519, + "smooth_loss": 5.3270672651706263e-05 }, { "epoch": 74.0, - "loss": 0.05544540658593178, - "image_loss": 0.05544362962245941, - "smooth_loss": 0.00017784371448215097 + "loss": 0.042093425989151, + "image_loss": 0.04209291934967041, + "smooth_loss": 5.081955896457657e-05 }, { "epoch": 75.0, - "loss": 0.05542201176285744, - "image_loss": 0.05542023852467537, - "smooth_loss": 0.00017743516946211457 + "loss": 0.04207289591431618, + "image_loss": 0.042072415351867676, + "smooth_loss": 4.802431067219004e-05 }, { "epoch": 76.0, - "loss": 0.0553855374455452, - "image_loss": 0.055383794009685516, - "smooth_loss": 0.0001741933374432847 + "loss": 0.04205827787518501, + "image_loss": 0.0420578271150589, + "smooth_loss": 4.526200791588053e-05 }, { "epoch": 77.0, - "loss": 0.05534776300191879, - "image_loss": 0.05534606799483299, - "smooth_loss": 0.00016937614418566227 + "loss": 0.042049601674079895, + "image_loss": 0.04204917326569557, + "smooth_loss": 4.2754021706059575e-05 }, { "epoch": 78.0, - "loss": 0.05531612038612366, - "image_loss": 0.055314477533102036, - "smooth_loss": 0.00016410183161497116 + "loss": 0.042045459151268005, + "image_loss": 0.04204505309462547, + "smooth_loss": 4.061676736455411e-05 }, { "epoch": 79.0, - "loss": 0.05529369041323662, - "image_loss": 0.055292099714279175, - "smooth_loss": 0.00015910626098047942 + "loss": 0.04204433411359787, + "image_loss": 0.04204394668340683, + "smooth_loss": 3.89082488254644e-05 }, { "epoch": 80.0, - "loss": 0.055279314517974854, - "image_loss": 0.05527777224779129, - "smooth_loss": 0.00015433150110766292 + "loss": 0.04204464331269264, + "image_loss": 0.0420442670583725, + "smooth_loss": 3.758826642297208e-05 } ] \ No newline at end of file diff --git a/outputs/checkpoints/vxm_head_ct_patient1.pt b/outputs/checkpoints/vxm_head_ct_patient1.pt index 39f402a..e4e97d0 100644 Binary files a/outputs/checkpoints/vxm_head_ct_patient1.pt and b/outputs/checkpoints/vxm_head_ct_patient1.pt differ diff --git a/outputs/preprocessed/patient1_neck_alignment.json b/outputs/preprocessed/patient1_neck_alignment.json new file mode 100644 index 0000000..6554c2d --- /dev/null +++ b/outputs/preprocessed/patient1_neck_alignment.json @@ -0,0 +1,57 @@ +{ + "fixed_input_path": "outputs/nifti/patient1_fixed.nii.gz", + "moving_input_path": "outputs/nifti/patient1_moving.nii.gz", + "fixed_output_path": "outputs/preprocessed/patient1_fixed_preprocessed.nii.gz", + "moving_output_path": "outputs/preprocessed/patient1_moving_preprocessed.nii.gz", + "fixed_resampled_shape_xyz": [ + 288, + 288, + 317 + ], + "moving_resampled_shape_xyz": [ + 331, + 331, + 334 + ], + "target_shape_xyz": [ + 320, + 320, + 352 + ], + "target_spacing_xyz": [ + 1.0, + 1.0, + 1.0 + ], + "fixed_neck_center_index_xyz": [ + 143.0, + 160.5, + 159.5 + ], + "moving_neck_center_index_xyz": [ + 164.5, + 195.0, + 170.0 + ], + "fixed_neck_center_world_mm": [ + 13.233566284179688, + -33.013214111328125, + -633.6527099609375 + ], + "moving_neck_center_world_mm": [ + 13.2132568359375, + 18.1754150390625, + -656.634033203125 + ], + "moving_translation_mm": [ + 0.0203094482421875, + -51.188629150390625, + 22.9813232421875 + ], + "body_threshold_hu": -500.0, + "foreground_percentiles": [ + 2.0, + 98.0 + ], + "method": "neck-mask-centroid-translation" +} \ No newline at end of file diff --git a/prealign.py b/prealign.py new file mode 100644 index 0000000..6605c19 --- /dev/null +++ b/prealign.py @@ -0,0 +1,297 @@ +"""颈部基准的 CT 预配准。 + +这一阶段发生在 VoxelMorph 训练/推理之前： +1. Fixed 与 Moving 原始 NIfTI 先重采样到同一 spacing。 +2. 基于 CT foreground mask 估计颈部/身体区域的稳健中心。 +3. 将 Moving 平移到 Fixed 的颈部中心。 +4. 在同一个物理网格中裁剪/填充并归一化，得到后续 VoxelMorph 输入。 + +当前实现是确定性的平移预配准，目标是先解决层数不同、扫描范围不同导致的 +初始空间不一致；细致非线性变形仍交给官方 VoxelMorph。 +""" + +from __future__ import annotations + +import argparse +import json +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Iterable, Sequence, Tuple + +import numpy as np +from scipy import ndimage + +from config import ( + DEFAULT_FIXED_NIFTI, + DEFAULT_MOVING_NIFTI, + DEFAULT_TARGET_SHAPE, + DEFAULT_TARGET_SPACING, + DEFAULT_WINDOW_LEVEL, + DEFAULT_WINDOW_WIDTH, + PREPROCESSED_DIR, +) +from preprocess import ( + affine_with_spacing, + ensure_multiple_of_16, + load_nifti, + resample_to_spacing, + save_nifti, + window_and_normalize, +) + + +@dataclass +class NeckAlignmentMeta: + fixed_input_path: str + moving_input_path: str + fixed_output_path: str + moving_output_path: str + fixed_resampled_shape_xyz: Tuple[int, int, int] + moving_resampled_shape_xyz: Tuple[int, int, int] + target_shape_xyz: Tuple[int, int, int] + target_spacing_xyz: Tuple[float, float, float] + fixed_neck_center_index_xyz: Tuple[float, float, float] + moving_neck_center_index_xyz: Tuple[float, float, float] + fixed_neck_center_world_mm: Tuple[float, float, float] + moving_neck_center_world_mm: Tuple[float, float, float] + moving_translation_mm: Tuple[float, float, float] + body_threshold_hu: float + foreground_percentiles: Tuple[float, float] + method: str = "neck-mask-centroid-translation" + + +def _tuple_float(values: Sequence[float]) -> Tuple[float, float, float]: + return tuple(float(v) for v in values) # type: ignore[return-value] + + +def _tuple_int(values: Sequence[int]) -> Tuple[int, int, int]: + return tuple(int(v) for v in values) # type: ignore[return-value] + + +def index_to_world(affine: np.ndarray, index_xyz: Sequence[float]) -> np.ndarray: + index = np.asarray(index_xyz, dtype=np.float64) + return affine[:3, :3] @ index + affine[:3, 3] + + +def largest_foreground_component(mask: np.ndarray) -> np.ndarray: + structure = ndimage.generate_binary_structure(rank=3, connectivity=1) + labeled, count = ndimage.label(mask, structure=structure) + if count < 1: + raise ValueError("未能从 CT 中提取颈部 foreground mask。") + sizes = np.bincount(labeled.ravel()) + sizes[0] = 0 + return labeled == int(np.argmax(sizes)) + + +def estimate_neck_center_index( + data_hu: np.ndarray, + body_threshold_hu: float = -500.0, + foreground_percentiles: Sequence[float] = (2.0, 98.0), +) -> Tuple[np.ndarray, Tuple[np.ndarray, np.ndarray]]: + """估计颈部 foreground 的稳健中心。 + + 先用 HU 阈值去掉空气，再保留最大连通域，最后用坐标百分位框的中心而不是 + 简单均值，减少局部噪声、少量外部物体和扫描边界的影响。 + """ + + data = np.asarray(data_hu, dtype=np.float32) + mask = np.isfinite(data) & (data > float(body_threshold_hu)) & (data < 3000.0) + mask = largest_foreground_component(mask) + coords = np.where(mask) + if not coords or coords[0].size == 0: + raise ValueError("颈部 foreground mask 为空。") + + p_low, p_high = (float(v) for v in foreground_percentiles) + lows = np.asarray([np.percentile(axis_coords, p_low) for axis_coords in coords], dtype=np.float64) + highs = np.asarray([np.percentile(axis_coords, p_high) for axis_coords in coords], dtype=np.float64) + center = (lows + highs) / 2.0 + return center, (lows, highs) + + +def target_affine_from_center( + reference_affine: np.ndarray, + center_world_mm: Sequence[float], + target_shape: Sequence[int], + target_spacing: Sequence[float], +) -> np.ndarray: + output_affine = affine_with_spacing(reference_affine, target_spacing) + center_index = (np.asarray(target_shape, dtype=np.float64) - 1.0) / 2.0 + center_world = np.asarray(center_world_mm, dtype=np.float64) + output_affine[:3, 3] = center_world - output_affine[:3, :3] @ center_index + return output_affine + + +def resample_to_reference_grid( + data: np.ndarray, + input_affine: np.ndarray, + output_affine: np.ndarray, + output_shape: Sequence[int], + order: int = 1, + cval: float = -1000.0, +) -> np.ndarray: + input_to_output = np.linalg.inv(input_affine) @ output_affine + matrix = input_to_output[:3, :3] + offset = input_to_output[:3, 3] + output = ndimage.affine_transform( + data.astype(np.float32, copy=False), + matrix=matrix, + offset=offset, + output_shape=tuple(int(v) for v in output_shape), + order=order, + mode="constant", + cval=float(cval), + prefilter=(order > 1), + output=np.float32, + ) + return output.astype(np.float32, copy=False) + + +def prealign_pair( + fixed_input_path: str | Path, + moving_input_path: str | Path, + fixed_output_path: str | Path = DEFAULT_FIXED_NIFTI, + moving_output_path: str | Path = DEFAULT_MOVING_NIFTI, + target_spacing: Sequence[float] = DEFAULT_TARGET_SPACING, + target_shape: Sequence[int] = DEFAULT_TARGET_SHAPE, + window_width: float = DEFAULT_WINDOW_WIDTH, + window_level: float = DEFAULT_WINDOW_LEVEL, + body_threshold_hu: float = -500.0, + foreground_percentiles: Sequence[float] = (2.0, 98.0), + max_memory_mb: int = 8192, + metadata_json: str | Path | None = None, +) -> NeckAlignmentMeta: + target_shape = ensure_multiple_of_16(target_shape) + fixed_input_path = Path(fixed_input_path) + moving_input_path = Path(moving_input_path) + fixed_output_path = Path(fixed_output_path) + moving_output_path = Path(moving_output_path) + + fixed_hu, fixed_affine, fixed_spacing = load_nifti(fixed_input_path) + moving_hu, moving_affine, moving_spacing = load_nifti(moving_input_path) + + fixed_resampled, fixed_resampled_affine = resample_to_spacing( + fixed_hu, + fixed_affine, + current_spacing=fixed_spacing, + target_spacing=target_spacing, + order=1, + max_memory_mb=max_memory_mb, + ) + moving_resampled, moving_resampled_affine = resample_to_spacing( + moving_hu, + moving_affine, + current_spacing=moving_spacing, + target_spacing=target_spacing, + order=1, + max_memory_mb=max_memory_mb, + ) + + fixed_center_idx, _ = estimate_neck_center_index( + fixed_resampled, + body_threshold_hu=body_threshold_hu, + foreground_percentiles=foreground_percentiles, + ) + moving_center_idx, _ = estimate_neck_center_index( + moving_resampled, + body_threshold_hu=body_threshold_hu, + foreground_percentiles=foreground_percentiles, + ) + fixed_center_world = index_to_world(fixed_resampled_affine, fixed_center_idx) + moving_center_world = index_to_world(moving_resampled_affine, moving_center_idx) + moving_translation = fixed_center_world - moving_center_world + + moving_aligned_affine = moving_resampled_affine.copy() + moving_aligned_affine[:3, 3] = moving_aligned_affine[:3, 3] + moving_translation + output_affine = target_affine_from_center( + fixed_resampled_affine, + center_world_mm=fixed_center_world, + target_shape=target_shape, + target_spacing=target_spacing, + ) + + fixed_grid_hu = resample_to_reference_grid( + fixed_resampled, + input_affine=fixed_resampled_affine, + output_affine=output_affine, + output_shape=target_shape, + order=1, + cval=-1000.0, + ) + moving_grid_hu = resample_to_reference_grid( + moving_resampled, + input_affine=moving_aligned_affine, + output_affine=output_affine, + output_shape=target_shape, + order=1, + cval=-1000.0, + ) + + fixed_norm = window_and_normalize(fixed_grid_hu, window_width=window_width, window_level=window_level) + moving_norm = window_and_normalize(moving_grid_hu, window_width=window_width, window_level=window_level) + save_nifti(fixed_norm, output_affine, fixed_output_path) + save_nifti(moving_norm, output_affine, moving_output_path) + + meta = NeckAlignmentMeta( + fixed_input_path=str(fixed_input_path), + moving_input_path=str(moving_input_path), + fixed_output_path=str(fixed_output_path), + moving_output_path=str(moving_output_path), + fixed_resampled_shape_xyz=_tuple_int(fixed_resampled.shape), + moving_resampled_shape_xyz=_tuple_int(moving_resampled.shape), + target_shape_xyz=_tuple_int(target_shape), + target_spacing_xyz=_tuple_float(target_spacing), + fixed_neck_center_index_xyz=_tuple_float(fixed_center_idx), + moving_neck_center_index_xyz=_tuple_float(moving_center_idx), + fixed_neck_center_world_mm=_tuple_float(fixed_center_world), + moving_neck_center_world_mm=_tuple_float(moving_center_world), + moving_translation_mm=_tuple_float(moving_translation), + body_threshold_hu=float(body_threshold_hu), + foreground_percentiles=(float(foreground_percentiles[0]), float(foreground_percentiles[1])), + ) + + if metadata_json is None: + metadata_json = PREPROCESSED_DIR / "patient1_neck_alignment.json" + metadata_path = Path(metadata_json) + metadata_path.parent.mkdir(parents=True, exist_ok=True) + metadata_path.write_text(json.dumps(asdict(meta), ensure_ascii=False, indent=2), encoding="utf-8") + return meta + + +def build_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description="以颈部 foreground 为基准的 CT 平移预配准。") + parser.add_argument("--fixed", required=True, help="Fixed 原始 NIfTI。") + parser.add_argument("--moving", required=True, help="Moving 原始 NIfTI。") + parser.add_argument("--fixed-output", default=str(DEFAULT_FIXED_NIFTI)) + parser.add_argument("--moving-output", default=str(DEFAULT_MOVING_NIFTI)) + parser.add_argument("--target-spacing", type=float, nargs=3, default=DEFAULT_TARGET_SPACING, metavar=("X", "Y", "Z")) + parser.add_argument("--target-shape", type=int, nargs=3, default=DEFAULT_TARGET_SHAPE, metavar=("X", "Y", "Z")) + parser.add_argument("--window-width", type=float, default=DEFAULT_WINDOW_WIDTH) + parser.add_argument("--window-level", type=float, default=DEFAULT_WINDOW_LEVEL) + parser.add_argument("--body-threshold-hu", type=float, default=-500.0) + parser.add_argument("--foreground-percentiles", type=float, nargs=2, default=(2.0, 98.0)) + parser.add_argument("--metadata-json", default=str(PREPROCESSED_DIR / "patient1_neck_alignment.json")) + parser.add_argument("--max-memory-mb", type=int, default=8192) + return parser + + +def main(argv: Iterable[str] | None = None) -> None: + args = build_arg_parser().parse_args(argv) + meta = prealign_pair( + fixed_input_path=args.fixed, + moving_input_path=args.moving, + fixed_output_path=args.fixed_output, + moving_output_path=args.moving_output, + target_spacing=args.target_spacing, + target_shape=args.target_shape, + window_width=args.window_width, + window_level=args.window_level, + body_threshold_hu=args.body_threshold_hu, + foreground_percentiles=args.foreground_percentiles, + metadata_json=args.metadata_json, + max_memory_mb=args.max_memory_mb, + ) + print(json.dumps(asdict(meta), ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/preprocess.py b/preprocess.py index 18db4c9..0c76237 100644 --- a/preprocess.py +++ b/preprocess.py @@ -53,9 +53,9 @@ def ensure_multiple_of_16(shape: Sequence[int]) -> Tuple[int, int, int]: shape = tuple(int(v) for v in shape) if len(shape) != 3: raise ValueError("target_shape 必须包含 3 个维度。") - bad = [v for v in shape if v <= 0 or v % 16 != 0] + bad = [v for v in shape if v <= 0 or v % 32 != 0] if bad: - raise ValueError(f"VoxelMorph 建议三维尺寸均为 16 的倍数，当前非法维度: {bad}") + raise ValueError(f"VoxelMorph 当前 U-Net 建议三维尺寸均为 32 的倍数，当前非法维度: {bad}") return shape # type: ignore[return-value] @@ -324,4 +324,3 @@ def main(argv: Iterable[str] | None = None) -> None: if __name__ == "__main__": main() -