admin/Voxelmorph_Head_CT
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add neck-based CT prealignment before VoxelMorph

Browse Source
This commit is contained in:
admin
2026-06-03 10:58:12 +08:00
parent 0a6a0ece00
commit 3d96bdb706
8 changed files with 665 additions and 266 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
README.md
View File

@@ -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` 特征配置。

58
app.py
View File

@@ -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)

4
config.py
View File

@@ -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。

480
outputs/checkpoints/vxm_head_ct_patient1.history.json
View File

@@ -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
}
]

BIN
outputs/checkpoints/vxm_head_ct_patient1.pt
View File

Binary file not shown.

57
outputs/preprocessed/patient1_neck_alignment.json Normal file
View File

@@ -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"
}

297
prealign.py Normal file
View File

@@ -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()

5
preprocess.py
View File

@@ -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()