Switch viewer and training to patient1 fixed flat CT

2026-06-03 10:48:14 +08:00
parent 972fb2435c
commit 0a6a0ece00
5 changed files with 566 additions and 46 deletions
--- a/README.md
+++ b/README.md
@@ -1,6 +1,6 @@
 # VoxelMorph Head CT Deformable Registration

-面向“患者平扫 CT（中立位）到仰头 CT（极度后仰位）”的 3D 形变配准工程。项目包含 DICOM 转 NIfTI、医学图像预处理、官方 VoxelMorph 训练适配器、独立推理，以及 Streamlit 交互式结果查看界面。
+面向“患者仰头 CT（Moving）到平扫 CT（Fixed）”的 3D 形变配准工程。项目包含 DICOM 转 NIfTI、医学图像预处理、官方 VoxelMorph 训练适配器、独立推理，以及 Streamlit 交互式结果查看界面。

 核心模型使用官方仓库 `voxelmorph/voxelmorph` 的 PyTorch 实现：

@@ -15,8 +15,8 @@
 ```text
 Voxelmorph_Head_CT/
 ├── Data/
-│   ├── 患者1-平扫CT/              # Moving 原始 DICOM
-│   ├── 患者1-仰头CT/              # Fixed 原始 DICOM
+│   ├── 患者1-平扫CT/              # Fixed 原始 DICOM
+│   ├── 患者1-仰头CT/              # Moving 原始 DICOM
 │   └── 患者2-平扫CT/
 ├── app.py                         # Streamlit Web 可视化界面
 ├── config.py                      # 默认路径与预处理参数
@@ -60,11 +60,11 @@ bash scripts/setup_env.sh
 ```bash
 python data_loader.py \
  --dicom-dir "Data/患者1-平扫CT" \
-  --output "outputs/nifti/patient1_moving.nii.gz"
+  --output "outputs/nifti/patient1_fixed.nii.gz"

 python data_loader.py \
  --dicom-dir "Data/患者1-仰头CT" \
-  --output "outputs/nifti/patient1_fixed.nii.gz"
+  --output "outputs/nifti/patient1_moving.nii.gz"
 ```

 `data_loader.py` 会优先按 `InstanceNumber` 排序，其次按 `SliceLocation` 排序，并保存 spacing、层厚、排序依据等元数据 JSON。
@@ -72,17 +72,17 @@ python data_loader.py \
 ## 2. 预处理

 ```bash
-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 160 192 224
-
 python preprocess.py \
  --input "outputs/nifti/patient1_fixed.nii.gz" \
  --output "outputs/preprocessed/patient1_fixed_preprocessed.nii.gz" \
  --target-spacing 1 1 1 \
-  --target-shape 160 192 224
+  --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
 ```

 默认窗口为 `W=400, L=40`，适合观察颈部软组织和气道。
@@ -93,10 +93,10 @@ python preprocess.py \
 python model_and_train.py \
  --moving "outputs/preprocessed/patient1_moving_preprocessed.nii.gz" \
  --fixed "outputs/preprocessed/patient1_fixed_preprocessed.nii.gz" \
-  --checkpoint "outputs/checkpoints/vxm_head_ct.pt" \
-  --epochs 200 \
-  --image-loss ncc \
-  --ncc-impl local \
+  --checkpoint "outputs/checkpoints/vxm_head_ct_patient1.pt" \
+  --epochs 80 \
+  --image-loss mse \
+  --nb-features 8 8 8 8 8 \
  --smooth-weight 0.01
 ```

@@ -109,7 +109,7 @@ python model_and_train.py \
 python infer.py \
  --moving "outputs/preprocessed/patient1_moving_preprocessed.nii.gz" \
  --fixed "outputs/preprocessed/patient1_fixed_preprocessed.nii.gz" \
-  --checkpoint "outputs/checkpoints/vxm_head_ct.pt" \
+  --checkpoint "outputs/checkpoints/vxm_head_ct_patient1.pt" \
  --out-dir "outputs/inference"
 ```

@@ -131,9 +131,8 @@ streamlit run app.py

 网页提供：

- Moving/Fixed/模型权重/输出目录输入。
- 自动发现 `outputs/` 与项目目录下的 NIfTI 和 checkpoint，也支持手动输入路径。
- “开始推理”按钮。
+- 患者1专用固定路径：Fixed 为 `患者1-平扫CT`，Moving 为 `患者1-仰头CT`。
+- “重新训练模型”和“开始推理”按钮。
 - Axial、Coronal、Sagittal 正交三视图；每个平面按行同时展示 Fixed、Moving、Warped。
 - Fixed 与 Warped 的 Alpha 融合或棋盘格对比。
 - DDF 位移强度热力图。
@@ -143,4 +142,4 @@ streamlit run app.py

 - DICOM 转换和重采样都有 `--max-memory-mb` 防护。
 - Web 界面对超大 NIfTI 会通过 nibabel proxy 按 stride 切片读取并下采样，只影响浏览器展示，不改变磁盘结果；侧栏可调整显示体素上限。
- 训练阶段的主要瓶颈是 3D U-Net 显存；`160x192x224` 是较重的 3D 输入，建议优先使用 CUDA GPU。
+- 训练阶段的主要瓶颈是 3D U-Net 显存；`256x256x352` 是较重的 3D 输入，建议优先使用 CUDA GPU。当前患者1默认使用较轻的 `8 8 8 8 8` 特征配置。
--- a/app.py
+++ b/app.py
@@ -31,7 +31,7 @@ from metrics import crop_to_common_shape, ddf_summary, registration_metrics, sli
 st.set_page_config(
    page_title="VoxelMorph 颈部 CT 配准工作台",
    layout="wide",
-    initial_sidebar_state="expanded",
+    initial_sidebar_state="collapsed",
 )


@@ -608,31 +608,70 @@ def run_inference_from_ui(moving_path: str, fixed_path: str, checkpoint_path: st
    )


+def run_training_from_ui(moving_path: str, fixed_path: str, checkpoint_path: str) -> None:
+    from model_and_train import train_pair
+
+    train_pair(
+        moving_path=moving_path,
+        fixed_path=fixed_path,
+        checkpoint_path=checkpoint_path,
+        epochs=80,
+        learning_rate=1e-4,
+        smooth_weight=0.01,
+        image_loss="mse",
+        nb_features=(8, 8, 8, 8, 8),
+        save_every=20,
+    )
+
+
 def main() -> None:
    st.markdown(CUSTOM_CSS, unsafe_allow_html=True)
    st.title("VoxelMorph 颈部 CT 配准工作台")

-    candidates = discover_nifti_files()
-    checkpoint_candidates = discover_checkpoint_files()
-    with st.sidebar:
-        st.header("输入")
-        moving_path = choose_path("Moving", DEFAULT_MOVING_NIFTI, candidates)
-        fixed_path = choose_path("Fixed", DEFAULT_FIXED_NIFTI, candidates)
-        checkpoint_path = choose_path("模型权重", DEFAULT_CHECKPOINT, checkpoint_candidates)
-        out_dir = st.text_input("输出目录", value=str(INFERENCE_DIR))
-        display_max_voxels = st.slider("Web显示体素上限", 2_000_000, 30_000_000, 14_000_000, 1_000_000)
+    moving_path = str(DEFAULT_MOVING_NIFTI)
+    fixed_path = str(DEFAULT_FIXED_NIFTI)
+    checkpoint_path = str(DEFAULT_CHECKPOINT)
+    out_dir = str(INFERENCE_DIR)
+    display_max_voxels = 30_000_000

-        start = st.button("开始推理", type="primary", width="stretch")
-        if start:
-            with st.spinner("推理运行中"):
-                try:
-                    load_nifti_cached.clear()
-                    result = run_inference_from_ui(moving_path, fixed_path, checkpoint_path, out_dir)
-                    load_nifti_cached.clear()
-                    st.session_state["last_result"] = result
-                    st.success("推理完成")
-                except Exception as exc:
-                    st.error(str(exc))
+    st.caption(
+        "患者1专用：固定图像 = 平扫CT；移动图像 = 仰头CT。"
+        "推理前会重采样、归一化并裁剪/填充到同一模型网格。"
+    )
+    info_cols = st.columns(4)
+    info_cols[0].metric("Fixed", "患者1-平扫CT")
+    info_cols[1].metric("Moving", "患者1-仰头CT")
+    info_cols[2].metric("模型", Path(checkpoint_path).name)
+    info_cols[3].metric("输出", Path(out_dir).name)
+
+    action_cols = st.columns([1, 1, 4])
+    with action_cols[0]:
+        train_now = st.button("重新训练模型", type="primary", width="stretch")
+    with action_cols[1]:
+        start = st.button("开始推理", width="stretch")
+
+    if train_now:
+        with st.spinner("患者1模型训练中"):
+            try:
+                load_nifti_cached.clear()
+                run_training_from_ui(moving_path, fixed_path, checkpoint_path)
+                result = run_inference_from_ui(moving_path, fixed_path, checkpoint_path, out_dir)
+                load_nifti_cached.clear()
+                st.session_state["last_result"] = result
+                st.success("训练和推理完成")
+            except Exception as exc:
+                st.error(str(exc))
+
+    if start:
+        with st.spinner("推理运行中"):
+            try:
+                load_nifti_cached.clear()
+                result = run_inference_from_ui(moving_path, fixed_path, checkpoint_path, out_dir)
+                load_nifti_cached.clear()
+                st.session_state["last_result"] = result
+                st.success("推理完成")
+            except Exception as exc:
+                st.error(str(exc))

    result_paths = load_result_paths(out_dir)
    last_result = st.session_state.get("last_result", {})
--- a/config.py
+++ b/config.py
@@ -10,8 +10,8 @@ from pathlib import Path
 PROJECT_ROOT = Path(__file__).resolve().parent

 DATA_ROOT = PROJECT_ROOT / "Data"
-DEFAULT_MOVING_DICOM_DIR = DATA_ROOT / "患者1-平扫CT"
-DEFAULT_FIXED_DICOM_DIR = DATA_ROOT / "患者1-仰头CT"
+DEFAULT_FIXED_DICOM_DIR = DATA_ROOT / "患者1-平扫CT"
+DEFAULT_MOVING_DICOM_DIR = DATA_ROOT / "患者1-仰头CT"

 OUTPUT_ROOT = PROJECT_ROOT / "outputs"
 NIFTI_DIR = OUTPUT_ROOT / "nifti"
@@ -21,10 +21,10 @@ INFERENCE_DIR = OUTPUT_ROOT / "inference"

 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.pt"
+DEFAULT_CHECKPOINT = CHECKPOINT_DIR / "vxm_head_ct_patient1.pt"

 # VoxelMorph 的 3D U-Net 多次下采样，三维尺寸建议均为 16 的倍数。
-DEFAULT_TARGET_SHAPE = (160, 192, 224)  # NIfTI 轴顺序: X, Y, Z
+DEFAULT_TARGET_SHAPE = (256, 256, 352)  # NIfTI 轴顺序: X, Y, Z
 DEFAULT_TARGET_SPACING = (1.0, 1.0, 1.0)  # mm, X/Y/Z

 # 颈部软组织/气道观察常用窗口：W=400, L=40。
--- a/outputs/checkpoints/vxm_head_ct_patient1.history.json
+++ b/outputs/checkpoints/vxm_head_ct_patient1.history.json
@@ -0,0 +1,482 @@
+[
+  {
+    "epoch": 1.0,
+    "loss": 0.058744531124830246,
+    "image_loss": 0.058744531124830246,
+    "smooth_loss": 1.5868747288427798e-12
+  },
+  {
+    "epoch": 2.0,
+    "loss": 0.058717112988233566,
+    "image_loss": 0.058717112988233566,
+    "smooth_loss": 2.6513908846226286e-09
+  },
+  {
+    "epoch": 3.0,
+    "loss": 0.058689698576927185,
+    "image_loss": 0.058689698576927185,
+    "smooth_loss": 1.0584165544003099e-08
+  },
+  {
+    "epoch": 4.0,
+    "loss": 0.058662042021751404,
+    "image_loss": 0.058662042021751404,
+    "smooth_loss": 2.401168686105848e-08
+  },
+  {
+    "epoch": 5.0,
+    "loss": 0.05863436311483383,
+    "image_loss": 0.05863436311483383,
+    "smooth_loss": 4.28750439596115e-08
+  },
+  {
+    "epoch": 6.0,
+    "loss": 0.05860644578933716,
+    "image_loss": 0.05860644578933716,
+    "smooth_loss": 6.765733928659756e-08
+  },
+  {
+    "epoch": 7.0,
+    "loss": 0.058578286319971085,
+    "image_loss": 0.058578286319971085,
+    "smooth_loss": 9.85552759402708e-08
+  },
+  {
+    "epoch": 8.0,
+    "loss": 0.05855008214712143,
+    "image_loss": 0.05855008214712143,
+    "smooth_loss": 1.3519149888452375e-07
+  },
+  {
+    "epoch": 9.0,
+    "loss": 0.05852168798446655,
+    "image_loss": 0.05852168798446655,
+    "smooth_loss": 1.7822365805386653e-07
+  },
+  {
+    "epoch": 10.0,
+    "loss": 0.05849294736981392,
+    "image_loss": 0.05849294364452362,
+    "smooth_loss": 2.282957041188638e-07
+  },
+  {
+    "epoch": 11.0,
+    "loss": 0.058464165776968,
+    "image_loss": 0.058464162051677704,
+    "smooth_loss": 2.843589470558072e-07
+  },
+  {
+    "epoch": 12.0,
+    "loss": 0.05843484401702881,
+    "image_loss": 0.05843484029173851,
+    "smooth_loss": 3.4935141002279124e-07
+  },
+  {
+    "epoch": 13.0,
+    "loss": 0.05840571969747543,
+    "image_loss": 0.058405715972185135,
+    "smooth_loss": 4.197248131276865e-07
+  },
+  {
+    "epoch": 14.0,
+    "loss": 0.05837589502334595,
+    "image_loss": 0.05837589129805565,
+    "smooth_loss": 5.000206897420867e-07
+  },
+  {
+    "epoch": 15.0,
+    "loss": 0.058346040546894073,
+    "image_loss": 0.05834603309631348,
+    "smooth_loss": 5.874838393538084e-07
+  },
+  {
+    "epoch": 16.0,
+    "loss": 0.058315880596637726,
+    "image_loss": 0.05831587314605713,
+    "smooth_loss": 6.829379231021449e-07
+  },
+  {
+    "epoch": 17.0,
+    "loss": 0.05828540772199631,
+    "image_loss": 0.05828540027141571,
+    "smooth_loss": 7.876523113736766e-07
+  },
+  {
+    "epoch": 18.0,
+    "loss": 0.05825469642877579,
+    "image_loss": 0.05825468897819519,
+    "smooth_loss": 9.008151096168149e-07
+  },
+  {
+    "epoch": 19.0,
+    "loss": 0.05822316184639931,
+    "image_loss": 0.05822315067052841,
+    "smooth_loss": 1.0293301784258801e-06
+  },
+  {
+    "epoch": 20.0,
+    "loss": 0.05819175764918327,
+    "image_loss": 0.05819174647331238,
+    "smooth_loss": 1.1623228601820301e-06
+  },
+  {
+    "epoch": 21.0,
+    "loss": 0.058159757405519485,
+    "image_loss": 0.05815974250435829,
+    "smooth_loss": 1.3093066399960662e-06
+  },
+  {
+    "epoch": 22.0,
+    "loss": 0.058127511292696,
+    "image_loss": 0.058127496391534805,
+    "smooth_loss": 1.465943910261558e-06
+  },
+  {
+    "epoch": 23.0,
+    "loss": 0.058094657957553864,
+    "image_loss": 0.05809464305639267,
+    "smooth_loss": 1.63818526743853e-06
+  },
+  {
+    "epoch": 24.0,
+    "loss": 0.0580613799393177,
+    "image_loss": 0.05806136131286621,
+    "smooth_loss": 1.824913169912179e-06
+  },
+  {
+    "epoch": 25.0,
+    "loss": 0.05802749842405319,
+    "image_loss": 0.0580274797976017,
+    "smooth_loss": 2.0261988993297564e-06
+  },
+  {
+    "epoch": 26.0,
+    "loss": 0.05799323692917824,
+    "image_loss": 0.05799321457743645,
+    "smooth_loss": 2.2423303107643733e-06
+  },
+  {
+    "epoch": 27.0,
+    "loss": 0.057958219200372696,
+    "image_loss": 0.05795819312334061,
+    "smooth_loss": 2.477491534591536e-06
+  },
+  {
+    "epoch": 28.0,
+    "loss": 0.05792278051376343,
+    "image_loss": 0.05792275443673134,
+    "smooth_loss": 2.7321416382619645e-06
+  },
+  {
+    "epoch": 29.0,
+    "loss": 0.05788649618625641,
+    "image_loss": 0.05788646638393402,
+    "smooth_loss": 3.0071300898271147e-06
+  },
+  {
+    "epoch": 30.0,
+    "loss": 0.057849571108818054,
+    "image_loss": 0.05784953758120537,
+    "smooth_loss": 3.304860001662746e-06
+  },
+  {
+    "epoch": 31.0,
+    "loss": 0.05781206861138344,
+    "image_loss": 0.057812031358480453,
+    "smooth_loss": 3.6249500681151403e-06
+  },
+  {
+    "epoch": 32.0,
+    "loss": 0.05777355656027794,
+    "image_loss": 0.057773515582084656,
+    "smooth_loss": 3.9738624764140695e-06
+  },
+  {
+    "epoch": 33.0,
+    "loss": 0.057734414935112,
+    "image_loss": 0.05773437023162842,
+    "smooth_loss": 4.350997187430039e-06
+  },
+  {
+    "epoch": 34.0,
+    "loss": 0.057694148272275925,
+    "image_loss": 0.057694099843502045,
+    "smooth_loss": 4.76802779303398e-06
+  },
+  {
+    "epoch": 35.0,
+    "loss": 0.05765282362699509,
+    "image_loss": 0.05765277147293091,
+    "smooth_loss": 5.221423634793609e-06
+  },
+  {
+    "epoch": 36.0,
+    "loss": 0.05761110037565231,
+    "image_loss": 0.057611044496297836,
+    "smooth_loss": 5.697434062312823e-06
+  },
+  {
+    "epoch": 37.0,
+    "loss": 0.05756793171167374,
+    "image_loss": 0.05756786838173866,
+    "smooth_loss": 6.224857315828558e-06
+  },
+  {
+    "epoch": 38.0,
+    "loss": 0.05752401426434517,
+    "image_loss": 0.0575239472091198,
+    "smooth_loss": 6.778077477065381e-06
+  },
+  {
+    "epoch": 39.0,
+    "loss": 0.05747854709625244,
+    "image_loss": 0.05747847259044647,
+    "smooth_loss": 7.391299732262269e-06
+  },
+  {
+    "epoch": 40.0,
+    "loss": 0.057432375848293304,
+    "image_loss": 0.05743229389190674,
+    "smooth_loss": 8.039884960453492e-06
+  },
+  {
+    "epoch": 41.0,
+    "loss": 0.0573841817677021,
+    "image_loss": 0.05738409236073494,
+    "smooth_loss": 8.764058293309063e-06
+  },
+  {
+    "epoch": 42.0,
+    "loss": 0.057334933429956436,
+    "image_loss": 0.057334836572408676,
+    "smooth_loss": 9.537392543279566e-06
+  },
+  {
+    "epoch": 43.0,
+    "loss": 0.057283785194158554,
+    "image_loss": 0.0572836808860302,
+    "smooth_loss": 1.0389785529696383e-05
+  },
+  {
+    "epoch": 44.0,
+    "loss": 0.05723174661397934,
+    "image_loss": 0.057231634855270386,
+    "smooth_loss": 1.1303329301881604e-05
+  },
+  {
+    "epoch": 45.0,
+    "loss": 0.057176779955625534,
+    "image_loss": 0.057176657021045685,
+    "smooth_loss": 1.2341002729954198e-05
+  },
+  {
+    "epoch": 46.0,
+    "loss": 0.057120226323604584,
+    "image_loss": 0.05712009221315384,
+    "smooth_loss": 1.3473189028445631e-05
+  },
+  {
+    "epoch": 47.0,
+    "loss": 0.057062555104494095,
+    "image_loss": 0.057062409818172455,
+    "smooth_loss": 1.4687300790683366e-05
+  },
+  {
+    "epoch": 48.0,
+    "loss": 0.05700160190463066,
+    "image_loss": 0.05700144171714783,
+    "smooth_loss": 1.6045180018409155e-05
+  },
+  {
+    "epoch": 49.0,
+    "loss": 0.056939497590065,
+    "image_loss": 0.056939322501420975,
+    "smooth_loss": 1.7514281353214756e-05
+  },
+  {
+    "epoch": 50.0,
+    "loss": 0.05687557905912399,
+    "image_loss": 0.05687538906931877,
+    "smooth_loss": 1.911029903567396e-05
+  },
+  {
+    "epoch": 51.0,
+    "loss": 0.056809235364198685,
+    "image_loss": 0.05680902674794197,
+    "smooth_loss": 2.0875198970315978e-05
+  },
+  {
+    "epoch": 52.0,
+    "loss": 0.05674029514193535,
+    "image_loss": 0.05674006789922714,
+    "smooth_loss": 2.2824997358839028e-05
+  },
+  {
+    "epoch": 53.0,
+    "loss": 0.056669920682907104,
+    "image_loss": 0.05666967108845711,
+    "smooth_loss": 2.4954817490652204e-05
+  },
+  {
+    "epoch": 54.0,
+    "loss": 0.0565965436398983,
+    "image_loss": 0.05659627169370651,
+    "smooth_loss": 2.7323461836203933e-05
+  },
+  {
+    "epoch": 55.0,
+    "loss": 0.056521233171224594,
+    "image_loss": 0.05652093514800072,
+    "smooth_loss": 2.9880573492846452e-05
+  },
+  {
+    "epoch": 56.0,
+    "loss": 0.05644303560256958,
+    "image_loss": 0.056442707777023315,
+    "smooth_loss": 3.275117705925368e-05
+  },
+  {
+    "epoch": 57.0,
+    "loss": 0.05636239051818848,
+    "image_loss": 0.056362032890319824,
+    "smooth_loss": 3.594645386328921e-05
+  },
+  {
+    "epoch": 58.0,
+    "loss": 0.056279003620147705,
+    "image_loss": 0.05627860873937607,
+    "smooth_loss": 3.945654316339642e-05
+  },
+  {
+    "epoch": 59.0,
+    "loss": 0.05619405582547188,
+    "image_loss": 0.056193623691797256,
+    "smooth_loss": 4.33080131188035e-05
+  },
+  {
+    "epoch": 60.0,
+    "loss": 0.056105729192495346,
+    "image_loss": 0.05610525235533714,
+    "smooth_loss": 4.769213410327211e-05
+  },
+  {
+    "epoch": 61.0,
+    "loss": 0.0560150220990181,
+    "image_loss": 0.056014496833086014,
+    "smooth_loss": 5.2607003453886136e-05
+  },
+  {
+    "epoch": 62.0,
+    "loss": 0.05592316761612892,
+    "image_loss": 0.05592258647084236,
+    "smooth_loss": 5.80552987230476e-05
+  },
+  {
+    "epoch": 63.0,
+    "loss": 0.05583030730485916,
+    "image_loss": 0.055829666554927826,
+    "smooth_loss": 6.423084414564073e-05
+  },
+  {
+    "epoch": 64.0,
+    "loss": 0.055736664682626724,
+    "image_loss": 0.05573595315217972,
+    "smooth_loss": 7.123746036086231e-05
+  },
+  {
+    "epoch": 65.0,
+    "loss": 0.0556456595659256,
+    "image_loss": 0.055644869804382324,
+    "smooth_loss": 7.909350097179413e-05
+  },
+  {
+    "epoch": 66.0,
+    "loss": 0.05555867403745651,
+    "image_loss": 0.055557794868946075,
+    "smooth_loss": 8.807046106085181e-05
+  },
+  {
+    "epoch": 67.0,
+    "loss": 0.055479418486356735,
+    "image_loss": 0.05547843500971794,
+    "smooth_loss": 9.828618203755468e-05
+  },
+  {
+    "epoch": 68.0,
+    "loss": 0.05541320890188217,
+    "image_loss": 0.055412109941244125,
+    "smooth_loss": 0.00010981244122376665
+  },
+  {
+    "epoch": 69.0,
+    "loss": 0.05536733567714691,
+    "image_loss": 0.05536610633134842,
+    "smooth_loss": 0.00012293846521060914
+  },
+  {
+    "epoch": 70.0,
+    "loss": 0.05535025894641876,
+    "image_loss": 0.05534888431429863,
+    "smooth_loss": 0.00013746831973548979
+  },
+  {
+    "epoch": 71.0,
+    "loss": 0.0553678423166275,
+    "image_loss": 0.05536631867289543,
+    "smooth_loss": 0.00015228186384774745
+  },
+  {
+    "epoch": 72.0,
+    "loss": 0.05540918558835983,
+    "image_loss": 0.055407531559467316,
+    "smooth_loss": 0.00016536397743038833
+  },
+  {
+    "epoch": 73.0,
+    "loss": 0.05544228106737137,
+    "image_loss": 0.05544053763151169,
+    "smooth_loss": 0.00017419658252038062
+  },
+  {
+    "epoch": 74.0,
+    "loss": 0.05544540658593178,
+    "image_loss": 0.05544362962245941,
+    "smooth_loss": 0.00017784371448215097
+  },
+  {
+    "epoch": 75.0,
+    "loss": 0.05542201176285744,
+    "image_loss": 0.05542023852467537,
+    "smooth_loss": 0.00017743516946211457
+  },
+  {
+    "epoch": 76.0,
+    "loss": 0.0553855374455452,
+    "image_loss": 0.055383794009685516,
+    "smooth_loss": 0.0001741933374432847
+  },
+  {
+    "epoch": 77.0,
+    "loss": 0.05534776300191879,
+    "image_loss": 0.05534606799483299,
+    "smooth_loss": 0.00016937614418566227
+  },
+  {
+    "epoch": 78.0,
+    "loss": 0.05531612038612366,
+    "image_loss": 0.055314477533102036,
+    "smooth_loss": 0.00016410183161497116
+  },
+  {
+    "epoch": 79.0,
+    "loss": 0.05529369041323662,
+    "image_loss": 0.055292099714279175,
+    "smooth_loss": 0.00015910626098047942
+  },
+  {
+    "epoch": 80.0,
+    "loss": 0.055279314517974854,
+    "image_loss": 0.05527777224779129,
+    "smooth_loss": 0.00015433150110766292
+  }
+]
--- a/outputs/checkpoints/vxm_head_ct_patient1.pt
+++ b/outputs/checkpoints/vxm_head_ct_patient1.pt