diff --git a/.agents/agents/testing.md b/.agents/agents/testing.md index 2ff0bd39..9ee28736 100644 --- a/.agents/agents/testing.md +++ b/.agents/agents/testing.md @@ -14,7 +14,7 @@ wherever practical. - `tests/baselines/` — stored via Git LFS; fetch with `git lfs pull` - `src/physiotwin4d/test_tools.py` — baseline comparison utilities (`TestTools`) - Markers (all opt-in via `--run-`): `slow`, `requires_gpu`, - `requires_simpleware`, `experiment`, `tutorial`. The `requires_data` marker + `requires_simpleware`, `tutorial`. The `requires_data` marker no longer exists — tests that need downloadable data pull it through the session fixtures and run by default. @@ -29,7 +29,8 @@ python -m pytest tests/ -v # fast, recomm python -m pytest tests/test_contour_tools.py -v # single file python -m pytest tests/test_contour_tools.py::TestContourTools -v # single class python -m pytest tests/ -v --run-slow # opt into slow tests -python -m pytest tests/ -v --run-gpu --run-slow # typical local GPU profile (CI runner adds --run-simpleware --run-experiments --run-tutorials) +# typical local GPU profile; CI adds --run-simpleware --run-tutorials +python -m pytest tests/ -v --run-gpu --run-slow python -m pytest tests/ --create-baselines # create missing baselines ``` @@ -59,8 +60,8 @@ python -m pytest tests/ --create-baselines # create missi 8. Prefer storing results in subdirectories under `./results/`. 9. Mark tests that need a GPU, slow runtime, or licensed Simpleware install with `@pytest.mark.requires_gpu`, `@pytest.mark.slow`, or - `@pytest.mark.requires_simpleware`. Mark experiment and tutorial tests - with `@pytest.mark.experiment` or `@pytest.mark.tutorial`. Tests that just + `@pytest.mark.requires_simpleware`. Mark tutorial tests with + `@pytest.mark.tutorial`. Tests that just need downloadable data need no marker. 10. Do not mock segmentation or registration models — test real outputs. 11. No emojis in test files (Windows cp1252 encoding has bitten this project). diff --git a/.github/workflows/README.md b/.github/workflows/README.md index 4e22ee18..ed81f499 100644 --- a/.github/workflows/README.md +++ b/.github/workflows/README.md @@ -226,7 +226,7 @@ uv pip install -e ".[test,cuda13,physicsnemo]" pytest tests/ --run-gpu # Enable every --run-* bucket at once (slow, GPU, simpleware, -# physicsnemo, experiments, tutorials) +# physicsnemo, tutorials) pytest tests/ --run-all ``` diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f2316677..ba25d1e5 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -13,7 +13,6 @@ name: CI # - requires_gpu -> --run-gpu # - requires_simpleware -> --run-simpleware (also implies GPU) # - requires_physicsnemo -> --run-physicsnemo (needs the [physicsnemo] extra) -# - experiment -> --run-experiments # - tutorial -> --run-tutorials # --run-all enables every bucket above at once. # Tests that need external data download it automatically via fixtures. @@ -107,13 +106,13 @@ jobs: run: | pip list - - name: Run unit tests (fast, no GPU/slow/experiment) - Ubuntu + - name: Run unit tests (fast, no GPU/slow/tutorial) - Ubuntu if: matrix.os == 'ubuntu-latest' run: | xvfb-run -a --server-args="-screen 0 1024x768x24" \ pytest tests/ -v --cov=physiotwin4d --cov-report=xml --cov-report=term --cov-report=html - - name: Run unit tests (fast, no GPU/slow/experiment) - Windows + - name: Run unit tests (fast, no GPU/slow/tutorial) - Windows if: matrix.os == 'windows-latest' run: | pytest tests/ -v --cov=physiotwin4d --cov-report=xml --cov-report=term --cov-report=html @@ -414,9 +413,9 @@ jobs: # - tests/test_transform_tools.py (depends on slow registration tests) # - tests/test_segment_chest_total_segmentator.py (requires CUDA for TotalSegmentator) # -# Experiment tests (EXTREMELY SLOW - hours to complete): -# - tests/test_experiments.py (runs all notebooks in experiments/ subdirectories) -# These tests are NEVER run in CI/CD and must be run manually +# Tutorial tests (SLOW - hours to complete): +# - tests/test_tutorials.py (runs every script in tutorials/ end-to-end) +# These tests are NEVER run in the PR CI and must be opted into # They execute end-to-end workflows that may take multiple hours # # To run locally: @@ -426,9 +425,9 @@ jobs: # pytest tests/test_register_images_ANTS.py -v --run-slow # # Self-hosted GPU runner enables ALL buckets via --run-all -# (--run-gpu --run-slow --run-simpleware --run-physicsnemo --run-experiments --run-tutorials). +# (--run-gpu --run-slow --run-simpleware --run-physicsnemo --run-tutorials). # That runner installs the [physicsnemo] extra in addition to [test,cuda13]. # -# To run experiment tests (manual only, extremely slow): -# pytest tests/test_experiments.py -v --run-experiments -# pytest tests/test_experiments.py::test_experiment_heart_gated_ct_to_usd -v --run-experiments +# To run tutorial tests (manual only, slow): +# pytest tests/test_tutorials.py -v --run-tutorials +# pytest tests/test_tutorials.py::TestTutorial01HeartGatedCTToUSD -v --run-tutorials diff --git a/.github/workflows/nightly-health.yml b/.github/workflows/nightly-health.yml index 7dd01918..923e69d6 100644 --- a/.github/workflows/nightly-health.yml +++ b/.github/workflows/nightly-health.yml @@ -1,6 +1,6 @@ name: Nightly Health -# Run the full test suite (including experiments) on the GPU runner every night. +# Run the full test suite (including tutorials) on the GPU runner every night. # # Schedule: 07:00 UTC daily ≈ 02:00 EST / 03:00 EDT # diff --git a/AGENTS.md b/AGENTS.md index 3e87f4dc..827335c7 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -81,7 +81,6 @@ python -m pytest tests/ -v --run-slow python -m pytest tests/ -v --run-gpu python -m pytest tests/ -v --run-simpleware python -m pytest tests/ -v --run-physicsnemo -python -m pytest tests/ -v --run-experiments python -m pytest tests/ -v --run-tutorials # Enable every bucket at once (equivalent to passing all --run-* flags) @@ -151,7 +150,7 @@ graphify update . # refresh after code changes (AST-only, no API c `X | None`. - Run `python -m pytest tests/ -v` from the active virtual venv to verify changes. Slow, GPU, Simpleware, - experiment, and tutorial tests are auto-skipped unless their opt-in flag is + and tutorial tests are auto-skipped unless their opt-in flag is passed. - Do not run pytest with `--run-slow` or `--run-all` unless the user explicitly asks; those runs take far too long for an interactive session. Default to the @@ -219,8 +218,7 @@ graphify update . # refresh after code changes (AST-only, no API c - Mark tests that need a GPU, a slow runtime, or a licensed Simpleware install with `@pytest.mark.requires_gpu`, `@pytest.mark.slow`, or `@pytest.mark.requires_simpleware`. -- Mark experiment and tutorial tests with `@pytest.mark.experiment` or - `@pytest.mark.tutorial`. +- Mark tutorial tests with `@pytest.mark.tutorial`. - Tests that just need downloadable data need no marker; the fixture chain handles it. - Prefer images from `ROOT/data/test/slicer_heart_small` for tests. diff --git a/CLAUDE.md b/CLAUDE.md index 177bf2d9..c8401783 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -127,15 +127,18 @@ came from. ## Testing -- Fast tests (recommended for development — slow/GPU/Simpleware/experiment - /tutorial tests are auto-skipped unless their opt-in flag is passed) +- Fast tests (recommended for development — slow/GPU/Simpleware/tutorial + tests are auto-skipped unless their opt-in flag is passed) py -m pytest tests/ -v - Baselines in `tests/baselines/` via Git LFS — run `git lfs pull` after cloning -- `tests/conftest.py`: session-scoped fixtures chaining download → convert → segment → register +- `tests/conftest.py`: session-scoped fixtures chaining + download → convert → segment → register - `src/physiotwin4d/test_tools.py`: baseline comparison utilities (`TestTools`, etc.) - Markers (all opt-in via `--run-`): `slow`, `requires_gpu`, - `requires_simpleware`, `experiment`, `tutorial`. Data-dependent tests no + `requires_simpleware`, `tutorial`. Data-dependent tests no longer use a marker — they pull data through fixtures and run by default. +- `experiments/` scripts are exploratory and are not run as tests; the + `tutorials/` scripts are the optional end-to-end suite - **Avoid `pytest --run-slow` (or `--run-all`) unless the user explicitly requests it.** Those runs take far too long for an interactive session. Default to the fast suite and let the user invoke the slow buckets. diff --git a/docs/cli_scripts/create_statistical_model.rst b/docs/cli_scripts/create_statistical_model.rst index 848b46b1..ceb044b3 100644 --- a/docs/cli_scripts/create_statistical_model.rst +++ b/docs/cli_scripts/create_statistical_model.rst @@ -71,7 +71,7 @@ With Custom Parameters --sample-meshes-dir ./meshes \ --reference-mesh average_mesh.vtk \ --output-dir ./pca_output \ - --pca-components 20 + --number-of-pca-components 20 Command-Line Arguments ====================== @@ -94,7 +94,7 @@ Required Arguments Optional Arguments ------------------ -``--pca-components N`` +``--number-of-pca-components N`` Number of PCA components to retain (default: 7). See :class:`physiotwin4d.WorkflowCreateStatisticalModel` for the full API and diff --git a/docs/cli_scripts/fit_statistical_model_to_patient.rst b/docs/cli_scripts/fit_statistical_model_to_patient.rst index 11ff2cfe..5f79ae8a 100644 --- a/docs/cli_scripts/fit_statistical_model_to_patient.rst +++ b/docs/cli_scripts/fit_statistical_model_to_patient.rst @@ -58,7 +58,7 @@ Include statistical shape model fitting: --patient-models lv.vtp rv.vtp myo.vtp \ --patient-image patient_ct.nii.gz \ --pca-json pca_model.json \ - --pca-number-of-modes 10 \ + --number-of-pca-components 10 \ --output-dir ./results Command-Line Arguments @@ -106,8 +106,8 @@ PCA Registration Options ``--pca-json PATH`` Path to PCA JSON file for shape-based registration (optional) -``--pca-number-of-modes NUM`` - Number of PCA modes to use (default: 0, uses all if PCA enabled) +``--number-of-pca-components NUM`` + Number of PCA components to use (default: 0, uses all if PCA enabled) Registration Configuration --------------------------- @@ -162,7 +162,7 @@ Example 2: PCA-Based Registration --patient-models lv.vtp rv.vtp \ --patient-image patient_ct.nii.gz \ --pca-json pca_model.json \ - --pca-number-of-modes 10 \ + --number-of-pca-components 10 \ --output-dir results/pca Output Files diff --git a/docs/cli_scripts/vtk_to_usd.rst b/docs/cli_scripts/vtk_to_usd.rst index 4bdea06f..50f407f8 100644 --- a/docs/cli_scripts/vtk_to_usd.rst +++ b/docs/cli_scripts/vtk_to_usd.rst @@ -36,7 +36,7 @@ Solid color: --appearance solid \ --color 1 0 0 -Anatomy material: +One anatomy material for every mesh: .. code-block:: bash @@ -45,6 +45,18 @@ Anatomy material: --appearance anatomy \ --anatomy-type heart +A material per structure. Omitting ``--anatomy-type`` picks each object's +material from its name, and with ``--static-merge`` the objects are named after +the structures recorded in each file's ``SegmentationLabelNames`` field data — +as written by the image-to-VTK workflow: + +.. code-block:: bash + + physiotwin4d-convert-vtk-to-usd patient_highres_*.vtp \ + --output heart_structures.usd \ + --appearance anatomy \ + --static-merge + Colormap from a VTK point data array: .. code-block:: bash diff --git a/docs/contributing.rst b/docs/contributing.rst index e915b9a6..a20c08bd 100644 --- a/docs/contributing.rst +++ b/docs/contributing.rst @@ -265,7 +265,7 @@ Run Tests # Run with coverage pytest tests/ --cov=src/physiotwin4d --cov-report=html - # Default invocation auto-skips slow/GPU/Simpleware/experiment/tutorial + # Default invocation auto-skips slow/GPU/Simpleware/tutorial pytest tests/ # Opt into specific buckets diff --git a/docs/developer/core.rst b/docs/developer/core.rst index 3e493b58..8b84f3c6 100644 --- a/docs/developer/core.rst +++ b/docs/developer/core.rst @@ -49,9 +49,9 @@ For most code changes, run: py -m pytest tests/ -v -(Slow / GPU / Simpleware / PhysicsNeMo / experiment / tutorial tests are +(Slow / GPU / Simpleware / PhysicsNeMo / tutorial tests are auto-skipped; opt in with ``--run-slow``, ``--run-gpu``, ``--run-simpleware``, -``--run-physicsnemo``, ``--run-experiments``, ``--run-tutorials``, or use +``--run-physicsnemo``, ``--run-tutorials``, or use ``--run-all`` to enable every bucket at once. Data-dependent tests download their data through the session fixtures and run by default.) diff --git a/docs/developer/registration_images.rst b/docs/developer/registration_images.rst index dd649dd3..e3fdf8d8 100644 --- a/docs/developer/registration_images.rst +++ b/docs/developer/registration_images.rst @@ -57,10 +57,10 @@ Workflows that accept a ``registration_method`` (e.g. :class:`WorkflowConvertImageToUSD`, :class:`RegisterTimeSeriesImages`) take any :class:`RegisterImagesBase` instance, including a composite chain that runs multiple backends in sequence. :class:`RegisterImagesChain` runs an -ordered list of registrars, feeding each stage's ``forward_transform`` as the -next stage's ``initial_forward_transform``. :class:`RegisterImagesGreedyICON` -is a named 2-stage convenience class for the common case of a fast Greedy -registration followed by ICON refinement: +ordered list of registrars, each stage refining the previous stage's +``forward_transform`` through ``register_from()`` (see `Seeding a registration`_ +below). :class:`RegisterImagesGreedyICON` is a named 2-stage convenience class +for the common case of a fast Greedy registration followed by ICON refinement: .. code-block:: python @@ -76,6 +76,23 @@ registration followed by ICON refinement: registrar.greedy.set_number_of_iterations([30, 15, 7, 3]) registrar.icon.set_number_of_iterations(20) +Seeding a registration +====================== + +To start from an alignment you already have, call ``register_from()`` instead of +``register()``: + +.. code-block:: python + + result = registrar.register_from(known_forward_transform, moving_image) + +It warps the moving image, mask and labelmap onto the fixed grid by that +transform, registers the residual, and composes the two, so the returned +transforms still map between the *original* moving image and the fixed image. +Every backend goes through this one implementation -- no registrar accepts an +initial transform of its own, which is what keeps the pre-warp, the composition +and the inversion identical for Greedy, ICON and ANTs. + Development Notes ================= diff --git a/docs/testing.rst b/docs/testing.rst index 2b00c34c..33890fe9 100644 --- a/docs/testing.rst +++ b/docs/testing.rst @@ -8,7 +8,7 @@ Run the fast test suite during development: pytest tests/ -v -Slow, GPU, Simpleware, experiment, and tutorial tests are auto-skipped unless +Slow, GPU, Simpleware, and tutorial tests are auto-skipped unless their opt-in flag is passed. Tests that depend on downloadable data fetch it automatically via the session fixtures, so no marker filter is needed for them. @@ -23,7 +23,6 @@ Each ``--run-`` flag enables one marker family: pytest tests/ -v --run-gpu # tests marked 'requires_gpu' pytest tests/ -v --run-simpleware # tests marked 'requires_simpleware' pytest tests/ -v --run-physicsnemo # tests marked 'requires_physicsnemo' - pytest tests/ -v --run-experiments # tests marked 'experiment' pytest tests/ -v --run-tutorials # tests marked 'tutorial' Flags compose. A typical local GPU profile is: diff --git a/docs/tutorials.rst b/docs/tutorials.rst index 5c328e28..8c25dc3f 100644 --- a/docs/tutorials.rst +++ b/docs/tutorials.rst @@ -11,7 +11,7 @@ Tutorials

PhysioTwin4D tutorials

From a CT scan to an animated digital twin

- Ten numbered stages across 15 runnable Python scripts. + Ten numbered stages across 16 runnable Python scripts. Each one drives the real workflow classes end-to-end on downloadable data, shows what it produced, and ends with the handful of constants to change so it runs on your own scans. @@ -219,6 +219,11 @@ Tutorial 2: Finetune ICON Registration Script ``tutorials/tutorial_02_lung_finetune_icon.py`` + ``tutorials/tutorial_02_lung_distancemap_finetune_icon.py`` — the + distance-map variant, which finetunes on distance maps rather than image + intensities so the labelmap-to-labelmap stage of Tutorials 7 and 8 has + in-distribution weights. + Workflow :class:`~physiotwin4d.WorkflowFinetuneICONRegistration`, then :class:`~physiotwin4d.RegisterImagesGreedy` and @@ -269,11 +274,13 @@ Run .. code-block:: bash python tutorials/tutorial_02_lung_finetune_icon.py + python tutorials/tutorial_02_lung_distancemap_finetune_icon.py Outputs The finetuned checkpoint under ``tutorials/network_weights/icon_dirlab_4dct/``, plus - ``registration_summary.csv``, the registered images, the fixed and warped + ``registration_summary.csv``, the fixed-minus-registered difference images + (residual structure is what separates the methods), the fixed and warped labelmaps, and before/after screenshots in ``tutorials/output/tutorial_02_lung/``. @@ -421,9 +428,10 @@ Run python tutorials/tutorial_04_lung_ct_to_vtk.py Outputs - ``patient_surfaces.vtp`` (all anatomy in one mesh), per-group and per-label - ``.vtp`` files, ``patient_labelmap.mha`` and two screenshots, under - ``tutorials/output/tutorial_04_{heart,lung}/``. + ``patient_surfaces.vtp`` (all anatomy in one mesh, with a per-cell + ``SegmentationLabelIds`` array so each cell still names the structure it came + from), per-group and per-label ``.vtp`` files, ``patient_labelmap.mha`` and + two screenshots, under ``tutorials/output/tutorial_04_{heart,lung}/``. Adapt to your data Change the input volume path, then choose the segmenter matching your scan: @@ -444,7 +452,8 @@ Workflow :class:`~physiotwin4d.WorkflowConvertVTKToUSD`. Dataset - Tutorial 4's ``patient_surfaces.vtp`` — no image data, no download. + Tutorial 4's per-structure ``patient_*.vtp`` surfaces — no image data, no + download. Requirements CPU only, seconds to run. The cheapest tutorial in the set. @@ -461,15 +470,22 @@ Inner API usage .. code-block:: python workflow = WorkflowConvertVTKToUSD( - input_meshes=[mesh], + input_meshes=meshes, usd_project_name=project_name, output_directory=output_dir, appearance="anatomy", - anatomy_type="heart", + static_merge=True, separate_by_connectivity=True, ) results = workflow.process() + Each input surface keeps the structure name that + :class:`~physiotwin4d.WorkflowConvertImageToVTK` wrote into its + ``field_data['SegmentationLabelNames']``. That name becomes the USD prim + name and, with ``anatomy_type`` left unset, selects the prim's material — + so the left chambers, right chambers, myocardium and great vessels each get + their own look rather than one shared heart material. + Run .. code-block:: bash @@ -481,11 +497,13 @@ Outputs Adapt to your data ``input_meshes`` takes any list of PyVista meshes — pass one per time point, - in order, for an animated scene instead of a static one, and set - ``frames_per_second`` to control playback. ``appearance="anatomy"`` binds - per-organ materials through :class:`~physiotwin4d.USDAnatomyTools`; use - ``anatomy_type`` to pick the palette. For file-in, file-out conversion - without Python, see :doc:`cli_scripts/vtk_to_usd`. + in order, for an animated scene instead of a static one (drop + ``static_merge``), and set ``frames_per_second`` to control playback. + ``appearance="anatomy"`` binds per-organ materials through + :class:`~physiotwin4d.USDAnatomyTools`; set ``anatomy_type`` to force one + palette onto every object, or ``object_names`` to name the prims yourself. + For file-in, file-out conversion without Python, see + :doc:`cli_scripts/vtk_to_usd`. Tutorial 6: Create a PCA Shape Model ==================================== @@ -533,7 +551,7 @@ Inner API usage workflow = WorkflowCreateStatisticalModel( sample_meshes=sample_surfaces, reference_mesh=reference_surface, - pca_number_of_components=pca_components, + number_of_pca_components=number_of_pca_components, ) result = workflow.process() @@ -553,7 +571,7 @@ Adapt to your data The workflow wants a population of meshes plus one reference; point ``sample_meshes`` at your own cohort and let :class:`~physiotwin4d.WorkflowCreateMeanSurface` build the reference when no - natural template exists. ``pca_number_of_components`` trades fidelity + natural template exists. ``number_of_pca_components`` trades fidelity against cohort size — you need more subjects than modes. The saved ``pca_model.json`` is the portable artifact: Tutorials 7 and 8 and :doc:`cli_scripts/create_statistical_model` all speak it. diff --git a/experiments/Convert_VTK_To_USD/convert_chop_alterra_valve_to_usd.py b/experiments/Convert_VTK_To_USD/convert_chop_alterra_valve_to_usd.py index f66b87d2..3fa3dd4b 100644 --- a/experiments/Convert_VTK_To_USD/convert_chop_alterra_valve_to_usd.py +++ b/experiments/Convert_VTK_To_USD/convert_chop_alterra_valve_to_usd.py @@ -30,9 +30,6 @@ from physiotwin4d import ConvertVTKToUSD -# Use as a test -from physiotwin4d.test_tools import TestTools - # Import USDTools for post-processing colormap from physiotwin4d.usd_tools import USDTools @@ -40,11 +37,6 @@ # ## 1. Discover and Organize Time-Series Files # %% -# Set to True to use as a test. Automatically done by -# TestTools.running_as_test() helper function. -test_mode = TestTools.running_as_test() -test_mode_step = 4 - # Define data directories (Alterra only). Anchored to the script's location # so the experiment runs from any working directory. script_dir = Path(__file__).resolve().parent @@ -52,10 +44,7 @@ alterra_dir = data_dir / "Alterra" output_dir = script_dir / "results" / "valve4d-alterra" -if test_mode: - output_usd = output_dir / "alterra_test.usd" -else: - output_usd = output_dir / "alterra_full.usd" +output_usd = output_dir / "alterra_full.usd" colormap_primvar_substrs = ["von_mises_stress"] colormap_name = "jet" # matplotlib colormap name @@ -137,10 +126,6 @@ alterra_files = [file_path for _, file_path in alterra_series] alterra_times = [float(time_step) for time_step, _ in alterra_series] -if test_mode: - alterra_files = alterra_files[::test_mode_step] - alterra_times = alterra_times[::test_mode_step] - print(f"\nConverting to: {output_usd}") print(f"Number of time steps: {len(alterra_times)}") print("\nThis may take several minutes...\n") diff --git a/experiments/Convert_VTK_To_USD/convert_chop_tpv25_valve_to_usd.py b/experiments/Convert_VTK_To_USD/convert_chop_tpv25_valve_to_usd.py index ccd8dff0..8b7d33ba 100644 --- a/experiments/Convert_VTK_To_USD/convert_chop_tpv25_valve_to_usd.py +++ b/experiments/Convert_VTK_To_USD/convert_chop_tpv25_valve_to_usd.py @@ -29,7 +29,6 @@ from pathlib import Path from physiotwin4d import ConvertVTKToUSD -from physiotwin4d.test_tools import TestTools # Import USDTools for post-processing colormap from physiotwin4d.usd_tools import USDTools @@ -38,11 +37,6 @@ # ## 1. Discover and Organize Time-Series Files # %% -# Set to True to use as a test. Automatically done by -# TestTools.running_as_test() helper function. -quick_run = TestTools.running_as_test() -quick_run_step = 4 - # Define data directories (TPV25 only). Anchored to the script's location # so the experiment runs from any working directory. script_dir = Path(__file__).resolve().parent @@ -50,10 +44,7 @@ tpv25_dir = data_dir / "TPV25" output_dir = script_dir / "results" / "valve4d-tpv25" -if quick_run: - output_usd = output_dir / "tpv25_quick.usd" -else: - output_usd = output_dir / "tpv25_full.usd" +output_usd = output_dir / "tpv25_full.usd" colormap_primvar_substrs = ["von_mises_stress"] colormap_name = "jet" # matplotlib colormap name @@ -134,10 +125,6 @@ tpv25_files = [file_path for _, file_path in tpv25_series] tpv25_times = [float(time_step) for time_step, _ in tpv25_series] -if quick_run: - tpv25_files = tpv25_files[::quick_run_step] - tpv25_times = tpv25_times[::quick_run_step] - print(f"\nConverting to: {output_usd}") print(f"Number of time steps: {len(tpv25_times)}") print("\nThis may take several minutes...\n") diff --git a/experiments/Heart-Create_Statistical_Model/1-input_meshes_to_input_surfaces.py b/experiments/Heart-Create_Statistical_Model/1-input_meshes_to_input_surfaces.py index c98fba56..89b27c77 100644 --- a/experiments/Heart-Create_Statistical_Model/1-input_meshes_to_input_surfaces.py +++ b/experiments/Heart-Create_Statistical_Model/1-input_meshes_to_input_surfaces.py @@ -12,8 +12,6 @@ import pyvista as pv -from physiotwin4d.test_tools import TestTools - _HERE = Path(__file__).parent # %% @@ -76,5 +74,4 @@ plotter = pv.Plotter() plotter.add_mesh(first_surface, color="lightblue", show_edges=True) plotter.add_axes() - if not TestTools.running_as_test(): - plotter.show() + plotter.show() diff --git a/experiments/Heart-Create_Statistical_Model/2-input_surfaces_to_surfaces_aligned.py b/experiments/Heart-Create_Statistical_Model/2-input_surfaces_to_surfaces_aligned.py index 0e3d4aec..bf59505d 100644 --- a/experiments/Heart-Create_Statistical_Model/2-input_surfaces_to_surfaces_aligned.py +++ b/experiments/Heart-Create_Statistical_Model/2-input_surfaces_to_surfaces_aligned.py @@ -22,7 +22,6 @@ import pyvista as pv from physiotwin4d.contour_tools import ContourTools -from physiotwin4d.test_tools import TestTools from physiotwin4d.register_models_icp import RegisterModelsICP _HERE = Path(__file__).parent @@ -168,8 +167,7 @@ plotter.show_axes() plotter.link_views() - if not TestTools.running_as_test(): - plotter.show() + plotter.show() # %% [markdown] # ## 6. Calculate Registration Statistics @@ -266,7 +264,6 @@ plt.tight_layout() plt.savefig(output_dir / "registration_statistics.png", dpi=150, bbox_inches="tight") -if not TestTools.running_as_test(): - plt.show() +plt.show() print(f"\nPlot saved to: {output_dir / 'registration_statistics.png'}") diff --git a/experiments/Heart-Create_Statistical_Model/3-registration_based_correspondence.py b/experiments/Heart-Create_Statistical_Model/3-registration_based_correspondence.py index ffd495ba..034c5c9b 100644 --- a/experiments/Heart-Create_Statistical_Model/3-registration_based_correspondence.py +++ b/experiments/Heart-Create_Statistical_Model/3-registration_based_correspondence.py @@ -27,7 +27,6 @@ from pathlib import Path from physiotwin4d.contour_tools import ContourTools -from physiotwin4d.test_tools import TestTools from physiotwin4d.register_models_distance_maps import RegisterModelsDistanceMaps _HERE = Path(__file__).parent @@ -219,8 +218,7 @@ # Link the camera views so they rotate together plotter.link_views() - if not TestTools.running_as_test(): - plotter.show() + plotter.show() # %% [markdown] # ## Visualize Deformation Magnitude @@ -274,8 +272,7 @@ plotter.show_axes() plotter.camera_position = "iso" - if not TestTools.running_as_test(): - plotter.show() + plotter.show() # %% # Save registration statistics @@ -318,8 +315,7 @@ plot_file = output_dir / "registration_statistics.png" plt.savefig(plot_file, dpi=150, bbox_inches="tight") print(f"\nPlot saved to: {plot_file}") - if not TestTools.running_as_test(): - plt.show() + plt.show() else: print("\nNo statistics to plot.") diff --git a/experiments/Heart-Create_Statistical_Model/4-surfaces_aligned_correspond_to_pca_inputs.py b/experiments/Heart-Create_Statistical_Model/4-surfaces_aligned_correspond_to_pca_inputs.py index 26e125e4..d7af7e98 100644 --- a/experiments/Heart-Create_Statistical_Model/4-surfaces_aligned_correspond_to_pca_inputs.py +++ b/experiments/Heart-Create_Statistical_Model/4-surfaces_aligned_correspond_to_pca_inputs.py @@ -6,7 +6,6 @@ import pyvista as pv from physiotwin4d.contour_tools import ContourTools -from physiotwin4d.test_tools import TestTools _HERE = Path(__file__).parent @@ -111,8 +110,7 @@ # Link the camera views so they rotate together plotter.link_views() - if not TestTools.running_as_test(): - plotter.show() + plotter.show() # %% [markdown] # ## Visualize Deformation Magnitude @@ -166,5 +164,4 @@ plotter.show_axes() plotter.camera_position = "iso" - if not TestTools.running_as_test(): - plotter.show() + plotter.show() diff --git a/experiments/Heart-Create_Statistical_Model/5-compute_pca_model.py b/experiments/Heart-Create_Statistical_Model/5-compute_pca_model.py index 365309c3..49b7325b 100644 --- a/experiments/Heart-Create_Statistical_Model/5-compute_pca_model.py +++ b/experiments/Heart-Create_Statistical_Model/5-compute_pca_model.py @@ -21,8 +21,6 @@ from sklearn.decomposition import PCA # SparsePCA, ... -from physiotwin4d.test_tools import TestTools - _HERE = Path(__file__).parent n_components = 15 @@ -242,8 +240,7 @@ def generate_pc_variation(pc_index, std_dev_multiplier=3.0): # Link all three views so camera movements are synchronized plotter.link_views() -if not TestTools.running_as_test(): - plotter.show() +plotter.show() # %% # Plot variance explained by each component @@ -266,8 +263,7 @@ def generate_pc_variation(pc_index, std_dev_multiplier=3.0): ax2.grid(True, alpha=0.3) plt.tight_layout() -if not TestTools.running_as_test(): - plt.show() +plt.show() print( f"\nTotal variance captured by {n_components} components: {cumulative_variance[-1] * 100:.2f}%" diff --git a/experiments/Heart-Create_Statistical_Model/README.md b/experiments/Heart-Create_Statistical_Model/README.md index 92727e9d..d450fe00 100644 --- a/experiments/Heart-Create_Statistical_Model/README.md +++ b/experiments/Heart-Create_Statistical_Model/README.md @@ -228,7 +228,8 @@ For ICON registration: - Check alignment quality from step 2 (ICP should produce good initial alignment) - Verify average surface looks reasonable before step 3 - If Greedy affine fails, check input mesh quality and topology -- If ICON deformable quality is poor, increase `icon_iterations` in the `register()` call +- If ICON deformable quality is poor, increase the iteration count on the + registrar's ICON instance: `registrar.registrar_ICON.set_number_of_iterations(...)` ### Import Errors - Ensure all PhysioTwin4D dependencies are installed diff --git a/experiments/Heart-GatedCT_To_USD/1-register_images.py b/experiments/Heart-GatedCT_To_USD/1-register_images.py index 4cf632ef..5cd95f2c 100644 --- a/experiments/Heart-GatedCT_To_USD/1-register_images.py +++ b/experiments/Heart-GatedCT_To_USD/1-register_images.py @@ -7,7 +7,6 @@ from physiotwin4d.segment_chest_total_segmentator_with_contrast import ( SegmentChestTotalSegmentatorWithContrast, ) -from physiotwin4d.test_tools import TestTools from physiotwin4d.transform_tools import TransformTools # nnUNetv2 (used by TotalSegmentator) spawns a multiprocessing.Pool. On Windows @@ -15,14 +14,12 @@ # __name__ == "__main__" guard around the top-level work, that re-import fires # segment() again and Python's spawn-cascade detector raises RuntimeError. if __name__ == "__main__": - test_mode = TestTools.running_as_test() - _HERE = Path(__file__).resolve().parent # %% # Number of cardiac frames and step size; downstream scripts must use the same values. N_FRAMES = 21 - FRAME_STEP = 21 if test_mode else 1 + FRAME_STEP = 1 data_dir = _HERE.parent.parent / "data" / "Slicer-Heart-CT" @@ -34,7 +31,6 @@ # %% seg = SegmentChestTotalSegmentatorWithContrast() - seg.fast_mode = test_mode result = seg.segment(fixed_image) # %% labelmap_mask = result["labelmap"] @@ -79,8 +75,7 @@ # %% reg = RegisterImagesANTS() reg.set_mask_dilation(5) - reg.fast_mode = test_mode - reg.set_number_of_iterations([2, 1, 1] if test_mode else [10, 5, 2]) + reg.set_number_of_iterations([10, 5, 2]) # %% for i in range(0, N_FRAMES, FRAME_STEP): @@ -108,8 +103,6 @@ print(f" Done registering whole image for slice {i:03d}.") inverse_transform = results["inverse_transform"] forward_transform = results["forward_transform"] - whole_image_forward_transform = forward_transform - whole_image_inverse_transform = inverse_transform moving_image_reg = TransformTools().transform_image( moving_image, forward_transform, fixed_image, "sinc" ) # Final resampling with sinc @@ -129,28 +122,20 @@ compression=True, ) - # Register the dynamic anatomy mask. In test mode, masked ANTs - # registration (mask applied at every pyramid level) is far more - # expensive than the unmasked whole-image case above; reuse the - # whole-image transform instead of re-registering with a mask. + # Register the dynamic anatomy mask. heart_arr = itk.GetArrayFromImage(heart_mask) contrast_arr = itk.GetArrayFromImage(contrast_mask) major_vessels_arr = itk.GetArrayFromImage(major_vessels_mask) dynamic_anatomy_arr = heart_arr + contrast_arr + major_vessels_arr moving_image_dynamic_anatomy_mask = itk.GetImageFromArray(dynamic_anatomy_arr) moving_image_dynamic_anatomy_mask.CopyInformation(moving_image) - if test_mode: - print(f" Reusing whole-image transform for slice {i:03d} (test mode).") - forward_transform = whole_image_forward_transform - inverse_transform = whole_image_inverse_transform - else: - print(f" Registering dynamic anatomy mask for slice {i:03d}...") - reg.set_fixed_image(fixed_image) - reg.set_fixed_mask(fixed_image_dynamic_anatomy_mask) - results = reg.register(moving_image, moving_image_dynamic_anatomy_mask) - print(f" Done registering dynamic anatomy mask for slice {i:03d}.") - inverse_transform = results["inverse_transform"] - forward_transform = results["forward_transform"] + print(f" Registering dynamic anatomy mask for slice {i:03d}...") + reg.set_fixed_image(fixed_image) + reg.set_fixed_mask(fixed_image_dynamic_anatomy_mask) + results = reg.register(moving_image, moving_image_dynamic_anatomy_mask) + print(f" Done registering dynamic anatomy mask for slice {i:03d}.") + inverse_transform = results["inverse_transform"] + forward_transform = results["forward_transform"] moving_image_reg_dynamic_anatomy = TransformTools().transform_image( moving_image, forward_transform, fixed_image, "sinc" ) # Final resampling with sinc @@ -175,7 +160,7 @@ compression=True, ) - # Register the static anatomy mask (same test-mode shortcut as above). + # Register the static anatomy mask. lung_arr = itk.GetArrayFromImage(lung_mask) bone_arr = itk.GetArrayFromImage(bone_mask) other_arr = itk.GetArrayFromImage(other_mask) @@ -183,18 +168,13 @@ lung_arr + bone_arr + other_arr ) moving_image_static_mask.CopyInformation(moving_image) - if test_mode: - print(f" Reusing whole-image transform for slice {i:03d} (test mode).") - forward_transform = whole_image_forward_transform - inverse_transform = whole_image_inverse_transform - else: - print(f" Registering static anatomy mask for slice {i:03d}...") - reg.set_fixed_image(fixed_image) - reg.set_fixed_mask(fixed_image_static_mask) - results = reg.register(moving_image, moving_image_static_mask) - print(f" Done registering static anatomy mask for slice {i:03d}.") - inverse_transform = results["inverse_transform"] - forward_transform = results["forward_transform"] + print(f" Registering static anatomy mask for slice {i:03d}...") + reg.set_fixed_image(fixed_image) + reg.set_fixed_mask(fixed_image_static_mask) + results = reg.register(moving_image, moving_image_static_mask) + print(f" Done registering static anatomy mask for slice {i:03d}.") + inverse_transform = results["inverse_transform"] + forward_transform = results["forward_transform"] moving_image_reg_static = TransformTools().transform_image( moving_image, forward_transform, fixed_image, "sinc" ) # Final resampling with sinc diff --git a/experiments/Heart-GatedCT_To_USD/2-generate_segmentation.py b/experiments/Heart-GatedCT_To_USD/2-generate_segmentation.py index 74824666..5adcdc67 100644 --- a/experiments/Heart-GatedCT_To_USD/2-generate_segmentation.py +++ b/experiments/Heart-GatedCT_To_USD/2-generate_segmentation.py @@ -9,7 +9,6 @@ from physiotwin4d.segment_chest_total_segmentator_with_contrast import ( SegmentChestTotalSegmentatorWithContrast, ) -from physiotwin4d.test_tools import TestTools # nnUNetv2 (used by TotalSegmentator) spawns a multiprocessing.Pool. On Windows # the spawn start method re-imports this script in each child; without the @@ -70,7 +69,6 @@ outname = "slice_max" seg = SegmentChestTotalSegmentatorWithContrast() - seg.fast_mode = TestTools.running_as_test() if re_run_image_segmentation: result = seg.segment(max_image) labelmap_image = result["labelmap"] @@ -162,5 +160,4 @@ ) pl.set_background("black") pl.camera_position = "xy" - if not TestTools.running_as_test(): - pl.show() + pl.show() diff --git a/experiments/Heart-GatedCT_To_USD/3-transform_dynamic_and_static_contours.py b/experiments/Heart-GatedCT_To_USD/3-transform_dynamic_and_static_contours.py index 166ddde7..0714bbb8 100644 --- a/experiments/Heart-GatedCT_To_USD/3-transform_dynamic_and_static_contours.py +++ b/experiments/Heart-GatedCT_To_USD/3-transform_dynamic_and_static_contours.py @@ -7,7 +7,6 @@ from physiotwin4d import ConvertVTKToUSD from physiotwin4d.contour_tools import ContourTools from physiotwin4d.segment_chest_total_segmentator import SegmentChestTotalSegmentator -from physiotwin4d.test_tools import TestTools from physiotwin4d.usd_anatomy_tools import USDAnatomyTools # Defensive: this script only reads `seg.all_mask_ids` today, but if anyone @@ -15,12 +14,11 @@ # multiprocessing.Pool which re-imports the script on Windows (spawn start # method) and crashes with a spawn-cascade RuntimeError. Guard pre-emptively. if __name__ == "__main__": - test_mode = TestTools.running_as_test() _HERE = os.path.dirname(os.path.abspath(__file__)) # Must match N_FRAMES / FRAME_STEP in 1-register_images.py N_FRAMES = 21 - FRAME_STEP = 21 if test_mode else 1 + FRAME_STEP = 1 # %% output_dir = os.path.join(_HERE, "results") diff --git a/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py b/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py index 37a66dfe..9ff7f4a0 100644 --- a/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py +++ b/experiments/Heart-Simpleware_Segmentation/simpleware_heart_segmentation.py @@ -23,6 +23,7 @@ # %% import logging import os +import sys import tkinter as tk from tkinter import filedialog @@ -33,7 +34,6 @@ from physiotwin4d.landmark_tools import LandmarkTools from physiotwin4d.segment_heart_simpleware import SegmentHeartSimpleware -from physiotwin4d.test_tools import TestTools _HERE = os.path.dirname(os.path.abspath(__file__)) @@ -59,18 +59,17 @@ # Load a cardiac CT image for segmentation. This should be a 3D volume containing the heart. # %% -if TestTools.running_as_test(): - input_image_path = os.path.join( - _HERE, "..", "..", "data", "CHOP-Valve4D", "CT", "RVOT28-Dias.nii.gz" - ) -else: - root = tk.Tk() - root.withdraw() - input_image_path = filedialog.askopenfilename( - title="Select a cardiac CT image", - filetypes=[("NIfTI", "*.nii.gz"), ("MetaIO", "*.mhd"), ("All files", "*.*")], - ) - root.destroy() +root = tk.Tk() +root.withdraw() +input_image_path = filedialog.askopenfilename( + title="Select a cardiac CT image", + filetypes=[("NIfTI", "*.nii.gz"), ("MetaIO", "*.mhd"), ("All files", "*.*")], +) +root.destroy() + +if not input_image_path: + print("No image selected; nothing to segment.") + sys.exit(0) # Load the image try: @@ -119,8 +118,7 @@ axes[2].axis("off") plt.tight_layout() - if not TestTools.running_as_test(): - plt.show() + plt.show() print(f"Image intensity range: [{image_array.min():.1f}, {image_array.max():.1f}]") else: @@ -362,8 +360,7 @@ plt.tight_layout() plt.savefig(os.path.join(output_dir, "segmentation_visualization.png"), dpi=150) - if not TestTools.running_as_test(): - plt.show() + plt.show() print( f"Visualization saved to: {os.path.join(output_dir, 'segmentation_visualization.png')}" @@ -417,8 +414,7 @@ # Save screenshot screenshot_path = os.path.join(output_dir, "3d_visualization.png") - if not TestTools.running_as_test(): - plotter.show(screenshot=screenshot_path) + plotter.show(screenshot=screenshot_path) print(f"3D visualization saved to: {screenshot_path}") else: diff --git a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_model_icp_itk.py b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_model_icp_itk.py index 73c29d26..03422961 100644 --- a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_model_icp_itk.py +++ b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_model_icp_itk.py @@ -29,7 +29,6 @@ TransformTools, ) from physiotwin4d.image_tools import ImageTools -from physiotwin4d.test_tools import TestTools # %% [markdown] # ## Define File Paths @@ -200,5 +199,4 @@ plotter.add_title("ICP Shape Fitting") plotter.add_axes() -if not TestTools.running_as_test(): - plotter.show() +plotter.show() diff --git a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_model_registration_pca.py b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_model_registration_pca.py index 19dc87c7..431b2f17 100644 --- a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_model_registration_pca.py +++ b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_model_registration_pca.py @@ -32,7 +32,6 @@ TransformTools, ) from physiotwin4d.image_tools import ImageTools -from physiotwin4d.test_tools import TestTools # %% [markdown] # ## Define File Paths @@ -161,12 +160,10 @@ icp_registrar = RegisterModelsICP(fixed_model=patient_surface) -# Use fewer iterations when run as test (pytest) for faster execution -max_iterations_icp = 100 if TestTools.running_as_test() else 2000 icp_result = icp_registrar.register( transform_type="Affine", moving_model=template_model_surface, - max_iterations=max_iterations_icp, + max_iterations=2000, ) # Get the aligned mesh and transform @@ -325,8 +322,7 @@ plotter.add_axes() plotter.link_views() -if not TestTools.running_as_test(): - plotter.show() +plotter.show() # %% [markdown] # ## Visualize PCA Displacement Magnitude @@ -385,8 +381,7 @@ ) plotter.add_title("PCA Signed Displacement on Registered Model") plotter.add_axes() -if not TestTools.running_as_test(): - plotter.show() +plotter.show() # Save the mesh with displacement data pca_registered_model_with_displacement.save( diff --git a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py index cacf30ed..bc4168d4 100644 --- a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py +++ b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient-CHOPValve.py @@ -14,7 +14,6 @@ SegmentHeartSimplewareTrimmedBranches, WorkflowFitStatisticalModelToPatient, ) -from physiotwin4d.test_tools import TestTools # %% [markdown] # ## Define File Paths @@ -55,7 +54,7 @@ ) registrar.set_use_pca_registration( - True, pca_model=model_pca_data, pca_number_of_modes=model_pca_n_modes + True, pca_model=model_pca_data, number_of_pca_components=model_pca_n_modes ) registrar.set_use_labelmap_to_labelmap_registration(True) @@ -110,8 +109,7 @@ plotter.add_title("Final Registration") plotter.link_views() -if not TestTools.running_as_test(): - plotter.show() +plotter.show() # %% [markdown] # ## Visualize Deformation Magnitude @@ -128,8 +126,7 @@ scalar_bar_args={"title": "Deformation (mm)"}, ) plotter.add_title("Deformation Magnitude") - if not TestTools.running_as_test(): - plotter.show() + plotter.show() # Print statistics deformation = registered_model_surface["DeformationMagnitude"] diff --git a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient.py b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient.py index 66a846ea..13fa11b2 100644 --- a/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient.py +++ b/experiments/Heart-Statistical_Model_To_Patient/heart_model_to_patient.py @@ -17,7 +17,6 @@ SegmentChestTotalSegmentator, WorkflowFitStatisticalModelToPatient, ) -from physiotwin4d.test_tools import TestTools # nnUNetv2 (used by TotalSegmentator) spawns a multiprocessing.Pool. On Windows # the spawn start method re-imports this script in each child; without the @@ -137,7 +136,7 @@ patient_image=patient_image, ) registrar.set_use_pca_registration( - True, pca_model=pca_model, pca_number_of_modes=pca_n_modes + True, pca_model=pca_model, number_of_pca_components=pca_n_modes ) registrar.set_use_labelmap_to_image_registration( True, @@ -278,8 +277,7 @@ plotter.add_title("Final Registration") plotter.link_views() - if not TestTools.running_as_test(): - plotter.show() + plotter.show() # %% [markdown] # ## Visualize Deformation Magnitude @@ -296,8 +294,7 @@ scalar_bar_args={"title": "Deformation (mm)"}, ) plotter.add_title("Deformation Magnitude") - if not TestTools.running_as_test(): - plotter.show() + plotter.show() # Print statistics deformation = registered_surface["DeformationMagnitude"] diff --git a/experiments/README.md b/experiments/README.md index 4e347fac..1c896c56 100644 --- a/experiments/README.md +++ b/experiments/README.md @@ -192,52 +192,30 @@ These experiments are **starting points for exploration**, not copy-paste soluti The **CLI commands and implementations in `src/physiotwin4d/cli/`** are the production-quality code you should use and extend for real-world digital twin projects. -## Automated Testing +## Testing -A comprehensive test suite is available to validate all experiment scripts: +Experiment scripts are **not** run as tests. They are exploratory research code: +they assume interactive display, full-resolution parameters, and data layouts +that only exist on the author's machine. -```bash -# Run all experiment tests (EXTREMELY SLOW - may take hours) -# NOTE: Requires --run-experiments flag! -pytest tests/test_experiments.py -v --run-experiments - -# Run a specific experiment subdirectory -pytest tests/test_experiments.py::test_experiment_heart_gated_ct_to_usd -v -s --run-experiments - -# List all scripts that would be run (without executing) -pytest tests/test_experiments.py::test_list_scripts_in_subdir -v -s --run-experiments +The optional end-to-end test suite runs the `tutorials/` scripts instead: -# Validate experiment directory structure -pytest tests/test_experiments.py::test_experiment_structure -v --run-experiments +```bash +pytest tests/test_tutorials.py -v --run-tutorials ``` -**IMPORTANT:** Experiment tests require the `--run-experiments` flag. Without this flag, they are automatically skipped, even if you run `pytest tests/` or target the test file directly. - -### Test Features - -- **Test-mode flag** - When run as tests (pytest with `--run-experiments`), the runner sets `PHYSIOTWIN_RUNNING_AS_TEST=1`. Scripts can read this (e.g. via `physiotwin4d.test_tools.TestTools.running_as_test()`) and use reduced parameters so test runs stay fast. See [tests/EXPERIMENT_TESTS_GUIDE.md](../tests/EXPERIMENT_TESTS_GUIDE.md#running-as-test-physiotwin_running_as_test). -- **One test per subdirectory** - Each experiment subdirectory gets its own test function -- **Alphanumeric ordering** - Scripts execute in alphanumeric order (e.g., `0-`, `1-`, `2-`) -- **Long timeouts** - Each script has up to 1 hour execution time, tests have multi-hour timeouts -- **Detailed output** - Progress reporting, execution summaries, and failure diagnostics -- **Opt-in only** - Requires `--run-experiments` flag; automatically skipped otherwise -- **Protected from CI/CD** - NEVER runs in automated workflows +Run experiment scripts manually, in the order their filename prefixes imply +(`0-`, `1-`, `2-`, ...), from within their own subdirectory: -### Requirements - -These tests require: -- All dependencies installed (see `pyproject.toml`) -- GPU/CUDA support for most experiments -- Large amounts of disk space and memory -- External data downloads (see individual experiment scripts) - -### Important Notes - -**These tests are extremely long-running** - Plan for multiple hours of execution time - -**Not part of CI/CD** - These tests are excluded from all automated workflows +```bash +cd experiments/Heart-GatedCT_To_USD +py 0-download_and_convert_4d_to_3d.py +py 1-register_images.py +``` -**Resource intensive** - Requires GPU, significant memory, and disk space +Each script requires all dependencies installed (see `pyproject.toml`), +GPU/CUDA support in most cases, large amounts of disk space and memory, and +external data downloads. ## Structure diff --git a/experiments/Reconstruct4DCT/reconstruct_4d_ct.py b/experiments/Reconstruct4DCT/reconstruct_4d_ct.py index a01c3e6b..a5c4ba90 100644 --- a/experiments/Reconstruct4DCT/reconstruct_4d_ct.py +++ b/experiments/Reconstruct4DCT/reconstruct_4d_ct.py @@ -4,7 +4,7 @@ import itk import numpy as np -from physiotwin4d import RegisterImagesGreedy, TestTools, TransformTools +from physiotwin4d import RegisterImagesGreedy, TransformTools _HERE = os.path.dirname(os.path.abspath(__file__)) @@ -16,10 +16,6 @@ if f.endswith(".mha") and f.startswith("slice_") ] -quick_run = TestTools.running_as_test() -if quick_run: - exit(0) - num_files = len(files) files_indx = list(range(num_files)) reference_image_num = 7 @@ -166,9 +162,7 @@ def register_slices( # Try identity as initial transform print(" Trying init with identity.") - results_init_identity = reg_tool.register( - img, initial_forward_transform=None - ) + results_init_identity = reg_tool.register(img) inverse_tranform_init_identity = results_init_identity["inverse_transform"] forward_transform_init_identity = results_init_identity["forward_transform"] loss_init_identity = results_init_identity["loss"] @@ -177,8 +171,8 @@ def register_slices( if portion_of_prior_to_use > 0.0: # Try with prior transform print(" Trying with init prior.") - results_init_prior = reg_tool.register( - img, initial_forward_transform=prior_forward_transform + results_init_prior = reg_tool.register_from( + prior_forward_transform, img ) inverse_transform_init_prior = results_init_prior["inverse_transform"] forward_transform_init_prior = results_init_prior["forward_transform"] diff --git a/experiments/Reconstruct4DCT/reconstruct_4d_ct_class.py b/experiments/Reconstruct4DCT/reconstruct_4d_ct_class.py index 931ca99c..a5f1ba97 100644 --- a/experiments/Reconstruct4DCT/reconstruct_4d_ct_class.py +++ b/experiments/Reconstruct4DCT/reconstruct_4d_ct_class.py @@ -24,7 +24,6 @@ RegisterTimeSeriesImages, TransformTools, ) -from physiotwin4d.test_tools import TestTools _HERE = os.path.dirname(os.path.abspath(__file__)) @@ -69,44 +68,19 @@ def _build_registrar(method_name: str, iterations=None) -> Optional[RegisterImag print(f"Found {len(files)} slice files") # %% -# Configuration: quick run when executed as test (pytest); full run when manual (set quick_run = True for interactive quick test) -quick_run = TestTools.running_as_test() - -# Select files and parameters based on mode -if quick_run: - print("=== QUICK RUN MODE ===") - total_num_files = len(files) - target_num_files = 2 - if total_num_files == 0: - raise FileNotFoundError(f"No slice_*.mha files found in {data_dir}") - target_num_files = min(target_num_files, total_num_files) - file_step = max(1, total_num_files // target_num_files) - files = files[0:total_num_files:file_step] - files_indx = list(range(0, total_num_files, file_step)) - num_files = len(files) - reference_image_num = num_files // 2 - - # Registration parameters - only Greedy for quick run. ICON and - # Greedy_ICON are exercised by dedicated registration tests elsewhere; - # this experiment validates the reconstruction pipeline, not every - # registration backend. - registration_method_names = ["Greedy"] - number_of_iterations_list = [[2, 1, 1]] # For Greedy -else: - print("=== FULL RUN MODE ===") - num_files = len(files) - files_indx = list(range(num_files)) - reference_image_num = 7 - - # Registration parameters - Greedy_ICON is the recommended method - registration_method_names = [ - "Default" - ] # Use default, or ["Greedy", "ICON", "Greedy_ICON"] - number_of_iterations_list = [None] # [ - # [30, 15, 7, 3], - # 20, # For ICON - # [[30, 15, 7, 3], 20], # For Greedy_ICON - # ] +num_files = len(files) +files_indx = list(range(num_files)) +reference_image_num = 7 + +# Registration parameters - Greedy_ICON is the recommended method +registration_method_names = [ + "Default" +] # Use default, or ["Greedy", "ICON", "Greedy_ICON"] +number_of_iterations_list = [None] # [ +# [30, 15, 7, 3], +# 20, # For ICON +# [[30, 15, 7, 3], 20], # For Greedy_ICON +# ] # Common parameters reference_image_file = os.path.join( @@ -254,41 +228,40 @@ def _build_registrar(method_name: str, iterations=None) -> Optional[RegisterImag print(f" Min loss: {np.min(losses):.6f}") print(f" Max loss: {np.max(losses):.6f}") - if not quick_run: - # Generate grid image for visualization - grid_image = tfm_tools.generate_grid_image(fixed_image, 30, 1) - - print(f"Generating {registration_method_name.upper()} grid visualizations...") - for i, img_indx in enumerate(files_indx): - print(f" Generating grid for slice {img_indx:03d}...") - - # Transform grid with inverse transform (FM) - inverse_grid_image = tfm_tools.transform_image( - grid_image, - inverse_transforms[i], - fixed_image, - ) - itk.imwrite( - inverse_grid_image, - os.path.join( - _RESULTS_DIR, - f"slice_fixed_{registration_method_name}_inverse_grid_{img_indx:03d}.mha", - ), - compression=True, - ) - - # Save displacement field as image - inverse_transform_image = tfm_tools.convert_transform_to_displacement_field( - inverse_transforms[i], - fixed_image, - np_component_type=np.float32, - ) - itk.imwrite( - inverse_transform_image, - os.path.join( - _RESULTS_DIR, - f"slice_{registration_method_name}_inverse_{img_indx:03d}_field.mha", - ), - compression=True, - ) - print(f"Grid visualizations saved for {registration_method_name.upper()}") + # Generate grid image for visualization + grid_image = tfm_tools.generate_grid_image(fixed_image, 30, 1) + + print(f"Generating {registration_method_name.upper()} grid visualizations...") + for i, img_indx in enumerate(files_indx): + print(f" Generating grid for slice {img_indx:03d}...") + + # Transform grid with inverse transform (FM) + inverse_grid_image = tfm_tools.transform_image( + grid_image, + inverse_transforms[i], + fixed_image, + ) + itk.imwrite( + inverse_grid_image, + os.path.join( + _RESULTS_DIR, + f"slice_fixed_{registration_method_name}_inverse_grid_{img_indx:03d}.mha", + ), + compression=True, + ) + + # Save displacement field as image + inverse_transform_image = tfm_tools.convert_transform_to_displacement_field( + inverse_transforms[i], + fixed_image, + np_component_type=np.float32, + ) + itk.imwrite( + inverse_transform_image, + os.path.join( + _RESULTS_DIR, + f"slice_{registration_method_name}_inverse_{img_indx:03d}_field.mha", + ), + compression=True, + ) + print(f"Grid visualizations saved for {registration_method_name.upper()}") diff --git a/pyproject.toml b/pyproject.toml index 8e7033ec..ed8dba5c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -327,6 +327,7 @@ module = [ # listed by name; add a line here when adding or renaming a tutorial. "tutorial_01_heart_gated_ct_to_usd", "tutorial_01_lung_gated_ct_to_usd", + "tutorial_02_lung_distancemap_finetune_icon", "tutorial_02_lung_finetune_icon", "tutorial_03_heart_reconstruct_highres_4d_ct", "tutorial_03_lung_reconstruct_highres_4d_ct", @@ -384,7 +385,6 @@ markers = [ "requires_gpu: marks tests that require GPU/CUDA support (opt-in via --run-gpu)", "requires_simpleware: marks tests that require a local Synopsys Simpleware Medical installation (opt-in via --run-simpleware)", "requires_physicsnemo: marks tests that require the optional [physicsnemo] extra (opt-in via --run-physicsnemo)", - "experiment: marks tests that run experiment scripts (extremely slow, manual only)", "tutorial: marks tests that run tutorial scripts (data/GPU gated, manual only)", "xdist_group: marks tests to run in the same pytest-xdist worker (prevents parallel conflicts)" ] @@ -395,6 +395,7 @@ source = ["src/physiotwin4d"] omit = [ "*/tests/*", "*/test_*", + "*/experiments/*", "*/__pycache__/*", "*/venv/*", "*/build/*", diff --git a/src/physiotwin4d/cli/convert_vtk_to_usd.py b/src/physiotwin4d/cli/convert_vtk_to_usd.py index a6d201b9..3971adfc 100644 --- a/src/physiotwin4d/cli/convert_vtk_to_usd.py +++ b/src/physiotwin4d/cli/convert_vtk_to_usd.py @@ -115,8 +115,11 @@ def main() -> int: parser.add_argument( "--anatomy-type", choices=ANATOMY_TYPES, - default="heart", - help="Anatomy material when --appearance anatomy (default: heart)", + default=None, + help=( + "Anatomy material applied to every mesh when --appearance anatomy. " + "Omit to pick a material per object from its name (default)." + ), ) parser.add_argument( "--primvar", diff --git a/src/physiotwin4d/cli/create_statistical_model.py b/src/physiotwin4d/cli/create_statistical_model.py index 863b83fd..547c11aa 100644 --- a/src/physiotwin4d/cli/create_statistical_model.py +++ b/src/physiotwin4d/cli/create_statistical_model.py @@ -39,7 +39,7 @@ def main() -> int: --sample-meshes-dir ./meshes \\ --reference-mesh average_mesh.vtk \\ --output-dir ./pca_model \\ - --pca-components 20 + --number-of-pca-components 20 """, ) @@ -75,7 +75,7 @@ def main() -> int: ) parser.add_argument( - "--pca-components", + "--number-of-pca-components", type=int, default=7, help="Number of PCA components to retain (default: 7)", @@ -142,7 +142,7 @@ def main() -> int: workflow = WorkflowCreateStatisticalModel( sample_meshes=sample_meshes, reference_mesh=reference_mesh, - pca_number_of_components=args.pca_components, + number_of_pca_components=args.number_of_pca_components, ) except (ValueError, RuntimeError) as e: print(f"Error initializing workflow: {e}") diff --git a/src/physiotwin4d/cli/fit_statistical_model_to_patient.py b/src/physiotwin4d/cli/fit_statistical_model_to_patient.py index adc16213..63df8ad8 100644 --- a/src/physiotwin4d/cli/fit_statistical_model_to_patient.py +++ b/src/physiotwin4d/cli/fit_statistical_model_to_patient.py @@ -32,7 +32,7 @@ def main() -> int: --patient-models lv.vtp rv.vtp myo.vtp \\ --patient-image patient_ct.nii.gz \\ --pca-json pca_model.json \\ - --pca-number-of-modes 10 \\ + --number-of-pca-components 10 \\ --output-dir ./results # Enable labelmap-to-image refinement (requires template labelmap and label IDs) @@ -110,10 +110,10 @@ def main() -> int: help="Path to PCA JSON file for shape-based registration (optional)", ) parser.add_argument( - "--pca-number-of-modes", + "--number-of-pca-components", type=int, default=0, - help="Number of PCA modes to use (default: 0, uses all if PCA enabled)", + help="Number of PCA components to use (default: 0, uses all if PCA enabled)", ) # Registration configuration @@ -237,7 +237,7 @@ def main() -> int: workflow.set_use_pca_registration( True, pca_model=pca_model, - pca_number_of_modes=args.pca_number_of_modes, + number_of_pca_components=args.number_of_pca_components, ) workflow.set_use_labelmap_to_labelmap_registration( diff --git a/src/physiotwin4d/cli/reconstruct_highres_4d_ct.py b/src/physiotwin4d/cli/reconstruct_highres_4d_ct.py index fbb13d8c..c44fc805 100644 --- a/src/physiotwin4d/cli/reconstruct_highres_4d_ct.py +++ b/src/physiotwin4d/cli/reconstruct_highres_4d_ct.py @@ -103,13 +103,6 @@ def main() -> int: "(default: enabled; use --no-register-reference for an identity transform)" ), ) - parser.add_argument( - "--prior-weight", - type=float, - default=0.0, - help="Weight for temporal smoothing with prior transforms [0.0-1.0] (default: 0.0)", - ) - # Registration iterations parser.add_argument( "--Greedy-iterations", @@ -221,11 +214,6 @@ def main() -> int: ) return 1 - # Validate prior weight - if not 0.0 <= args.prior_weight <= 1.0: - print(f"Error: Prior weight must be in [0.0, 1.0], got {args.prior_weight}") - return 1 - # Create output directory os.makedirs(args.output_dir, exist_ok=True) @@ -296,7 +284,6 @@ def main() -> int: # Configure registration parameters workflow.set_modality(args.modality) - workflow.set_prior_weight(args.prior_weight) workflow.set_mask_dilation(args.mask_dilation_mm) if fixed_mask is not None: diff --git a/src/physiotwin4d/contour_tools.py b/src/physiotwin4d/contour_tools.py index 67558d1f..8ec8c0c1 100644 --- a/src/physiotwin4d/contour_tools.py +++ b/src/physiotwin4d/contour_tools.py @@ -318,6 +318,54 @@ def create_labelmap_from_meshes( return labelmap_image + @staticmethod + def sample_mesh_faces(mesh: pv.DataSet, max_spacing: float) -> np.ndarray: + """Return mesh points supplemented by samples across the triangle faces. + + Rasterizing vertices alone leaves gaps between them on meshes that are + coarse relative to the voxel size, which makes a distance map built from + them ripple. Adding barycentric samples dense enough that consecutive + samples are closer than ``max_spacing`` closes those gaps. + + Args: + mesh: Source mesh; its surface is triangulated if needed. + max_spacing: Target spacing between samples, in mm. + + Returns: + (n, 3) array of sample points, starting with the mesh's own points. + """ + points = np.asarray(mesh.points, dtype=np.float64) + surface = mesh.extract_surface() if not isinstance(mesh, pv.PolyData) else mesh + surface = surface.triangulate() + if surface.faces.size == 0: + return points + faces = surface.faces.reshape(-1, 4)[:, 1:] + + vertices = np.asarray(surface.points, dtype=np.float64) + corners = vertices[faces] # (n_faces, 3, 3) + edge_lengths = np.linalg.norm( + corners - np.roll(corners, 1, axis=1), axis=2 + ).max(axis=1) + + # Group faces by how finely they need to be subdivided so each division + # level is generated as one vectorized batch. + divisions = np.maximum(1, np.ceil(edge_lengths / max(max_spacing, 1e-6))) + divisions = np.minimum(divisions, 64).astype(np.int64) + + samples = [points] + for level in np.unique(divisions): + if level < 2: + continue + selected = corners[divisions == level] + # Barycentric lattice with `level` divisions per edge. + steps = np.arange(level + 1, dtype=np.float64) / level + u, v = np.meshgrid(steps, steps, indexing="ij") + mask = (u + v) <= 1.0 + weights = np.column_stack([1.0 - u[mask] - v[mask], u[mask], v[mask]]) + samples.append(np.einsum("fca,kc->fka", selected, weights).reshape(-1, 3)) + + return np.concatenate(samples, axis=0) + def create_distance_map( self, mesh: pv.DataSet | pv.UnstructuredGrid, @@ -326,37 +374,62 @@ def create_distance_map( negative_inside: bool = True, zero_inside: bool = False, norm_to_max_distance: float = 0.0, + sample_faces: bool = True, ) -> itk.Image: - self.log_info("Computing signed distance map...") + """Compute a distance map of a mesh on the reference image's grid. + + Args: + mesh: Mesh whose surface the distances are measured to. + reference_image: Image defining the output grid. + squared_distance: Sign-preserving square of the result. Default: False + negative_inside: Keep the signed output. Default: True + zero_inside: Clip negative values to zero before anything else. + Default: False + norm_to_max_distance: If non-zero, divide by this value and clip to + [-1, 1]. Default: 0.0 (distances stay in mm) + sample_faces: Rasterize samples across the triangle faces as well as + the vertices, so that coarse meshes do not leave gaps in the + rasterized surface. Default: True - # Convert mask to binary - points = mesh.points + Returns: + ITK image of distances on the reference grid. + """ + self.log_info("Computing signed distance map...") size = reference_image.GetLargestPossibleRegion().GetSize() + if sample_faces: + points = self.sample_mesh_faces( + mesh, 0.5 * float(min(reference_image.GetSpacing())) + ) + self.log_debug( + "Distance map: %d face samples from %d mesh points", + len(points), + mesh.n_points, + ) + else: + points = np.asarray(mesh.points, dtype=np.float64) + # NumPy convention is (z, y, x); ITK GetSize() returns (x, y, z) - tmp_arr = np.zeros((size[2], size[1], size[0]), dtype=np.int32) - itk_point = itk.Point[itk.D, 3]() - point_count = 0 - for point in points: - itk_point[0] = float(point[0]) - itk_point[1] = float(point[1]) - itk_point[2] = float(point[2]) - indx = reference_image.TransformPhysicalPointToIndex(itk_point) - if ( - indx[0] < 0 - or indx[1] < 0 - or indx[2] < 0 - or indx[0] >= size[0] - or indx[1] >= size[1] - or indx[2] >= size[2] - ): - continue - tmp_arr[indx[2], indx[1], indx[0]] = 1 - point_count += 1 + tmp_arr = np.zeros((size[2], size[1], size[0]), dtype=np.uint8) + + # Bulk equivalent of TransformPhysicalPointToIndex, which rounds half up. + index_to_world = itk.array_from_matrix( + reference_image.GetDirection() + ) @ np.diag(np.asarray(reference_image.GetSpacing())) + origin = np.asarray(reference_image.GetOrigin(), dtype=np.float64) + indices = np.floor( + (points - origin) @ np.linalg.inv(index_to_world).T + 0.5 + ).astype(np.int64) + size_arr = np.array([size[0], size[1], size[2]], dtype=np.int64) + inside = np.all((indices >= 0) & (indices < size_arr), axis=1) + indices = indices[inside] + point_count = len(indices) + if point_count: + tmp_arr[indices[:, 2], indices[:, 1], indices[:, 0]] = 1 self.log_info( - "Distance map: %d/%d surface points within reference image", + "Distance map: %d/%d surface samples within reference image", point_count, len(points), ) @@ -370,8 +443,17 @@ def create_distance_map( str(size), str(reference_image.GetSpacing()), ) + elif not inside.all(): + # Distances near the dropped region are measured to whatever samples + # remain in the grid, so they are larger than the true distance. + self.log_warning( + "%d of %d surface samples fall outside the reference image; " + "distances near that boundary are overestimated.", + len(points) - point_count, + len(points), + ) - tmp_binary_image = itk.GetImageFromArray(tmp_arr.astype(np.uint8)) + tmp_binary_image = itk.GetImageFromArray(tmp_arr) tmp_binary_image.CopyInformation(reference_image) assert ( tmp_binary_image.GetLargestPossibleRegion().GetSize() @@ -545,17 +627,27 @@ def save_combined_surfaces( surface's annotation, enabling colour-by-anatomy rendering in Paraview, PyVista, etc. + It also gains a per-cell ``SegmentationLabelIds`` (int32) array, which + carries each cell's originating label ID so structure identity survives + the merge. Downstream, :class:`ConvertVTKToUSD` splits on this array + when given ``mask_ids``, giving one prim (and one anatomy material) per + structure. A surface whose ``field_data['SegmentationLabelIds']`` does + not hold exactly one ID has no per-cell attribution — that is the case + for the per-group surfaces of :class:`WorkflowConvertImageToVTK`, which + are contoured from a merged binary mask — so its cells are tagged ``0``. + Pass the per-label surfaces (``'label_surfaces'``) to get real IDs. + Per-object ``field_data`` is *not* preserved: it is per-object, so a - single merged mesh cannot carry one value per input surface. The keys - set by :meth:`WorkflowConvertImageToVTK._annotate` are therefore lost: + single merged mesh cannot carry one value per input surface. The + remaining keys set by :meth:`WorkflowConvertImageToVTK._annotate` are + therefore lost: - ``AnatomyGroup`` — group name, e.g. ``'heart'``. - ``SegmentationLabelNames`` — structure names within the group. - - ``SegmentationLabelIds`` — corresponding integer label IDs. - ``AnatomyColor`` — RGB float color (survives indirectly as the per-cell ``Color`` array). - Use :meth:`save_surfaces` instead when structure identity must be + Use :meth:`save_surfaces` instead when structure *names* must be recoverable from the saved files. Args: @@ -574,8 +666,20 @@ def save_combined_surfaces( output_dir = os.path.dirname(output_filename) if output_dir: os.makedirs(output_dir, exist_ok=True) - merged = cast( - pv.PolyData, pv.merge(list(surfaces.values()), merge_points=False) - ) + # Shallow copies so tagging does not add an array to the caller's + # surfaces; the point/cell arrays themselves stay shared. + tagged: list[pv.PolyData] = [] + for surface in surfaces.values(): + label_ids = surface.field_data.get("SegmentationLabelIds") + if label_ids is not None and len(label_ids) == 1: + label_id = int(label_ids[0]) + else: + label_id = 0 + tagged_surface = surface.copy(deep=False) + tagged_surface.cell_data["SegmentationLabelIds"] = np.full( + tagged_surface.n_cells, label_id, dtype=np.int32 + ) + tagged.append(tagged_surface) + merged = cast(pv.PolyData, pv.merge(tagged, merge_points=False)) merged.save(output_filename) return output_filename diff --git a/src/physiotwin4d/convert_vtk_to_usd.py b/src/physiotwin4d/convert_vtk_to_usd.py index ff35ba99..6ccc1831 100644 --- a/src/physiotwin4d/convert_vtk_to_usd.py +++ b/src/physiotwin4d/convert_vtk_to_usd.py @@ -14,6 +14,7 @@ from __future__ import annotations import logging +from collections import Counter from collections.abc import Sequence from pathlib import Path from typing import Any, Literal, Optional, Union, cast @@ -41,20 +42,6 @@ validate_time_series_topology, ) -_USD_EXTENSIONS = {".usd", ".usda", ".usdc"} - - -def _split_usd_extension(name: str) -> tuple[str, str]: - """Split a trailing USD extension off ``name``. - - Returns ``(name_without_extension, extension)``. ``extension`` is ``".usd"`` - when ``name`` has no recognized USD extension (``.usd``, ``.usda``, ``.usdc``). - """ - suffix = Path(name).suffix - if suffix.lower() in _USD_EXTENSIONS: - return name[: -len(suffix)], suffix - return name, ".usd" - class ConvertVTKToUSD(PhysioTwin4DBase): """ @@ -98,6 +85,7 @@ def __init__( solid_color: tuple[float, float, float] = (0.8, 0.8, 0.8), static_merge: bool = False, time_codes: Optional[list[float]] = None, + object_names: Optional[Sequence[str]] = None, segmenter: Optional[SegmentAnatomyBase] = None, log_level: int | str = logging.INFO, ) -> None: @@ -125,6 +113,12 @@ def __init__( time_codes: Explicit time codes aligned to input_polydata, used when static_merge is False. If None, uses sequential integers [0, 1, 2, ...]. + object_names: Optional prim names aligned to input_polydata, used + when static_merge is True. If None, objects are named + ``{data_basename}_{index}``. Naming objects after the + structure they hold (e.g. "heart_ventricle_left") makes + the stage self-describing and lets downstream material + assignment key off the prim name. segmenter: Optional SegmentAnatomyBase instance used to classify each mask_ids label into an anatomy group (heart / lung / bone / major_vessels / contrast / soft_tissue / other) so labeled @@ -134,12 +128,18 @@ def __init__( log_level: Logging level Raises: - ValueError: If time_codes is not None and its length does not match - input_polydata, or its values are not non-decreasing. + ValueError: If time_codes or object_names is not None and its length + does not match input_polydata, or if time_codes values are not + non-decreasing. """ super().__init__(class_name=self.__class__.__name__, log_level=log_level) - self.data_basename, _ = _split_usd_extension(data_basename) + suffix = Path(data_basename).suffix + self.data_basename = ( + data_basename[: -len(suffix)] + if suffix.lower() in {".usd", ".usda", ".usdc"} + else data_basename + ) self.input_polydata = list(input_polydata) self.mask_ids = mask_ids self.compute_normals = compute_normals @@ -166,8 +166,28 @@ def __init__( "time_codes must be in non-decreasing order; " "got values that decrease between consecutive frames" ) + if object_names is not None: + if len(object_names) != len(self.input_polydata): + raise ValueError( + f"object_names length ({len(object_names)}) must match " + f"input_polydata length ({len(self.input_polydata)})" + ) + # Each name becomes a prim path component, so it has to be a legal + # USD identifier and unique or prims silently collide. + invalid = [n for n in object_names if not Sdf.Path.IsValidIdentifier(n)] + counts = Counter(object_names) + duplicated = sorted(n for n, count in counts.items() if count > 1) + if invalid or duplicated: + raise ValueError( + "object_names must be unique and valid USD prim names " + "(letter or underscore followed by letters, digits or " + f"underscores); invalid: {invalid}, duplicated: {duplicated}" + ) self._is_static_merge: bool = static_merge self._time_codes: Optional[list[float]] = time_codes + self.object_names: Optional[list[str]] = ( + list(object_names) if object_names is not None else None + ) # Pre-converted MeshData for each time step; populated by from_files() so # _convert_unified() can reuse the topology-validation work instead of # calling _vtk_to_mesh_data() a second time. @@ -734,7 +754,11 @@ def _convert_static_merge( ) for i, vtk_mesh in enumerate(self.input_polydata): mesh_data = self._vtk_to_mesh_data(vtk_mesh, i) - frame_name = f"{self.data_basename}_{i}" + frame_name = ( + self.object_names[i] + if self.object_names is not None + else f"{self.data_basename}_{i}" + ) if self.separate_by == "none": parts = [(mesh_data, frame_name)] @@ -935,13 +959,20 @@ def _split_by_labels( if isinstance(vtk_mesh, pv.UnstructuredGrid) and self.convert_to_surface: vtk_mesh = vtk_mesh.extract_surface(algorithm="dataset_surface") - # Get boundary labels - if "boundary_labels" not in vtk_mesh.cell_data: - self.logger.warning("No 'boundary_labels' array found - using unified mesh") + # Get per-cell label IDs. 'SegmentationLabelIds' is written by + # ContourTools.save_combined_surfaces when merging per-label surfaces; + # 'boundary_labels' comes from contouring a multi-label labelmap. + if "SegmentationLabelIds" in vtk_mesh.cell_data: + label_array = vtk_mesh.cell_data["SegmentationLabelIds"] + elif "boundary_labels" in vtk_mesh.cell_data: + label_array = vtk_mesh.cell_data["boundary_labels"] + else: + self.log_warning( + "No 'SegmentationLabelIds' or 'boundary_labels' array found " + "- using unified mesh" + ) return {"default": self._vtk_to_mesh_data(vtk_mesh, time_idx)} - label_array = vtk_mesh.cell_data["boundary_labels"] - # Create submeshes for each label labeled_meshes = {} for label_id, label_name in mask_ids.items(): diff --git a/src/physiotwin4d/image_tools.py b/src/physiotwin4d/image_tools.py index 0e5c4e22..284c71fb 100644 --- a/src/physiotwin4d/image_tools.py +++ b/src/physiotwin4d/image_tools.py @@ -6,7 +6,7 @@ """ import logging -from typing import Any, Optional, Union, overload +from typing import Any, Optional, Union, cast, overload import itk import numpy as np @@ -270,6 +270,148 @@ def make_isotropic_image(self, image: itk.Image) -> itk.Image: result.DisconnectPipeline() return result + @staticmethod + def _per_axis_values( + value: Union[float, int, list, tuple, NDArray[Any]], + dimension: int, + name: str, + ) -> list[float]: + """Broadcast a scalar to every axis, or validate a per-axis sequence. + + Args: + value: Scalar applied to every axis, or one value per axis. + dimension: Number of image dimensions expected. + name: Parameter name, used in error messages. + + Returns: + One value per axis. + + Raises: + ValueError: If a sequence has the wrong length, or any value is + negative. + """ + if np.isscalar(value): + values = [float(cast(float, value))] * dimension + else: + values = [float(v) for v in value] # type: ignore[union-attr] + if len(values) != dimension: + raise ValueError( + f"{name} needs a scalar or one value per image dimension " + f"({dimension}), got {len(values)}." + ) + if any(v < 0.0 for v in values): + raise ValueError(f"{name} must be >= 0, got {values}") + return values + + def pad_image( + self, + image: itk.Image, + pad_portion: Optional[ + Union[float, list[float], tuple[float, ...], NDArray[Any]] + ] = None, + pad_voxels: Optional[ + Union[int, list[int], tuple[int, ...], NDArray[Any]] + ] = None, + background_value: float = 0.0, + ) -> itk.Image: + """Pad *image* on every side with a constant-valued margin. + + The margin is given either as a portion of each axis' physical extent + (*pad_portion*) or directly in voxels (*pad_voxels*); exactly one of the + two must be supplied. Either accepts a scalar, applied to every axis, or + one value per image dimension. Both pad the lower and the upper end of + every axis, so ``pad_voxels=10`` grows all six faces of a 3-D image by + ten voxels. Spacing and direction are untouched. + + The origin and size are updated together so the original voxels keep + their physical positions: the padded image's index ``(0, 0, 0)`` sits one + margin below the input's, and the input data occupies the interior. + (``itk.ConstantPadImageFilter`` alone reports the margin as a negative + start index instead, which most file formats drop on write — shifting the + data. The region-of-interest pass here folds that index back into the + origin.) + + Args: + image: ITK image to pad. + pad_portion: Portion of an axis' physical extent (``size * spacing``) + to add at both ends, as a scalar for every axis or one value per + dimension: ``0.1`` grows an axis spanning 200 mm by 20 mm per + side. Rounded up to whole voxels. Mutually exclusive with + *pad_voxels*. + pad_voxels: Margin in voxels, as a scalar for every axis or one value + per dimension, applied at both ends of each axis. Mutually + exclusive with *pad_portion*. + background_value: Pixel value written into the new margin + (default: 0.0). + + Returns: + Padded image with the same pixel type, spacing and direction. + + Raises: + ValueError: If neither or both of *pad_portion* and *pad_voxels* are + given, if either is negative, or if a sequence does not have one + entry per image dimension. + """ + if (pad_portion is None) == (pad_voxels is None): + raise ValueError( + "Specify exactly one of pad_portion or pad_voxels; got " + f"pad_portion={pad_portion}, pad_voxels={pad_voxels}." + ) + + size = [int(s) for s in image.GetLargestPossibleRegion().GetSize()] + if pad_portion is not None: + portions = self._per_axis_values(pad_portion, len(size), "pad_portion") + # extent_i = size_i * spacing_i, and pad_portion * extent_i of margin + # is that distance divided by spacing_i, so the spacing cancels. + margin = [int(np.ceil(p * s)) for p, s in zip(portions, size)] + self.log_info( + "Padding by %s voxels per side (%s of extent); size %s -> %s", + margin, + [f"{p * 100.0:.1f}%" for p in portions], + size, + [s + 2 * p for s, p in zip(size, margin)], + ) + else: + assert pad_voxels is not None # guaranteed by the check above + margin = [ + int(v) + for v in self._per_axis_values(pad_voxels, len(size), "pad_voxels") + ] + self.log_info( + "Padding by %s voxels per side; size %s -> %s", + margin, + size, + [s + 2 * p for s, p in zip(size, margin)], + ) + + ImageType = type(image) + # SetConstant is typed to the pixel type, so an integer image rejects a + # Python float. + pixel_type = itk.template(image)[1][0] + constant = ( + float(background_value) + if pixel_type in (itk.F, itk.D) + else int(round(background_value)) + ) + + pad_filter = itk.ConstantPadImageFilter[ImageType, ImageType].New() + pad_filter.SetInput(image) + pad_filter.SetPadLowerBound(margin) + pad_filter.SetPadUpperBound(margin) + pad_filter.SetConstant(constant) + pad_filter.Update() + padded = pad_filter.GetOutput() + + # Re-anchor the padded region at index 0, moving the margin into the + # origin so the physical position of the original data is preserved. + roi_filter = itk.RegionOfInterestImageFilter[ImageType, ImageType].New() + roi_filter.SetInput(padded) + roi_filter.SetRegionOfInterest(padded.GetLargestPossibleRegion()) + roi_filter.Update() + result = roi_filter.GetOutput() + result.DisconnectPipeline() + return result + def binary_dilate_image( self, image: itk.Image, diff --git a/src/physiotwin4d/physicsnemo_tools.py b/src/physiotwin4d/physicsnemo_tools.py index 19a1c3f1..346d9865 100644 --- a/src/physiotwin4d/physicsnemo_tools.py +++ b/src/physiotwin4d/physicsnemo_tools.py @@ -405,8 +405,8 @@ def _target_values(self, path: Path) -> np.ndarray: f"{path} has {values.shape[0]} points, expected {self._n_points}." ) + self._cache[path] = values if self._cache_max_samples != 0: - self._cache[path] = values while len(self._cache) > self._cache_max_samples: self._cache.popitem(last=False) return values diff --git a/src/physiotwin4d/register_images_ants.py b/src/physiotwin4d/register_images_ants.py index 5c2f013e..cd32134d 100644 --- a/src/physiotwin4d/register_images_ants.py +++ b/src/physiotwin4d/register_images_ants.py @@ -510,7 +510,6 @@ def registration_method( moving_mask: Optional[itk.Image] = None, moving_labelmap: Optional[itk.Image] = None, moving_image_pre: Optional[itk.Image] = None, - initial_forward_transform: Optional[itk.Transform] = None, ) -> dict[str, Union[itk.Transform, float]]: """Register moving image to fixed image using ANTs registration algorithm. @@ -524,14 +523,6 @@ def registration_method( region of interest in the moving image moving_image_pre (itk.Image, optional): Pre-processed moving image. If None, preprocessing is performed automatically - initial_forward_transform (itk.Transform, optional): Initial - forward transform (same convention as the returned - forward_transform: used to warp the moving image onto the fixed - grid). Can be any ITK transform type (Affine, Rigid, - DisplacementField, Composite, etc.). It is applied by pre-warping - the moving image onto the fixed grid before registration; the - returned transforms compose this initial alignment with the - registration refinement. Returns: dict: Dictionary containing: @@ -548,20 +539,15 @@ def registration_method( consistent. The forward and inverse transforms are stored separately by ANTs. - IMPORTANT: the initial transform is applied by pre-warping the - moving image onto the fixed grid (the same approach as - RegisterImagesICON) rather than via ants.registration's - initial_transform argument, which mishandles matrix (affine/ - translation) initials. This method composes the initial transform - with the registration result, so the returned transforms include - both the initial alignment and the registration refinement. + To seed the registration with a known alignment, use + RegisterImagesBase.register_from(), which handles the pre-warp and + the composition. Implementation details: - Uses ANTs registration with configurable transform types - Supports multi-resolution optimization - Handles masked and unmasked registration - Returns ITK-compatible displacement field transforms - - Initial transforms are converted from ITK to ANTs format automatically Example: >>> # Basic registration @@ -573,9 +559,9 @@ def registration_method( >>> registrar.set_fixed_mask(heart_mask_fixed) >>> result = registrar.register(moving_image, moving_mask=heart_mask_moving) >>> - >>> # Registration with initial transform + >>> # Registration seeded with a known alignment >>> initial_tfm = itk.AffineTransform[itk.D, 3].New() - >>> result = registrar.register(moving_image, initial_forward_transform=initial_tfm) + >>> result = registrar.register_from(initial_tfm, moving_image) """ if moving_image is not None: self.moving_image = moving_image @@ -591,22 +577,6 @@ def registration_method( if self.fixed_image_pre is None: self.fixed_image_pre = self.preprocess(self.fixed_image, self.modality) - if initial_forward_transform is not None: - self.log_info("Pre-warping moving image with initial transform...") - transform_tools = TransformTools() - self.moving_image_pre = transform_tools.transform_image( - self.moving_image_pre, - initial_forward_transform, - self.fixed_image, - ) - if self.moving_mask is not None: - self.moving_mask = transform_tools.transform_image( - self.moving_mask, - initial_forward_transform, - self.fixed_image, - interpolation_method="nearest", - ) - transform_type = None if self.transform_type == "Deformable": transform_type = "antsRegistrationSyNQuick[so]" @@ -706,46 +676,8 @@ def registration_method( reference_image=self.moving_image, ) - # Important: ANTs does NOT include the initial_transform in the output transforms - # We need to manually compose them - if initial_forward_transform is not None: - self.log_info("Composing initial transform with registration result...") - - # For forward_transform (Moving -> Fixed): Apply initial_forward_transform first, then registration - # Transform order: point -> initial_forward_transform -> forward_reg - forward_transform = itk.CompositeTransform[itk.D, 3].New() - forward_transform.AddTransform(initial_forward_transform) - # Add transforms from registration result (may be composite) - if isinstance(forward_reg, itk.CompositeTransform[itk.D, 3]): - for i in range(forward_reg.GetNumberOfTransforms()): - forward_transform.AddTransform(forward_reg.GetNthTransform(i)) - else: - forward_transform.AddTransform(forward_reg) - - # For inverse_transform (Fixed -> Moving): Apply registration inverse first, then initial inverse - # Transform order: point -> inverse_reg -> initial_forward_transform^(-1) - inverse_transform = itk.CompositeTransform[itk.D, 3].New() - # Add registration inverse transforms - if isinstance(inverse_reg, itk.CompositeTransform[itk.D, 3]): - for i in range(inverse_reg.GetNumberOfTransforms()): - inverse_transform.AddTransform(inverse_reg.GetNthTransform(i)) - else: - inverse_transform.AddTransform(inverse_reg) - # Invert and add initial transform - # For displacement field transforms, we need to invert properly - transform_tools = TransformTools() - initial_inverse = transform_tools.invert_displacement_field_transform( - transform_tools.convert_transform_to_displacement_field_transform( - initial_forward_transform, self.moving_image - ) - ) - inverse_transform.AddTransform(initial_inverse) - - self.log_info("Transforms composed successfully") - else: - # No initial transform, use registration results directly - forward_transform = forward_reg - inverse_transform = inverse_reg + forward_transform = forward_reg + inverse_transform = inverse_reg moving_image_reg = registration_result["warpedmovout"] loss = ants.image_similarity( self._itk_to_ants_image(self.fixed_image), diff --git a/src/physiotwin4d/register_images_base.py b/src/physiotwin4d/register_images_base.py index 64e0d7c1..7f3c1971 100644 --- a/src/physiotwin4d/register_images_base.py +++ b/src/physiotwin4d/register_images_base.py @@ -236,7 +236,6 @@ def registration_method( moving_mask: Optional[itk.Image] = None, moving_labelmap: Optional[itk.Image] = None, moving_image_pre: Optional[itk.Image] = None, - initial_forward_transform: Optional[itk.Transform] = None, ) -> dict[str, Union[itk.Transform, float]]: """Main registration method to align moving image to fixed image. @@ -252,7 +251,6 @@ def registration_method( moving_mask (itk.image, optional): Binary mask for moving image ROI moving_labelmap (itk.image, optional): Multi-label segmentation for the moving image moving_image_pre (itk.image, optional): Preprocessed moving image - initial_forward_transform (itk.Transform, optional): Initial transformation from moving to fixed Returns: dict: Dictionary containing: @@ -274,7 +272,6 @@ def register( moving_mask: Optional[itk.Image] = None, moving_labelmap: Optional[itk.Image] = None, moving_image_pre: Optional[itk.Image] = None, - initial_forward_transform: Optional[itk.Transform] = None, ) -> dict[str, Union[itk.Transform, float]]: """Register a moving image to the fixed image. @@ -282,12 +279,14 @@ def register( concrete subclasses. It should align the moving image to the fixed image using the specific algorithm implemented by the subclass. + To start from a known alignment, use :meth:`register_from` rather than + seeding the backend directly. + Args: moving_image (itk.image): The 3D image to be registered to the fixed image moving_mask (itk.image, optional): Binary mask for moving image ROI moving_labelmap (itk.image, optional): Multi-label segmentation for the moving image moving_image_pre (itk.image, optional): Preprocessed moving image - initial_forward_transform (itk.Transform, optional): Initial transformation from moving to fixed Returns: dict: Dictionary containing transformation results: @@ -347,7 +346,6 @@ def register( moving_mask=new_moving_mask, moving_labelmap=moving_labelmap, moving_image_pre=moving_image_pre, - initial_forward_transform=initial_forward_transform, ) self.forward_transform = result["forward_transform"] @@ -360,6 +358,190 @@ def register( "loss": self.loss, } + def register_from( + self, + initial_forward_transform: itk.Transform, + moving_image: itk.Image, + moving_mask: Optional[itk.Image] = None, + moving_labelmap: Optional[itk.Image] = None, + ) -> dict[str, Union[itk.Transform, float]]: + """Register starting from a known alignment. + + The moving data is warped onto the fixed grid by + ``initial_forward_transform`` first, :meth:`register` then measures only + the residual misalignment, and the two are composed. This is the single + supported way to seed a registration: doing it here rather than inside + each backend keeps the pre-warp, the composition and the inversion + identical for every algorithm. + + The image, the mask and the labelmap are all pre-warped, so they stay in + the same frame as each other; the mask and labelmap use nearest-neighbor + interpolation to preserve their discrete values. + + Args: + initial_forward_transform: Starting alignment, in the same + convention as the returned ``forward_transform`` -- it warps the + moving image onto the fixed grid. + moving_image: The 3D image to be registered to the fixed image. + moving_mask: Binary mask for the moving image ROI. + moving_labelmap: Multi-label segmentation for the moving image. + + Returns: + dict: Same keys as :meth:`register`, with the transforms composed so + they map between the *original* moving image and the fixed image. + + Raises: + ValueError: If the fixed image has not been set. + """ + warped_image, warped_mask, warped_labelmap = self._prewarp_moving( + initial_forward_transform, moving_image, moving_mask, moving_labelmap + ) + result = self.register( + warped_image, + moving_mask=warped_mask, + moving_labelmap=warped_labelmap, + ) + composed = self._compose_with_initial( + initial_forward_transform, result, moving_image + ) + + # register() left the pre-warped image on self; the composed transforms + # are defined against the original, so restore it and drop any + # registered-image cache built for the pre-warped one. + self.moving_image = moving_image + self.moving_image_registered = None + + self.forward_transform = composed["forward_transform"] + self.inverse_transform = composed["inverse_transform"] + self.loss = composed["loss"] + return composed + + def _prewarp_moving( + self, + initial_forward_transform: itk.Transform, + moving_image: itk.Image, + moving_mask: Optional[itk.Image], + moving_labelmap: Optional[itk.Image], + ) -> tuple[itk.Image, Optional[itk.Image], Optional[itk.Image]]: + """Warp the moving image, mask and labelmap onto the fixed grid. + + Args: + initial_forward_transform: Alignment to apply, in the image-warp + convention. + moving_image: Raw moving image. + moving_mask: Moving mask, or None. + moving_labelmap: Moving labelmap, or None. + + Returns: + Tuple of the warped ``(image, mask, labelmap)``, the latter two None + when not supplied. The mask and labelmap are warped with + nearest-neighbor interpolation to keep their discrete values. + + Raises: + ValueError: If the fixed image has not been set. + """ + if self.fixed_image is None: + raise ValueError("Fixed image must be set before registration.") + + transform_tools = TransformTools() + self.log_info("Pre-warping moving data with the initial transform...") + + def _warp(image: Optional[itk.Image], nearest: bool) -> Optional[itk.Image]: + if image is None: + return None + return transform_tools.transform_image( + image, + initial_forward_transform, + self.fixed_image, + interpolation_method="nearest" if nearest else "linear", + ) + + return ( + _warp(moving_image, nearest=False), + _warp(moving_mask, nearest=True), + _warp(moving_labelmap, nearest=True), + ) + + def _compose_with_initial( + self, + initial_forward_transform: itk.Transform, + result: dict[str, Union[itk.Transform, float]], + moving_image: itk.Image, + ) -> dict[str, Union[itk.Transform, float]]: + """Compose a residual registration result onto its initial transform. + + Args: + initial_forward_transform: The alignment the moving data was + pre-warped by. + result: Result of registering the pre-warped data. + moving_image: Raw moving image, whose grid defines the domain the + initial transform is inverted over. + + Returns: + The result dict with both transforms mapping between the *original* + moving image and the fixed image. + """ + transform_tools = TransformTools() + + # The registration measured the residual from the pre-warped position, + # so the total is the initial transform followed by that residual. An + # itk.CompositeTransform applies its transforms in reverse order of + # addition, so adding the initial first makes the residual apply first -- + # which is what the image-warp direction needs: a fixed-grid sample is + # mapped by the residual, then by the initial transform, to land in the + # original moving image. + forward_transform = itk.CompositeTransform[itk.D, 3].New() + self._add_transform_flattened(forward_transform, initial_forward_transform) + self._add_transform_flattened( + forward_transform, cast(itk.Transform, result["forward_transform"]) + ) + + # The inverse runs the other way -- a moving-grid sample is mapped by the + # initial transform's inverse into the pre-warped frame, then by the + # residual's inverse into the fixed image -- so the additions are + # reversed too. + initial_inverse = transform_tools.invert_transform( + initial_forward_transform, moving_image + ) + inverse_transform = itk.CompositeTransform[itk.D, 3].New() + self._add_transform_flattened( + inverse_transform, cast(itk.Transform, result["inverse_transform"]) + ) + self._add_transform_flattened(inverse_transform, initial_inverse) + + return { + "forward_transform": forward_transform, + "inverse_transform": inverse_transform, + "loss": result["loss"], + } + + @staticmethod + def _add_transform_flattened( + composite: itk.CompositeTransform, transform: itk.Transform + ) -> None: + """Append a transform to a composite, splicing in nested composites. + + itk.HDF5TransformIO refuses to write a CompositeTransform that holds + another CompositeTransform ("Composite Transform can only be 1st + transform in a file"), which every multi-stage registration would + otherwise produce: RegisterImagesGreedy already returns an affine+warp + composite, and composing a residual onto it would nest that composite. + + Splicing the sub-transforms in at the position their composite occupied + leaves the mapping unchanged, since itk.CompositeTransform applies its + queue back to front either way. + + The down_cast is required: ITK hands back base-typed ``itkTransformD33`` + Python objects from ``GetInverseTransform()`` and ``GetNthTransform()``, + which carry none of CompositeTransform's methods. + """ + transform = itk.down_cast(transform) + if isinstance(transform, itk.CompositeTransform[itk.D, 3]): + for i in range(transform.GetNumberOfTransforms()): + composite.AddTransform(transform.GetNthTransform(i)) + else: + composite.AddTransform(transform) + def _delegate_to( self, other: "RegisterImagesBase", diff --git a/src/physiotwin4d/register_images_chain.py b/src/physiotwin4d/register_images_chain.py index 0e5fab50..d05a719a 100644 --- a/src/physiotwin4d/register_images_chain.py +++ b/src/physiotwin4d/register_images_chain.py @@ -14,8 +14,9 @@ class RegisterImagesChain(RegisterImagesBase): - """Run an ordered list of registrars in sequence, feeding each stage's - forward_transform as the next stage's initial_forward_transform. + """Run an ordered list of registrars in sequence, each stage refining the + previous stage's forward_transform via + :meth:`RegisterImagesBase.register_from`. Use this to combine independent registration backends into a multi-stage pipeline (e.g. a fast coarse registrar followed by a refinement stage). @@ -67,12 +68,14 @@ def registration_method( moving_mask: Optional[itk.Image] = None, moving_labelmap: Optional[itk.Image] = None, moving_image_pre: Optional[itk.Image] = None, - initial_forward_transform: Optional[itk.Transform] = None, ) -> dict[str, Union[itk.Transform, float]]: """Run each registrar in ``self.registrars`` in order. - Each stage's ``forward_transform`` becomes the next stage's - ``initial_forward_transform``. + The first stage registers the raw moving image; every later stage sees + the moving data pre-warped by the running result and contributes only a + refinement, which is composed back on -- the same mechanics as + :meth:`RegisterImagesBase.register_from`, run through the delegated + ``registration_method`` path so masks are not re-converted per stage. Note: ``moving_image_pre`` is ignored: each stage may need different @@ -87,25 +90,40 @@ def registration_method( moving_labelmap (itk.image, optional): Multi-label segmentation for the moving image moving_image_pre (itk.image, optional): Ignored - see Note above - initial_forward_transform (itk.Transform, optional): Initial - transformation from moving to fixed, used to initialize the - first stage Returns: dict: The last stage's result dict (see :meth:`RegisterImagesBase.register`) """ - current_initial = initial_forward_transform + current_forward: Optional[itk.Transform] = None result: dict[str, Union[itk.Transform, float]] = {} for registrar in self.registrars: - self._delegate_to(registrar, moving_image, moving_mask, moving_labelmap) - result = registrar.registration_method( - moving_image, - moving_mask=moving_mask, - moving_labelmap=moving_labelmap, + if current_forward is None: + stage_image, stage_mask, stage_labelmap = ( + moving_image, + moving_mask, + moving_labelmap, + ) + else: + stage_image, stage_mask, stage_labelmap = self._prewarp_moving( + current_forward, moving_image, moving_mask, moving_labelmap + ) + + self._delegate_to(registrar, stage_image, stage_mask, stage_labelmap) + stage_result = registrar.registration_method( + stage_image, + moving_mask=stage_mask, + moving_labelmap=stage_labelmap, moving_image_pre=None, - initial_forward_transform=current_initial, ) - self._capture_delegate_result(registrar, result) - current_initial = cast(itk.Transform, result["forward_transform"]) + self._capture_delegate_result(registrar, stage_result) + + result = ( + stage_result + if current_forward is None + else self._compose_with_initial( + current_forward, stage_result, moving_image + ) + ) + current_forward = cast(itk.Transform, result["forward_transform"]) return result diff --git a/src/physiotwin4d/register_images_greedy.py b/src/physiotwin4d/register_images_greedy.py index d5af2108..bf745ce0 100644 --- a/src/physiotwin4d/register_images_greedy.py +++ b/src/physiotwin4d/register_images_greedy.py @@ -237,13 +237,38 @@ def _write_affine_matrix_file(self, mat_4x4: NDArray[np.float64]) -> str: self.log_debug("Wrote Greedy affine init matrix to %s", path) return path + # RAS <-> LPS basis change: negate x and y, leave z. Its own inverse. + _RAS_TO_LPS = np.diag([-1.0, -1.0, 1.0]) + def _matrix_to_itk_affine(self, mat_4x4: NDArray[np.float64]) -> itk.Transform: - """Convert 4x4 affine matrix to ITK AffineTransform.""" + """Convert Greedy's 4x4 RAS affine matrix to an ITK AffineTransform. + + Greedy works in RAS (its logs print "Final RAS Transform" and its + ``.mat`` files are RAS), while ITK -- and every transform this project + stores or applies -- is LPS. The matrix therefore has to change basis on + the way across: ``M_lps = F M_ras F`` and ``t_lps = F t_ras`` with + ``F = diag(-1, -1, 1)``. Skipping this negates the x and y components of + every Greedy result, which for a near-identity registration is a + sub-millimetre error that is easy to miss and impossible to correct + downstream. + + Only the affine needs this. Greedy's displacement fields come back as + images carrying the input's own LPS metadata, so + :meth:`_sitk_warp_to_itk_displacement_transform` passes them through + unchanged. See docs/developer/transform_conventions. + + Args: + mat_4x4: 4x4 affine matrix in Greedy's RAS convention. + + Returns: + The equivalent ITK affine transform, in LPS. + """ mat_4x4 = np.asarray(mat_4x4, dtype=np.float64) if mat_4x4.shape != (4, 4): raise ValueError(f"Expected 4x4 matrix, got shape {mat_4x4.shape}") - M = mat_4x4[:3, :3] - t = mat_4x4[:3, 3] + flip = self._RAS_TO_LPS + M = flip @ mat_4x4[:3, :3] @ flip + t = flip @ mat_4x4[:3, 3] center = itk.Point[itk.D, 3]() for i in range(3): center[i] = 0.0 @@ -283,7 +308,6 @@ def _registration_method_affine_or_rigid( moving_mask_sitk: Optional[Any] = None, fixed_labelmap_sitk: Optional[Any] = None, moving_labelmap_sitk: Optional[Any] = None, - initial_affine: Optional[NDArray[np.float64]] = None, ) -> tuple[NDArray[np.float64], float]: """Run Greedy affine or rigid registration. Returns (4x4 matrix, loss).""" Greedy3D = _try_import_greedy() @@ -307,19 +331,8 @@ def _registration_method_affine_or_rigid( cmd += " -gm fixed_mask -mm moving_mask" kwargs["fixed_mask"] = fixed_mask_sitk kwargs["moving_mask"] = moving_mask_sitk - # Greedy crashes (heap corruption) when an initial affine is passed as an - # in-memory matrix; write it to a temp file and pass the path instead. - initial_affine_file: Optional[str] = None - if initial_affine is not None: - initial_affine_file = self._write_affine_matrix_file(initial_affine) - cmd += f" -ia {initial_affine_file}" - self.log_debug("Greedy affine/rigid command: %s", cmd) - try: - g.execute(cmd, **kwargs) - finally: - if initial_affine_file is not None: - os.remove(initial_affine_file) + g.execute(cmd, **kwargs) mat = np.array(g["aff_out"], dtype=np.float64) try: ml = g.metric_log() @@ -339,36 +352,35 @@ def _registration_method_deformable( moving_mask_sitk: Optional[Any] = None, fixed_labelmap_sitk: Optional[Any] = None, moving_labelmap_sitk: Optional[Any] = None, - initial_affine: Optional[NDArray[np.float64]] = None, ) -> tuple[Optional[NDArray[np.float64]], Any, float]: """Run Greedy deformable registration. Returns (affine 4x4 or None, warp_sitk, loss).""" Greedy3D = _try_import_greedy() g = Greedy3D() - # Optional affine init (uses configured metric) - if initial_affine is None: - cmd_aff = "-d 3" - if fixed_labelmap_sitk is not None and moving_labelmap_sitk is not None: - cmd_aff += " -w 0.60" - cmd_aff += " -i fixed moving" - kwargs_aff = { - "fixed": fixed_sitk, - "moving": moving_sitk, - } - if fixed_labelmap_sitk is not None and moving_labelmap_sitk is not None: - cmd_aff += " -w 0.40 -i fixed_labelmap moving_labelmap" - kwargs_aff["fixed_labelmap"] = fixed_labelmap_sitk - kwargs_aff["moving_labelmap"] = moving_labelmap_sitk - cmd_aff += f" -a -dof 12 -n {iterations_str} -m {metric_str} -o aff_init" - kwargs_aff["aff_init"] = None - if fixed_mask_sitk is not None and moving_mask_sitk is not None: - cmd_aff += " -gm fixed_mask -mm moving_mask" - kwargs_aff["fixed_mask"] = fixed_mask_sitk - kwargs_aff["moving_mask"] = moving_mask_sitk - self.log_debug("Greedy deformable affine-init command: %s", cmd_aff) - g.execute(cmd_aff, **kwargs_aff) - initial_affine = np.array(g["aff_init"], dtype=np.float64) - self.log_info("Greedy deformable affine init complete") + # Greedy seeds its own deformable stage with an affine pass, using the + # configured metric. + cmd_aff = "-d 3" + if fixed_labelmap_sitk is not None and moving_labelmap_sitk is not None: + cmd_aff += " -w 0.60" + cmd_aff += " -i fixed moving" + kwargs_aff: dict[str, Any] = { + "fixed": fixed_sitk, + "moving": moving_sitk, + } + if fixed_labelmap_sitk is not None and moving_labelmap_sitk is not None: + cmd_aff += " -w 0.40 -i fixed_labelmap moving_labelmap" + kwargs_aff["fixed_labelmap"] = fixed_labelmap_sitk + kwargs_aff["moving_labelmap"] = moving_labelmap_sitk + cmd_aff += f" -a -dof 12 -n {iterations_str} -m {metric_str} -o aff_init" + kwargs_aff["aff_init"] = None + if fixed_mask_sitk is not None and moving_mask_sitk is not None: + cmd_aff += " -gm fixed_mask -mm moving_mask" + kwargs_aff["fixed_mask"] = fixed_mask_sitk + kwargs_aff["moving_mask"] = moving_mask_sitk + self.log_debug("Greedy deformable affine-init command: %s", cmd_aff) + g.execute(cmd_aff, **kwargs_aff) + initial_affine = np.array(g["aff_init"], dtype=np.float64) + self.log_info("Greedy deformable affine init complete") # Greedy crashes (heap corruption) when the affine init is passed as an # in-memory matrix via -it; write it to a temp file and pass the path. @@ -415,13 +427,11 @@ def registration_method( moving_mask: Optional[itk.Image] = None, moving_labelmap: Optional[itk.Image] = None, moving_image_pre: Optional[itk.Image] = None, - initial_forward_transform: Optional[itk.Transform] = None, ) -> dict[str, Union[itk.Transform, float]]: """Register moving image to fixed image using Greedy. Converts ITK images to SimpleITK, runs Greedy (affine and/or deformable), - then converts outputs back to ITK transforms. Composes with - initial_forward_transform when provided. + then converts outputs back to ITK transforms. Returns a dict with "forward_transform", "inverse_transform", and "loss". As with the other image-registration backends, @@ -487,24 +497,6 @@ def registration_method( iterations_str = self._greedy_iterations_str() metric_str = self._greedy_metric() - # Optional initial transform: convert ITK -> 4x4 for Greedy - initial_affine: Optional[NDArray[np.float64]] = None - if initial_forward_transform is not None: - # If it's affine-like, extract 4x4; else convert to displacement and skip for Greedy init - if hasattr(initial_forward_transform, "GetMatrix"): - M = np.eye(4, dtype=np.float64) - M[:3, :3] = np.asarray(initial_forward_transform.GetMatrix()).reshape( - 3, 3 - ) - if hasattr(initial_forward_transform, "GetTranslation"): - M[:3, 3] = np.asarray(initial_forward_transform.GetTranslation()) - if hasattr(initial_forward_transform, "GetCenter"): - c = np.asarray(initial_forward_transform.GetCenter()) - M[:3, 3] += c - M[:3, :3] @ c - initial_affine = M - # Non-affine initial: we could convert to disp field and pass; for simplicity we skip Greedy init - # and compose at the end (same as ANTs). - forward_transform: itk.Transform inverse_transform: itk.Transform loss_val: float @@ -520,7 +512,6 @@ def registration_method( iterations_str=iterations_str, metric_str=metric_str, dof=6, - initial_affine=initial_affine, ) forward_transform = self._matrix_to_itk_affine(mat) inverse_affine = itk.AffineTransform[itk.D, 3].New() @@ -537,7 +528,6 @@ def registration_method( iterations_str=iterations_str, metric_str=metric_str, dof=12, - initial_affine=initial_affine, ) forward_transform = self._matrix_to_itk_affine(mat) inverse_affine = itk.AffineTransform[itk.D, 3].New() @@ -554,7 +544,6 @@ def registration_method( moving_labelmap_sitk=moving_labelmap_sitk, iterations_str=iterations_str, metric_str=metric_str, - initial_affine=initial_affine, ) aff_tfm = ( self._matrix_to_itk_affine(aff_mat) if aff_mat is not None else None @@ -597,26 +586,6 @@ def registration_method( inverse_composite.AddTransform(inv_aff) inverse_transform = inverse_composite - # Compose with user-provided initial transform (same semantics as ANTs) - if initial_forward_transform is not None: - transform_tools = TransformTools() - forward_composite = itk.CompositeTransform[itk.D, 3].New() - forward_composite.AddTransform(initial_forward_transform) - forward_composite.AddTransform(forward_transform) - initial_disp = ( - transform_tools.convert_transform_to_displacement_field_transform( - initial_forward_transform, self.moving_image - ) - ) - inv_initial = transform_tools.invert_displacement_field_transform( - initial_disp - ) - inverse_composite = itk.CompositeTransform[itk.D, 3].New() - inverse_composite.AddTransform(inverse_transform) - inverse_composite.AddTransform(inv_initial) - forward_transform = forward_composite - inverse_transform = inverse_composite - return { "forward_transform": forward_transform, "inverse_transform": inverse_transform, diff --git a/src/physiotwin4d/register_images_icon.py b/src/physiotwin4d/register_images_icon.py index 1c964081..6b71a627 100644 --- a/src/physiotwin4d/register_images_icon.py +++ b/src/physiotwin4d/register_images_icon.py @@ -18,7 +18,6 @@ import numpy as np from .register_images_base import RegisterImagesBase -from .transform_tools import TransformTools DEFAULT_FINETUNE_LEARNING_RATE = 2e-5 @@ -196,7 +195,6 @@ def registration_method( moving_mask: Optional[itk.Image] = None, moving_labelmap: Optional[itk.Image] = None, moving_image_pre: Optional[itk.Image] = None, - initial_forward_transform: Optional[itk.Transform] = None, ) -> dict[str, Union[itk.Transform, float]]: """Register moving image to fixed image using ICON registration algorithm. @@ -212,9 +210,6 @@ def registration_method( fixed_mask, enables mask-constrained registration moving_image_pre (itk.image, optional): Pre-processed moving image. If None, preprocessing is performed automatically - initial_forward_transform (itk.Transform, optional): Initial transformation from moving - to fixed. If provided, it is used to transform the moving image before - registration. Returns: dict: Dictionary containing: @@ -252,18 +247,10 @@ def registration_method( >>> result = registrar.register(moving_image, moving_mask=heart_mask_moving) """ - tfm_tools = TransformTools() - if moving_image_pre is None: moving_image_pre = self.preprocess(moving_image, self.modality) new_moving_image_pre = moving_image_pre - if initial_forward_transform is not None: - new_moving_image_pre = tfm_tools.transform_image( - moving_image_pre, - initial_forward_transform, - self.fixed_image, - ) # Prefer labelmap over binary mask when both sides have a labelmap. use_labelmaps = moving_labelmap is not None and self.fixed_labelmap is not None @@ -301,22 +288,6 @@ def registration_method( loss = loss_artifacts[0] - if initial_forward_transform is not None: - forward_transform = tfm_tools.combine_displacement_field_transforms( - initial_forward_transform, - forward_transform, - self.fixed_image, - tfm1_weight=1.0, - tfm2_weight=1.0, - mode="compose", - ) - - dftfm = tfm_tools.convert_transform_to_displacement_field_transform( - forward_transform, - self.fixed_image, - ) - inverse_transform = tfm_tools.invert_displacement_field_transform(dftfm) - return { "forward_transform": forward_transform, "inverse_transform": inverse_transform, diff --git a/src/physiotwin4d/register_models_distance_maps.py b/src/physiotwin4d/register_models_distance_maps.py index f7d94d83..43998b54 100644 --- a/src/physiotwin4d/register_models_distance_maps.py +++ b/src/physiotwin4d/register_models_distance_maps.py @@ -35,7 +35,7 @@ ... reference_image=reference_image, ... mask_dilation_mm=20, ... ) - >>> result = registrar.register(transform_type='Deformable', icon_iterations=50) + >>> result = registrar.register(transform_type='Deformable') >>> >>> # Access results >>> aligned_model = result['registered_model'] @@ -86,6 +86,7 @@ class RegisterModelsDistanceMaps(PhysioTwin4DBase): fixed_model (pv.PolyData): Target surface model reference_image (itk.Image): Reference image for coordinate frame mask_dilation_mm (float): Dilation amount in mm for binary registration masks + distance_squared_max (float): Maximum squared distance for distance map normalization transform_tools (TransformTools): Transform utility instance contour_tools (ContourTools): Model utility instance registrar_Greedy (RegisterImagesGreedy): Greedy registration instance @@ -110,7 +111,7 @@ class RegisterModelsDistanceMaps(PhysioTwin4DBase): >>> result = registrar.register(transform_type='Affine') >>> >>> # Or run deformable (Greedy affine + ICON) - >>> result = registrar.register(transform_type='Deformable', icon_iterations=50) + >>> result = registrar.register(transform_type='Deformable') >>> >>> # Get aligned model and transforms >>> aligned_model = result['registered_model'] @@ -122,6 +123,7 @@ def __init__( moving_model: pv.PolyData, fixed_model: pv.PolyData, reference_image: itk.Image, + distance_squared_max: float = 50.0, mask_dilation_mm: float = 20, log_level: int | str = logging.INFO, ): @@ -132,6 +134,8 @@ def __init__( fixed_model: PyVista target surface model reference_image: ITK image providing coordinate frame (origin, spacing, direction) for mask generation. Typically the patient CT/MRI image. + distance_squared_max: Maximum squared distance, in squared millimeters, + that the distance maps are normalized against. Default: 50.0 mask_dilation_mm: Dilation amount in millimeters for binary registration mask generation. Default: 20mm log_level: Logging level (default: logging.INFO) @@ -145,6 +149,7 @@ def __init__( self.moving_model = moving_model self.fixed_model = fixed_model self.reference_image = reference_image + self.distance_squared_max = distance_squared_max self.mask_dilation_mm = mask_dilation_mm # Utilities @@ -169,6 +174,23 @@ def __init__( self.inverse_transform: Optional[itk.CompositeTransform] = None # Fixed→moving self.registered_model: Optional[pv.PolyData] = None + def set_icon_weights_path(self, weights_path: str) -> None: + """Use a finetuned uniGradICON checkpoint for the deformable stage. + + The distance maps this class registers are not CT intensities, so stock + uniGradICON weights are out of distribution for them. Weights finetuned + on distance maps, e.g. by + ``tutorials/tutorial_02_lung_distancemap_finetune_icon.py``, are + supplied here. + + Args: + weights_path: Path to an existing uniGradICON checkpoint. + + Raises: + FileNotFoundError: If weights_path does not exist. + """ + self.registrar_ICON.set_weights_path(weights_path) + def _create_masks_from_models(self) -> None: """Generate distance maps and binary registration masks from moving and fixed models. @@ -189,20 +211,23 @@ def _create_masks_from_models(self) -> None: squared_distance=True, negative_inside=True, zero_inside=False, - norm_to_max_distance=50.0, + norm_to_max_distance=self.distance_squared_max, ) - # Create fixed binary registration mask with dilation - self.log_info( - "Dilating fixed mask by %.1fmm for registration mask...", - self.mask_dilation_mm, - ) - binary_mask = self.contour_tools.create_mask_from_mesh( - self.fixed_model, self.reference_image - ) - self.fixed_mask_image = self.labelmap_tools.convert_labelmap_to_mask( - binary_mask, dilation_in_mm=self.mask_dilation_mm - ) + if self.mask_dilation_mm > 0: + # Create fixed binary registration mask with dilation + self.log_info( + "Dilating fixed mask by %.1fmm for registration mask...", + self.mask_dilation_mm, + ) + binary_mask = self.contour_tools.create_mask_from_mesh( + self.fixed_model, self.reference_image + ) + self.fixed_mask_image = self.labelmap_tools.convert_labelmap_to_mask( + binary_mask, dilation_in_mm=self.mask_dilation_mm + ) + else: + self.fixed_mask_image = None # Create moving distance map self.moving_distance_map_image = self.contour_tools.create_distance_map( @@ -211,7 +236,7 @@ def _create_masks_from_models(self) -> None: squared_distance=True, negative_inside=True, zero_inside=False, - norm_to_max_distance=50.0, + norm_to_max_distance=self.distance_squared_max, ) # Emulate CT intensity range by multiplying by 1000 @@ -221,24 +246,26 @@ def _create_masks_from_models(self) -> None: tmp_arr = itk.GetArrayViewFromImage(self.moving_distance_map_image) tmp_arr *= 1000 - # Create moving binary registration mask with dilation - self.log_info( - "Dilating moving mask by %.1fmm for registration mask...", - self.mask_dilation_mm, - ) - binary_mask = self.contour_tools.create_mask_from_mesh( - self.moving_model, self.reference_image - ) - self.moving_mask_image = self.labelmap_tools.convert_labelmap_to_mask( - binary_mask, dilation_in_mm=self.mask_dilation_mm - ) + if self.mask_dilation_mm > 0: + # Create moving binary registration mask with dilation + self.log_info( + "Dilating moving mask by %.1fmm for registration mask...", + self.mask_dilation_mm, + ) + binary_mask = self.contour_tools.create_mask_from_mesh( + self.moving_model, self.reference_image + ) + self.moving_mask_image = self.labelmap_tools.convert_labelmap_to_mask( + binary_mask, dilation_in_mm=self.mask_dilation_mm + ) + else: + self.moving_mask_image = None self.log_info("Distance map and mask generation complete") def register( self, transform_type: str = "Deformable", - icon_iterations: int = 50, ) -> dict: """Perform mask-based registration of moving model to fixed model. @@ -259,7 +286,6 @@ def register( Args: transform_type: Registration transform type - 'None', 'Rigid', 'Affine', or 'Deformable'. Default: 'Deformable' - icon_iterations: Number of ICON optimization iterations for 'Deformable' mode. Default: 50 Returns: Dictionary containing: @@ -278,7 +304,7 @@ def register( >>> result = registrar.register(transform_type='Affine') >>> >>> # Deformable registration (Greedy affine + ICON) - >>> result = registrar.register(transform_type='Deformable', icon_iterations=100) + >>> result = registrar.register(transform_type='Deformable') """ if transform_type not in ["None", "Rigid", "Affine", "Deformable"]: raise ValueError( @@ -319,10 +345,7 @@ def register( # Step 3: ICON deformable stage (only for Deformable mode) if transform_type == "Deformable": - self.log_info( - "Performing ICON deformable registration (%d iterations)...", - icon_iterations, - ) + self.log_info("Performing ICON deformable registration...") # Pre-align moving distance map and binary mask into the fixed grid using the Greedy affine result moving_distance_map_affine_transformed = ( @@ -340,8 +363,8 @@ def register( # interpolation_method="nearest", # ) - # Configure and run ICON - self.registrar_ICON.set_number_of_iterations(icon_iterations) + # Configure and run ICON. Iteration count and any other ICON tuning + # come from registrar_ICON itself, configured by the caller. self.registrar_ICON.set_fixed_image(self.fixed_distance_map_image) # self.registrar_ICON.set_fixed_mask(self.fixed_mask_image) @@ -359,18 +382,20 @@ def register( # (patient-space δ), then Greedy (patient→ICP-template). # Inverse (moving→fixed for point push-forward): apply Greedy first # (ICP-template→patient), then ICON (patient-space refinement). + # combine_displacement_field_transforms(a, b) evaluates b then a, so + # the stage that runs first is the second argument. self.forward_transform = ( self.transform_tools.combine_displacement_field_transforms( - forward_transform_ICON, forward_transform_Greedy, + forward_transform_ICON, reference_image=self.reference_image, mode="compose", ) ) self.inverse_transform = ( self.transform_tools.combine_displacement_field_transforms( - inverse_transform_Greedy, inverse_transform_ICON, + inverse_transform_Greedy, reference_image=self.reference_image, mode="compose", ) diff --git a/src/physiotwin4d/register_models_icp.py b/src/physiotwin4d/register_models_icp.py index c845a658..e0602086 100644 --- a/src/physiotwin4d/register_models_icp.py +++ b/src/physiotwin4d/register_models_icp.py @@ -3,15 +3,17 @@ This module provides the RegisterModelsICP class for aligning anatomical models using Iterative Closest Point (ICP) algorithm. The workflow includes: 1. Initial centroid alignment -2. Rigid or affine ICP alignment +2. Isotropic bounding-box scaling (Similarity and Affine transform types) +3. Rigid, similarity or affine ICP alignment The registration is particularly useful for initial rough alignment of generic models to patient-specific anatomical data. Key Features: - Centroid-based initial alignment - - VTK ICP algorithm with rigid or affine transformation modes - - Three-stage affine pipeline: centroid → rigid ICP → affine ICP + - Bounding-box size matching before any ICP iteration + - VTK ICP with rigid, similarity or affine transformation modes + - Four-stage affine pipeline: centroid → scale → similarity ICP → affine ICP - Support for PyVista models - Automatic transform composition @@ -57,8 +59,15 @@ class RegisterModelsICP(PhysioTwin4DBase): centroid alignment for initialization followed by VTK's ICP algorithm. **Registration Pipelines:** - - **Rigid transform type**: Centroid alignment → Rigid ICP - - **Affine transform type**: Centroid alignment → Rigid ICP → Affine ICP + - **Rigid**: Centroid alignment → Rigid ICP + - **Similarity**: Centroid alignment → bounding-box scaling → + Similarity ICP + - **Affine**: Centroid alignment → bounding-box scaling → + Similarity ICP → Affine ICP + + The bounding-box scaling estimates the size ratio between the models from + their bounding-box diagonals before ICP starts. It is skipped for 'Rigid', + whose transform has no scale degree of freedom. **Transform Convention:** These are POINT transforms (applied with TransformPoint, e.g. via @@ -133,6 +142,118 @@ def __init__( self.inverse_point_transform: Optional[itk.AffineTransform] = None self.registered_model: Optional[pv.PolyData] = None + # ICP stages run for each transform type, in order. + _ICP_STAGES = { + "Rigid": ("Rigid",), + "Similarity": ("Similarity",), + "Affine": ("Similarity", "Affine"), + } + + def _icp_stage( + self, model: pv.PolyData, mode: str, max_iterations: int + ) -> tuple[pv.PolyData, itk.AffineTransform]: + """Run one VTK ICP stage against the fixed model. + + Args: + model: Current state of the moving model. + mode: Landmark-transform mode, one of ``"Rigid"``, ``"Similarity"`` + or ``"Affine"``. + max_iterations: Maximum ICP iterations for this stage. + + Returns: + Tuple of the transformed model and this stage's moving→fixed point + transform. + """ + icp = vtk.vtkIterativeClosestPointTransform() + icp.SetSource(model) + icp.SetTarget(self.fixed_model) + landmark = icp.GetLandmarkTransform() + if mode == "Rigid": + landmark.SetModeToRigidBody() + elif mode == "Similarity": + landmark.SetModeToSimilarity() + else: + landmark.SetModeToAffine() + icp.SetMaximumNumberOfIterations(max_iterations) + icp.Update() + + stage_transform = self.transform_tools.convert_vtk_matrix_to_itk_transform( + icp.GetMatrix() + ) + transformed = self.transform_tools.transform_pvcontour( + model, + stage_transform, + with_deformation_magnitude=False, + ) + return transformed, stage_transform + + def _bounding_box_scale( + self, moving_model: pv.PolyData, fixed_model: pv.PolyData + ) -> float: + """Return the isotropic scale matching the models' bounding-box diagonals. + + The diagonal is used rather than the three side lengths separately: the + bounding boxes are axis-aligned, so per-axis ratios would fold any + residual rotation between the models into an anisotropic stretch. A + single factor changes size without distorting shape, leaving the + remaining anisotropy to the affine ICP stage. + + Args: + moving_model: Model whose size is being matched. + fixed_model: Model supplying the target size. + + Returns: + The scale factor, or ``1.0`` when either model is degenerate (a + single point or a zero-extent bounding box), where no meaningful + ratio exists. + """ + + def _diagonal(model: pv.PolyData) -> float: + x_min, x_max, y_min, y_max, z_min, z_max = model.bounds + extents = np.array( + [x_max - x_min, y_max - y_min, z_max - z_min], dtype=np.float64 + ) + return float(np.linalg.norm(extents)) + + moving_diagonal = _diagonal(moving_model) + fixed_diagonal = _diagonal(fixed_model) + self.log_debug( + "Bounding-box diagonals - moving: %.4f, fixed: %.4f", + moving_diagonal, + fixed_diagonal, + ) + if moving_diagonal <= 0.0 or fixed_diagonal <= 0.0: + self.log_warning( + "Degenerate bounding box (moving diagonal %.4f, fixed diagonal " + "%.4f); skipping the scaling step.", + moving_diagonal, + fixed_diagonal, + ) + return 1.0 + return fixed_diagonal / moving_diagonal + + def _scale_transform(self, scale: float, center: np.ndarray) -> itk.AffineTransform: + """Build the transform scaling isotropically about ``center``. + + Args: + scale: Isotropic scale factor. + center: Fixed point of the scaling, in world coordinates. + + Returns: + An ITK affine point transform mapping ``p`` to + ``center + scale * (p - center)``. + """ + matrix = np.eye(3, dtype=np.float64) * scale + offset = itk.Vector[itk.D, 3]() + for i in range(3): + offset[i] = float(center[i]) * (1.0 - scale) + + transform = itk.AffineTransform[itk.D, 3].New() + transform.SetIdentity() + transform.SetMatrix(itk.Matrix[itk.D, 3, 3](itk.GetVnlMatrixFromArray(matrix))) + transform.SetOffset(offset) + return transform + def register( self, moving_model: pv.PolyData, @@ -141,22 +262,31 @@ def register( ) -> dict: """Perform ICP alignment of moving model to fixed model. - This method executes alignment with either rigid or affine transformations: + **Rigid transform type** (rotation + translation): + 1. Centroid alignment: Translate moving model to align mass centers + 2. Rigid ICP: Refine with rigid-body transformation - **Rigid transform type:** + **Similarity transform type** (rotation + translation + one uniform scale): 1. Centroid alignment: Translate moving model to align mass centers - 2. Rigid ICP: Refine with rigid-body transformation (rotation + translation) + 2. Bounding-box scaling: Scale isotropically about the fixed centroid + so the two bounding-box diagonals match, which keeps ICP's + closest-point search from locking onto a size mismatch + 3. Similarity ICP: Refine rotation, translation and uniform scale - **Affine transform type:** + **Affine transform type** (adds anisotropic scale and shear): 1. Centroid alignment: Translate moving model to align mass centers - 2. Rigid ICP: Refine with rigid-body transformation - 3. Affine ICP: Further refine with affine transformation (includes - scaling/shearing) + 2. Bounding-box scaling: As above + 3. Similarity ICP: Refine rotation, translation and uniform scale + 4. Affine ICP: Further refine with affine transformation + + 'Rigid' skips the bounding-box scaling so its result stays a pure + rigid-body transform; use 'Similarity' when the models differ in size but + the shape should not be distorted. Args: moving_model: PyVista surface model to be aligned to fixed model - transform_type: Registration transform type, either 'Rigid' or 'Affine'. - Default: 'Affine' + transform_type: Registration transform type, one of 'Rigid', + 'Similarity' or 'Affine'. Default: 'Affine' max_iterations: Maximum number of ICP iterations per stage. Default: 2000 Returns: @@ -168,7 +298,7 @@ def register( (ITK AffineTransform) Raises: - ValueError: If transform_type is not 'Rigid' or 'Affine' + ValueError: If transform_type is not 'Rigid', 'Similarity' or 'Affine' Example: >>> # Rigid registration @@ -178,6 +308,13 @@ def register( ... moving_model=moving_model, ... ) >>> + >>> # Similarity registration (rigid plus one uniform scale) + >>> result = registrar.register( + ... transform_type='Similarity', + ... max_iterations=2000, + ... moving_model=moving_model, + ... ) + >>> >>> # Affine registration >>> result = registrar.register( ... transform_type='Affine', @@ -185,9 +322,10 @@ def register( ... moving_model=moving_model, ... ) """ - if transform_type not in ["Rigid", "Affine"]: + if transform_type not in self._ICP_STAGES: raise ValueError( - f"Invalid transform '{transform_type}'. Must be 'Rigid' or 'Affine'." + f"Invalid transform '{transform_type}'. Must be one of " + f"{sorted(self._ICP_STAGES)}." ) self.log_section("%s ICP Alignment", transform_type.upper()) @@ -195,15 +333,15 @@ def register( self.moving_model = moving_model self.transform_type = transform_type - # Step 1: Centroid alignment (common to both modes) - self.registered_model = self.moving_model.copy(deep=True) + # Centroid alignment (common to every mode) + registered_model = self.moving_model.copy(deep=True) - moving_centroid = np.array(self.registered_model.center) + moving_centroid = np.array(registered_model.center) self.log_debug("Moving model centroid: %s", moving_centroid) fixed_centroid = np.array(self.fixed_model.center) self.log_debug("Fixed model centroid: %s", fixed_centroid) translation = fixed_centroid - moving_centroid - self.log_info("Step 1: Translating by %s to align centroids...", translation) + self.log_info("Translating by %s to align centroids...", translation) # Create ITK affine transform with translation forward_point_transform = itk.AffineTransform[itk.D, 3].New() @@ -211,70 +349,55 @@ def register( forward_point_transform.SetOffset(translation) # Apply centroid alignment to model - self.registered_model = self.transform_tools.transform_pvcontour( - self.registered_model, + registered_model = self.transform_tools.transform_pvcontour( + registered_model, forward_point_transform, with_deformation_magnitude=False, ) - self.log_debug("Center after Step 1: %s", self.registered_model.center) - - # Step 2: Rigid ICP (common to both modes) - self.log_info( - "Step 2: Performing rigid ICP (max iterations: %d)...", max_iterations - ) - icp_rigid = vtk.vtkIterativeClosestPointTransform() - icp_rigid.SetSource(self.registered_model) - icp_rigid.SetTarget(self.fixed_model) - icp_rigid.GetLandmarkTransform().SetModeToRigidBody() # Rigid mode - icp_rigid.SetMaximumNumberOfIterations(max_iterations) - icp_rigid.Update() - - # Convert VTK transform to ITK and compose with centroid transform - rigid_transform = self.transform_tools.convert_vtk_matrix_to_itk_transform( - icp_rigid.GetMatrix() - ) - forward_point_transform.Compose(rigid_transform) - - # Apply rigid ICP transform to model - self.registered_model = self.transform_tools.transform_pvcontour( - self.registered_model, - rigid_transform, - with_deformation_magnitude=False, - ) - - self.log_debug("Center after Step 2: %s", self.registered_model.center) - - # Step 3: Affine ICP (only if affine mode) - if transform_type == "Affine": + self.log_debug("Center after centroid alignment: %s", registered_model.center) + + # Bounding-box scaling, for the modes whose transform admits a scale. ICP + # only searches for correspondences among nearest points, so a template + # that differs from the patient in overall size drags the closest-point + # matching into a local minimum. Matching the bounding-box diagonals first + # puts the two models on the same scale before any ICP iteration runs. + # 'Rigid' skips it: a size estimate there would make the result a + # similarity transform, which is what 'Similarity' is for. + if transform_type != "Rigid": + scale = self._bounding_box_scale(registered_model, self.fixed_model) self.log_info( - "Step 3: Performing affine ICP (max iterations: %d)...", max_iterations - ) - icp_affine = vtk.vtkIterativeClosestPointTransform() - icp_affine.SetSource(self.registered_model) - icp_affine.SetTarget(self.fixed_model) - icp_affine.GetLandmarkTransform().SetModeToAffine() # Affine mode - icp_affine.SetMaximumNumberOfIterations(max_iterations) - icp_affine.Update() - - # Convert VTK transform to ITK and compose - affine_transform = self.transform_tools.convert_vtk_matrix_to_itk_transform( - icp_affine.GetMatrix() + "Scaling by %.4f about the fixed centroid to match bounding boxes...", + scale, ) - forward_point_transform.Compose(affine_transform) - - # Apply affine ICP transform to model - self.registered_model = self.transform_tools.transform_pvcontour( - self.registered_model, - affine_transform, + scale_transform = self._scale_transform(scale, fixed_centroid) + forward_point_transform.Compose(scale_transform) + registered_model = self.transform_tools.transform_pvcontour( + registered_model, + scale_transform, with_deformation_magnitude=False, ) + self.log_debug("Bounds after scaling: %s", registered_model.bounds) - self.log_debug("Center after Step 3: %s", self.registered_model.center) + # ICP stages. Affine runs similarity first so the shear and anisotropic + # scale degrees of freedom refine an already-oriented, already-sized model + # rather than absorbing rotation and overall scale themselves. + for stage in self._ICP_STAGES[transform_type]: + self.log_info( + "Performing %s ICP (max iterations: %d)...", + stage.lower(), + max_iterations, + ) + registered_model, stage_transform = self._icp_stage( + registered_model, stage, max_iterations + ) + forward_point_transform.Compose(stage_transform) + self.log_debug("Center after %s ICP: %s", stage, registered_model.center) # Compute inverse transform # Ths forward transform for ICP is consistent with the transform convention # used with images-to-images registration. + self.registered_model = registered_model self.forward_point_transform = forward_point_transform self.inverse_point_transform = forward_point_transform.GetInverseTransform() diff --git a/src/physiotwin4d/register_models_pca.py b/src/physiotwin4d/register_models_pca.py index ef93d97c..fd42293a 100644 --- a/src/physiotwin4d/register_models_pca.py +++ b/src/physiotwin4d/register_models_pca.py @@ -8,7 +8,9 @@ import itk import numpy as np import pyvista as pv +from scipy.ndimage import map_coordinates from scipy.optimize import minimize +from scipy.spatial import cKDTree from typing_extensions import Self from .contour_tools import ContourTools @@ -17,29 +19,44 @@ class RegisterModelsPCA(PhysioTwin4DBase): - """Register PCA-based shape models to medical images using mean distance optimization. + """Register PCA-based shape models to images by minimizing a distance metric. This class implements a registration pipeline for fitting statistical shape models to patient-specific medical images: **PCA Deformable Registration** - Optimizes PCA coefficients - - Model equation: P = mean + Σ(b_i * std_i * pca_eigenvector_i) - - Maximizes mean distance at deformed model points P + - Model equation: P = template + Σ(b_i * std_i * pca_eigenvector_i) + - Minimizes the mean distance-to-target at the deformed model points P **Optimization Objective:** - Maximize the mean distance of the image sampled at model points using - ITK's LinearInterpolateImageFunction. This aligns the model with bright - regions in contrast-enhanced images (e.g., blood pool in cardiac CT). + ``fixed_distance_map`` is zero on the target surface and grows with + distance away from it (in mm), so the objective is *minimized*:: + + E(b) = (1 - w) * mean_i D(P_i(b)) # model -> target + + w * mean_j ||Q_j - P_nn(j)|| # target -> model + + lambda * Σ b_i² # Mahalanobis prior + + ``w`` is ``symmetric_weight`` and ``lambda`` is ``pca_prior_weight``. + Because the deformation is linear in ``b``, the gradient is analytic and + is supplied to the optimizer directly. + + **Coordinate frames:** + The eigenvectors are directions in the statistical model's own training + frame, so they are only valid when added to a template in that same + frame. Any rigid/affine alignment to the target must be supplied as + ``post_pca_transform``, which is applied *after* the deformation, rather + than pre-applied to ``pca_template_model``. Attributes: pca_template_model (pv.DataSet): Mean shape model pca_eigenvectors (np.ndarray): PCA eigenvectors/components (modes × n_points*3) pca_std_deviations (np.ndarray): Standard deviations per mode (modes,) - fixed_distance_map (itk.Image): Patient image providing distance data - n_points (int): Number of points in the model + fixed_distance_map (itk.Image): Distance map of the target, in mm + fixed_model (pv.DataSet): Target model, when one was supplied. Required + for the symmetric (target-to-model) term. pca_number_of_modes (int): Number of PCA modes available - pca_coefficients (np.ndarray): Optimized PCA coefficients + registered_model_pca_coefficients (np.ndarray): Optimized PCA coefficients registered_model (pv.DataSet): Final registered and deformed model post_pca_transform (itk.Transform): Transform to apply after PCA registration forward_point_transform (itk.DisplacementFieldTransform): POINT transform @@ -90,63 +107,99 @@ def __init__( fixed_distance_map: Optional[itk.Image] = None, fixed_model: Optional[pv.DataSet] = None, reference_image: Optional[itk.Image] = None, + pca_prior_weight: float = 0.0, + symmetric_weight: float = 0.5, log_level: int | str = logging.INFO, ): """Initialize the PCA-based model-to-image registration. Args: pca_template_model: PyVista model containing the mean 3D shape model - (unstructured grid or polydata) + (unstructured grid or polydata). It must be in the same frame + the PCA modes were trained in; supply any alignment to the + target as post_pca_transform instead of pre-applying it here. pca_eigenvectors: Numpy array of PCA eigenvectors/components. Shape: (modes, n_points*3) Each row is a flattened eigenmode with 3D displacements: [x1,y1,z1, x2,y2,z2, ...] pca_std_deviations: Numpy array of standard deviations per PCA mode. Shape: (modes,) These are the square roots of pca_eigenvalues - pca_number_of_modes: Number of PCA modes to use. Default: -1 (use all) + pca_number_of_modes: Number of PCA modes to use. Default: 0 (use all) pca_template_model_point_subsample: Step size for subsampling model points. Default: 4 post_pca_transform: Optional ITK transform to apply after PCA registration. Default: None - fixed_distance_map: ITK image providing the distance map. + fixed_distance_map: ITK image providing the distance map, in mm. Default: None fixed_model: PyVista model used to compute the distance map, if one isn't provided. + Also supplies the target points for the symmetric term. reference_image: ITK image providing coordinate frame for computing the distance map. + pca_prior_weight: Weight (in mm) of the Mahalanobis shape prior + ``lambda * sum(b_i**2)``. Because b is expressed in standard + deviations, this term is the squared Mahalanobis distance in + shape space and makes the fit a MAP estimate rather than a pure + data fit constrained only by the coefficient bounds. + Default: 0.0 (prior disabled). + symmetric_weight: Weight in [0, 1] of the target-to-model distance + term. 0.0 measures model-to-target only, which lets the model + satisfy the metric while covering just part of the target. + Requires fixed_model; ignored with a warning when only a + distance map is available. Default: 0.5 log_level: Logging level (logging.DEBUG, logging.INFO, logging.WARNING). Default: logging.INFO Raises: - ValueError: If pca_eigenvector dimensions don't match model points + ValueError: If pca_eigenvector dimensions don't match model points, + if the mode counts disagree, or if neither a distance map nor a + fixed model plus reference image is provided. """ # Initialize base class with logging super().__init__(class_name="RegisterModelsPCA", log_level=log_level) # Store model data self.pca_template_model: pv.DataSet = pca_template_model - self.pca_eigenvectors: np.ndarray = pca_eigenvectors - self.pca_std_deviations: np.ndarray = pca_std_deviations + self.pca_eigenvectors: np.ndarray = np.asarray( + pca_eigenvectors, dtype=np.float64 + ) + self.pca_std_deviations: np.ndarray = np.asarray( + pca_std_deviations, dtype=np.float64 + ) + + if self.pca_eigenvectors.ndim != 2: + raise ValueError( + f"pca_eigenvectors must be 2D (modes, n_points*3), got shape " + f"{self.pca_eigenvectors.shape}" + ) + expected_size = pca_template_model.n_points * 3 + if self.pca_eigenvectors.shape[1] != expected_size: + raise ValueError( + f"Component dimension mismatch: expected {expected_size} " + f"(3 × {pca_template_model.n_points} points), got " + f"{self.pca_eigenvectors.shape[1]}" + ) + if self.pca_eigenvectors.shape[0] != self.pca_std_deviations.shape[0]: + raise ValueError( + f"Mode count mismatch: {self.pca_eigenvectors.shape[0]} eigenvectors " + f"but {self.pca_std_deviations.shape[0]} standard deviations" + ) self.post_pca_transform = post_pca_transform self._contour_tools = ContourTools() + self.fixed_model: Optional[pv.DataSet] = fixed_model self.fixed_distance_map = fixed_distance_map if ( self.fixed_distance_map is None and fixed_model is not None and reference_image is not None ): - self.fixed_model = fixed_model - self.fixed_distance_map = self._contour_tools.create_distance_map( - fixed_model, - reference_image, - squared_distance=False, - negative_inside=False, - zero_inside=True, - norm_to_max_distance=200.0, + self.fixed_distance_map = self._create_distance_map( + fixed_model, reference_image ) elif self.fixed_distance_map is not None and ( fixed_model is not None or reference_image is not None ): self.log_warning( - "Fixed model and reference image will be ignored because a distance map is provided." + "A distance map was provided, so the reference image is ignored; " + "the fixed model is retained only for the symmetric metric term." ) elif self.fixed_distance_map is None and ( fixed_model is None or reference_image is None @@ -160,9 +213,15 @@ def __init__( self.pca_number_of_modes: int = pca_number_of_modes if self.pca_number_of_modes <= 0: - self.pca_number_of_modes = len(pca_std_deviations) + self.pca_number_of_modes = len(self.pca_std_deviations) self.pca_template_model_point_subsample = pca_template_model_point_subsample + self.pca_prior_weight = pca_prior_weight + if not 0.0 <= symmetric_weight <= 1.0: + raise ValueError( + f"symmetric_weight must be in [0, 1]; got {symmetric_weight}" + ) + self.symmetric_weight = symmetric_weight # outputs self.registered_model_pca_coefficients: Optional[np.ndarray] = None @@ -172,17 +231,26 @@ def __init__( self.forward_point_transform: Optional[itk.DisplacementFieldTransform] = None self.inverse_point_transform: Optional[itk.DisplacementFieldTransform] = None - # Image interpolator (created when needed) - self._fixed_distance_map_interpolator: Optional[ - itk.LinearInterpolateImageFunction - ] = None + # Sampling caches, built lazily by _prepare_sampling() + self._sampling_ready: bool = False + self._analytic_gradient: bool = True self._fixed_distance_map_max_distance: float = 0.0 + self._post_pca_affine_key: Optional[itk.Transform] = None + self._post_pca_affine: Optional[tuple[np.ndarray, np.ndarray]] = None self._metric_call_count: int = 0 - # Pre-convert mean shape points to ITK format - self._pca_template_model_points_itk: Optional[list[itk.Point]] = None - self._create_itk_points() + def _create_distance_map( + self, fixed_model: pv.DataSet, reference_image: itk.Image + ) -> itk.Image: + """Build the unsigned, un-normalized (mm) distance map of the target.""" + return self._contour_tools.create_distance_map( + fixed_model, + reference_image, + squared_distance=False, + negative_inside=False, + zero_inside=True, + ) @classmethod def from_json( @@ -195,6 +263,8 @@ def from_json( fixed_distance_map: Optional[itk.Image] = None, fixed_model: Optional[pv.DataSet] = None, reference_image: Optional[itk.Image] = None, + pca_prior_weight: float = 0.0, + symmetric_weight: float = 0.5, log_level: int | str = logging.INFO, ) -> Self: """Create RegisterModelsPCA from PCA model JSON file. @@ -217,6 +287,8 @@ def from_json( for registration. If None, must be set later before registration. fixed_model: Target surface mesh to register to. Default: None reference_image: Reference image defining coordinate space. Default: None + pca_prior_weight: Weight (mm) of the Mahalanobis shape prior. Default: 0.0 + symmetric_weight: Weight of the target-to-model term. Default: 0.5 log_level: Logging level (logging.DEBUG, logging.INFO, logging.WARNING). Default: logging.INFO @@ -288,6 +360,8 @@ def from_json( fixed_distance_map=fixed_distance_map, fixed_model=fixed_model, reference_image=reference_image, + pca_prior_weight=pca_prior_weight, + symmetric_weight=symmetric_weight, log_level=log_level, ) @@ -302,6 +376,8 @@ def from_pca_model( fixed_distance_map: Optional[itk.Image] = None, fixed_model: Optional[pv.DataSet] = None, reference_image: Optional[itk.Image] = None, + pca_prior_weight: float = 0.0, + symmetric_weight: float = 0.5, log_level: int | str = logging.INFO, ) -> Self: """Create RegisterModelsPCA from a PCA model dictionary. @@ -320,6 +396,8 @@ def from_pca_model( fixed_distance_map: ITK image providing the distance values for registration. fixed_model: Target surface mesh to register to. reference_image: Reference image defining coordinate space. + pca_prior_weight: Weight (mm) of the Mahalanobis shape prior. Default: 0.0 + symmetric_weight: Weight of the target-to-model term. Default: 0.5 log_level: Logging level. Returns: @@ -334,12 +412,6 @@ def from_pca_model( if "components" not in pca_model: raise ValueError("'components' field not found in pca_model") pca_eigenvectors = np.array(pca_model["components"], dtype=np.float64) - expected_size = pca_template_model.n_points * 3 - if pca_eigenvectors.shape[1] != expected_size: - raise ValueError( - f"Component dimension mismatch: expected {expected_size} " - f"(3 × {pca_template_model.n_points} points), got {pca_eigenvectors.shape[1]}" - ) return cls( pca_template_model=pca_template_model, pca_eigenvectors=pca_eigenvectors, @@ -350,38 +422,18 @@ def from_pca_model( fixed_distance_map=fixed_distance_map, fixed_model=fixed_model, reference_image=reference_image, + pca_prior_weight=pca_prior_weight, + symmetric_weight=symmetric_weight, log_level=log_level, ) - def _create_itk_points(self) -> None: - """Pre-convert mean shape points to ITK Point format for efficiency. - - This method creates ITK Point objects once at initialization, avoiding - repeated conversions during optimization iterations. - """ - self.log_info("Converting mean shape points to ITK format...") - - self._pca_template_model_points_itk = [] - for point in self.pca_template_model.points: - itk_point = itk.Point[itk.D, 3]() - itk_point[0] = float(point[0]) - itk_point[1] = float(point[1]) - itk_point[2] = float(point[2]) - self._pca_template_model_points_itk.append(itk_point) - - self.log_info( - f" Converted {len(self._pca_template_model_points_itk)} points to ITK format" - ) - def set_fixed_model( self, fixed_model: pv.UnstructuredGrid, reference_image: Optional[itk.Image] ) -> None: - """Set the fixed model for registration. - - If this is set, the fixed distance map will be set to None. + """Set the fixed model for registration and rebuild its distance map. Args: - fixed_model: PyVista model used to compute the distance map, if one isn't provided. + fixed_model: PyVista model used to compute the distance map. reference_image: ITK image providing coordinate frame for computing the distance map. """ if reference_image is None: @@ -389,26 +441,20 @@ def set_fixed_model( "reference_image must not be None when setting a fixed model" ) - self.fixed_distance_map = self._contour_tools.create_distance_map( - fixed_model, - reference_image, - squared_distance=False, - negative_inside=False, - zero_inside=True, - norm_to_max_distance=200.0, + self.fixed_model = fixed_model + self.fixed_distance_map = self._create_distance_map( + fixed_model, reference_image ) - self._fixed_distance_map_interpolator = None + self._sampling_ready = False def set_fixed_distance_map(self, fixed_distance_map: Optional[itk.Image]) -> None: - """Set the reference image for registration. - - If this is set, the fixed model will be set to None. + """Set the distance map used as the registration target. Args: - fixed_distance_map: ITK image providing distance data + fixed_distance_map: ITK image providing distance data, in mm """ self.fixed_distance_map = fixed_distance_map - self._fixed_distance_map_interpolator = None + self._sampling_ready = False def set_pca_template_model(self, pca_template_model: pv.UnstructuredGrid) -> None: """Set the average model for registration. @@ -418,132 +464,310 @@ def set_pca_template_model(self, pca_template_model: pv.UnstructuredGrid) -> Non (unstructured grid or polydata) """ self.pca_template_model = pca_template_model - - self._pca_template_model_points_itk = None - - self._create_itk_points() + self._sampling_ready = False self.log_info(" Average model set successfully!") - def _mean_distance_metric( - self, - params: np.ndarray, - ) -> float: - """Evaluate the optimization metric (mean intensity) at model points. - - This is the objective function to be MAXIMIZED during optimization. - Higher values indicate better alignment with bright regions. + def _affine_of_transform( + self, transform: itk.Transform + ) -> Optional[tuple[np.ndarray, np.ndarray]]: + """Recover (matrix, offset) if ``transform`` acts affinely, else None. - Args: - pca_deformation: Nx3 numpy array of PCA deformation vectors to add to points. - If None, no deformation is applied. + Probing the transform rather than querying ``GetMatrix()`` works for any + ITK transform type and correctly rejects the non-affine ones (such as a + displacement field), for which no constant Jacobian exists. + """ - Returns: - Mean distance value across all points + def apply(vector: np.ndarray) -> np.ndarray: + point = itk.Point[itk.D, 3]() + point[0], point[1], point[2] = (float(v) for v in vector) + result = transform.TransformPoint(point) + return np.array([result[0], result[1], result[2]], dtype=np.float64) + + # Probe inside the model's own extent: a displacement field evaluated + # outside its grid returns no displacement, so probing at the origin and + # the unit cube would report such a transform as the identity affine. + bounds = np.asarray(self.pca_template_model.bounds, dtype=np.float64) + low, high = bounds[0::2], bounds[1::2] + center = 0.5 * (low + high) + step = 0.25 * np.maximum(high - low, 1.0) + + base = apply(center) + matrix = np.column_stack( + [ + (apply(center + step[i] * basis) - base) / step[i] + for i, basis in enumerate(np.eye(3, dtype=np.float64)) + ] + ) + offset = base - matrix @ center + probe = center + step * np.array([0.37, -0.61, 0.83], dtype=np.float64) + scale = max(1.0, float(np.abs(base).max()), float(np.abs(matrix).max())) + if not np.allclose(apply(probe), matrix @ probe + offset, atol=1e-9 * scale): + return None + return matrix, offset + + def _get_post_pca_affine(self) -> Optional[tuple[np.ndarray, np.ndarray]]: + """Return the cached (matrix, offset) of post_pca_transform, or None. + + Keyed on the transform object so that reassigning post_pca_transform + invalidates the cache. """ - pca_deformation = self._compute_pca_deformation(params) - - # Create interpolator if not already cached (inline creation) - if self._fixed_distance_map_interpolator is None: - if self.fixed_distance_map is None: - self.log_error("Distance map is not set.") - raise ValueError("Distance map must be set before registering.") - ImageType = type(self.fixed_distance_map) - self._fixed_distance_map_interpolator = itk.LinearInterpolateImageFunction[ - ImageType, itk.D - ].New() - self._fixed_distance_map_interpolator.SetInputImage(self.fixed_distance_map) - fixed_distance_map_array = itk.GetArrayFromImage(self.fixed_distance_map) - self._fixed_distance_map_max_distance = fixed_distance_map_array.max() - self.log_debug("Interpolator created") - self.log_debug( - " Max distance = %s", self._fixed_distance_map_max_distance - ) + if self.post_pca_transform is None: + return None + if self._post_pca_affine_key is not self.post_pca_transform: + self._post_pca_affine_key = self.post_pca_transform + self._post_pca_affine = self._affine_of_transform(self.post_pca_transform) + return self._post_pca_affine + + def _prepare_sampling(self) -> None: + """Build the cached arrays the objective and its gradient are made of. + + Everything that does not depend on the PCA coefficients is computed once + here: the subsampled template points, the per-mode displacement vectors + already scaled by their standard deviation and mapped through the + post-PCA transform, the distance map and its gradient, and the + physical-to-index mapping used to sample them. + """ + if self.fixed_distance_map is None: + self.log_error("Distance map is not set.") + raise ValueError("Distance map must be set before registering.") + + template_points = np.asarray(self.pca_template_model.points, dtype=np.float64) + step = max(1, self.pca_template_model_point_subsample) + self._sample_slice = slice(None, None, step) + self._sample_points = template_points[self._sample_slice] + + # (modes, m, 3) displacement per unit coefficient, in template space. + modes = self.pca_eigenvectors.reshape(self.pca_eigenvectors.shape[0], -1, 3) + self._sample_modes = ( + modes[:, self._sample_slice, :] * self.pca_std_deviations[:, None, None] + ) - self.log_debug("Evaluating params = %s", params) - self.log_debug(" Max displacement = %s", pca_deformation.max(axis=0)) + # Fold the post-PCA transform into the mode directions so the gradient + # is expressed directly in world space. A non-affine post-PCA transform + # has no constant Jacobian, so the analytic gradient is disabled. + affine = self._get_post_pca_affine() + if self.post_pca_transform is not None and affine is None: + self.log_warning( + "post_pca_transform is not affine; falling back to a " + "finite-difference gradient." + ) + self._sample_modes_world = ( + self._sample_modes if affine is None else self._sample_modes @ affine[0].T + ) + self._analytic_gradient = self.post_pca_transform is None or affine is not None + + # Physical point -> continuous index: index = affine_inv @ (p - origin). + image = self.fixed_distance_map + direction = itk.array_from_matrix(image.GetDirection()) + index_to_world = direction @ np.diag(np.asarray(image.GetSpacing())) + self._index_to_world = index_to_world + self._world_to_index = np.linalg.inv(index_to_world) + self._image_origin = np.asarray(image.GetOrigin(), dtype=np.float64) + size = image.GetLargestPossibleRegion().GetSize() + self._image_size = np.array([size[0], size[1], size[2]], dtype=np.float64) + + # Array axes are (k, j, i), so index component a is array axis 2 - a. + # The forward differences are the exact derivative of the trilinear + # interpolant used to sample the map, which keeps the objective and its + # gradient consistent; a central difference would not. + self._distance_array = np.asarray( + itk.array_view_from_image(image), dtype=np.float64 + ) + self._fixed_distance_map_max_distance = float(self._distance_array.max()) + self._distance_forward_diff = tuple( + np.diff(self._distance_array, axis=2 - axis) + if self._distance_array.shape[2 - axis] > 1 + else None + for axis in range(3) + ) - # Sample distance at each point - n_valid_points = 0 - n_invalid_points = 0 - total_distance = 0.0 - center = np.zeros(3) - point = itk.Point[itk.D, 3]() - assert self.fixed_distance_map is not None, "fixed_distance_map must be set" - assert self._pca_template_model_points_itk is not None, ( - "ITK points must be initialized" + # Target points for the symmetric term. + self._target_points: Optional[np.ndarray] = None + if self.symmetric_weight > 0.0: + if self.fixed_model is None: + self.log_warning( + "symmetric_weight is %.3g but no fixed_model is available; " + "the target-to-model term is disabled.", + self.symmetric_weight, + ) + else: + self._target_points = np.asarray( + self.fixed_model.points, dtype=np.float64 + )[self._sample_slice] + + self._sampling_ready = True + self.log_debug( + "Sampling prepared: %d model points, %s target points, max distance %.3f mm", + self._sample_points.shape[0], + "no" if self._target_points is None else str(len(self._target_points)), + self._fixed_distance_map_max_distance, ) - image_size = self.fixed_distance_map.GetBufferedRegion().GetSize() - for i, base_point in enumerate(self._pca_template_model_points_itk): - if i % self.pca_template_model_point_subsample != 0: - continue - # Start with base point - point[0] = base_point[0] - point[1] = base_point[1] - point[2] = base_point[2] + def _sample_distance( + self, points: np.ndarray + ) -> tuple[np.ndarray, np.ndarray, int]: + """Sample the distance map and its spatial gradient at world points. - # Add PCA deformation if provided - point[0] += pca_deformation[i, 0] - point[1] += pca_deformation[i, 1] - point[2] += pca_deformation[i, 2] + Points outside the image are clamped to the grid and charged the extra + travel from the clamped location, which keeps the metric continuous and + gives the optimizer a gradient that pushes such points back inside -- + unlike a constant out-of-bounds penalty, which is flat. - if self.post_pca_transform is not None: - point = self.post_pca_transform.TransformPoint(point) + Args: + points: (n, 3) array of world-space points - # Check if point is inside image bounds + Returns: + Tuple of (distances (n,), world-space gradients (n, 3), n_outside) + """ + index = (points - self._image_origin) @ self._world_to_index.T + clamped = np.clip(index, 0.0, self._image_size - 1.0) + outside = index - clamped + is_outside = np.any(outside != 0.0, axis=1) + n_outside = int(np.count_nonzero(is_outside)) + + # map_coordinates indexes the array as (k, j, i). + array_coordinates = clamped[:, ::-1] + distances = map_coordinates( + self._distance_array, array_coordinates.T, order=1, mode="nearest" + ) - coord_index = ( - self.fixed_distance_map.TransformPhysicalPointToContinuousIndex(point) + # d/dc_a of the trilinear interpolant is the forward difference along a, + # interpolated linearly across the other two axes within the same cell. + gradient_index = np.zeros_like(clamped) + for axis in range(3): + differences = self._distance_forward_diff[axis] + if differences is None: + continue + coordinates = array_coordinates.copy() + coordinates[:, 2 - axis] = np.clip( + np.floor(clamped[:, axis]), 0.0, self._image_size[axis] - 2.0 ) - if ( - 0 <= coord_index[0] < image_size[0] - and 0 <= coord_index[1] < image_size[1] - and 0 <= coord_index[2] < image_size[2] - ): - center[0] += point[0] - center[1] += point[1] - center[2] += point[2] - distance = ( - self._fixed_distance_map_interpolator.EvaluateAtContinuousIndex( - coord_index - ) - ) - total_distance += distance - n_valid_points += 1 - else: - n_invalid_points += 1 - - if n_invalid_points >= 0.05 * n_valid_points: - self.log_warning( - "%d of %d mapped outside of image. Rejecting.", - n_invalid_points, - n_valid_points + n_invalid_points, + gradient_index[:, axis] = map_coordinates( + differences, coordinates.T, order=1, mode="nearest" ) - return self._fixed_distance_map_max_distance + # A clamped axis cannot change the sampled value, so it carries no + # gradient from the map; the out-of-bounds term supplies it instead. + gradient_index[outside != 0.0] = 0.0 + gradients = gradient_index @ self._world_to_index + + if n_outside: + outside_world = outside @ self._index_to_world.T + outside_distance = np.linalg.norm(outside_world, axis=1) + safe = np.where(outside_distance > 0.0, outside_distance, 1.0) + distances = distances + outside_distance + gradients = gradients + outside_world / safe[:, None] + + return distances, gradients, n_outside + + def _deform(self, pca_coefficients: np.ndarray) -> np.ndarray: + """Deform the subsampled template points into world space.""" + n_modes = len(pca_coefficients) + points = self._sample_points + np.tensordot( + pca_coefficients, self._sample_modes[:n_modes], axes=(0, 0) + ) + return self._apply_post_pca_transform(points) - # Compute mean distance - mean_distance = total_distance / n_valid_points - center /= n_valid_points + def _objective_and_gradient(self, params: np.ndarray) -> tuple[float, np.ndarray]: + """Evaluate the registration objective and its gradient. - log_level_int = ( - self.log_level - if isinstance(self.log_level, int) - else logging.getLevelName(self.log_level) - ) - if log_level_int <= logging.DEBUG or self._metric_call_count % 100 == 0: - self.log_info( - " Metric %d: %s -> %f", - (self._metric_call_count + 1), - center, - mean_distance, + The objective is MINIMIZED: the distance map is zero on the target + surface and grows away from it. + + Args: + params: PCA coefficients b, in units of standard deviations + + Returns: + Tuple of (objective value in mm, gradient with respect to params) + """ + if not self._sampling_ready: + self._prepare_sampling() + + n_modes = len(params) + modes_world = self._sample_modes_world[:n_modes] + points = self._deform(params) + + distances, gradients, n_outside = self._sample_distance(points) + forward_distance = float(distances.mean()) + # d/db_j of mean_i D(p_i) = mean_i grad_D(p_i) . (sigma_j * v_ij) + forward_gradient = np.einsum("ia,jia->j", gradients, modes_world) / len(points) + + weight = self.symmetric_weight if self._target_points is not None else 0.0 + reverse_distance = 0.0 + reverse_gradient = np.zeros(n_modes, dtype=np.float64) + if weight > 0.0: + assert self._target_points is not None, "target points must be set" + # Target -> model: each target point is charged its distance to the + # nearest deformed model point, so a model that covers only part of + # the target scores badly. The forward term alone cannot see this. + nearest, nearest_index = cKDTree(points).query(self._target_points) + nearest_distance = np.atleast_1d(np.asarray(nearest, dtype=np.float64)) + reverse_distance = float(nearest_distance.mean()) + safe = np.where(nearest_distance > 0.0, nearest_distance, 1.0) + direction = (points[nearest_index] - self._target_points) / safe[:, None] + # Accumulate each target's pull onto the model point it selected, + # then contract once against the modes. + pull = np.zeros_like(points) + np.add.at(pull, nearest_index, direction) + pull /= len(self._target_points) + reverse_gradient = np.einsum("ia,jia->j", pull, modes_world) + + objective = (1.0 - weight) * forward_distance + weight * reverse_distance + gradient = (1.0 - weight) * forward_gradient + weight * reverse_gradient + + prior = 0.0 + if self.pca_prior_weight > 0.0: + prior = self.pca_prior_weight * float(np.dot(params, params)) + objective += prior + gradient = gradient + 2.0 * self.pca_prior_weight * params + + if n_outside > 0.25 * len(points): + self.log_warning( + "%d of %d model points mapped outside the distance map.", + n_outside, + len(points), ) + + if self.log_level <= logging.DEBUG or self._metric_call_count % 25 == 0: self.log_info( - " Params %s", - params, + " Metric %d: %.4f mm (model->target %.4f, target->model %.4f, " + "prior %.4f, outside %d)", + self._metric_call_count + 1, + objective, + forward_distance, + reverse_distance, + prior, + n_outside, ) + self.log_debug(" Params %s", params) self._metric_call_count += 1 - return mean_distance + return objective, gradient + + def _mean_distance_metric(self, params: np.ndarray) -> float: + """Evaluate the registration objective at the given PCA coefficients. + + Args: + params: PCA coefficients b, in units of standard deviations + + Returns: + Objective value, in mm. Lower is better. + """ + return self._objective_and_gradient(np.asarray(params, dtype=np.float64))[0] + + def _apply_post_pca_transform(self, points: np.ndarray) -> np.ndarray: + """Apply post_pca_transform to an (n, 3) array of world points.""" + affine = self._get_post_pca_affine() + if affine is not None: + return np.asarray(points @ affine[0].T + affine[1], dtype=np.float64) + if self.post_pca_transform is None: + return points + transformed = np.empty_like(points) + point = itk.Point[itk.D, 3]() + for i, source in enumerate(points): + point[0], point[1], point[2] = (float(v) for v in source) + result = self.post_pca_transform.TransformPoint(point) + transformed[i] = (result[0], result[1], result[2]) + return transformed def _compute_pca_deformation(self, pca_coefficients: np.ndarray) -> np.ndarray: """Compute PCA deformation vectors for all points. @@ -552,28 +776,18 @@ def _compute_pca_deformation(self, pca_coefficients: np.ndarray) -> np.ndarray: displacement = Σ(b_i * std_i * pca_eigenvector_i) Args: - pca_coefficients: Array of PCA coefficients b_i (one per mode) - pca_number_of_modes: Number of PCA modes to use. Default: use all available modes + pca_coefficients: Array of PCA coefficients b_i. Only as many modes + as there are coefficients contribute. Returns: - Nx3 array of deformation vectors (displacement from mean shape) + Nx3 array of deformation vectors (displacement from the template) """ - # Initialize deformation to zero - deformation = np.zeros((self.pca_template_model.n_points, 3), dtype=np.float64) - - # Add contribution from each PCA mode - for i in range(self.pca_number_of_modes): - pca_eigenvector_flat = self.pca_eigenvectors[i, :] - - # Reshape to (N, 3) - pca_eigenvector_3d = pca_eigenvector_flat.reshape(-1, 3) - - # Add weighted deformation: b_i * std_i * pca_eigenvector_i - deformation += ( - pca_coefficients[i] * self.pca_std_deviations[i] * pca_eigenvector_3d - ) - - return deformation + n_modes = len(pca_coefficients) + scaled = pca_coefficients * self.pca_std_deviations[:n_modes] + deformation = np.asarray( + scaled @ self.pca_eigenvectors[:n_modes], dtype=np.float64 + ) + return deformation.reshape(-1, 3) def _optimize_pca_coefficients( self, @@ -584,12 +798,12 @@ def _optimize_pca_coefficients( ) -> tuple[np.ndarray, float]: """Optimize PCA coefficients - This method optimizes PCA mode coefficients to deform the model to better match - low values in the distance map. + Minimizes the mean distance between the deformed model and the target, + supplying the analytic gradient of the objective to the optimizer. Args: pca_number_of_modes: Number of PCA modes to use in optimization. Using fewer - modes provides smoother deformations. Default: 10 + modes provides smoother deformations. Default: 0 (use all) pca_coefficient_bounds: Bound on PCA coefficients in units of std deviations. Default: 3.0 (±3 std deviations per mode) method: Optimization method for scipy.optimize.minimize. @@ -600,52 +814,65 @@ def _optimize_pca_coefficients( Returns: Tuple of (pca_coefficients, mean_distance): - pca_coefficients: Optimized PCA coefficients - - mean_distance: Final mean distance metric value + - mean_distance: Final objective value, in mm Raises: ValueError: If number of PCA modes to use exceeds available modes """ + n_available = len(self.pca_eigenvectors) if pca_number_of_modes <= 0: - pca_number_of_modes = len(self.pca_eigenvectors) - if pca_number_of_modes > len(self.pca_eigenvectors): + pca_number_of_modes = n_available + if pca_number_of_modes > n_available: raise ValueError( - f"Number of PCA modes to use ({pca_number_of_modes}) exceeds available modes ({len(self.pca_std_deviations)})" + f"Number of PCA modes to use ({pca_number_of_modes}) exceeds " + f"available modes ({n_available})" ) self.pca_number_of_modes = pca_number_of_modes + self._prepare_sampling() + self.log_info(f"Number of PCA modes: {pca_number_of_modes}") self.log_info( f"PCA coefficient bounds: ±{pca_coefficient_bounds} std deviations" ) self.log_info(f"Optimization method: {method}") self.log_info(f"Max iterations: {max_iterations}") + self.log_info(f"Shape prior weight: {self.pca_prior_weight}") + self.log_info(f"Symmetric weight: {self.symmetric_weight}") - bounds = [] - for _ in range(pca_number_of_modes): - bounds.append((-pca_coefficient_bounds, pca_coefficient_bounds)) + bounds = [ + (-pca_coefficient_bounds, pca_coefficient_bounds) + for _ in range(pca_number_of_modes) + ] - log_level_int = ( - self.log_level - if isinstance(self.log_level, int) - else logging.getLevelName(self.log_level) - ) - disp = log_level_int <= logging.INFO + disp = self.log_level <= logging.INFO + + # The metric is in mm, so the default gradient tolerance is meaningful. + # Without an analytic gradient the finite-difference step must be large + # enough to move sample points by a useful fraction of a voxel. + options: dict = {"maxiter": max_iterations, "disp": disp, "gtol": 1e-6} + if not self._analytic_gradient: + options["eps"] = 1e-2 self.log_info("Running optimization...") result_pca = minimize( # type: ignore[call-overload] - lambda params: self._mean_distance_metric(params), - np.zeros(self.pca_number_of_modes), + self._objective_and_gradient + if self._analytic_gradient + else self._mean_distance_metric, + np.zeros(pca_number_of_modes), method=method, + jac=self._analytic_gradient, bounds=bounds, - options={"maxiter": max_iterations, "disp": disp}, + options=options, ) optimized_pca_coefficients = result_pca.x - optimized_mean_distance = result_pca.fun + optimized_mean_distance = float(result_pca.fun) self.log_info("Optimization completed!") self.log_info(f"Optimized PCA coefficients: {optimized_pca_coefficients}") - self.log_info(f"Final mean intensity: {optimized_mean_distance:.2f}") + self.log_info(f"Metric evaluations: {self._metric_call_count}") + self.log_info(f"Final mean distance: {optimized_mean_distance:.4f} mm") return optimized_pca_coefficients, optimized_mean_distance @@ -672,35 +899,12 @@ def transform_template_model(self) -> pv.DataSet: self.registered_model_pca_coefficients, ) - # Apply deformation and affine transform to each point - final_points = np.zeros((self.pca_template_model.n_points, 3), dtype=np.float64) - - n_points = self.pca_template_model.n_points - progress_interval = max(1, n_points // 10) # Report progress every 10% - - point = itk.Point[itk.D, 3]() - for i in range(n_points): - # Report progress - if i % progress_interval == 0 or i == n_points - 1: - self.log_progress(i + 1, n_points, prefix="Transforming points") - - # Start with mean shape point - point[0] = float(self.pca_template_model.points[i][0]) - point[1] = float(self.pca_template_model.points[i][1]) - point[2] = float(self.pca_template_model.points[i][2]) - - # Add PCA deformation - point[0] += self.registered_model_pca_deformation[i, 0] - point[1] += self.registered_model_pca_deformation[i, 1] - point[2] += self.registered_model_pca_deformation[i, 2] - - if self.post_pca_transform is not None: - point = self.post_pca_transform.TransformPoint(point) - - # Store result - final_points[i, 0] = point[0] - final_points[i, 1] = point[1] - final_points[i, 2] = point[2] + # Deform in the template frame, then map into target space. + deformed_points = ( + np.asarray(self.pca_template_model.points, dtype=np.float64) + + self.registered_model_pca_deformation + ) + final_points = self._apply_post_pca_transform(deformed_points) # Create new model with transformed points self.registered_model = self.pca_template_model.copy(deep=True) @@ -717,30 +921,36 @@ def transform_point( point: itk.Point, include_post_pca_transform: bool = True, ) -> itk.Point: - """Transform an arbitrary point using nearest neighbor interpolation. + """Transform an arbitrary point through the PCA deformation field. Args: point: ITK point to transform (itk.Point[itk.D, 3]) + include_post_pca_transform: Also apply post_pca_transform. Default: True Returns: Transformed ITK point + Raises: + ValueError: If compute_pca_transforms() has not been called yet + Notes: - 1) if the point is outside the image bounds, the point is not transformed. - 2) if the forward point transform is set, it is applied. - 3) if the post_pca_transform is set and enabled, it is applied. - 4) if the forward point transform is not set, no errors are raised. + This samples the *approximated* deformation field built by + compute_pca_transforms(), which is splatted and blurred, so it does + not reproduce transform_template_model() exactly; the RMS of that + difference is logged when the field is built. Points outside the + field's reference image are not displaced. Example: >>> p = itk.Point[itk.D, 3]() >>> p[0], p[1], p[2] = 10.0, 20.0, 30.0 >>> transformed_p = registrar.transform_point(p) """ - - if self.forward_point_transform is not None: - transformed_point = self.forward_point_transform.TransformPoint(point) - else: - transformed_point = point + if self.forward_point_transform is None: + self.log_error("Forward point transform is not set.") + raise ValueError( + "compute_pca_transforms() must be called before transform_point()" + ) + transformed_point = self.forward_point_transform.TransformPoint(point) if include_post_pca_transform and self.post_pca_transform is not None: transformed_point = self.post_pca_transform.TransformPoint( @@ -749,9 +959,21 @@ def transform_point( return transformed_point - def compute_pca_transforms(self, reference_image: itk.Image) -> dict: + def compute_pca_transforms( + self, reference_image: itk.Image, blur_sigma: float = 2.5 + ) -> dict: """Compute PCA transforms. + The field is built by splatting the per-point PCA displacements onto the + reference grid and blurring them, so it only approximates the exact + per-point deformation. The RMS of that approximation error, and of the + forward/inverse round trip, are both logged. + + Args: + reference_image: ITK image providing the coordinate frame for the field. + blur_sigma: Sigma for Gaussian blurring of the deformation field. + Default: 2.5 + Returns: Dictionary containing: - 'forward_point_transform': POINT transform mapping template @@ -761,17 +983,19 @@ def compute_pca_transforms(self, reference_image: itk.Image) -> dict: Note: These are point transforms, oriented opposite to image-registration - transforms; see docs/developer/transform_conventions. + transforms; see docs/developer/transform_conventions. Neither + includes post_pca_transform. """ assert self.registered_model_pca_deformation is not None, ( "PCA deformation must be computed" ) + template_points = np.asarray(self.pca_template_model.points, dtype=np.float64) template_model_pca_deformation_field_image = ( self._contour_tools.create_deformation_field( - np.array(self.pca_template_model.points), + template_points, self.registered_model_pca_deformation, reference_image=reference_image, - blur_sigma=2.5, + blur_sigma=blur_sigma, ptype=itk.D, ) ) @@ -787,11 +1011,50 @@ def compute_pca_transforms(self, reference_image: itk.Image) -> dict: self.forward_point_transform ) ) + + self._log_transform_fidelity(template_points) + return { "forward_point_transform": self.forward_point_transform, "inverse_point_transform": self.inverse_point_transform, } + def _log_transform_fidelity(self, template_points: np.ndarray) -> None: + """Report how well the field reproduces the deformation and inverts.""" + assert self.forward_point_transform is not None, "forward transform must be set" + assert self.inverse_point_transform is not None, "inverse transform must be set" + assert self.registered_model_pca_deformation is not None, ( + "PCA deformation must be computed" + ) + + # Two TransformPoint calls per point is costly on dense templates, and a + # strided subset reports the same RMS to within sampling noise. + stride = max(1, len(template_points) // 5000) + sampled = template_points[::stride] + deformation = self.registered_model_pca_deformation[::stride] + + point = itk.Point[itk.D, 3]() + forward = np.empty_like(sampled) + round_trip = np.empty_like(sampled) + for i, source in enumerate(sampled): + point[0], point[1], point[2] = (float(v) for v in source) + mapped = self.forward_point_transform.TransformPoint(point) + forward[i] = (mapped[0], mapped[1], mapped[2]) + back = self.inverse_point_transform.TransformPoint(mapped) + round_trip[i] = (back[0], back[1], back[2]) + + expected = sampled + deformation + field_rms = float(np.sqrt(np.mean(np.sum((forward - expected) ** 2, axis=1)))) + inverse_rms = float( + np.sqrt(np.mean(np.sum((round_trip - sampled) ** 2, axis=1))) + ) + self.log_info( + "Deformation field RMS error: %.4f mm (approximation of the " + "per-point deformation)", + field_rms, + ) + self.log_info("Forward/inverse round-trip RMS error: %.4f mm", inverse_rms) + def register( self, pca_number_of_modes: int = 0, @@ -799,40 +1062,40 @@ def register( method: str = "L-BFGS-B", max_iterations: int = 100, ) -> dict: - """Optimize PCA coefficients to deform the model to better match - low values in the distance map. + """Optimize PCA coefficients to deform the model onto the target. Args: pca_number_of_modes: Number of PCA modes to use. Default: 0 (use all available modes) - pca_coefficient_bounds: PCA coefficient bounds (±std devs). Default: 3.0 + pca_coefficient_bounds: PCA coefficient bounds (±std devs). Default: 3.5 method: Optimization method for scipy.optimize.minimize. Default: 'L-BFGS-B' (supports bounds) max_iterations: Maximum number of optimization iterations. - Default: 50 + Default: 100 Returns: Dictionary containing: - 'registered_model': Final registered PyVista model - 'pca_coefficients': Optimized PCA coefficients - - 'mean_distance': Final mean distance metric value + - 'mean_distance': Final objective value, in mm Raises: - ValueError: If reference image is not set + ValueError: If the distance map is not set Example: >>> result = registrar.register(pca_number_of_modes=10) >>> result['registered_model'].save('registered_heart.vtk') """ if self.fixed_distance_map is None: - raise ValueError("Reference image must be set before registration") + raise ValueError("A distance map must be set before registration") if pca_number_of_modes <= 0: pca_number_of_modes = self.pca_number_of_modes - self.log_section("PCA-BASED MODEL-TO-IMAGE REGISTRATION", width=70) + self.log_section("PCA-BASED MODEL-TO-MODEL REGISTRATION", width=70) self.log_info(f"Number of points: {self.pca_template_model.n_points}") self.log_info(f"Modes to use: {pca_number_of_modes}") + self._metric_call_count = 0 self.registered_model_pca_coefficients, self.registered_model_mean_distance = ( self._optimize_pca_coefficients( pca_number_of_modes=pca_number_of_modes, diff --git a/src/physiotwin4d/register_time_series_images.py b/src/physiotwin4d/register_time_series_images.py index 074fb320..4da6f2e7 100644 --- a/src/physiotwin4d/register_time_series_images.py +++ b/src/physiotwin4d/register_time_series_images.py @@ -23,25 +23,15 @@ class RegisterTimeSeriesImages(RegisterImagesBase): """Register a time series of images to a fixed image. - This class extends RegisterImagesBase to provide sequential registration - of multiple images (time series) to a fixed image, using a - caller-supplied registration backend. It can propagate information from - prior registrations to initialize subsequent ones. - - The registration proceeds in two passes from a reference frame: - - 1. Forward pass: from reference_frame to the end of the series - 2. Backward pass: from reference_frame-1 to the beginning - - This bidirectional approach helps maintain temporal coherence in the - registration results. + This class extends RegisterImagesBase to provide registration of multiple + images (time series) to a fixed image, using a caller-supplied registration + backend. Every frame is registered to the fixed image independently. Key features: - Sequential registration of ordered image lists - Supports any RegisterImagesBase backend, including RegisterImagesChain / RegisterImagesGreedyICON for multi-stage registration - - Optional use of prior transforms to initialize next registration - Configurable starting point in the time series - Returns all transforms and loss values for the entire series @@ -60,7 +50,6 @@ class RegisterTimeSeriesImages(RegisterImagesBase): ... moving_images=time_series_images, ... reference_frame=5, # Start from middle of cardiac cycle ... register_reference=True, - ... prior_weight=0.5, ... ) >>> >>> forward_tfms = result['forward_transforms'] # warp moving images -> fixed grid @@ -105,18 +94,6 @@ def __init__( self.transform_tools: TransformTools = TransformTools() - self.smooth_prior_transform_sigma: float = 0.5 - - def set_smooth_prior_transform_sigma( - self, smooth_prior_transform_sigma: float - ) -> None: - """Set the sigma for smoothing the prior transform. - - Args: - smooth_prior_transform_sigma (float): Sigma for smoothing the prior transform. - """ - self.smooth_prior_transform_sigma = smooth_prior_transform_sigma - def set_mask_dilation(self, mask_dilation_mm: float) -> None: """Set the dilation of the fixed and moving image masks. @@ -175,17 +152,12 @@ def register_time_series( moving_labelmaps: Optional[list[Optional[itk.Image]]] = None, reference_frame: int = 0, register_reference: bool = True, - prior_weight: float = 0.0, ) -> dict[str, list[itk.Transform] | list[float]]: """Register a time series of images to the fixed image. This method registers an ordered sequence of images to a common fixed - frame. Registration proceeds bidirectionally from a reference frame: - forward to the end and backward to the beginning. - - For each image after the reference image, the method can optionally use - the transform from the previous image to initialize the registration, - which can improve convergence and temporal coherence. + frame. The reference frame is registered first, then every other frame, + each independently of the others. Args: moving_images (list[itk.Image]): List of 3D images to register @@ -195,18 +167,12 @@ def register_time_series( moving_labelmaps (list[itk.Image], optional): Per-frame multi-label segmentations, one for each moving image. If None, no labelmaps are used. If provided, must have the same length as moving_images. Default: None - reference_frame (int, optional): Index of the reference image to register first. - Registration proceeds forward from this index to the end, then - backward from this index to the beginning. Default: 0 + reference_frame (int, optional): Index of the reference image, which + is registered first. Default: 0 register_reference (bool, optional): If True, register the reference image to the fixed image. If False, use identity transform for the reference image. Default: True - prior_weight (float, optional): - Weight (0.0 to 1.0) for using the prior image's transform to - initialize the next registration. 0.0 means no prior information - is used (each registration starts from identity). Higher values - provide more temporal smoothness but may propagate errors. - Default: 0.0 + Returns: dict: Dictionary containing results: - "forward_transforms" (list[itk.Transform]): one per image; @@ -222,13 +188,11 @@ def register_time_series( Raises: ValueError: If fixed_image is not set ValueError: If reference_frame is out of range - ValueError: If prior_weight not in [0, 1] ValueError: If moving_masks length doesn't match moving_images length Note: - The method compares registration with identity initialization versus - prior transform initialization and selects the result with lower loss. - This helps prevent error propagation in the temporal sequence. + Every frame is registered independently, so an error in one frame + cannot propagate along the series. The fixed image mask can be set using set_fixed_mask() before calling this method. @@ -245,7 +209,6 @@ def register_time_series( ... moving_labelmaps=labelmap_list, # Optional ... reference_frame=5, ... register_reference=True, - ... prior_weight=0.5, ... ) >>> >>> # Access results using new intuitive names @@ -273,9 +236,6 @@ def register_time_series( f"reference_frame {reference_frame} out of range [0, {num_images - 1}]" ) - if not 0.0 <= prior_weight <= 1.0: - raise ValueError("prior_weight must be in [0.0, 1.0]") - if moving_masks is not None and len(moving_masks) != num_images: raise ValueError( f"moving_masks length ({len(moving_masks)}) must match " @@ -329,97 +289,25 @@ def register_time_series( inverse_transforms[reference_frame] = inverse_transform losses[reference_frame] = loss - # Compute prior transform for reference frame if needed - prior_forward_ref = None - if prior_weight > 0.0: - prior_forward_ref = ( - self.transform_tools.combine_displacement_field_transforms( - identity_tfm, - forward_transform, - self.fixed_image, - tfm1_weight=1.0, - tfm2_weight=prior_weight, - tfm1_blur_sigma=0.0, - tfm2_blur_sigma=0.5, - mode="add", - ) + # Register every remaining frame; each is independent of the others. + for img_idx in range(num_images): + if img_idx == reference_frame: + continue + moving_image = moving_images[img_idx] + moving_mask = moving_masks[img_idx] if moving_masks is not None else None + moving_labelmap = ( + moving_labelmaps[img_idx] if moving_labelmaps is not None else None ) - # Register forward and backward from reference frame - for step, start_idx, end_idx in [ - (1, reference_frame + 1, num_images), # Forward pass - (-1, reference_frame - 1, -1), # Backward pass - ]: - prior_forward = prior_forward_ref - - for img_idx in range(start_idx, end_idx, step): - moving_image = moving_images[img_idx] - moving_mask = ( - moving_masks[img_idx] if moving_masks is not None else None - ) - moving_labelmap = ( - moving_labelmaps[img_idx] if moving_labelmaps is not None else None - ) + result = self.registrar.register( + moving_image=moving_image, + moving_mask=moving_mask, + moving_labelmap=moving_labelmap, + ) - # Try registration with identity initialization - result_init_identity = self.registrar.register( - moving_image=moving_image, - moving_mask=moving_mask, - moving_labelmap=moving_labelmap, - ) - forward_init_identity = result_init_identity["forward_transform"] - inverse_init_identity = result_init_identity["inverse_transform"] - loss_init_identity = result_init_identity["loss"] - - # Select best result based on prior usage - if prior_weight > 0.0: - # Try with prior transform initialization - result_init_prior = self.registrar.register( - moving_image=moving_image, - moving_mask=moving_mask, - moving_labelmap=moving_labelmap, - initial_forward_transform=prior_forward, - ) - forward_init_prior = result_init_prior["forward_transform"] - inverse_init_prior = result_init_prior["inverse_transform"] - loss_init_prior = result_init_prior["loss"] - - # Select result with lower loss - if loss_init_identity < loss_init_prior: - # Identity initialization was better - prior_forward = identity_tfm - forward_transform = forward_init_identity - inverse_transform = inverse_init_identity - loss = loss_init_identity - else: - # Prior initialization was better - forward_transform = forward_init_prior - inverse_transform = inverse_init_prior - loss = loss_init_prior - - # Update prior for next iteration - prior_forward = ( - self.transform_tools.combine_displacement_field_transforms( - identity_tfm, - forward_transform, - self.fixed_image, - tfm1_weight=1.0, - tfm2_weight=prior_weight, - tfm1_blur_sigma=0.0, - tfm2_blur_sigma=self.smooth_prior_transform_sigma, - mode="add", - ) - ) - else: - # No prior usage, just use identity result - forward_transform = forward_init_identity - inverse_transform = inverse_init_identity - loss = loss_init_identity - - # Store results - forward_transforms[img_idx] = forward_transform - inverse_transforms[img_idx] = inverse_transform - losses[img_idx] = loss + forward_transforms[img_idx] = result["forward_transform"] + inverse_transforms[img_idx] = result["inverse_transform"] + losses[img_idx] = cast(float, result["loss"]) assert all(t is not None for t in forward_transforms) assert all(t is not None for t in inverse_transforms) @@ -567,7 +455,6 @@ def registration_method( moving_mask: Optional[itk.Image] = None, moving_labelmap: Optional[itk.Image] = None, moving_image_pre: Optional[itk.Image] = None, - initial_forward_transform: Optional[itk.Transform] = None, ) -> dict[str, Union[itk.Transform, float]]: """Registration method required by RegisterImagesBase. @@ -581,7 +468,6 @@ def registration_method( moving_labelmap (itk.Image, optional): Multi-label segmentation moving_image_pre (itk.Image, optional): Ignored - the registrar computes its own preprocessing from the raw moving_image - initial_forward_transform (itk.Transform, optional): Initial transform Returns: dict: Registration result with forward_transform, inverse_transform, and loss @@ -592,7 +478,6 @@ def registration_method( moving_mask=moving_mask, moving_labelmap=moving_labelmap, moving_image_pre=None, - initial_forward_transform=initial_forward_transform, ) self._capture_delegate_result(self.registrar, result) return { diff --git a/src/physiotwin4d/train_physicsnemo_mgn.py b/src/physiotwin4d/train_physicsnemo_mgn.py index f50dcd74..82338eb4 100644 --- a/src/physiotwin4d/train_physicsnemo_mgn.py +++ b/src/physiotwin4d/train_physicsnemo_mgn.py @@ -66,6 +66,22 @@ def set_num_layers(self, num_layers: int) -> None: raise ValueError(f"num_layers must be >= 1, got {num_layers}") self.num_layers = num_layers + def set_num_processor_checkpoint_segments(self, num_segments: int) -> None: + """Set the gradient-checkpointing segment count for the processor. + + Gradient checkpointing recomputes processor activations during the + backward pass instead of storing them, trading compute for GPU memory. + A large mesh graph makes that trade worthwhile: a 179k-point, 1.07M-edge + template peaks near 43 GiB at ``batch_size`` 4 without it. + + Args: + num_segments: Number of checkpointed segments; ``0`` (the default) + disables checkpointing and stores every activation. + """ + if num_segments < 0: + raise ValueError(f"num_segments must be >= 0, got {num_segments}") + self.num_processor_checkpoint_segments = num_segments + def build_model(self, in_features: int, out_features: int) -> "torch.nn.Module": try: import torch_geometric # noqa: F401 - needed by the graph seams diff --git a/src/physiotwin4d/transform_tools.py b/src/physiotwin4d/transform_tools.py index 75f86efd..b891156a 100644 --- a/src/physiotwin4d/transform_tools.py +++ b/src/physiotwin4d/transform_tools.py @@ -82,6 +82,10 @@ def combine_displacement_field_transforms( In ``add`` mode, returns a single displacement field transform with weighted summed vectors. In ``compose`` mode, returns a composite transform containing both weighted displacement field transforms. + + ``compose`` follows ITK's CompositeTransform convention, where the + last-added transform is applied first: the result evaluates + ``tfm1(tfm2(x))``, so ``tfm2`` is the stage that runs first. """ assert mode in ["add", "compose"], "Invalid mode" @@ -255,9 +259,30 @@ def convert_transform_to_displacement_field_transform( new_tfm.SetDisplacementField(field) return new_tfm - def invert_displacement_field_transform(self, tfm: itk.Transform) -> itk.Transform: + def invert_displacement_field_transform( + self, + tfm: itk.Transform, + max_iterations: int = 20, + max_error_tolerance: float = 0.05, + mean_error_tolerance: float = 0.0005, + ) -> itk.Transform: """ Invert a displacement field transform. + + Uses SimpleITK's fixed-point iterative inversion on the input field's + own grid. The defaults are tighter than SimpleITK's own (10 iterations, + 0.1 mm max error) because the fields produced here are not smooth + everywhere and converge slowly near their support boundary. + + Args: + tfm: Displacement field transform to invert. + max_iterations: Fixed-point iterations per voxel. + max_error_tolerance: Convergence threshold on the maximum error, in + the field's units. + mean_error_tolerance: Convergence threshold on the mean error. + + Returns: + The inverted displacement field transform. """ assert "DisplacementFieldTransform" in str(type(tfm)), ( "Input transform must be a displacement field transform" @@ -268,7 +293,12 @@ def invert_displacement_field_transform(self, tfm: itk.Transform) -> itk.Transfo field_sitk = image_tools.convert_itk_image_to_sitk(field_itk) - field_sitk_inv = sitk.InvertDisplacementField(field_sitk) + field_sitk_inv = sitk.InvertDisplacementField( + field_sitk, + maximumNumberOfIterations=max_iterations, + maxErrorToleranceThreshold=max_error_tolerance, + meanErrorToleranceThreshold=mean_error_tolerance, + ) field_itk_inv = image_tools.convert_sitk_image_to_itk(field_sitk_inv) @@ -277,6 +307,38 @@ def invert_displacement_field_transform(self, tfm: itk.Transform) -> itk.Transfo return new_tfm + def invert_transform( + self, tfm: itk.Transform, reference_image: itk.Image + ) -> itk.Transform: + """Invert any transform, analytically when the type supports it. + + Prefers ITK's analytic inverse (available for translation, rigid, affine + and composites of them) and falls back to rasterizing a displacement + field over ``reference_image`` and inverting that numerically. + + The analytic inverse is preferred because the fallback is only defined on + ``reference_image``'s grid: outside it the field is zero, so the inverse + silently degrades to the identity there. + + Args: + tfm (itk.Transform): Transform to invert. + reference_image (itk.Image): Grid used by the displacement-field + fallback. + + Returns: + itk.Transform: The inverse transform. + """ + try: + analytic = tfm.GetInverseTransform() + except Exception: # pragma: no cover - transform type dependent + analytic = None + if analytic is not None: + return cast(itk.Transform, analytic) + + return self.invert_displacement_field_transform( + self.convert_transform_to_displacement_field_transform(tfm, reference_image) + ) + def transform_pvcontour( self, contour: pv.PolyData, diff --git a/src/physiotwin4d/usd_anatomy_tools.py b/src/physiotwin4d/usd_anatomy_tools.py index e99e8f99..7bc73fc5 100644 --- a/src/physiotwin4d/usd_anatomy_tools.py +++ b/src/physiotwin4d/usd_anatomy_tools.py @@ -844,6 +844,25 @@ def _resolve_render_params(self, anatomy_type: str) -> Optional[dict[str, Any]]: return self.render_params[key] return None + def resolve_anatomy_type(self, anatomy_type: str) -> Optional[str]: + """Return the material name *anatomy_type* selects, or ``None``. + + Lets callers test a group/organ name before applying it, instead of + catching the :exc:`ValueError` raised by + :meth:`apply_anatomy_material_to_mesh`. Matching is the same as that + method's (see :meth:`_resolve_render_params`), so ``"kidney_left"`` + returns ``"Kidney"``. + + Args: + anatomy_type: A group/organ name or registered render-params key. + + Returns: + The matching material name (e.g. ``"Heart"``), or ``None`` when + nothing matches. + """ + params = self._resolve_render_params(anatomy_type) + return str(params["name"]) if params is not None else None + def get_anatomy_diffuse_color( self, anatomy_type: str ) -> tuple[float, float, float]: diff --git a/src/physiotwin4d/workflow_convert_vtk_to_usd.py b/src/physiotwin4d/workflow_convert_vtk_to_usd.py index c3346ba9..a920a7f5 100644 --- a/src/physiotwin4d/workflow_convert_vtk_to_usd.py +++ b/src/physiotwin4d/workflow_convert_vtk_to_usd.py @@ -8,13 +8,14 @@ """ import logging +import re from pathlib import Path -from typing import Any, Literal, Optional, Sequence, Union +from typing import Any, Literal, Mapping, Optional, Sequence, Union import pyvista as pv import vtk -from .convert_vtk_to_usd import ConvertVTKToUSD, _split_usd_extension +from .convert_vtk_to_usd import ConvertVTKToUSD from .physiotwin4d_base import PhysioTwin4DBase from .usd_anatomy_tools import USDAnatomyTools from .usd_tools import USDTools @@ -42,7 +43,8 @@ def __init__( time_codes: Optional[list[float]] = None, appearance: AppearanceKind = "solid", solid_color: tuple[float, float, float] = (0.8, 0.8, 0.8), - anatomy_type: str = "heart", + anatomy_type: Optional[str] = None, + object_names: Optional[Sequence[str]] = None, colormap_primvar: Optional[str] = None, colormap_name: str = "viridis", colormap_intensity_range: Optional[tuple[float, float]] = None, @@ -74,8 +76,21 @@ def __init__( static_merge is False. If None, uses sequential integers [0, 1, 2, ...]. appearance: "solid" | "anatomy" | "colormap". solid_color: RGB in [0,1] when appearance == "solid". - anatomy_type: Anatomy material name when appearance == "anatomy" - (e.g. heart, lung, bone, soft_tissue). + anatomy_type: Anatomy material name applied to every mesh when + appearance == "anatomy" (e.g. heart, lung, bone, soft_tissue). + None (default) instead resolves a material per mesh prim from + that prim's name, so a stage whose objects are named after the + structures they hold gets per-structure materials (e.g. + ventricle_left vs. myocardium). A name matching no material + falls back to the object's ``field_data['AnatomyGroup']`` + (so "rib_left_3" still reaches the bone material) and then to + the "other" material. + object_names: Prim names aligned to input_meshes, used when + static_merge is True. None (default) derives them from each + mesh's ``field_data['SegmentationLabelNames']`` when that holds + exactly one name (as written by + :class:`WorkflowConvertImageToVTK`), and falls back to + ``{usd_project_name}_{index}`` otherwise. colormap_primvar: Primvar name for coloring when appearance == "colormap" (e.g. vtk_point_stress_c0). If None, a candidate is auto-picked when possible. colormap_name: Matplotlib colormap name when appearance == "colormap". @@ -84,9 +99,13 @@ def __init__( """ super().__init__(class_name=self.__class__.__name__, log_level=log_level) self.input_meshes = list(input_meshes) - self.usd_project_name, self._usd_extension = _split_usd_extension( - usd_project_name - ) + suffix = Path(usd_project_name).suffix + if suffix.lower() in {".usd", ".usda", ".usdc"}: + self.usd_project_name = usd_project_name[: -len(suffix)] + self._usd_extension = suffix + else: + self.usd_project_name = usd_project_name + self._usd_extension = ".usd" self.output_directory = Path(output_directory) self.separate_by_connectivity = separate_by_connectivity self.separate_by_cell_type = separate_by_cell_type @@ -97,6 +116,7 @@ def __init__( self.appearance = appearance self.solid_color = solid_color self.anatomy_type = anatomy_type + self.object_names = list(object_names) if object_names is not None else None self.colormap_primvar = colormap_primvar self.colormap_name = colormap_name self.colormap_intensity_range = colormap_intensity_range @@ -106,6 +126,51 @@ def __init__( "separate_by_connectivity and separate_by_cell_type cannot both be True" ) + def _read_object_annotations(self) -> list[tuple[Optional[str], Optional[str]]]: + """Return ``(structure name, anatomy group)`` per input mesh. + + Both come from the annotation :class:`WorkflowConvertImageToVTK` writes + onto each surface: the name from ``field_data['SegmentationLabelNames']`` + when it holds exactly one entry, the group from + ``field_data['AnatomyGroup']``. Either is ``None`` when absent. + """ + annotations: list[tuple[Optional[str], Optional[str]]] = [] + for mesh in self.input_meshes: + if not isinstance(mesh, pv.DataSet) and isinstance(mesh, vtk.vtkDataSet): + mesh = pv.wrap(mesh) + if not isinstance(mesh, pv.DataSet): + annotations.append((None, None)) + continue + label_names = mesh.field_data.get("SegmentationLabelNames") + groups = mesh.field_data.get("AnatomyGroup") + name = ( + str(label_names[0]) + if label_names is not None and len(label_names) == 1 + else None + ) + group = str(groups[0]) if groups is not None and len(groups) else None + annotations.append((name, group)) + return annotations + + def _anatomy_candidates( + self, mesh_path: str, object_groups: Mapping[str, str] + ) -> list[str]: + """Return the anatomy names to try for *mesh_path*, best match first. + + With ``anatomy_type`` set, that one name is the only candidate. Without + it, the prim's own name is tried first, then the anatomy group of the + object it came from — so ``"rib_left_3"``, which matches no material of + its own, still lands on the bone material through its group. + """ + if self.anatomy_type is not None: + return [self.anatomy_type] + # Connectivity/cell-type splitting appends "_objectN" to the object + # name; strip it to recover the name object_groups is keyed by. + leaf = mesh_path.rsplit("/", 1)[-1] + object_name = re.sub(r"_object\d+$", "", leaf) + group = object_groups.get(object_name) + return [object_name] if group is None else [object_name, group] + def process(self) -> dict[str, Any]: """ Run the full workflow: convert meshes to USD, then apply the chosen appearance. @@ -147,6 +212,33 @@ def process(self) -> dict[str, Any]: else "none" ) + # Object names only name prims in the static-merge layout; a time + # series writes one prim per part across all frames instead. + annotations = self._read_object_annotations() + object_names = None + if self.static_merge: + object_names = self.object_names + if object_names is None and any(name for name, _ in annotations): + object_names = [ + name or f"{self.usd_project_name}_{index}" + for index, (name, _) in enumerate(annotations) + ] + if object_names is not None: + self.log_info("Naming objects: %s", ", ".join(object_names)) + + # Anatomy group per object name, used as the fallback when the name + # itself matches no material (e.g. "rib_left_3" -> the bone group). + # Keyed by the prim names ConvertVTKToUSD will actually emit, which fall + # back to "_" when no object_names were derived. + object_groups: dict[str, str] = {} + if self.static_merge: + group_keys = object_names or [ + f"{self.usd_project_name}_{index}" for index in range(len(annotations)) + ] + for object_name, (_, group) in zip(group_keys, annotations): + if group is not None: + object_groups[object_name] = group + converter = ConvertVTKToUSD( data_basename=self.usd_project_name, input_polydata=self.input_meshes, @@ -156,6 +248,7 @@ def process(self) -> dict[str, Any]: solid_color=self.solid_color, static_merge=self.static_merge, time_codes=time_codes, + object_names=object_names, log_level=self.log_level, ) stage = converter.convert(str(output_usd)) @@ -191,9 +284,22 @@ def process(self) -> dict[str, Any]: elif self.appearance == "anatomy": anatomy_tools = USDAnatomyTools(stage, log_level=self.log_level) for mesh_path in mesh_paths: - anatomy_tools.apply_anatomy_material_to_mesh( - mesh_path, self.anatomy_type + candidates = self._anatomy_candidates(mesh_path, object_groups) + selected = next( + ( + candidate + for candidate in candidates + if anatomy_tools.resolve_anatomy_type(candidate) is not None + ), + None, ) + if selected is None: + self.log_warning( + "No anatomy material matches %s; using 'other'", + " or ".join(candidates), + ) + selected = "other" + anatomy_tools.apply_anatomy_material_to_mesh(mesh_path, selected) stage.Save() elif self.appearance == "colormap": diff --git a/src/physiotwin4d/workflow_create_statistical_model.py b/src/physiotwin4d/workflow_create_statistical_model.py index 0be75ea9..bf5769d2 100644 --- a/src/physiotwin4d/workflow_create_statistical_model.py +++ b/src/physiotwin4d/workflow_create_statistical_model.py @@ -40,7 +40,7 @@ class WorkflowCreateStatisticalModel(PhysioTwin4DBase): Attributes: sample_meshes (list): List of sample mesh DataSets (.vtk/.vtu/.vtp geometry) reference_mesh (pv.DataSet): Reference mesh; its surface is used for alignment - pca_number_of_components (int): Number of PCA components to retain + number_of_pca_components (int): Number of PCA components to retain reference_spatial_resolution (float): Resolution for reference image from mesh reference_buffer_factor (float): Buffer around mesh for reference image """ @@ -49,7 +49,7 @@ def __init__( self, sample_meshes: list[pv.DataSet], reference_mesh: pv.DataSet, - pca_number_of_components: int = 7, + number_of_pca_components: int = 7, reference_spatial_resolution: float = 1.0, reference_buffer_factor: float = 0.25, solve_for_surface_pca: bool = True, @@ -60,7 +60,7 @@ def __init__( Args: sample_meshes: List of sample mesh DataSets (PyVista PolyData or UnstructuredGrid). reference_mesh: Reference mesh; its surface is used to align all samples. - pca_number_of_components: Number of PCA components. Default 7. + number_of_pca_components: Number of PCA components. Default 7. reference_spatial_resolution: Isotropic resolution (mm) for reference image. Default 1.0. reference_buffer_factor: Buffer factor around mesh for reference image. Default 0.25. solve_for_surface_pca: Whether to reduce the reference mesh to a surface. Default True. @@ -71,7 +71,7 @@ def __init__( ) self.sample_meshes = list(sample_meshes) self.reference_mesh = reference_mesh - self.pca_number_of_components = pca_number_of_components + self.number_of_pca_components = number_of_pca_components self.reference_spatial_resolution = reference_spatial_resolution self.reference_buffer_factor = reference_buffer_factor self.solve_for_surface_pca = solve_for_surface_pca @@ -91,9 +91,9 @@ def __init__( self.pca_mean_surface: Optional[pv.PolyData] = None self.pca_mean_mesh: Optional[pv.DataSet] = None - def set_pca_number_of_components(self, n: int) -> None: + def set_number_of_pca_components(self, n: int) -> None: """Set number of PCA components to retain.""" - self.pca_number_of_components = n + self.number_of_pca_components = n def _step1_extract_surfaces(self) -> None: """Extract reference surface and all sample surfaces (notebook 1).""" @@ -239,11 +239,11 @@ def _step5_compute_pca(self) -> None: raise ValueError( f"At least 2 samples are required for PCA. Got {data_matrix.shape[0]} samples." ) - n_comp = min(self.pca_number_of_components, data_matrix.shape[0] - 1) - if n_comp < self.pca_number_of_components: + n_comp = min(self.number_of_pca_components, data_matrix.shape[0] - 1) + if n_comp < self.number_of_pca_components: self.log_warning( "Reducing PCA components from %d to %d (n_samples=%d)", - self.pca_number_of_components, + self.number_of_pca_components, n_comp, data_matrix.shape[0], ) diff --git a/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py b/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py index 44041aaf..30ec9e33 100644 --- a/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py +++ b/src/physiotwin4d/workflow_fit_statistical_model_to_patient.py @@ -23,6 +23,7 @@ """ import logging +from pathlib import Path from typing import Any, Optional, cast import itk @@ -84,7 +85,7 @@ class WorkflowFitStatisticalModelToPatient(PhysioTwin4DBase): registrar_Greedy (RegisterImagesGreedy): Greedy registration instance use_pca_registration (bool): Whether PCA registration is enabled (set via set_use_pca_registration) pca_model (dict): PCA model dict when PCA enabled; same structure as WorkflowCreateStatisticalModel output - pca_number_of_modes (int): Number of PCA modes when PCA enabled + number_of_pca_components (int): Number of PCA components when PCA enabled labelmap_interior_object_ids (list): List of labelmap IDs corresponding to interior objects that should not be used when computing a distance map. icp_forward_point_transform : ICP transforms @@ -116,7 +117,7 @@ class WorkflowFitStatisticalModelToPatient(PhysioTwin4DBase): ... ) >>> registrar.set_mask_dilation_mm(20) >>> # To enable PCA registration, call before process(): - >>> # registrar.set_use_pca_registration(True, pca_model=pca_model_dict, pca_number_of_modes=10) + >>> # registrar.set_use_pca_registration(True, pca_model=pca_model_dict, number_of_pca_components=10) >>> # To enable labelmap-to-image refinement: >>> # registrar.set_use_labelmap_to_image_registration(True, template_labelmap, organ_mesh_ids, organ_extra_ids, background_ids) >>> result = registrar.process() @@ -245,6 +246,9 @@ def __init__( # Parameters for labelmap and mask generation self.mask_dilation_mm: float = 10.0 # For binary registration mask generation + # Optional finetuned ICON checkpoint for the labelmap-to-labelmap stage + self.l2l_icon_weights_path: Optional[str] = None + # Stage 1: ICP alignment results self.icp_registrar: Optional[RegisterModelsICP] = None self.icp_inverse_point_transform: Optional[itk.Transform] = None @@ -253,13 +257,13 @@ def __init__( self.icp_template_model_surface: Optional[pv.PolyData] = None self.icp_template_labelmap: Optional[itk.Image] = None - # Stage 1.5: PCA registration results (optional; enable via set_use_pca_registration(True, pca_model, pca_number_of_modes)) + # Stage 1.5: PCA registration results (optional; enable via set_use_pca_registration(True, pca_model, number_of_pca_components)) self.use_pca_registration = False self.pca_registrar: Optional[RegisterModelsPCA] = None self.pca_forward_point_transform: Optional[itk.Transform] = None self.pca_inverse_point_transform: Optional[itk.Transform] = None self.pca_model: Optional[dict[str, Any]] = None - self.pca_number_of_modes: int = 0 + self.number_of_pca_components: int = 0 self.pca_coefficients: Optional[np.ndarray] = None self.pca_template_model: Optional[pv.DataSet] = None self.pca_template_model_surface: Optional[pv.PolyData] = None @@ -296,24 +300,44 @@ def set_mask_dilation_mm(self, mask_dilation_mm: float) -> None: """ self.mask_dilation_mm = mask_dilation_mm + def set_labelmap_to_labelmap_icon_weights_path(self, weights_path: str) -> None: + """Set a finetuned ICON checkpoint for the labelmap-to-labelmap stage. + + That stage (:meth:`register_labelmap_to_labelmap`) registers distance + maps rather than image intensities, so it benefits from weights + finetuned on distance maps -- e.g. by + ``tutorials/tutorial_02_lung_distancemap_finetune_icon.py``. The + labelmap-to-image stage keeps the stock weights: it registers the + patient image itself. + + Args: + weights_path: Path to an existing uniGradICON checkpoint. + + Raises: + FileNotFoundError: If weights_path does not exist. + """ + if not Path(weights_path).exists(): + raise FileNotFoundError(f"ICON weights not found: {weights_path}") + self.l2l_icon_weights_path = weights_path + def set_use_pca_registration( self, use_pca_registration: bool, pca_model: Optional[dict[str, Any]] = None, - pca_number_of_modes: int = 0, + number_of_pca_components: int = 0, use_surface: bool = False, ) -> None: """Set whether to use PCA-based registration and provide the PCA model. - When enabling (True), pca_model and pca_number_of_modes must be provided. + When enabling (True), pca_model and number_of_pca_components must be provided. Args: use_pca_registration: Whether to use PCA registration after ICP. pca_model: Required when use is True. PCA model dict (e.g. from WorkflowCreateStatisticalModel result["pca_model"]) with keys "eigenvalues" and "components". - pca_number_of_modes: Required when use is True. Number of PCA modes to use. - Default 0 means use all modes. + number_of_pca_components: Required when use is True. Number of PCA + components to use. Default 0 means use all components. use_surface: Whether to use the surface of the patient model for PCA registration. Raises: ValueError: If use is True and pca_model is None. @@ -324,10 +348,10 @@ def set_use_pca_registration( "When enabling PCA registration, pca_model must be provided." ) self.pca_model = pca_model - self.pca_number_of_modes = pca_number_of_modes + self.number_of_pca_components = number_of_pca_components else: self.pca_model = None - self.pca_number_of_modes = 0 + self.number_of_pca_components = 0 self.use_surface = use_surface self.use_pca_registration = use_pca_registration @@ -480,12 +504,17 @@ def register_model_to_model_icp(self) -> dict: def register_model_to_model_pca(self) -> dict: """Perform PCA-based registration after ICP alignment. - Uses RegisterModelsPCA class for intensity-based PCA registration. - This method requires PCA data to be set via set_pca_data(). + Uses RegisterModelsPCA to optimize shape coefficients against a distance + map of the patient. The statistical model's modes are defined in the + un-aligned template frame, so the registrar is given the raw template + and the ICP alignment is passed as its ``post_pca_transform``. Returns: dict: Dictionary containing: - - 'forward_point_transform': Rigid transform from PCA registration + - 'forward_point_transform': DisplacementFieldTransform mapping + un-aligned template points to their PCA-deformed positions. + It excludes the ICP alignment, which is applied separately. + - 'inverse_point_transform': its inverse - 'pca_coefficients': PCA shape coefficients - 'registered_template_model_surface': PCA-registered model surface @@ -514,13 +543,18 @@ def register_model_to_model_pca(self) -> dict: "inverse_point_transform": self.pca_inverse_point_transform, } + # PCA modes are directions in the statistical model's own training + # frame, so they must be added to the un-aligned template and the ICP + # alignment applied afterwards. Deforming the ICP-aligned template + # instead would yield A*mean + sum(b*sigma*v) rather than + # A*(mean + sum(b*sigma*v)), mis-rotating and mis-scaling every mode. pca_template_model: Optional[pv.DataSet] if self.use_surface: - pca_template_model = self.icp_template_model_surface + pca_template_model = self.template_model_surface fixed_model = self.patient_model_surface fixed_distance_map = None else: - pca_template_model = self.icp_template_model + pca_template_model = self.template_model fixed_model = self.combined_patient_model if self.patient_labelmap is not None: fixed_distance_map = self.labelmap_tools.create_distance_map( @@ -538,7 +572,8 @@ def register_model_to_model_pca(self) -> dict: self.pca_registrar = RegisterModelsPCA.from_pca_model( pca_template_model=pca_template_model, pca_model=self.pca_model, - pca_number_of_modes=self.pca_number_of_modes, + pca_number_of_modes=self.number_of_pca_components, + post_pca_transform=self.icp_forward_point_transform, fixed_model=fixed_model, fixed_distance_map=fixed_distance_map, reference_image=self.patient_image, @@ -556,35 +591,48 @@ def register_model_to_model_pca(self) -> dict: algorithm="dataset_surface" ) + # The PCA field is splatted at the *un-aligned* template's points, which + # generally fall outside the patient image, so grid it in the template's + # own frame; create_deformation_field drops samples that land off-grid. + # The ICP alignment is applied separately, after this field. + pca_field_reference_image = self.contour_tools.create_reference_image( + pca_template_model, + spatial_resolution=float(min(self.patient_image.GetSpacing())), + ) pca_transforms = self.pca_registrar.compute_pca_transforms( - reference_image=self.patient_image, + reference_image=pca_field_reference_image, ) self.pca_forward_point_transform = pca_transforms["forward_point_transform"] self.pca_inverse_point_transform = pca_transforms["inverse_point_transform"] if self.log_level == logging.DEBUG: - tfm_arr = itk.GetArrayFromImage( - self.pca_forward_point_transform.GetDisplacementField() - ) + tfm_field = self.pca_forward_point_transform.GetDisplacementField() + tfm_arr = itk.GetArrayFromImage(tfm_field) tfm_x_arr = tfm_arr[:, :, :, 0] tfm_y_arr = tfm_arr[:, :, :, 1] tfm_z_arr = tfm_arr[:, :, :, 2] tfm_x_img = itk.GetImageFromArray(tfm_x_arr) tfm_y_img = itk.GetImageFromArray(tfm_y_arr) tfm_z_img = itk.GetImageFromArray(tfm_z_arr) - tfm_x_img.CopyInformation(self.patient_image) - tfm_y_img.CopyInformation(self.patient_image) - tfm_z_img.CopyInformation(self.patient_image) - itk.imwrite(tfm_x_img, "pca_forward_point_transform_x.nii.gz") - itk.imwrite(tfm_y_img, "pca_forward_point_transform_y.nii.gz") - itk.imwrite(tfm_z_img, "pca_forward_point_transform_z.nii.gz") + tfm_x_img.CopyInformation(tfm_field) + tfm_y_img.CopyInformation(tfm_field) + tfm_z_img.CopyInformation(tfm_field) if self.use_surface: - assert self.icp_template_model is not None, "ICP template model must be set" - self.pca_template_model = self._transform_model_dataset( - self.icp_template_model, + # forward_point_transform excludes the post-PCA step and is defined + # in the un-aligned template frame, so warp the raw volumetric + # template with it and then apply the ICP alignment. + assert self.icp_forward_point_transform is not None, ( + "ICP forward transform must be set" + ) + deformed_template_model = self._transform_model_dataset( + self.template_model, self.pca_forward_point_transform, ) + self.pca_template_model = self._transform_model_dataset( + deformed_template_model, + self.icp_forward_point_transform, + ) else: self.pca_template_model = registered_model @@ -593,10 +641,22 @@ def register_model_to_model_pca(self) -> dict: self.registered_template_model = self.pca_template_model self.registered_template_model_surface = self.pca_template_model_surface - if self.icp_template_labelmap is not None: + if self.template_labelmap is not None: + # Resampling pulls back: a patient-grid sample is mapped by the ICP + # inverse into the deformed-template frame and then by the PCA + # inverse onto the un-deformed template. itk.CompositeTransform + # applies its transforms in reverse order of addition, so the ICP + # inverse is added last. Resampling the raw template labelmap in one + # step also avoids a second round of nearest-neighbor sampling. + assert self.icp_inverse_point_transform is not None, ( + "ICP inverse transform must be set" + ) + pca_image_transform = itk.CompositeTransform[itk.D, 3].New() + pca_image_transform.AddTransform(self.pca_inverse_point_transform) + pca_image_transform.AddTransform(self.icp_inverse_point_transform) self.pca_template_labelmap = self.transform_tools.transform_image( - self.icp_template_labelmap, - self.pca_inverse_point_transform, + self.template_labelmap, + pca_image_transform, self.patient_image, interpolation_method="nearest", ) @@ -642,12 +702,27 @@ def register_labelmap_to_labelmap(self) -> Optional[dict]: assert self.pca_template_model_surface is not None, ( "PCA template model surface must be set" ) + + # Create a padded patient image since often the surface of interest + # is not fully contained within the original image, which causes trouble + # with distance map registration. The margin is physical -- it has to + # hold the dilated masks -- so it is converted per axis rather than + # padding a fixed voxel count on grids of any spacing. + margin_mm = 2.5 * self.mask_dilation_mm + spacing = np.asarray(self.patient_image.GetSpacing(), dtype=np.float64) + pad_voxels = np.maximum(1, np.ceil(margin_mm / spacing)).astype(int).tolist() + padded_patient_image = ImageTools().pad_image( + self.patient_image, pad_voxels=pad_voxels, background_value=-1000 + ) labelmap_registrar = RegisterModelsDistanceMaps( moving_model=self.pca_template_model_surface, fixed_model=self.patient_model_surface, - reference_image=self.patient_image, + reference_image=padded_patient_image, mask_dilation_mm=self.mask_dilation_mm, + distance_squared_max=(1.25 * self.mask_dilation_mm) ** 2, ) + if self.l2l_icon_weights_path is not None: + labelmap_registrar.set_icon_weights_path(self.l2l_icon_weights_path) # Run deformable registration l2l_result = labelmap_registrar.register( @@ -772,10 +847,10 @@ def register_labelmap_to_image( self.registrar_ICON.set_fixed_image(self.patient_image) self.registrar_ICON.set_fixed_mask(patient_mask) - # Perform Icon registration - result = self.registrar_ICON.register( - initial_forward_transform=self.l2i_forward_transform, - moving_image=template_labelmap, + # Perform Icon registration, refining the alignment found so far + result = self.registrar_ICON.register_from( + self.l2i_forward_transform, + template_labelmap, moving_mask=template_mask, ) self.l2i_inverse_transform = result["inverse_transform"] @@ -847,9 +922,10 @@ def transform_model( transformed_model = base_model.copy(deep=True) transform_steps: list[tuple[str, itk.Transform]] = [] - if self.icp_forward_point_transform is not None: - transform_steps.append(("ICP", self.icp_forward_point_transform)) if self.pca_coefficients is not None: + # PCA registration runs in the un-aligned template frame and carries + # the ICP alignment in its post-PCA transform, so ICP must not be + # applied again here. assert self.pca_registrar is not None, "PCA registrar must be set" pca_transform = ( self.pca_forward_point_transform @@ -861,6 +937,8 @@ def transform_model( transform_steps.append( ("PCA post-transform", self.pca_registrar.post_pca_transform) ) + elif self.icp_forward_point_transform is not None: + transform_steps.append(("ICP", self.icp_forward_point_transform)) if self.use_l2l_registration and self.l2l_inverse_transform is not None: transform_steps.append(("Labelmap-to-labelmap", self.l2l_inverse_transform)) if self.use_l2i_registration and self.l2i_inverse_transform is not None: @@ -908,9 +986,7 @@ def process( Returns: dict with registered_template_model and registered_template_model_surface """ - self.log_section( - "STARTING COMPLETE MODEL-TO-IMAGE-AND-MODEL REGISTRATION WORKFLOW", width=70 - ) + self.log_section("STARTING COMPLETE MODEL REGISTRATION WORKFLOW", width=70) self.use_ICON_registration_refinement = use_ICON_registration_refinement diff --git a/src/physiotwin4d/workflow_reconstruct_highres_4d_ct.py b/src/physiotwin4d/workflow_reconstruct_highres_4d_ct.py index a6c96951..3d1ed2db 100644 --- a/src/physiotwin4d/workflow_reconstruct_highres_4d_ct.py +++ b/src/physiotwin4d/workflow_reconstruct_highres_4d_ct.py @@ -64,7 +64,6 @@ class WorkflowReconstructHighres4DCT(PhysioTwin4DBase): reference_time_frame (int): Index of reference time frame in time series register_reference_time_frame_to_reference_image (bool): Whether to register the reference time frame to the reference image - prior_weight (float): Weight for temporal smoothing (0.0-1.0) upsample_to_fixed_resolution (bool): Whether to upsample reconstruction registrar (RegisterTimeSeriesImages): Internal registration object forward_transforms (list[itk.Transform]): one per frame; each warps its @@ -152,7 +151,6 @@ def __init__( ) # Initialize parameters with defaults - self.prior_weight: float = 0.0 self.upsample_to_fixed_resolution: bool = True self.modality: str = "ct" self.mask_dilation_mm: float = 0.0 @@ -170,22 +168,6 @@ def __init__( self.losses: Optional[list[float]] = None self.reconstructed_images: Optional[list[itk.Image]] = None - def set_prior_weight(self, prior_weight: float) -> None: - """Set the weight for temporal smoothing with prior transforms. - - Args: - prior_weight (float): Weight (0.0 to 1.0) for using the prior image's - transform to initialize the next registration. 0.0 means no prior - information is used (each registration starts from identity). - Higher values provide more temporal smoothness but may propagate errors. - - Raises: - ValueError: If prior_weight not in [0.0, 1.0] - """ - if not 0.0 <= prior_weight <= 1.0: - raise ValueError(f"prior_weight must be in [0.0, 1.0], got {prior_weight}") - self.prior_weight = prior_weight - def set_modality(self, modality: str) -> None: """Set the imaging modality for registration optimization. @@ -267,7 +249,6 @@ def register_time_series(self) -> dict: "Register reference time frame to reference image: " f"{self.register_reference_time_frame_to_reference_image}" ) - self.log_info(f"Prior weight: {self.prior_weight}") # Perform registration result = self.registrar.register_time_series( @@ -275,7 +256,6 @@ def register_time_series(self) -> dict: moving_masks=self.moving_masks, reference_frame=self.reference_time_frame, register_reference=self.register_reference_time_frame_to_reference_image, - prior_weight=self.prior_weight, ) # Store results @@ -378,7 +358,6 @@ def process(self) -> dict: registrar_type = type(self.registrar.registrar).__name__ self.log_info(f" Registration method: {registrar_type}") self.log_info(f" Reference time frame: {self.reference_time_frame}") - self.log_info(f" Prior weight: {self.prior_weight}") self.log_info( f" Upsample to fixed resolution: {self.upsample_to_fixed_resolution}" ) diff --git a/statistics.md b/statistics.md index 82d262ed..a5a85941 100644 --- a/statistics.md +++ b/statistics.md @@ -152,8 +152,8 @@ PhysioTwin4D operates across several technically demanding domains: - `requires_gpu` - GPU/CUDA-dependent tests (opt-in via `--run-gpu`) - `requires_simpleware` - tests needing a local Synopsys Simpleware Medical install (opt-in via `--run-simpleware`) - `requires_physicsnemo` - tests needing the optional `[physicsnemo]` extra (opt-in via `--run-physicsnemo`) -- `experiment` - runs experiment scripts end-to-end (opt-in via `--run-experiments`; multi-hour) -- `tutorial` - runs tutorial scripts end-to-end (opt-in via `--run-tutorials`) +- `tutorial` - runs tutorial scripts end-to-end + (opt-in via `--run-tutorials`; multi-hour) --- diff --git a/tests/README.md b/tests/README.md index 0798ced7..58eb2511 100644 --- a/tests/README.md +++ b/tests/README.md @@ -6,9 +6,7 @@ This directory contains comprehensive test suites for the PhysioTwin4D package, - **[TESTING_GUIDE.md](TESTING_GUIDE.md)** - Comprehensive testing guide with setup, troubleshooting, and best practices - **[GITHUB_WORKFLOWS.md](GITHUB_WORKFLOWS.md)** - CI/CD documentation and GitHub Actions workflow details -- **[EXPERIMENT_TESTS_GUIDE.md](EXPERIMENT_TESTS_GUIDE.md)** - Guide for running experiment script tests -- **[PARALLEL_EXECUTION_GUIDE.md](PARALLEL_EXECUTION_GUIDE.md)** - How parallel execution works with experiment tests -- **[EXPERIMENT_FLAG_USAGE.md](EXPERIMENT_FLAG_USAGE.md)** - Details on the --run-experiments flag +- **[PARALLEL_EXECUTION_GUIDE.md](PARALLEL_EXECUTION_GUIDE.md)** - How parallel execution works - **[TEST_FIXES_SUMMARY.md](TEST_FIXES_SUMMARY.md)** - Recent bug fixes and known issues ## Test Categories @@ -33,16 +31,14 @@ This directory contains comprehensive test suites for the PhysioTwin4D package, - **`test_usd_merge.py`** - USD file merging with material preservation - **`test_usd_time_preservation.py`** - Time-varying data validation -### Experiment Tests (EXTREMELY SLOW - Manual Only) -- **`test_experiments.py`** - End-to-end experiment script execution (hours to complete) - - **Opt-in only** - Requires `--run-experiments` flag to run - - **NOT included in CI/CD** - Never runs in automated workflows - - **Automatically skipped** - Won't run with `pytest tests/` unless flag is set - - Runs every `*.py` script in each `experiments/` subdirectory - - Each subdirectory gets its own test - - Scripts run in alphanumeric order - - Requires GPU, CUDA, and all dependencies installed - - See [EXPERIMENT_TESTS_GUIDE.md](EXPERIMENT_TESTS_GUIDE.md) for detailed usage instructions +### Tutorial Tests (SLOW - Opt-in) +- **`test_tutorials.py`** - End-to-end execution of each `tutorials/*.py` script, + comparing the screenshots it writes against stored baselines + - **Opt-in only** - Requires `--run-tutorials` flag to run + - Requires GPU, CUDA, all dependencies, and the tutorial datasets + - Scripts in `experiments/` are exploratory and are **not** run as tests + + ## Directory Structure @@ -82,7 +78,7 @@ uv pip install -e ".[test]" ### Run Tests The fast path is the default. Heavy buckets (slow tests, GPU tests, Simpleware -tests, experiment notebooks, tutorial scripts) are **auto-skipped** unless you +tests, tutorial scripts) are **auto-skipped** unless you pass their `--run-` flag. Tests that need downloadable data fetch it through the session fixtures and run by default — there is no `requires_data` marker any more. @@ -105,8 +101,7 @@ Each flag enables one marker family. Flags compose, so you can stack them. | `--run-gpu` | `requires_gpu` | CUDA-dependent tests (ICON, Simpleware, etc.) | | `--run-simpleware` | `requires_simpleware` | Need a licensed Synopsys Simpleware Medical install (also marked `requires_gpu`) | | `--run-physicsnemo` | `requires_physicsnemo` | Need the optional `[physicsnemo]` extra installed | -| `--run-experiments` | `experiment` | End-to-end experiment notebooks (hours to run) | -| `--run-tutorials` | `tutorial` | Tutorial scripts run end-to-end | +| `--run-tutorials` | `tutorial` | Tutorial scripts run end-to-end (hours to run) | | `--run-all` | every bucket above | Equivalent to passing all `--run-*` flags at once | ```bash @@ -122,11 +117,11 @@ pytest tests/ -v --run-all # Full Simpleware coverage (requires Simpleware Medical installed locally) pytest tests/ -v --run-simpleware --run-gpu --run-slow -# Experiment tests (EXTREMELY SLOW — hours to complete) -pytest tests/test_experiments.py -v --run-experiments +# Tutorial tests (SLOW — hours to complete) +pytest tests/test_tutorials.py -v --run-tutorials -# A single experiment by name -pytest tests/test_experiments.py::test_experiment_heart_gated_ct_to_usd -v -s --run-experiments +# A single tutorial by name +pytest tests/test_tutorials.py::TestTutorial01HeartGatedCTToUSD -v -s --run-tutorials ``` ### Common Test Commands @@ -145,7 +140,10 @@ pytest tests/ --create-baselines ## Test Timing Reports -All test runs automatically generate a comprehensive timing report at the end showing individual test durations, session time, and pass/fail/skip counts. The report separates regular tests from experiment tests and highlights the slowest tests. +All test runs automatically generate a comprehensive timing report at the end +showing individual test durations, session time, and pass/fail/skip counts. The +report separates regular tests from tutorial tests and highlights the slowest +tests. ## Test Configuration @@ -163,10 +161,8 @@ All test runs automatically generate a comprehensive timing report at the end sh - `@pytest.mark.requires_physicsnemo` — Tests needing the optional `[physicsnemo]` extra (`pip install "physiotwin4d[physicsnemo]"`, requires Python >= 3.11). Opt in: `--run-physicsnemo`. -- `@pytest.mark.experiment` — End-to-end experiment notebooks (EXTREMELY - SLOW, never in CI). Opt in: `--run-experiments`. -- `@pytest.mark.tutorial` — Tutorial scripts run end-to-end. Opt in: - `--run-tutorials`. +- `@pytest.mark.tutorial` — Tutorial scripts run end-to-end (SLOW, never in + CI). Opt in: `--run-tutorials`. `--run-all` is a convenience flag that turns on every `--run-*` bucket at once. - `@pytest.mark.integration` — Integration tests vs unit tests (filter-only). @@ -201,7 +197,7 @@ Tests automatically run on pull requests via GitHub Actions. The CI workflow: - **Runs fast tests** - USD utilities, data conversion, basic validation - **Skips slow tests** - Registration and segmentation (too slow for CI) -- **Automatically skips experiment tests** - Protected by `--run-experiments` flag requirement +- **Automatically skips tutorial tests** - Protected by `--run-tutorials` flag requirement - **Caches test data** - Speeds up subsequent runs - **Generates coverage** - Reports uploaded to Codecov @@ -209,7 +205,7 @@ Tests automatically run on pull requests via GitHub Actions. The CI workflow: - Platforms: Ubuntu, Windows, macOS - Python versions: 3.11, 3.12 - Target coverage: >70% -- Protection: Experiment tests require `--run-experiments` flag (never used in CI/CD) +- Protection: Tutorial tests require `--run-tutorials` flag (never used in CI/CD) **For detailed CI/CD information**, see [GITHUB_WORKFLOWS.md](GITHUB_WORKFLOWS.md) @@ -269,9 +265,7 @@ Tests automatically: ## Additional Resources - **Detailed Testing Guide**: [TESTING_GUIDE.md](TESTING_GUIDE.md) -- **Experiment Tests Guide**: [EXPERIMENT_TESTS_GUIDE.md](EXPERIMENT_TESTS_GUIDE.md) - **Parallel Execution Guide**: [PARALLEL_EXECUTION_GUIDE.md](PARALLEL_EXECUTION_GUIDE.md) -- **Experiment Flag Usage**: [EXPERIMENT_FLAG_USAGE.md](EXPERIMENT_FLAG_USAGE.md) - **CI/CD Documentation**: [GITHUB_WORKFLOWS.md](GITHUB_WORKFLOWS.md) - **Recent Fixes**: [TEST_FIXES_SUMMARY.md](TEST_FIXES_SUMMARY.md) - **Main Project**: [../README.md](../README.md) diff --git a/tests/baselines/registration_time_series_images/basic_forward_transform_0.hdf b/tests/baselines/registration_time_series_images/basic_forward_transform_0.hdf index 75823466..229d4d4f 100644 --- a/tests/baselines/registration_time_series_images/basic_forward_transform_0.hdf +++ b/tests/baselines/registration_time_series_images/basic_forward_transform_0.hdf @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:24697b25fa57cc7d83ad221b42789e75409ef96470a1ad2d89981b90386ccfc6 -size 40878719 +oid sha256:3fdb6bab6c362e6d8b2f7f23db540e24eb6b732e7f6cb82ce713433daa0670bd +size 4512579 diff --git a/tests/baselines/registration_time_series_images/basic_time_series_registered_0.mha b/tests/baselines/registration_time_series_images/basic_time_series_registered_0.mha index 3098be93..73f512a1 100644 --- a/tests/baselines/registration_time_series_images/basic_time_series_registered_0.mha +++ b/tests/baselines/registration_time_series_images/basic_time_series_registered_0.mha @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:4deb7987fcaa77c9f77660ff8ccbe6446e19051c954912de6d7514a14aa8927a -size 4814459 +oid sha256:fa8492f316a3c9c78b75a06ef45152fc8f5a8db36b0931386ef1ca22a6f6197e +size 4767196 diff --git a/tests/baselines/registration_time_series_images/middle_frame_forward_transform_0.hdf b/tests/baselines/registration_time_series_images/middle_frame_forward_transform_0.hdf new file mode 100644 index 00000000..afb51242 --- /dev/null +++ b/tests/baselines/registration_time_series_images/middle_frame_forward_transform_0.hdf @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:8b6d8ed46f2654a646111b6c6aad58641bf35573277d339d8d02d5e27e6737c4 +size 4158463 diff --git a/tests/baselines/registration_time_series_images/prior_forward_transform_0.hdf b/tests/baselines/registration_time_series_images/prior_forward_transform_0.hdf deleted file mode 100644 index 664096cd..00000000 --- a/tests/baselines/registration_time_series_images/prior_forward_transform_0.hdf +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:fd4bb595e748dcc777ffe290b18f8b1a6436bc2e579b298bc529157326e7b40f -size 40878719 diff --git a/tests/baselines/registration_time_series_images/prior_time_series_registered_0.mha b/tests/baselines/registration_time_series_images/prior_time_series_registered_0.mha deleted file mode 100644 index 3098be93..00000000 --- a/tests/baselines/registration_time_series_images/prior_time_series_registered_0.mha +++ /dev/null @@ -1,3 +0,0 @@ -version https://git-lfs.github.com/spec/v1 -oid sha256:4deb7987fcaa77c9f77660ff8ccbe6446e19051c954912de6d7514a14aa8927a -size 4814459 diff --git a/tests/baselines/registration_time_series_images/transform_application_time_series_0.mha b/tests/baselines/registration_time_series_images/transform_application_time_series_0.mha index 3098be93..e2ec568b 100644 --- a/tests/baselines/registration_time_series_images/transform_application_time_series_0.mha +++ b/tests/baselines/registration_time_series_images/transform_application_time_series_0.mha @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:4deb7987fcaa77c9f77660ff8ccbe6446e19051c954912de6d7514a14aa8927a -size 4814459 +oid sha256:5615af55c808458fc9e9e7c999ad4ddfc85a6082f2db1bfb6ec5bc770c03d892 +size 4779024 diff --git a/tests/conftest.py b/tests/conftest.py index 31b5a686..4befdafc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -12,6 +12,7 @@ from typing import Any, Optional import itk +import numpy as np import pytest from physiotwin4d.contour_tools import ContourTools @@ -38,7 +39,6 @@ _RUN_BUCKET_FLAGS = ( - "--run-experiments", "--run-tutorials", "--run-simpleware", "--run-slow", @@ -57,12 +57,6 @@ def _run_bucket_enabled(config: pytest.Config, flag: str) -> bool: def pytest_addoption(parser: pytest.Parser) -> None: """Add custom command-line options for pytest.""" - parser.addoption( - "--run-experiments", - action="store_true", - default=False, - help="Run experiment tests (extremely long-running notebook tests)", - ) parser.addoption( "--run-tutorials", action="store_true", @@ -124,10 +118,6 @@ def pytest_configure(config: pytest.Config) -> None: config.getoption("--create-baselines", default=False) ) - config.addinivalue_line( - "markers", - "experiment: marks tests that run experiment notebooks (extremely slow, manual only)", - ) config.addinivalue_line( "markers", "tutorial: marks tests that run tutorial scripts (data/GPU gated, manual only)", @@ -154,21 +144,12 @@ def pytest_collection_modifyitems( config: pytest.Config, items: list[pytest.Item] ) -> None: """ - Automatically skip experiment and tutorial tests unless their opt-in flags - are passed. + Automatically skip bucketed tests unless their opt-in flags are passed. - This ensures that experiment tests are opt-in only and won't run + This ensures that long-running tests are opt-in only and won't run accidentally when running the normal test suite. """ for item in items: - if "experiment" in item.keywords and not _run_bucket_enabled( - config, "--run-experiments" - ): - item.add_marker( - pytest.mark.skip( - reason="Experiment tests require --run-experiments (or --run-all) to run" - ) - ) if "tutorial" in item.keywords and not _run_bucket_enabled( config, "--run-tutorials" ): @@ -232,7 +213,6 @@ def pytest_runtest_logreport(report: pytest.TestReport) -> None: "nodeid": report.nodeid, "duration": report.duration, "outcome": report.outcome, - "is_experiment": "experiment" in report.keywords, "is_tutorial": "tutorial" in report.keywords, } @@ -249,7 +229,7 @@ def pytest_terminal_summary( Print comprehensive test timing report after all tests complete. This hook is called at the end of the test session to display - timing statistics for all tests, including experiment tests. + timing statistics for all tests, including tutorial tests. """ timings = config._test_timings # type: ignore[attr-defined] tests = timings["tests"] @@ -260,12 +240,9 @@ def pytest_terminal_summary( # Calculate session duration session_duration = datetime.now() - timings["start_time"] - # Separate regular, tutorial, and experiment tests - regular_tests = [ - t for t in tests if not t["is_experiment"] and not t["is_tutorial"] - ] + # Separate regular and tutorial tests + regular_tests = [t for t in tests if not t["is_tutorial"]] tutorial_tests = [t for t in tests if t["is_tutorial"]] - experiment_tests = [t for t in tests if t["is_experiment"]] # Write the timing report terminalreporter.write_sep("=", "TEST TIMING REPORT", bold=True) @@ -322,33 +299,6 @@ def pytest_terminal_summary( ) terminalreporter.write_line("") - # Experiment tests section - if experiment_tests: - terminalreporter.write_sep("-", "Experiment Tests", bold=True) - terminalreporter.write_line(f"Count: {len(experiment_tests)}") - - # Sort by duration (longest first) - sorted_experiments = sorted( - experiment_tests, key=lambda x: x["duration"], reverse=True - ) - - # Calculate total time - experiment_total = sum(t["duration"] for t in experiment_tests) - terminalreporter.write_line( - f"Total Time: {timedelta(seconds=int(experiment_total))}" - ) - terminalreporter.write_line("") - - # Show all experiment tests with timing - terminalreporter.write_line("Individual Test Times:") - for test in sorted_experiments: - outcome_symbol = "+" if test["outcome"] == "passed" else "x" - duration_str = _format_duration(test["duration"]) - terminalreporter.write_line( - f" {outcome_symbol} {duration_str:>10s} {test['nodeid']}" - ) - terminalreporter.write_line("") - # Top 10 slowest tests overall if len(tests) > 10: terminalreporter.write_sep("-", "Top 10 Slowest Tests", bold=True) @@ -357,12 +307,7 @@ def pytest_terminal_summary( for i, test in enumerate(sorted_all, 1): outcome_symbol = "+" if test["outcome"] == "passed" else "x" duration_str = _format_duration(test["duration"]) - if test["is_experiment"]: - test_type = "[EXP]" - elif test["is_tutorial"]: - test_type = "[TUT]" - else: - test_type = "[REG]" + test_type = "[TUT]" if test["is_tutorial"] else "[REG]" terminalreporter.write_line( f" {i:2d}. {outcome_symbol} {duration_str:>10s} {test_type} {test['nodeid']}" ) @@ -662,3 +607,74 @@ def registrar_ICON() -> RegisterImagesICON: def transform_tools() -> TransformTools: """Create a TransformTools instance.""" return TransformTools() + + +class KnownShiftCase: + """A registration case whose correct answer is known exactly. + + ``moving`` is built by resampling ``fixed`` through a translation of + ``shift_mm``, so ``moving(q) == fixed(q + shift_mm)``. Warping ``moving`` + back onto the fixed grid therefore requires a ``forward_transform`` of + ``-shift_mm``, which gives an absolute accuracy target instead of the + "did it return something" checks that let a sign error go unnoticed. + """ + + def __init__(self, fixed_image: itk.Image, shift_mm: tuple[float, float, float]): + """Build the shifted pair. + + Args: + fixed_image: Image used as the registration target. + shift_mm: Content displacement applied to build the moving image. + Use a different magnitude and sign per axis so an axis swap or a + sign flip cannot pass. + """ + self.transform_tools = TransformTools() + self.fixed = fixed_image + self.shift_mm = shift_mm + self.expected_displacement = np.array([-v for v in shift_mm]) + + shift = itk.TranslationTransform[itk.D, 3].New() + shift.SetOffset(list(shift_mm)) + self.moving = self.transform_tools.transform_image( + fixed_image, shift, fixed_image, interpolation_method="linear" + ) + + size = itk.size(fixed_image) + self._center = list( + fixed_image.TransformIndexToPhysicalPoint( + [int(size[i]) // 2 for i in range(3)] + ) + ) + # Score over the brightest 30% of the fixed image (tissue and blood + # pool); background air correlates trivially and would mask errors. + self._fixed_array = itk.array_from_image(fixed_image) + self._foreground = self._fixed_array >= np.percentile(self._fixed_array, 70) + + def center_error_mm(self, forward_transform: itk.Transform) -> float: + """Distance, in mm, between the recovered and true displacement.""" + displacement = np.array( + list(forward_transform.TransformPoint(self._center)) + ) - np.array(self._center) + return float(np.linalg.norm(displacement - self.expected_displacement)) + + def foreground_ncc(self, forward_transform: itk.Transform) -> float: + """Normalized cross-correlation after warping moving onto the fixed grid.""" + warped = self.transform_tools.transform_image( + self.moving, forward_transform, self.fixed, interpolation_method="linear" + ) + moved = itk.array_from_image(warped)[self._foreground] + target = self._fixed_array[self._foreground] + moved = moved - moved.mean() + target = target - target.mean() + denominator = np.sqrt((moved**2).sum() * (target**2).sum()) + return float((moved * target).sum() / denominator) if denominator else 0.0 + + def unregistered_ncc(self) -> float: + """Baseline score with no registration, for a floor to beat.""" + return self.foreground_ncc(itk.TranslationTransform[itk.D, 3].New()) + + +@pytest.fixture(scope="session") +def known_shift_case(test_images: list[Any]) -> KnownShiftCase: + """A moving/fixed pair separated by a known (6, -4, 3) mm shift.""" + return KnownShiftCase(test_images[0], (6.0, -4.0, 3.0)) diff --git a/tests/test_contour_tools.py b/tests/test_contour_tools.py index d01bdbd0..c5f81b03 100644 --- a/tests/test_contour_tools.py +++ b/tests/test_contour_tools.py @@ -337,5 +337,56 @@ def test_contours_from_both_time_points( print(f"Extracted contours from {len(test_labelmaps)} time points") +class TestSaveCombinedSurfaces: + """Structure identity must survive merging into one file.""" + + @staticmethod + def _annotated_sphere( + center: tuple[float, float, float], label_id: int + ) -> pv.PolyData: + surface = pv.Sphere(radius=1.0, center=center) + surface.field_data["SegmentationLabelIds"] = np.array( + [label_id], dtype=np.int32 + ) + return surface + + def test_per_cell_label_ids_from_single_label_surfaces( + self, tmp_path: Path + ) -> None: + """Each merged cell carries the label ID of the surface it came from.""" + surfaces = { + "atrium_left": self._annotated_sphere((0.0, 0.0, 0.0), 141), + "ventricle_left": self._annotated_sphere((3.0, 0.0, 0.0), 142), + } + + output_file = tmp_path / "combined.vtp" + ContourTools.save_combined_surfaces(surfaces, str(output_file)) + + merged = pv.read(str(output_file)) + label_ids = merged.cell_data["SegmentationLabelIds"] + assert set(np.unique(label_ids)) == {141, 142} + for name, surface in surfaces.items(): + expected = int(surface.field_data["SegmentationLabelIds"][0]) + assert int(np.sum(label_ids == expected)) == surface.n_cells, ( + f"Cell count for {name} did not survive the merge" + ) + # Tagging must not leak back into the caller's surfaces. + for surface in surfaces.values(): + assert "SegmentationLabelIds" not in surface.cell_data + + def test_multi_label_surfaces_tagged_unknown(self, tmp_path: Path) -> None: + """Per-group surfaces list several labels, so no cell can be attributed.""" + group_surface = pv.Sphere(radius=1.0) + group_surface.field_data["SegmentationLabelIds"] = np.array( + [141, 142], dtype=np.int32 + ) + + output_file = tmp_path / "combined_group.vtp" + ContourTools.save_combined_surfaces({"heart": group_surface}, str(output_file)) + + merged = pv.read(str(output_file)) + assert set(np.unique(merged.cell_data["SegmentationLabelIds"])) == {0} + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/test_convert_vtk_to_usd.py b/tests/test_convert_vtk_to_usd.py index 35c1ec68..2ab2c68d 100644 --- a/tests/test_convert_vtk_to_usd.py +++ b/tests/test_convert_vtk_to_usd.py @@ -477,6 +477,49 @@ def test_static_merge_prim_names_use_data_basename(self, tmp_path: Path) -> None f"{prim_path} should have no time samples but got {samples}" ) + def test_mask_ids_split_on_segmentation_label_ids(self, tmp_path: Path) -> None: + """A merged surface file splits on the array save_combined_surfaces writes.""" + mesh = _make_poly() + mesh.cell_data["SegmentationLabelIds"] = np.array( + [141] * 4 + [142] * 5, dtype=np.int32 + ) + converter = ConvertVTKToUSD( + data_basename="P", + input_polydata=[mesh], + mask_ids={141: "atrium_left", 142: "ventricle_left"}, + ) + stage = converter.convert(str(tmp_path / "out.usd")) + + assert stage.GetPrimAtPath("/World/P/Anatomy/atrium_left").IsValid() + assert stage.GetPrimAtPath("/World/P/Anatomy/ventricle_left").IsValid() + + def test_static_merge_object_names_name_prims(self, tmp_path: Path) -> None: + """object_names replaces the positional {data_basename}_{i} naming.""" + mesh_a, mesh_b = _make_poly(), _make_poly() + converter = ConvertVTKToUSD( + data_basename="Organ", + input_polydata=[mesh_a, mesh_b], + static_merge=True, + object_names=["myocardium", "ventricle_left"], + ) + stage = converter.convert(str(tmp_path / "out.usd")) + + assert stage.GetPrimAtPath("/World/Organ/myocardium").IsValid() + assert stage.GetPrimAtPath("/World/Organ/ventricle_left").IsValid() + assert not stage.GetPrimAtPath("/World/Organ/Organ_0").IsValid(), ( + "Positional naming still present" + ) + + def test_object_names_length_mismatch_raises(self) -> None: + """A short object_names list would silently mis-name prims.""" + with pytest.raises(ValueError, match="object_names length"): + ConvertVTKToUSD( + data_basename="Organ", + input_polydata=[_make_poly(), _make_poly()], + static_merge=True, + object_names=["myocardium"], + ) + # ------------------------------------------------------------------ # Gap D — mask_ids / _convert_with_labels # ------------------------------------------------------------------ diff --git a/tests/test_experiments.py b/tests/test_experiments.py deleted file mode 100644 index 9d73368e..00000000 --- a/tests/test_experiments.py +++ /dev/null @@ -1,574 +0,0 @@ -""" -Test suite for running experiment scripts. - -These tests execute Python scripts in the experiments/ directory. Each subdirectory -in experiments/ gets its own test that runs all scripts in that subdirectory in -alphanumeric order. - -Scripts are Jupytext percent-format files (# %% cell separators), converted from the -original Jupyter notebooks while preserving git history via git mv. - -WARNING: These are EXTREMELY long-running tests that may take hours to complete. -They are NOT part of CI/CD and should only be run manually. - -Usage: - # Run all experiment tests - pytest tests/test_experiments.py -v -m experiment - - # Run a specific experiment subdirectory - pytest tests/test_experiments.py::test_experiment_colormap_vtk_to_usd -v - - # Run with detailed output - pytest tests/test_experiments.py -v -s -m experiment - -Note: These tests require all dependencies installed and GPU/CUDA support for -many of the experiments. -""" - -import os -import subprocess -import sys -from pathlib import Path -from typing import Any - -import pytest - -# Base directories -REPO_ROOT = Path(__file__).parent.parent -EXPERIMENTS_DIR = REPO_ROOT / "experiments" - -# Experiment subdirectories to test (in order of complexity/dependencies) -EXPERIMENT_SUBDIRS = [ - "Colormap-VTK_To_USD", - "Convert_VTK_To_USD", - # 'DisplacementField_To_USD', # Disabled - scripts not ready - "Reconstruct4DCT", - "Heart-VTKSeries_To_USD", - "Heart-GatedCT_To_USD", - "Heart-Create_Statistical_Model", - "Heart-Statistical_Model_To_Patient", - "Lung-GatedCT_To_USD", - # 'Lung-VesselsAirways', # Disabled - scripts not ready -] - - -def get_scripts_in_subdir(subdir_name: str) -> list[Path]: - """ - Get all Python scripts in a subdirectory, sorted alphanumerically. - - Args: - subdir_name: Name of the subdirectory in experiments/ - - Returns: - List of Path objects for .py script files, sorted alphanumerically - """ - subdir = EXPERIMENTS_DIR / subdir_name - if not subdir.exists(): - return [] - - scripts = sorted(subdir.glob("*.py")) - return scripts - - -def execute_script(script_path: Path, timeout: int = 3600) -> dict[str, Any]: - """ - Execute a Python experiment script. - - Args: - script_path: Path to the .py script file - timeout: Maximum execution time in seconds (default: 1 hour) - - Returns: - Dictionary with execution results: - - success: bool - - stdout: str - - stderr: str - - returncode: int - - Raises: - subprocess.TimeoutExpired: If script execution exceeds timeout - """ - print(f"\n{'=' * 80}") - print(f"Executing script: {script_path.name}") - print(f"Path: {script_path}") - print(f"Timeout: {timeout} seconds ({timeout // 60} minutes)") - print(f"{'=' * 80}\n") - - cmd = [sys.executable, str(script_path)] - - # So scripts can use reduced parameters when run as tests - env = os.environ.copy() - env["PHYSIOTWIN_RUNNING_AS_TEST"] = "1" - - try: - result = subprocess.run( - cmd, - capture_output=True, - text=True, - timeout=timeout * 1.5, # Give extra time for startup overhead - cwd=script_path.parent, # Run in script's directory - env=env, - check=False, - ) - - success = result.returncode == 0 - - if success: - print(f"OK: {script_path.name}") - else: - print(f"FAILED: {script_path.name}") - print(f"Return code: {result.returncode}") - if result.stderr: - print(f"Error output:\n{result.stderr}") - - return { - "success": success, - "stdout": result.stdout, - "stderr": result.stderr, - "returncode": result.returncode, - } - - except subprocess.TimeoutExpired: - print(f"TIMEOUT: {script_path.name}") - print(f"Exceeded: {timeout} seconds") - raise - - -def _heart_statistical_model_pca_prerequisites_met() -> tuple[bool, str]: - """ - Check whether PCA model outputs from Heart-Create_Statistical_Model exist. - - Heart-Statistical_Model_To_Patient scripts expect these artifacts from - the Heart-Create_Statistical_Model experiment (notably 5-compute_pca_model.py). - - Returns: - (True, "") if all required paths exist, else (False, reason_message). - """ - pca_output_dir = ( - EXPERIMENTS_DIR / "Heart-Create_Statistical_Model" / "kcl-heart-model" - ) - pca_json = pca_output_dir / "pca_model.json" - pca_mean_vtp = pca_output_dir / "pca_mean.vtp" - - if not pca_output_dir.is_dir(): - return ( - False, - f"PCA model output directory not found: {pca_output_dir}. " - "Run the Heart-Create_Statistical_Model experiment first " - "(e.g. pytest tests/test_experiments.py::test_experiment_create_statistical_model -v -s --run-experiments).", - ) - if not pca_json.is_file(): - return ( - False, - f"PCA model JSON not found: {pca_json}. " - "Complete Heart-Create_Statistical_Model (including 5-compute_pca_model.py) first.", - ) - if not pca_mean_vtp.is_file(): - return ( - False, - f"PCA mean surface not found: {pca_mean_vtp}. " - "Complete Heart-Create_Statistical_Model (including 5-compute_pca_model.py) first.", - ) - return (True, "") - - -def run_experiment_scripts(subdir_name: str, timeout_per_script: int = 3600) -> None: - """ - Run all Python scripts in an experiment subdirectory in alphanumeric order. - - IMPORTANT: Scripts are executed SEQUENTIALLY in alphanumeric order within - this function. This ensures proper dependency handling even when running - tests with multiple pytest workers (e.g., pytest -n 2). - - The sequential execution is enforced by: - 1. Using a standard Python for loop (not parallelized) - 2. Each script must complete before the next begins - 3. Failures in earlier scripts prevent later ones from running - - Args: - subdir_name: Name of the subdirectory in experiments/ - timeout_per_script: Timeout in seconds for each script (default: 1 hour) - - Raises: - AssertionError: If any script fails to execute successfully - """ - scripts = get_scripts_in_subdir(subdir_name) - - if not scripts: - pytest.skip(f"No scripts found in experiments/{subdir_name}") - - print(f"\n{'#' * 80}") - print(f"# Experiment: {subdir_name}") - print(f"# Found {len(scripts)} script(s)") - print("# Sequential execution enforced (scripts run in order)") - print(f"{'#' * 80}\n") - - failed_scripts: list[dict[str, Any]] = [] - successful_scripts: list[str] = [] - - for i, script in enumerate(scripts, 1): - print(f"\n--- Script {i}/{len(scripts)} ---") - print(f"Sequential execution: script {i} must complete before {i + 1} starts") - - try: - result = execute_script(script, timeout=timeout_per_script) - - if result["success"]: - successful_scripts.append(script.name) - else: - failed_scripts.append( - { - "name": script.name, - "returncode": result["returncode"], - "stderr": result["stderr"], - } - ) - # Stop execution on first failure to maintain dependencies - print(f"\nStopping execution: {script.name} failed") - print("Remaining scripts in this experiment will not run.") - break - - except subprocess.TimeoutExpired: - failed_scripts.append( - { - "name": script.name, - "returncode": -1, - "stderr": f"Timeout after {timeout_per_script} seconds", - } - ) - # Stop execution on timeout - print(f"\nStopping execution: {script.name} timed out") - print("Remaining scripts in this experiment will not run.") - break - - except Exception as e: - failed_scripts.append( - {"name": script.name, "returncode": -2, "stderr": str(e)} - ) - # Stop execution on exception - print(f"\nStopping execution: {script.name} raised exception") - print("Remaining scripts in this experiment will not run.") - break - - # Print summary - print(f"\n{'=' * 80}") - print(f"Experiment Summary: {subdir_name}") - print(f"{'=' * 80}") - print(f"Total scripts: {len(scripts)}") - print(f"Successful: {len(successful_scripts)}") - print(f"Failed: {len(failed_scripts)}") - - if successful_scripts: - print("\nSuccessful scripts:") - for name in successful_scripts: - print(f" - {name}") - - if failed_scripts: - print("\nFailed scripts:") - for failure in failed_scripts: - print(f" - {failure['name']}") - print(f" Return code: {failure['returncode']}") - if failure["stderr"]: - # Print first few lines of error - error_lines = failure["stderr"].split("\n")[:10] - for line in error_lines: - print(f" {line}") - - print(f"{'=' * 80}\n") - - # Assert all scripts succeeded - assert not failed_scripts, ( - f"{len(failed_scripts)} script(s) failed in {subdir_name}: " - f"{[f['name'] for f in failed_scripts]}" - ) - - -# ============================================================================ -# Test Functions - One per Experiment Subdirectory -# ============================================================================ - - -@pytest.mark.experiment -@pytest.mark.slow -@pytest.mark.xdist_group( - name="experiment_colormap" -) # Prevent parallel execution within group -def test_experiment_colormap_vtk_to_usd() -> None: - """ - Test Colormap-VTK_To_USD experiment scripts. - - This experiment demonstrates converting VTK files with colormaps to USD format. - - EXECUTION MODEL: - - Scripts run SEQUENTIALLY in alphanumeric order within this test - - This test function is atomic - pytest-xdist treats it as a single unit - - Multiple experiment tests CAN run in parallel (different subdirectories) - - Scripts within THIS experiment CANNOT run in parallel or out of order - """ - run_experiment_scripts("Colormap-VTK_To_USD", timeout_per_script=3600) - - -# DISABLED - Scripts not ready -# @pytest.mark.experiment -# @pytest.mark.slow -# def test_experiment_displacement_field_to_usd(): -# """ -# Test DisplacementField_To_USD experiment scripts. -# -# This experiment demonstrates converting registration displacement fields to USD -# format for visualization in PhysicsNeMo and Omniverse. -# """ -# run_experiment_scripts('DisplacementField_To_USD', timeout_per_script=3600) - - -@pytest.mark.experiment -@pytest.mark.slow -@pytest.mark.requires_gpu -@pytest.mark.xdist_group(name="experiment_reconstruct4dct") -def test_experiment_reconstruct_4dct() -> None: - """ - Test Reconstruct4DCT experiment scripts. - - This experiment demonstrates 4D CT reconstruction techniques. - - EXECUTION MODEL: - - Scripts run SEQUENTIALLY in alphanumeric order within this test - - Each script must complete before the next begins - - Failure in one script stops execution of remaining scripts - """ - run_experiment_scripts("Reconstruct4DCT", timeout_per_script=1200) - - -@pytest.mark.experiment -@pytest.mark.slow -@pytest.mark.xdist_group(name="experiment_heart_vtk") -def test_experiment_heart_vtk_series_to_usd() -> None: - """ - Test Heart-VTKSeries_To_USD experiment scripts. - - This experiment converts heart VTK time series data to USD format. - - EXECUTION ORDER (ENFORCED): - 1. 0-download_and_convert_4d_to_3d.py (downloads data) - 2. 1-heart_vtkseries_to_usd.py (uses downloaded data) - - Each script must complete successfully before the next begins. - """ - run_experiment_scripts("Heart-VTKSeries_To_USD", timeout_per_script=5400) - - -@pytest.mark.experiment -@pytest.mark.slow -@pytest.mark.requires_gpu -@pytest.mark.xdist_group(name="experiment_heart_gated_ct") -def test_experiment_heart_gated_ct_to_usd() -> None: - """ - Test Heart-GatedCT_To_USD experiment scripts. - - This is the main cardiac imaging pipeline experiment with strict dependencies. - - EXECUTION ORDER (STRICTLY ENFORCED): - 1. 0-download_and_convert_4d_to_3d.py (downloads and converts data) - 2. 1-register_images.py (registers converted images) - 3. 2-generate_segmentation.py (segments registered images) - 4. 3-transform_dynamic_and_static_contours.py (transforms segmentations) - 5. 4-merge_dynamic_and_static_usd.py (merges into final USD) - - Each script depends on outputs from previous scripts. - Execution stops on first failure to prevent cascading errors. - """ - run_experiment_scripts("Heart-GatedCT_To_USD", timeout_per_script=5400) - - -@pytest.mark.experiment -@pytest.mark.slow -@pytest.mark.xdist_group(name="experiment_convert_vtk_to_usd") -def test_experiment_convert_vtk_to_usd() -> None: - """ - Test Convert_VTK_To_USD experiment scripts. - - This experiment demonstrates VTK to USD conversion using the library classes. - - EXECUTION ORDER (ENFORCED): - 1. convert_chop_valve_to_usd.py (converts CHOP valve data) - 2. convert_vtk_to_usd_using_class.py (demonstrates library usage) - - Sequential execution ensures examples build on each other. - """ - run_experiment_scripts("Convert_VTK_To_USD", timeout_per_script=3600) - - -@pytest.mark.experiment -@pytest.mark.slow -@pytest.mark.xdist_group(name="experiment_create_statistical_model") -def test_experiment_create_statistical_model() -> None: - """ - Test Heart-Create_Statistical_Model experiment scripts. - - This experiment demonstrates creating a PCA statistical shape model from the - KCL Heart Model dataset. - - EXECUTION ORDER (ENFORCED): - 1. 1-input_meshes_to_input_surfaces.py (convert meshes to surfaces) - 2. 2-input_surfaces_to_surfaces_aligned.py (align surfaces) - 3. 3-registration_based_correspondence.py (establish point correspondence) - 4. 4-surfaces_aligned_correspond_to_pca_inputs.py (prepare PCA inputs) - 5. 5-compute_pca_model.py (compute PCA model using sklearn) - - Sequential execution ensures data dependencies are met. - """ - run_experiment_scripts("Heart-Create_Statistical_Model", timeout_per_script=5400) - - -@pytest.mark.experiment -@pytest.mark.slow -@pytest.mark.requires_gpu -@pytest.mark.xdist_group(name="experiment_heart_statistical_model") -def test_experiment_heart_statistical_model_to_patient() -> None: - """ - Test Heart-Statistical_Model_To_Patient experiment scripts. - - This experiment demonstrates heart model to patient registration using - statistical shape models (PCA). - - PREREQUISITE: Complete Heart-Create_Statistical_Model experiment first to generate - the PCA model data required for this experiment. - - If PCA outputs (kcl-heart-model/pca_model.json, pca_mean.vtp) are missing, this test - is skipped with a clear message so it can be run in isolation after generating them. - - EXECUTION ORDER (ENFORCED): - 1. heart_model_to_model_icp_itk.py (ICP registration) - 2. heart_model_to_model_registration_pca.py (PCA-based registration) - 3. heart_model_to_patient.py (applies registration to patient) - - Sequential execution ensures registration results are available for subsequent steps. - """ - prereq_met, skip_reason = _heart_statistical_model_pca_prerequisites_met() - if not prereq_met: - pytest.skip(skip_reason) - - run_experiment_scripts( - "Heart-Statistical_Model_To_Patient", timeout_per_script=7200 - ) - - -@pytest.mark.experiment -@pytest.mark.slow -@pytest.mark.requires_gpu -@pytest.mark.xdist_group(name="experiment_lung_gated_ct") -def test_experiment_lung_gated_ct_to_usd() -> None: - """ - Test Lung-GatedCT_To_USD experiment scripts. - - This is the lung imaging pipeline experiment using DirLab 4DCT data. - - EXECUTION ORDER (STRICTLY ENFORCED): - 1. 0-register_dirlab_4dct.py (registers lung 4DCT data) - 2. 1-make_dirlab_models.py (creates 3D models from registered data) - 3. 2-paint_dirlab_models.py (applies textures/materials to models) - 4. Experiment_ArrangeOnStage.py (arranges models in USD scene) - 5. Experiment_CombineModels.py (combines models into single USD) - 6. Experiment_SegReg.py (segmentation and registration experiments) - 7. Experiment_SubSurfaceScatter.py (applies advanced materials) - - Each script depends on outputs from previous scripts. - Execution is sequential and stops on first failure. - """ - run_experiment_scripts("Lung-GatedCT_To_USD", timeout_per_script=5400) - - -# DISABLED - Scripts not ready -# @pytest.mark.experiment -# @pytest.mark.slow -# @pytest.mark.requires_gpu -# @pytest.mark.requires_data -# def test_experiment_lung_vessels_airways(): -# """ -# Test Lung-VesselsAirways experiment scripts. -# -# This experiment demonstrates specialized vessel and airway segmentation -# using deep learning models. -# Expected scripts (in order): -# - 0-GenData.py -# """ -# run_experiment_scripts('Lung-VesselsAirways', timeout_per_script=3600) - - -# ============================================================================ -# Discovery Test - Validate Experiment Structure -# ============================================================================ - - -@pytest.mark.experiment -def test_experiment_structure() -> None: - """ - Validate the structure of the experiments directory. - - This test checks that: - 1. The experiments directory exists - 2. Each expected subdirectory exists - 3. Each subdirectory contains at least one .py script - """ - assert EXPERIMENTS_DIR.exists(), ( - f"Experiments directory not found: {EXPERIMENTS_DIR}" - ) - - missing_subdirs = [] - empty_subdirs = [] - - for subdir_name in EXPERIMENT_SUBDIRS: - subdir = EXPERIMENTS_DIR / subdir_name - - if not subdir.exists(): - missing_subdirs.append(subdir_name) - continue - - scripts = list(subdir.glob("*.py")) - if not scripts: - empty_subdirs.append(subdir_name) - - # Report findings - if missing_subdirs: - print(f"\nWARNING: Missing subdirectories: {missing_subdirs}") - - if empty_subdirs: - print(f"\nWARNING: Empty subdirectories (no scripts): {empty_subdirs}") - - # Print discovered scripts - print("\nDiscovered Scripts:") - for subdir_name in EXPERIMENT_SUBDIRS: - scripts = get_scripts_in_subdir(subdir_name) - if scripts: - print(f"\n{subdir_name}/ ({len(scripts)} script(s)):") - for s in scripts: - print(f" - {s.name}") - - assert not missing_subdirs, f"Missing subdirectories: {missing_subdirs}" - assert not empty_subdirs, f"Empty subdirectories: {empty_subdirs}" - - -# ============================================================================ -# Helper Test - Script Discovery -# ============================================================================ - - -@pytest.mark.experiment -@pytest.mark.parametrize("subdir_name", EXPERIMENT_SUBDIRS) -def test_list_scripts_in_subdir(subdir_name: str) -> None: - """ - List all scripts in each experiment subdirectory. - - This helper test can be used to preview what scripts will be run - without actually executing them. - - Usage: - pytest tests/test_experiments.py::test_list_scripts_in_subdir -v -s - """ - scripts = get_scripts_in_subdir(subdir_name) - - print(f"\n{subdir_name}/ - {len(scripts)} script(s):") - for i, s in enumerate(scripts, 1): - print(f" {i}. {s.name}") - - assert scripts, f"No scripts found in {subdir_name}" diff --git a/tests/test_register_images_ants.py b/tests/test_register_images_ants.py index a40a5237..862512cb 100644 --- a/tests/test_register_images_ants.py +++ b/tests/test_register_images_ants.py @@ -18,6 +18,8 @@ from physiotwin4d.register_images_ants import RegisterImagesANTS from physiotwin4d.transform_tools import TransformTools +from .conftest import KnownShiftCase + def _foreground_ncc( reference_arr: np.ndarray, warped_arr: np.ndarray, foreground: np.ndarray @@ -294,6 +296,35 @@ def test_preprocess_images( print("Image preprocessing complete") print(f" Preprocessed spacing: {preprocessed_spacing}") + def test_recovers_known_shift(self, known_shift_case: KnownShiftCase) -> None: + """ANTs must recover a known shift, in the right direction. + + The companion of the Greedy check in test_register_images_greedy.py: + ANTs works in ITK's LPS frame throughout, so this asserts no equivalent + RAS/LPS conversion is missing here. + """ + registrar = RegisterImagesANTS() + registrar.set_modality("ct") + registrar.set_transform_type("Affine") + registrar.set_fixed_image(known_shift_case.fixed) + + result = registrar.register(moving_image=known_shift_case.moving) + forward_transform = result["forward_transform"] + + error_mm = known_shift_case.center_error_mm(forward_transform) + ncc = known_shift_case.foreground_ncc(forward_transform) + unregistered_ncc = known_shift_case.unregistered_ncc() + + print("\nANTs known-shift recovery:") + print(f" error: {error_mm:.2f} mm") + print(f" foreground NCC: {ncc:.4f} (unregistered {unregistered_ncc:.4f})") + + assert error_mm < 2.0, f"ANTs recovered the shift {error_mm:.2f} mm off" + assert ncc > unregistered_ncc, ( + f"ANTs left the images less aligned than they started " + f"({ncc:.4f} vs {unregistered_ncc:.4f})" + ) + def test_registration_with_initial_transform( self, registrar_ANTS: RegisterImagesANTS, @@ -318,9 +349,9 @@ def test_registration_with_initial_transform( registrar_ANTS.set_modality("ct") registrar_ANTS.set_fixed_image(fixed_image) - result = registrar_ANTS.register( - moving_image=moving_image, - initial_forward_transform=initial_tfm_forward, + result = registrar_ANTS.register_from( + initial_tfm_forward, + moving_image, ) assert isinstance(result, dict), "Result should be a dictionary" @@ -335,7 +366,7 @@ def test_initial_transform_composition_metrics( test_images: list[Any], test_directories: dict[str, Path], ) -> None: - """Verify the initial_forward_transform composition path with metrics. + """Verify the register_from() composition path with metrics. Exercises the two initial-transform inputs the platform actually uses (identity and a prior deformable forward_transform, as in prior-based @@ -396,18 +427,16 @@ def warp_score(forward_transform: Any) -> float: registrar_identity = RegisterImagesANTS() registrar_identity.set_modality("ct") registrar_identity.set_fixed_image(fixed_image) - identity_result = registrar_identity.register( - moving_image=moving_image, initial_forward_transform=identity - ) + identity_result = registrar_identity.register_from(identity, moving_image) ncc_identity = warp_score(identity_result["forward_transform"]) # Prior deformable initial: the realistic time-series prior use case. registrar_prior = RegisterImagesANTS() registrar_prior.set_modality("ct") registrar_prior.set_fixed_image(fixed_image) - prior_result = registrar_prior.register( - moving_image=moving_image, - initial_forward_transform=baseline["forward_transform"], + prior_result = registrar_prior.register_from( + baseline["forward_transform"], + moving_image, ) ncc_prior = warp_score(prior_result["forward_transform"]) @@ -475,9 +504,7 @@ def test_initial_transform_matrix_composition( registrar_ANTS.set_modality("ct") registrar_ANTS.set_fixed_image(fixed_image) - result = registrar_ANTS.register( - moving_image=moving_image, initial_forward_transform=translation - ) + result = registrar_ANTS.register_from(translation, moving_image) transform_tools = TransformTools() warped = transform_tools.transform_image( diff --git a/tests/test_register_images_chain.py b/tests/test_register_images_chain.py index 26e999c7..627245e1 100644 --- a/tests/test_register_images_chain.py +++ b/tests/test_register_images_chain.py @@ -10,7 +10,7 @@ from __future__ import annotations -from typing import Any, Optional, Union +from typing import Any, Optional, Union, cast import itk import numpy as np @@ -41,7 +41,6 @@ def __init__(self, name: str, sentinel_value: float) -> None: self.seen_moving_image: Optional[itk.Image] = None self.seen_moving_image_pre: Optional[itk.Image] = None self.seen_moving_mask: Optional[itk.Image] = None - self.seen_initial_forward_transform: Optional[object] = None self.preprocess_call_count = 0 def preprocess(self, image: itk.Image, modality: str = "ct") -> itk.Image: @@ -55,37 +54,99 @@ def registration_method( moving_mask: Optional[itk.Image] = None, moving_labelmap: Optional[itk.Image] = None, moving_image_pre: Optional[itk.Image] = None, - initial_forward_transform: Optional[object] = None, ) -> dict[str, Union[object, float]]: - """Record the state visible at call time; return a sentinel result.""" + """Record the state visible at call time; return a sentinel result. + + The sentinel transforms are real translations along x of + ``sentinel_value`` mm, so a chained result can be checked numerically. + """ self.seen_fixed_image_pre = self.fixed_image_pre self.seen_moving_image = self.moving_image self.seen_moving_image_pre = moving_image_pre self.seen_moving_mask = moving_mask - self.seen_initial_forward_transform = initial_forward_transform + forward = itk.TranslationTransform[itk.D, 3].New() + forward.SetOffset([self.sentinel_value, 0.0, 0.0]) + inverse = itk.TranslationTransform[itk.D, 3].New() + inverse.SetOffset([-self.sentinel_value, 0.0, 0.0]) return { - "forward_transform": f"forward_{self.name}", - "inverse_transform": f"inverse_{self.name}", + "forward_transform": forward, + "inverse_transform": inverse, "loss": self.sentinel_value, } -def test_chain_feeds_forward_transform_to_next_stage() -> None: - """Stage 2 must receive stage 1's forward_transform as its - initial_forward_transform.""" +def test_chain_refines_previous_stage_result() -> None: + """Stage 2 must refine stage 1's alignment rather than start over. + + Stage 1 is called on the raw moving image; stage 2 sees a moving image + pre-warped by stage 1's forward transform, and the chain's composed + forward transform is the sum of both stages' translations. + """ stage1 = _RecordingRegistrar("stage1", 1.0) stage2 = _RecordingRegistrar("stage2", 2.0) chain = RegisterImagesChain([stage1, stage2]) chain.set_fixed_image(_small_image()) + moving = _small_image() - result = chain.register(_small_image()) + result = chain.register(moving) - assert stage1.seen_initial_forward_transform is None - assert stage2.seen_initial_forward_transform == "forward_stage1" - assert result["forward_transform"] == "forward_stage2" + assert stage1.seen_moving_image is moving + # Stage 2 registers the pre-warped image, not the caller's image. + assert stage2.seen_moving_image is not moving + composed = cast(itk.Transform, result["forward_transform"]) + assert list(composed.TransformPoint([0.0, 0.0, 0.0])) == [3.0, 0.0, 0.0] assert result["loss"] == 2.0 +class _CompositeRegistrar(_RecordingRegistrar): + """Stub registrar returning a CompositeTransform, as RegisterImagesGreedy + does (its result is an affine plus a displacement field).""" + + def registration_method( + self, + moving_image: itk.Image, + moving_mask: Optional[itk.Image] = None, + moving_labelmap: Optional[itk.Image] = None, + moving_image_pre: Optional[itk.Image] = None, + ) -> dict[str, Union[object, float]]: + """Wrap the parent's translations in single-entry composites.""" + result = super().registration_method( + moving_image, moving_mask, moving_labelmap, moving_image_pre + ) + for key in ("forward_transform", "inverse_transform"): + composite = itk.CompositeTransform[itk.D, 3].New() + composite.AddTransform(cast(itk.Transform, result[key])) + result[key] = composite + return result + + +def test_chain_result_holds_no_nested_composite(tmp_path: Any) -> None: + """A stage returning a CompositeTransform must not end up nested. + + itk.HDF5TransformIO refuses to write a CompositeTransform that holds + another one, so a chain over Greedy (which returns affine+warp composites) + would produce transforms that cannot be saved. + """ + stage1 = _CompositeRegistrar("stage1", 1.0) + stage2 = _CompositeRegistrar("stage2", 2.0) + chain = RegisterImagesChain([stage1, stage2]) + chain.set_fixed_image(_small_image()) + + result = chain.register(_small_image()) + + for key in ("forward_transform", "inverse_transform"): + composed = cast(itk.Transform, result[key]) + assert isinstance(composed, itk.CompositeTransform[itk.D, 3]) + for i in range(composed.GetNumberOfTransforms()): + sub = composed.GetNthTransform(i) + assert "Composite" not in sub.GetNameOfClass() + itk.transformwrite(composed, str(tmp_path / f"{key}.hdf")) + + # Splicing the sub-transforms in must leave the mapping unchanged. + forward = cast(itk.Transform, result["forward_transform"]) + assert list(forward.TransformPoint([0.0, 0.0, 0.0])) == [3.0, 0.0, 0.0] + + def test_chain_propagates_fixed_and_moving_state_to_each_child() -> None: """Every child must see a non-None fixed_image_pre (computed via its own preprocess(), not copied from the chain's no-op preprocess()) and the @@ -101,12 +162,16 @@ def test_chain_propagates_fixed_and_moving_state_to_each_child() -> None: for stage in (stage1, stage2): assert stage.seen_fixed_image_pre is not None - assert stage.seen_moving_image is moving + assert stage.seen_moving_image is not None # Each stage must compute its own preprocessing (moving_image_pre is # not inherited from the chain, which has no meaningful preprocess() # of its own). assert stage.seen_moving_image_pre is None assert stage.preprocess_call_count == 1 + # The first stage gets the caller's image; later stages get it pre-warped + # by the running result. + assert stage1.seen_moving_image is moving + assert stage2.seen_moving_image is not moving def test_chain_recomputes_fixed_image_pre_when_fixed_image_changes() -> None: diff --git a/tests/test_register_images_greedy.py b/tests/test_register_images_greedy.py index d079c709..efdbf065 100644 --- a/tests/test_register_images_greedy.py +++ b/tests/test_register_images_greedy.py @@ -15,6 +15,8 @@ from physiotwin4d.register_images_greedy import RegisterImagesGreedy from physiotwin4d.transform_tools import TransformTools +from .conftest import KnownShiftCase + @pytest.mark.slow class TestRegisterImagesGreedy: @@ -183,6 +185,48 @@ def test_register_affine_with_mask( print("Greedy affine registration complete with masks") + @pytest.mark.parametrize("transform_type", ["Rigid", "Affine", "Deformable"]) + def test_recovers_known_shift( + self, + known_shift_case: KnownShiftCase, + transform_type: str, + ) -> None: + """Greedy must recover a known shift, in the right direction. + + Regression guard for the RAS/LPS conversion: Greedy reports its affine + in RAS while ITK is LPS, so omitting the basis change negates x and y. + Before the fix this recovered (+6, -4, -3) mm instead of (-6, +4, -3) + and scored *below* the unregistered pair; the sign error passed every + other test in this file, which only check that transforms exist. + """ + registrar = RegisterImagesGreedy() + registrar.set_modality("ct") + registrar.set_transform_type(transform_type) + registrar.set_number_of_iterations([60, 30, 10]) + registrar.set_fixed_image(known_shift_case.fixed) + + result = registrar.register(moving_image=known_shift_case.moving) + forward_transform = result["forward_transform"] + + error_mm = known_shift_case.center_error_mm(forward_transform) + ncc = known_shift_case.foreground_ncc(forward_transform) + unregistered_ncc = known_shift_case.unregistered_ncc() + + print(f"\nGreedy {transform_type} known-shift recovery:") + print(f" expected displacement: {known_shift_case.expected_displacement}") + print(f" error: {error_mm:.2f} mm") + print(f" foreground NCC: {ncc:.4f} (unregistered {unregistered_ncc:.4f})") + + assert error_mm < 2.0, ( + f"Greedy {transform_type} recovered the shift {error_mm:.2f} mm off; " + "a sign or axis error inverts the transform (see the RAS/LPS " + "conversion in RegisterImagesGreedy._matrix_to_itk_affine)" + ) + assert ncc > unregistered_ncc, ( + f"Greedy {transform_type} left the images less aligned than they " + f"started ({ncc:.4f} vs {unregistered_ncc:.4f})" + ) + def test_transform_application( self, registrar_greedy: RegisterImagesGreedy, diff --git a/tests/test_register_images_icon.py b/tests/test_register_images_icon.py index 7ecde324..f3e369aa 100644 --- a/tests/test_register_images_icon.py +++ b/tests/test_register_images_icon.py @@ -16,12 +16,43 @@ from physiotwin4d.register_images_icon import RegisterImagesICON from physiotwin4d.transform_tools import TransformTools +from .conftest import KnownShiftCase + @pytest.mark.requires_gpu @pytest.mark.slow class TestRegisterImagesICON: """Test suite for ICON-based image registration.""" + def test_recovers_known_shift(self, known_shift_case: KnownShiftCase) -> None: + """ICON must recover a known shift, in the right direction. + + The companion of the Greedy check in test_register_images_greedy.py: + ICON returns ITK transforms in the LPS frame throughout, so this asserts + no equivalent RAS/LPS conversion is missing here. + """ + registrar = RegisterImagesICON() + registrar.set_modality("ct") + registrar.set_number_of_iterations(5) + registrar.set_fixed_image(known_shift_case.fixed) + + result = registrar.register(moving_image=known_shift_case.moving) + forward_transform = result["forward_transform"] + + error_mm = known_shift_case.center_error_mm(forward_transform) + ncc = known_shift_case.foreground_ncc(forward_transform) + unregistered_ncc = known_shift_case.unregistered_ncc() + + print("\nICON known-shift recovery:") + print(f" error: {error_mm:.2f} mm") + print(f" foreground NCC: {ncc:.4f} (unregistered {unregistered_ncc:.4f})") + + assert error_mm < 2.0, f"ICON recovered the shift {error_mm:.2f} mm off" + assert ncc > unregistered_ncc, ( + f"ICON left the images less aligned than they started " + f"({ncc:.4f} vs {unregistered_ncc:.4f})" + ) + def test_registrar_initialization(self, registrar_ICON: RegisterImagesICON) -> None: """Test that RegisterImagesICON initializes correctly.""" assert registrar_ICON is not None, "Registrar not initialized" @@ -386,9 +417,9 @@ def test_registration_with_initial_transform( registrar_ICON.set_fixed_image(fixed_image) registrar_ICON.set_number_of_iterations(2) - result = registrar_ICON.register( - moving_image=moving_image, - initial_forward_transform=initial_tfm_forward, + result = registrar_ICON.register_from( + initial_tfm_forward, + moving_image, ) assert isinstance(result, dict), "Result should be a dictionary" diff --git a/tests/test_register_models_pca.py b/tests/test_register_models_pca.py index 66b19746..884bc103 100644 --- a/tests/test_register_models_pca.py +++ b/tests/test_register_models_pca.py @@ -8,11 +8,12 @@ import numpy as np import pytest import pyvista as pv +from scipy.optimize import approx_fprime from physiotwin4d.register_models_pca import RegisterModelsPCA -def _make_registrar() -> RegisterModelsPCA: +def _make_registrar(**kwargs: Any) -> RegisterModelsPCA: """Create a small PCA registrar with a three-point template surface.""" template_model = pv.PolyData( np.array( @@ -27,27 +28,52 @@ def _make_registrar() -> RegisterModelsPCA: pca_eigenvectors = np.zeros((1, template_model.n_points * 3), dtype=np.float64) pca_std_deviations = np.ones(1, dtype=np.float64) fixed_distance_map = itk.image_from_array(np.zeros((4, 4, 4), dtype=np.float32)) + kwargs.setdefault("symmetric_weight", 0.0) return RegisterModelsPCA( pca_template_model=template_model, pca_eigenvectors=pca_eigenvectors, pca_std_deviations=pca_std_deviations, pca_number_of_modes=1, fixed_distance_map=fixed_distance_map, + **kwargs, ) -def test_itk_template_points_are_distinct_objects() -> None: - """Cached ITK points are distinct per template vertex.""" - registrar = _make_registrar() +def _sphere_registrar( + radius: float, + modes: int = 1, + **kwargs: Any, +) -> tuple[RegisterModelsPCA, np.ndarray]: + """Build a registrar whose single mode inflates a sphere radially. + + Returns the registrar and the per-point unit-norm eigenvector it was given. + """ + template = pv.Sphere(radius=radius, theta_resolution=24, phi_resolution=24) + directions = np.asarray(template.points, dtype=np.float64) + directions /= np.linalg.norm(directions, axis=1, keepdims=True) + + eigenvectors = np.zeros((modes, template.n_points * 3), dtype=np.float64) + eigenvectors[0] = directions.reshape(-1) / np.linalg.norm(directions) + for mode in range(1, modes): + # Orthogonal filler modes: displace a disjoint slab of points along x. + filler = np.zeros((template.n_points, 3), dtype=np.float64) + filler[mode::modes, 0] = 1.0 + eigenvectors[mode] = filler.reshape(-1) / np.linalg.norm(filler) - points = registrar._pca_template_model_points_itk - assert points is not None - assert len({id(point) for point in points}) == len(points) - assert [float(points[0][0]), float(points[1][0]), float(points[2][1])] == [ - 0.0, - 1.0, - 1.0, - ] + reference_image = itk.image_from_array(np.zeros((48, 48, 48), dtype=np.float32)) + reference_image.SetSpacing([1.0, 1.0, 1.0]) + reference_image.SetOrigin([-24.0, -24.0, -24.0]) + + registrar = RegisterModelsPCA( + pca_template_model=template, + pca_eigenvectors=eigenvectors, + pca_std_deviations=np.full(modes, 5.0), + pca_template_model_point_subsample=1, + fixed_model=pv.Sphere(radius=radius, theta_resolution=24, phi_resolution=24), + reference_image=reference_image, + **kwargs, + ) + return registrar, eigenvectors[0].reshape(-1, 3) def test_set_fixed_model_requires_reference_image() -> None: @@ -60,6 +86,48 @@ def test_set_fixed_model_requires_reference_image() -> None: ) +def test_mode_count_mismatch_is_rejected() -> None: + """Eigenvector and standard-deviation counts must agree.""" + template_model = pv.PolyData(np.zeros((3, 3), dtype=np.float64)) + with pytest.raises(ValueError, match="Mode count mismatch"): + RegisterModelsPCA( + pca_template_model=template_model, + pca_eigenvectors=np.zeros((2, 9), dtype=np.float64), + pca_std_deviations=np.ones(3, dtype=np.float64), + fixed_distance_map=itk.image_from_array( + np.zeros((4, 4, 4), dtype=np.float32) + ), + ) + + +def test_compute_pca_deformation_scales_eigenvectors_by_std() -> None: + """Deformation is exactly sum(b_i * std_i * eigenvector_i), reshaped (N, 3).""" + template_model = pv.PolyData(np.zeros((2, 3), dtype=np.float64)) + eigenvectors = np.array( + [ + [1.0, 0.0, 0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 1.0, 0.0, 0.0, 0.0], + ], + dtype=np.float64, + ) + std_deviations = np.array([2.0, 5.0], dtype=np.float64) + registrar = RegisterModelsPCA( + pca_template_model=template_model, + pca_eigenvectors=eigenvectors, + pca_std_deviations=std_deviations, + fixed_distance_map=itk.image_from_array(np.zeros((4, 4, 4), dtype=np.float32)), + ) + + deformation = registrar._compute_pca_deformation(np.array([1.5, -1.0])) + + # Point 0: 1.5*2*[1,0,0] + (-1)*5*[0,0,1]; point 1: 1.5*2*[0,1,0]. + assert np.allclose(deformation, [[3.0, 0.0, -5.0], [0.0, 3.0, 0.0]]) + + # A shorter coefficient vector uses only the leading modes. + leading = registrar._compute_pca_deformation(np.array([1.5])) + assert np.allclose(leading, [[3.0, 0.0, 0.0], [0.0, 3.0, 0.0]]) + + def test_transform_template_model_applies_post_pca_transform_after_deformation() -> ( None ): @@ -78,3 +146,190 @@ def test_transform_template_model_applies_post_pca_transform_after_deformation() assert np.allclose(registered_model.points[0], [2.0, 0.0, 0.0]) assert np.allclose(registered_model.points[1], [4.0, 0.0, 0.0]) + + +def test_modes_are_deformed_in_the_template_frame_then_transformed() -> None: + """Regression: modes must be rotated with the template, not added after it. + + The registered model must equal ``A @ (template + deformation)``, never + ``A @ template + deformation``. The two differ whenever the post-PCA + transform contains a rotation, which is the case for every ICP alignment. + """ + registrar = _make_registrar() + registrar.registered_model_pca_coefficients = np.array([1.0], dtype=np.float64) + deformation = np.tile( + np.array([1.0, 0.0, 0.0], dtype=np.float64), + (registrar.pca_template_model.n_points, 1), + ) + registrar.registered_model_pca_deformation = deformation + + # 90 degrees about z, so x-displacements must come out along y. + matrix = np.array( + [[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]], dtype=np.float64 + ) + offset = np.array([10.0, -3.0, 2.0], dtype=np.float64) + transform = itk.AffineTransform[itk.D, 3].New() + transform.SetMatrix(itk.matrix_from_array(matrix)) + transform.SetTranslation(offset) + registrar.post_pca_transform = transform + + registered_model: Any = registrar.transform_template_model() + + template_points = np.asarray(registrar.pca_template_model.points, dtype=np.float64) + correct = (template_points + deformation) @ matrix.T + offset + wrong = template_points @ matrix.T + offset + deformation + + assert np.allclose(registered_model.points, correct) + assert not np.allclose(registered_model.points, wrong) + + +def test_analytic_gradient_matches_finite_differences() -> None: + """The supplied Jacobian agrees with a finite-difference gradient.""" + registrar, _ = _sphere_registrar(radius=8.0, modes=3, symmetric_weight=0.5) + registrar.pca_prior_weight = 0.1 + registrar._prepare_sampling() + + params = np.array([0.4, -0.25, 0.15], dtype=np.float64) + _, analytic = registrar._objective_and_gradient(params) + numeric = approx_fprime(params, registrar._mean_distance_metric, 1e-5) + + assert np.allclose(analytic, numeric, atol=2e-3) + + +def test_register_recovers_known_coefficients() -> None: + """Fitting to a target built from known coefficients recovers them.""" + registrar, mode = _sphere_registrar(radius=8.0, modes=1, symmetric_weight=0.0) + + # Target = template inflated by b = 1.0 along the single radial mode. + truth = 1.0 + target = registrar.pca_template_model.copy(deep=True) + target.points = ( + np.asarray(registrar.pca_template_model.points, dtype=np.float64) + + truth * registrar.pca_std_deviations[0] * mode + ) + + reference_image = itk.image_from_array(np.zeros((64, 64, 64), dtype=np.float32)) + reference_image.SetSpacing([0.75, 0.75, 0.75]) + reference_image.SetOrigin([-24.0, -24.0, -24.0]) + registrar.set_fixed_model(cast(pv.UnstructuredGrid, target), reference_image) + + result = registrar.register(pca_number_of_modes=1, max_iterations=60) + + assert result["pca_coefficients"][0] == pytest.approx(truth, abs=0.1) + assert result["mean_distance"] < 0.5 + + +def test_symmetric_term_penalizes_partial_coverage() -> None: + """The target-to-model term sees coverage the model-to-target term misses. + + A hemisphere sitting on a full sphere scores near-perfectly one-way: every + model point lies on the target surface. Only the target-to-model term + notices that half the target has no model near it. + """ + target = pv.Sphere(radius=8.0, theta_resolution=24, phi_resolution=24) + points = np.asarray(target.points, dtype=np.float64) + hemisphere = pv.PolyData(points[points[:, 2] > 0.0]) + + reference_image = itk.image_from_array(np.zeros((64, 64, 64), dtype=np.float32)) + reference_image.SetSpacing([0.75, 0.75, 0.75]) + reference_image.SetOrigin([-24.0, -24.0, -24.0]) + + registrar = RegisterModelsPCA( + pca_template_model=hemisphere, + pca_eigenvectors=np.zeros((1, hemisphere.n_points * 3), dtype=np.float64), + pca_std_deviations=np.ones(1, dtype=np.float64), + pca_template_model_point_subsample=1, + fixed_model=target, + reference_image=reference_image, + symmetric_weight=0.0, + ) + + registrar._prepare_sampling() + one_way = registrar._mean_distance_metric(np.zeros(1)) + + registrar.symmetric_weight = 0.5 + registrar._prepare_sampling() + symmetric = registrar._mean_distance_metric(np.zeros(1)) + + # Model points all lie on the target surface, so the one-way term is small. + assert one_way < 0.5 + assert symmetric > 1.0 + + +def test_prior_shrinks_coefficients() -> None: + """Raising pca_prior_weight pulls the solution toward the mean shape.""" + registrar, mode = _sphere_registrar(radius=8.0, modes=1, symmetric_weight=0.0) + + target = registrar.pca_template_model.copy(deep=True) + target.points = ( + np.asarray(registrar.pca_template_model.points, dtype=np.float64) + + 1.0 * registrar.pca_std_deviations[0] * mode + ) + + reference_image = itk.image_from_array(np.zeros((64, 64, 64), dtype=np.float32)) + reference_image.SetSpacing([0.75, 0.75, 0.75]) + reference_image.SetOrigin([-24.0, -24.0, -24.0]) + registrar.set_fixed_model(cast(pv.UnstructuredGrid, target), reference_image) + + unregularized = registrar.register(pca_number_of_modes=1, max_iterations=60) + + registrar.pca_prior_weight = 5.0 + registrar._sampling_ready = False + regularized = registrar.register(pca_number_of_modes=1, max_iterations=60) + + assert abs(regularized["pca_coefficients"][0]) < abs( + unregularized["pca_coefficients"][0] + ) + + +def test_transform_point_requires_computed_transforms() -> None: + """transform_point raises instead of silently returning the input.""" + registrar = _make_registrar() + point = itk.Point[itk.D, 3]() + point[0], point[1], point[2] = 1.0, 2.0, 3.0 + + with pytest.raises(ValueError, match="compute_pca_transforms"): + registrar.transform_point(point) + + +def test_pca_transforms_round_trip() -> None: + """forward reproduces the deformation and inverse undoes it.""" + registrar, mode = _sphere_registrar(radius=8.0, modes=1, symmetric_weight=0.0) + registrar.registered_model_pca_coefficients = np.array([1.0], dtype=np.float64) + registrar.registered_model_pca_deformation = ( + 1.0 * registrar.pca_std_deviations[0] * mode + ) + + reference_image = itk.image_from_array(np.zeros((64, 64, 64), dtype=np.float32)) + reference_image.SetSpacing([0.75, 0.75, 0.75]) + reference_image.SetOrigin([-24.0, -24.0, -24.0]) + + transforms = registrar.compute_pca_transforms(reference_image, blur_sigma=1.5) + forward = transforms["forward_point_transform"] + inverse = transforms["inverse_point_transform"] + + template_points = np.asarray(registrar.pca_template_model.points, dtype=np.float64) + expected = template_points + registrar.registered_model_pca_deformation + + point = itk.Point[itk.D, 3]() + mapped = np.empty_like(template_points) + back = np.empty_like(template_points) + for i, source in enumerate(template_points): + point[0], point[1], point[2] = (float(v) for v in source) + forward_point = forward.TransformPoint(point) + mapped[i] = (forward_point[0], forward_point[1], forward_point[2]) + inverse_point = inverse.TransformPoint(forward_point) + back[i] = (inverse_point[0], inverse_point[1], inverse_point[2]) + + # The field is splatted and blurred, so it only approximates the deformation. + field_rms = np.sqrt(np.mean(np.sum((mapped - expected) ** 2, axis=1))) + round_trip_rms = np.sqrt(np.mean(np.sum((back - template_points) ** 2, axis=1))) + + assert field_rms < 1.5 + assert round_trip_rms < 1.0 + + # An identity field would satisfy both bounds above, so require that the + # forward transform actually moves the points a comparable distance. + displacement_rms = np.sqrt(np.mean(np.sum((mapped - template_points) ** 2, axis=1))) + expected_rms = np.sqrt(np.mean(np.sum((expected - template_points) ** 2, axis=1))) + assert displacement_rms > 0.5 * expected_rms diff --git a/tests/test_register_time_series_images.py b/tests/test_register_time_series_images.py index e7268d9b..10396c4b 100644 --- a/tests/test_register_time_series_images.py +++ b/tests/test_register_time_series_images.py @@ -145,7 +145,6 @@ def test_register_time_series_basic( moving_images=moving_images, reference_frame=0, register_reference=True, - prior_weight=0.0, ) # Verify result structure @@ -207,16 +206,15 @@ def test_register_time_series_basic( assert (results_dir / "basic_forward_transform_0.hdf").exists() assert (results_dir / "basic_time_series_registered_0.mha").exists() - def test_register_time_series_with_prior( + def test_register_time_series_from_middle_frame( self, test_images: list[Any], test_directories: dict[str, Path] ) -> None: - """Test time series registration with prior transform usage.""" + """Test time series registration starting from a middle reference frame.""" fixed_image = test_images[0] moving_images = test_images[1:4] - print("\nRegistering time series (with prior)...") + print("\nRegistering time series (middle reference frame)...") print(f" Number of moving images: {len(moving_images)}") - print(" Using prior transform weight: 0.5") greedy = RegisterImagesGreedy() greedy.set_number_of_iterations([20, 10, 2]) @@ -228,7 +226,6 @@ def test_register_time_series_with_prior( moving_images=moving_images, reference_frame=1, # Start from middle register_reference=True, - prior_weight=0.5, ) forward_transforms = result["forward_transforms"] @@ -246,7 +243,7 @@ def test_register_time_series_with_prior( for i, forward_transform in enumerate(forward_transforms): assert forward_transform is not None, f"forward_transform[{i}] is None" - print("Time series registration with prior complete") + print("Time series registration from the middle frame complete") print(f" Losses: {[f'{loss:.6f}' for loss in losses]}") test_tools = TestTools( @@ -259,14 +256,14 @@ def test_register_time_series_with_prior( # bit-reproducible across runs, so we save artifacts without # asserting an exact baseline match. test_tools.write_result_transform( - forward_transforms[0], "prior_forward_transform_0.hdf" + forward_transforms[0], "middle_frame_forward_transform_0.hdf" ) test_tools.write_result_image( - moving_image, "prior_time_series_registered_0.mha" + moving_image, "middle_frame_time_series_registered_0.mha" ) results_dir = test_directories["output"] / self._class_name - assert (results_dir / "prior_forward_transform_0.hdf").exists() - assert (results_dir / "prior_time_series_registered_0.mha").exists() + assert (results_dir / "middle_frame_forward_transform_0.hdf").exists() + assert (results_dir / "middle_frame_time_series_registered_0.mha").exists() def test_register_time_series_identity_start(self, test_images: list[Any]) -> None: """Test time series registration with identity for starting image.""" @@ -285,7 +282,6 @@ def test_register_time_series_identity_start(self, test_images: list[Any]) -> No moving_images=moving_images, reference_frame=0, register_reference=False, # Use identity - prior_weight=0.0, ) # Starting image should have very low/zero loss @@ -317,7 +313,6 @@ def test_register_time_series_different_starting_indices( moving_images=moving_images, reference_frame=starting_index, register_reference=True, - prior_weight=0.0, ) assert len(result["forward_transforms"]) == len(moving_images), ( @@ -360,31 +355,6 @@ def test_register_time_series_error_invalid_starting_index( print("\nInvalid starting index correctly rejected") - def test_register_time_series_error_invalid_prior_portion( - self, test_images: list[Any] - ) -> None: - """Test that error is raised for invalid prior portion value.""" - registrar = RegisterTimeSeriesImages(registration_method=RegisterImagesGreedy()) - registrar.set_fixed_image(test_images[0]) - - moving_images = test_images[1:4] - - # Test negative value - with pytest.raises(ValueError, match="must be in"): - registrar.register_time_series( - moving_images=moving_images, - prior_weight=-0.1, - ) - - # Test value > 1 - with pytest.raises(ValueError, match="must be in"): - registrar.register_time_series( - moving_images=moving_images, - prior_weight=1.5, - ) - - print("\nInvalid prior portion correctly rejected") - def test_transform_application_time_series( self, test_images: list[Any], test_directories: dict[str, Path] ) -> None: @@ -404,7 +374,6 @@ def test_transform_application_time_series( moving_images=moving_images, reference_frame=0, register_reference=True, - prior_weight=0.0, ) forward_transforms = result["forward_transforms"] @@ -457,7 +426,6 @@ def test_register_time_series_ICON(self, test_images: list[Any]) -> None: moving_images=moving_images, reference_frame=0, register_reference=True, - prior_weight=0.0, ) assert len(result["forward_transforms"]) == len(moving_images) @@ -505,7 +473,6 @@ def test_register_time_series_with_mask( moving_images=moving_images, reference_frame=0, register_reference=True, - prior_weight=0.0, ) assert len(result["forward_transforms"]) == len(moving_images) @@ -531,7 +498,6 @@ def test_bidirectional_registration(self, test_images: list[Any]) -> None: moving_images=moving_images, reference_frame=2, # Middle image register_reference=True, - prior_weight=0.0, ) forward_transforms = result["forward_transforms"] diff --git a/tests/test_tutorials.py b/tests/test_tutorials.py index d5e2fb64..9f15ca1b 100644 --- a/tests/test_tutorials.py +++ b/tests/test_tutorials.py @@ -202,25 +202,22 @@ class TestTutorial05HeartVTKToUSD: def test_run( self, test_directories: dict[str, Path], test_images: list[Any] ) -> None: - # The script reads this exact path and offers no input override, so - # bootstrap Tutorial 4 rather than pointing it at another surface. - vtk_file = ( - _REPO_ROOT - / "tutorials" - / "output" - / "tutorial_04_heart" - / "patient_surfaces.vtp" - ) - if not vtk_file.exists(): + # The script reads this exact directory and offers no input override, + # so bootstrap Tutorial 4 rather than pointing it at other surfaces. + input_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_04_heart" + if not list(input_dir.glob("patient_*.vtp")): _run_tutorial_script("tutorial_04_heart_ct_to_vtk.py") - assert vtk_file.exists(), ( - f"Tutorial 4 bootstrap did not create the expected surface: {vtk_file}" + assert list(input_dir.glob("patient_*.vtp")), ( + f"Tutorial 4 bootstrap did not create surfaces in: {input_dir}" ) out_dir = _REPO_ROOT / "tutorials" / "output" / "tutorial_05_heart" results = _run_tutorial_script("tutorial_05_heart_vtk_to_usd.py") assert results["usd_file"], "USD file path should not be empty" assert Path(results["usd_file"]).exists(), "USD file should exist" + assert len(results["structures"]) > 1, ( + "Per-structure surfaces expected, so that each becomes its own prim" + ) tt = TestTools( class_name=self._class_name, diff --git a/tests/test_workflow_convert_vtk_to_usd.py b/tests/test_workflow_convert_vtk_to_usd.py new file mode 100644 index 00000000..1fc6ce2a --- /dev/null +++ b/tests/test_workflow_convert_vtk_to_usd.py @@ -0,0 +1,169 @@ +"""Tests for the appearance and object-naming behavior of the VTK-to-USD workflow. + +Synthetic meshes only - no segmentation or image data required. +""" + +from pathlib import Path +from typing import Optional + +import numpy as np +import pyvista as pv +from pxr import Usd, UsdShade + +from physiotwin4d import WorkflowConvertVTKToUSD + + +def _labeled_sphere( + center: tuple[float, float, float], + label_name: str, + group: Optional[str] = None, +) -> pv.PolyData: + """Return a sphere annotated the way WorkflowConvertImageToVTK annotates one.""" + surface = pv.Sphere(radius=1.0, center=center, theta_resolution=8, phi_resolution=8) + surface.field_data["SegmentationLabelNames"] = np.array([label_name]) + if group is not None: + surface.field_data["AnatomyGroup"] = np.array([group]) + return surface + + +def _bound_material_path(stage: Usd.Stage, mesh_path: str) -> str: + prim = stage.GetPrimAtPath(mesh_path) + assert prim.IsValid(), f"Missing prim: {mesh_path}" + binding = UsdShade.MaterialBindingAPI(prim).GetDirectBinding() + return str(binding.GetMaterialPath()) + + +class TestAnatomyAppearance: + """Per-structure materials must follow the structure names on the meshes.""" + + def test_label_names_drive_prim_names_and_materials(self, tmp_path: Path) -> None: + """Each labeled mesh becomes its own prim with its own anatomy material.""" + meshes = [ + _labeled_sphere((0.0, 0.0, 0.0), "highres_myocardium"), + _labeled_sphere((3.0, 0.0, 0.0), "highres_ventricle_left"), + ] + + workflow = WorkflowConvertVTKToUSD( + input_meshes=meshes, + usd_project_name="heart", + output_directory=tmp_path, + appearance="anatomy", + static_merge=True, + ) + result = workflow.process() + + stage = Usd.Stage.Open(result["usd_file"]) + # separate_by_connectivity defaults to True, and each sphere is a + # single connected component, so every object gains one "_object1" part. + myocardium = _bound_material_path( + stage, "/World/heart/highres_myocardium_object1" + ) + ventricle = _bound_material_path( + stage, "/World/heart/highres_ventricle_left_object1" + ) + assert myocardium.endswith("OmniSurface_Myocardium") + assert ventricle.endswith("OmniSurface_Ventricle_Left") + + def test_explicit_anatomy_type_overrides_names(self, tmp_path: Path) -> None: + """A caller-supplied anatomy_type still paints every object the same.""" + meshes = [ + _labeled_sphere((0.0, 0.0, 0.0), "highres_myocardium"), + _labeled_sphere((3.0, 0.0, 0.0), "highres_ventricle_left"), + ] + + workflow = WorkflowConvertVTKToUSD( + input_meshes=meshes, + usd_project_name="heart", + output_directory=tmp_path, + appearance="anatomy", + anatomy_type="heart", + static_merge=True, + ) + result = workflow.process() + + stage = Usd.Stage.Open(result["usd_file"]) + for name in ("highres_myocardium", "highres_ventricle_left"): + material = _bound_material_path(stage, f"/World/heart/{name}_object1") + assert material.endswith("OmniSurface_Heart") + + def test_unmatched_name_falls_back_to_group(self, tmp_path: Path) -> None: + """No material is named "rib_left_3", so its anatomy group decides.""" + meshes = [ + _labeled_sphere((0.0, 0.0, 0.0), "rib_left_3", group="bone"), + _labeled_sphere((3.0, 0.0, 0.0), "vertebrae_T7", group="bone"), + ] + + workflow = WorkflowConvertVTKToUSD( + input_meshes=meshes, + usd_project_name="chest", + output_directory=tmp_path, + appearance="anatomy", + static_merge=True, + ) + result = workflow.process() + + stage = Usd.Stage.Open(result["usd_file"]) + for name in ("rib_left_3", "vertebrae_T7"): + material = _bound_material_path(stage, f"/World/chest/{name}_object1") + assert material.endswith("OmniSurface_Bone") + + def test_structure_name_wins_over_group(self, tmp_path: Path) -> None: + """A structure with its own material must not collapse onto its group.""" + mesh = _labeled_sphere((0.0, 0.0, 0.0), "highres_myocardium", group="heart") + + workflow = WorkflowConvertVTKToUSD( + input_meshes=[mesh], + usd_project_name="heart", + output_directory=tmp_path, + appearance="anatomy", + static_merge=True, + ) + result = workflow.process() + + stage = Usd.Stage.Open(result["usd_file"]) + material = _bound_material_path( + stage, "/World/heart/highres_myocardium_object1" + ) + assert material.endswith("OmniSurface_Myocardium") + + def test_unmatched_name_falls_back_to_other(self, tmp_path: Path) -> None: + """A mesh whose name matches no anatomy still gets a material.""" + mesh = _labeled_sphere((0.0, 0.0, 0.0), "calibration_phantom") + + workflow = WorkflowConvertVTKToUSD( + input_meshes=[mesh], + usd_project_name="scan", + output_directory=tmp_path, + appearance="anatomy", + static_merge=True, + ) + result = workflow.process() + + stage = Usd.Stage.Open(result["usd_file"]) + material = _bound_material_path( + stage, "/World/scan/calibration_phantom_object1" + ) + assert material.endswith("OmniSurface_Other") + + def test_unlabeled_meshes_keep_positional_names(self, tmp_path: Path) -> None: + """Without SegmentationLabelNames, naming stays {project}_{index}.""" + meshes = [ + pv.Sphere(radius=1.0, theta_resolution=8, phi_resolution=8), + pv.Sphere( + radius=1.0, center=(3.0, 0.0, 0.0), theta_resolution=8, phi_resolution=8 + ), + ] + + workflow = WorkflowConvertVTKToUSD( + input_meshes=meshes, + usd_project_name="scan", + output_directory=tmp_path, + appearance="anatomy", + anatomy_type="heart", + static_merge=True, + ) + result = workflow.process() + + stage = Usd.Stage.Open(result["usd_file"]) + assert stage.GetPrimAtPath("/World/scan/scan_0_object1").IsValid() + assert stage.GetPrimAtPath("/World/scan/scan_1_object1").IsValid() diff --git a/tutorials/README.md b/tutorials/README.md index ba92cf11..02c37c7f 100644 --- a/tutorials/README.md +++ b/tutorials/README.md @@ -28,6 +28,7 @@ current working directory. | 1 | [tutorial_01_heart_gated_ct_to_usd.py](tutorial_01_heart_gated_ct_to_usd.py) | `WorkflowConvertImageToUSD` | Slicer-Heart-CT (prepare first) | | 1 | [tutorial_01_lung_gated_ct_to_usd.py](tutorial_01_lung_gated_ct_to_usd.py) | `WorkflowConvertImageToUSD` | Lung gated 4D CT (prepare first) | | 2 | [tutorial_02_lung_finetune_icon.py](tutorial_02_lung_finetune_icon.py) | `WorkflowFinetuneICONRegistration` | DirLab-4DCT (manual) | +| 2 | [distancemap variant](tutorial_02_lung_distancemap_finetune_icon.py) | `WorkflowFinetuneICONRegistration` on distance maps | DirLab-4DCT (manual) | | 3 | [tutorial_03_heart_reconstruct_highres_4d_ct.py](tutorial_03_heart_reconstruct_highres_4d_ct.py) | `WorkflowReconstructHighres4DCT` | Slicer-Heart-CT (prepare first) | | 3 | [tutorial_03_lung_reconstruct_highres_4d_ct.py](tutorial_03_lung_reconstruct_highres_4d_ct.py) | `WorkflowReconstructHighres4DCT` | DirLab-4DCT (manual) | | 4 | [tutorial_04_heart_ct_to_vtk.py](tutorial_04_heart_ct_to_vtk.py) | `WorkflowConvertImageToVTK` | Slicer-Heart-CT (prepare first) | diff --git a/tutorials/tutorial_02_lung_distancemap_finetune_icon.py b/tutorials/tutorial_02_lung_distancemap_finetune_icon.py new file mode 100644 index 00000000..1d754140 --- /dev/null +++ b/tutorials/tutorial_02_lung_distancemap_finetune_icon.py @@ -0,0 +1,612 @@ +""" +Tutorial 2 (Lung, distance maps): Finetune uniGradICON on lung distance maps + +Purpose +------- +``RegisterModelsDistanceMaps`` -- the labelmap-to-labelmap stage of +``WorkflowFitStatisticalModelToPatient`` -- does not register CT intensities. +It rasterizes a signed squared distance map from each surface, normalizes it to +[-1, 1] and scales it by 1000 so that it fills the CT window uniGradICON +preprocesses with, then hands that pair to ICON. Stock uniGradICON has never +seen such an image, so this tutorial finetunes it on exactly that +representation: distance maps rasterized from the lung surfaces segmented out of +the DIR-Lab 4D CT cases. + +The finetuning cohort is every DIR-Lab case except Case 1, and within each case +only every other respiratory phase (``T00``, ``T20``, ``T40``, ``T60``, ``T80``) +-- half the time points, spanning the full breathing cycle at half the +segmentation cost. Each selected phase is segmented once with +``SegmentNVSegmentCTMRI``; the lung labelmap is kept for uniGradICON's Dice +loss and the lung surfaces are combined and rasterized into the distance map +that serves as the training "image". Segmentation outputs are cached on disk, +so a second run of this tutorial only re-runs the finetuning. + +Accuracy is measured on Case 1, which is never seen during finetuning, by +registering ``Case1Pack_T00.mha`` (moving) to ``Case1Pack_T50.mha`` (fixed) five +ways: ``RegisterImagesGreedy`` deformable on the distance maps, +``RegisterImagesICON`` on the distance maps with the stock and with the +finetuned weights, ``RegisterImagesGreedyICON`` on the distance maps with the +finetuned weights in its ICON stage -- which separates what the finetuned +network adds on its own from what it adds on top of the Greedy result -- and, +as the intensity-based reference point, stock ``RegisterImagesICON`` on the CT +images themselves. No registration masks are used, so the comparison isolates +the method and the weights. + +The two metrics match ``tutorial_02_lung_finetune_icon.py``. The primary metric +is target registration error: DIR-Lab ships 300 expert landmarks for the extreme +phases (T00 and T50) of every case, so each fixed-image landmark is mapped +through the registration transform and compared, in millimeters, against its +moving-image counterpart. The secondary metric is label overlap: the moving +lung labelmap is warped onto the fixed grid by every transform, so the Dice +scores reflect the transform rather than segmentation variability. The moving +image and labelmap resampled onto the fixed grid without registration supply the +"before registration" reference row for both metrics. + +Reported per method: the mean, standard deviation, 95th percentile and maximum +landmark error in millimeters; the mean, 5th percentile, median, 95th +percentile, minimum and maximum of the per-class Dice scores; the number of +mislabeled voxels; and the wall-clock registration time. + +Finetuning artifacts (dataset JSON, YAML config, checkpoint tree) are written +under ``tutorials/network_weights/icon_dirlab_4dct_distancemap``. The final +checkpoint is ``tutorials/network_weights/icon_dirlab_4dct_distancemap/ +icon_dirlab_4dct_distancemap_model/checkpoints/network_weights_final.trch``, the +path returned by ``WorkflowFinetuneICONRegistration.expected_weights_path()`` +and the path ``tutorial_07_lung_fit_statistical_model_to_patient.py`` and +``tutorial_08_lung_fit_model_to_4d_patients.py`` look for. That directory is +deleted at the start of every run, so each run finetunes from scratch; see the +comment above the ``shutil.rmtree`` call for how to reuse a previous run. + +Data Required +------------- +Full data: ``data/DirLab-4DCT`` (all 10 cases, converted to HU ``.mha`` by +``data/DirLab-4DCT/fix_downloaded_data.py``), including the raw +``downloaded_data/Case1Pack/ExtremePhases`` landmark files +Test data: ``data/test/DirLab-4DCT`` +""" + +# Imports +from __future__ import annotations + +import csv +import logging +import shutil +import time +from pathlib import Path +from typing import Any, Optional, cast + +import itk +import numpy as np +import pyvista as pv + +from physiotwin4d import ( + ContourTools, + PhysioTwin4DBase, + RegisterImagesBase, + RegisterImagesGreedy, + RegisterImagesGreedyICON, + RegisterImagesICON, + SegmentNVSegmentCTMRI, + TestTools, + TransformTools, + WorkflowConvertImageToVTK, + WorkflowFinetuneICONRegistration, +) + +# Only run if this script is not imported as a module + +# nnUNetv2 (used inside the segmenter) spawns a multiprocessing.Pool and +# unigradicon finetuning is launched as a subprocess that spawns torch workers; +# on Windows the spawn start method re-imports this script in each child, so all +# top-level work stays under the __name__ == "__main__" guard. +if __name__ == "__main__": + # Data directory specification + repo_root = Path(__file__).resolve().parent.parent + tutorials_dir = Path(__file__).resolve().parent + + class_name = "tutorial_02_lung_distancemap_finetune_icon" + + output_dir = tutorials_dir / "output" / "tutorial_02_lung_distancemap" + # Segmented labelmaps, lung surfaces, and rasterized distance maps. Cached + # so re-runs skip the segmentation, which dominates this tutorial's runtime. + derived_dir = output_dir / "distance_maps" + baselines_dir = repo_root / "tests" / "baselines" + + # The workflow writes its dataset JSON, YAML config, and checkpoint tree + # under ``weights_dir / finetune_name``. + weights_dir = tutorials_dir / "network_weights" + finetune_name = "icon_dirlab_4dct_distancemap" + + # Distance-map normalization. WorkflowFitStatisticalModelToPatient passes + # (1.25 * mask_dilation_mm) ** 2 as RegisterModelsDistanceMaps' + # distance_squared_max, so the value below fixes the saturation radius of + # every distance map the finetuned weights will ever see. The lung fitting + # tutorials set the same mask_dilation_mm; changing one without the other + # trains on a different image distribution than it infers on. + mask_dilation_mm = 40.0 + distance_squared_max = (1.25 * mask_dilation_mm) ** 2 + + run_finetuning = True + + # Half the respiratory phases per case, taken as every other phase so the + # selection still spans inhale through exhale. + phase_stride = 2 + + test_mode = TestTools.running_as_test() + if test_mode: + data_dir = repo_root / "data" / "test" / "DirLab-4DCT" + number_of_iterations_greedy: Optional[list[int]] = [1, 0] + number_of_iterations_icon = 1 + epochs = 1 + else: + data_dir = repo_root / "data" / "DirLab-4DCT" + number_of_iterations_greedy = [60, 30, 20] + number_of_iterations_icon = 10 + epochs = 200 + + log_level = logging.INFO + reporter = PhysioTwin4DBase(class_name=class_name, log_level=log_level) + + derived_dir.mkdir(parents=True, exist_ok=True) + + # Held-out evaluation pair (Case 1 is excluded from finetuning). T00 and + # T50 are the extreme inhale/exhale phases, the only pair DIR-Lab supplies + # expert landmarks for. + fixed_file = data_dir / "Case1Pack_T50.mha" + moving_file = data_dir / "Case1Pack_T00.mha" + landmark_dir = data_dir / "downloaded_data" / "Case1Pack" / "ExtremePhases" + fixed_landmark_file = landmark_dir / "Case1_300_T50_xyz.txt" + moving_landmark_file = landmark_dir / "Case1_300_T00_xyz.txt" + missing = [ + str(p) + for p in (fixed_file, moving_file, fixed_landmark_file, moving_landmark_file) + if not p.exists() + ] + if missing: + raise FileNotFoundError( + f"Missing DirLab phase images or landmarks: {missing}.\n" + "See data/DirLab-4DCT/README.md for download instructions." + ) + + # Segmentation and distance-map generation + segmentation_workflow = WorkflowConvertImageToVTK( + segmentation_method=SegmentNVSegmentCTMRI(log_level=log_level), + log_level=log_level, + ) + contour_tools = ContourTools(log_level=log_level) + transform_tools = TransformTools() + + def segment_phase(image_file: Path) -> tuple[Path, Path]: + """Segment one phase's lungs and rasterize their distance map. + + Returns the distance map and lung labelmap paths. Both, plus the + combined lung surface, are cached under ``derived_dir``; an existing + pair short-circuits the segmentation. + """ + distance_map_file = derived_dir / f"{image_file.stem}_distance_map.mha" + labelmap_file = derived_dir / f"{image_file.stem}_lung_labelmap.nii.gz" + surface_file = derived_dir / f"{image_file.stem}_lung_surface.vtp" + if distance_map_file.exists() and labelmap_file.exists(): + return distance_map_file, labelmap_file + + reporter.log_info("Segmenting lungs in %s", image_file.name) + image = itk.imread(str(image_file), pixel_type=itk.F) + segmentation_result = segmentation_workflow.process( + input_image=image, + anatomy_groups=["lung"], + extract_label_surfaces=True, + ) + contour_tools.save_combined_surfaces( + segmentation_result["label_surfaces"], str(surface_file) + ) + itk.imwrite( + segmentation_result["labelmap"], str(labelmap_file), compression=True + ) + surface = cast(pv.PolyData, pv.read(str(surface_file))) + # Rasterize the lung surface into ICON's distance-map representation, + # mirroring ``RegisterModelsDistanceMaps._create_masks_from_models`` so + # the finetuning inputs match what that class feeds ICON at inference: + # a signed squared distance normalized to [-1, 1] by + # ``distance_squared_max``, then multiplied by 1000 to fill the + # [-1000, 1000] HU window uniGradICON's CT preprocessing expects. + distance_map = contour_tools.create_distance_map( + surface, + image, + squared_distance=True, + negative_inside=True, + zero_inside=False, + norm_to_max_distance=distance_squared_max, + ) + itk.GetArrayViewFromImage(distance_map)[...] *= 1000 + itk.imwrite(distance_map, str(distance_map_file), compression=True) + return distance_map_file, labelmap_file + + # Finetuning cohort: every case except Case1Pack, every other phase. + # ``Case10Pack_*`` is kept because only the exact ``Case1Pack_`` prefix is + # excluded. + case_phase_files: dict[str, list[Path]] = {} + for path in sorted(data_dir.glob("Case*_T??.mha")): + if path.name.startswith("Case1Pack_"): + continue + case_phase_files.setdefault(path.name.split("_")[0], []).append(path) + if not case_phase_files: + raise FileNotFoundError( + f"No non-Case1 DirLab phase images found under {data_dir}.\n" + "See data/DirLab-4DCT/README.md for download instructions." + ) + case_phase_files = { + case_id: files[::phase_stride] for case_id, files in case_phase_files.items() + } + reporter.log_info( + "Finetuning cohort: %d cases, %d frames (every %d-th phase)", + len(case_phase_files), + sum(len(files) for files in case_phase_files.values()), + phase_stride, + ) + + subject_distance_map_files: list[list[str]] = [] + subject_labelmap_files: list[list[Optional[str]]] = [] + for case_id, phase_files in case_phase_files.items(): + segmented = [segment_phase(path) for path in phase_files] + subject_distance_map_files.append([str(pair[0]) for pair in segmented]) + subject_labelmap_files.append([str(pair[1]) for pair in segmented]) + + # Always finetune from scratch. uniGradICON refuses to overwrite an + # existing experiment directory: it appends "-N" to the name instead + # (``icon_dirlab_4dct_distancemap_model-5``, ...), while + # expected_weights_path() keeps pointing at the original, never-written + # path. Deleting the tree up front keeps the two in agreement. + # + # To reuse a previous run instead, delete the shutil.rmtree call below and + # guard the process() call: + # weights_path = workflow.expected_weights_path() + # if not weights_path.exists(): + # weights_path = workflow.process() + experiment_dir = weights_dir / finetune_name + if run_finetuning: + if experiment_dir.exists(): + reporter.log_info( + "Removing previous finetuning outputs: %s", experiment_dir + ) + shutil.rmtree(experiment_dir) + + # Unlike tutorial_02_lung_finetune_icon.py, the lungs are segmented here + # anyway, so the labelmaps are supplied and uniGradICON's Dice loss stays + # enabled at its default weight. + # + # lncc_sigma matches the sigma RegisterImagesICON uses at inference, so + # finetuning optimizes the similarity this comparison scores. The + # distance maps are already scaled into [-1000, 1000], so the default CT + # window passes them through unclipped. + workflow = WorkflowFinetuneICONRegistration( + subject_image_files=subject_distance_map_files, + output_dir=weights_dir, + finetune_name=finetune_name, + subject_ids=list(case_phase_files.keys()), + subject_labelmap_files=subject_labelmap_files, + epochs=epochs, + lncc_sigma=5, + log_level=log_level, + ) + weights_path = workflow.process() + else: + weights_path = ( + experiment_dir + / f"{finetune_name}_model" + / "checkpoints" + / "network_weights_final.trch" + ) + # Checked here rather than at the first set_weights_path() call, which + # only happens after the greedy and stock-ICON rows have already run. + if not weights_path.exists(): + raise FileNotFoundError( + f"run_finetuning is False but no checkpoint at {weights_path}. " + "Set run_finetuning = True to finetune from scratch." + ) + + # Registration comparison + fixed_image = itk.imread(str(fixed_file), pixel_type=itk.F) + moving_image = itk.imread(str(moving_file), pixel_type=itk.F) + + fixed_distance_map_file, fixed_labelmap_file = segment_phase(fixed_file) + moving_distance_map_file, moving_labelmap_file = segment_phase(moving_file) + fixed_distance_map = itk.imread(str(fixed_distance_map_file), pixel_type=itk.F) + moving_distance_map = itk.imread(str(moving_distance_map_file), pixel_type=itk.F) + fixed_labelmap = itk.imread(str(fixed_labelmap_file)) + moving_labelmap = itk.imread(str(moving_labelmap_file)) + fixed_labels = itk.array_from_image(fixed_labelmap) + + def read_landmarks(landmark_file: Path, image: itk.Image) -> np.ndarray: + """Read a DIR-Lab landmark file as an (N, 3) array of world points. + + Each line holds one 1-based voxel index as ``x y z``. + """ + indices = np.loadtxt(landmark_file, dtype=int) - 1 + return np.array( + [ + image.TransformIndexToPhysicalPoint([int(v) for v in index]) + for index in indices + ] + ) + + fixed_landmarks = read_landmarks(fixed_landmark_file, fixed_image) + moving_landmarks = read_landmarks(moving_landmark_file, moving_image) + + def landmark_metrics(errors_mm: np.ndarray) -> dict[str, Any]: + """Summarize per-landmark target registration errors, in millimeters.""" + return { + "tre_mean": float(errors_mm.mean()), + "tre_std": float(errors_mm.std()), + "tre_p95": float(np.percentile(errors_mm, 95)), + "tre_max": float(errors_mm.max()), + } + + def landmark_errors(transform: itk.Transform) -> np.ndarray: + """Distance from each mapped fixed landmark to its moving counterpart. + + ``forward_transform`` is the resampling transform: it maps points on the + fixed grid back into moving space, which is the direction the landmark + correspondences are defined in. + """ + mapped = np.array( + [transform.TransformPoint(tuple(point)) for point in fixed_landmarks] + ) + return np.asarray(np.linalg.norm(mapped - moving_landmarks, axis=1)) + + def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: + """Per-class Dice summary against the fixed lung labelmap. + + Classes are the union of the two labelmaps' non-zero ids, so a class + found in only one of them scores 0 rather than being dropped. + """ + labels = itk.array_from_image(labelmap) + classes = np.union1d(np.unique(fixed_labels), np.unique(labels)) + classes = classes[classes != 0] + dice = np.array( + [ + 2.0 + * np.count_nonzero((fixed_labels == c) & (labels == c)) + / (np.count_nonzero(fixed_labels == c) + np.count_nonzero(labels == c)) + for c in classes + ] + ) + return { + "n_classes": int(dice.size), + "dice_mean": float(dice.mean()), + "dice_p5": float(np.percentile(dice, 5)), + "dice_median": float(np.median(dice)), + "dice_p95": float(np.percentile(dice, 95)), + "dice_min": float(dice.min()), + "dice_max": float(dice.max()), + "mislabeled_voxels": int(np.count_nonzero(fixed_labels != labels)), + } + + # Reference row: the moving distance map and labelmap on the fixed grid, + # unregistered. + unregistered_distance_map = itk.resample_image_filter( + moving_distance_map, + ReferenceImage=fixed_distance_map, + UseReferenceImage=True, + ) + unregistered_labelmap = itk.resample_image_filter( + moving_labelmap, + Interpolator=itk.NearestNeighborInterpolateImageFunction.New(moving_labelmap), + ReferenceImage=fixed_labelmap, + UseReferenceImage=True, + ) + + registered_distance_maps: dict[str, itk.Image] = { + "unregistered": unregistered_distance_map + } + labelmaps: dict[str, itk.Image] = {"unregistered": unregistered_labelmap} + rows: list[dict[str, Any]] = [ + { + "method": "unregistered", + "input": "-", + "weights": "-", + "registration_time_s": None, + "loss": None, + **landmark_metrics( + np.linalg.norm(fixed_landmarks - moving_landmarks, axis=1) + ), + **overlap_metrics(unregistered_labelmap), + } + ] + # (method, registration input, ICON weights). The last row registers the CT + # images themselves, as the intensity-based reference point for what the + # distance-map methods achieve from surfaces alone. + methods: list[tuple[str, str, Optional[Path]]] = [ + ("greedy_dmap", "distance_map", None), + ("icon_stock_dmap", "distance_map", None), + ("icon_finetuned_dmap", "distance_map", weights_path), + ("greedy_icon_finetuned_dmap", "distance_map", weights_path), + ("icon_stock_ct", "ct", None), + ] + for method_name, method_input, method_weights in methods: + if method_input == "distance_map": + method_fixed, method_moving = fixed_distance_map, moving_distance_map + else: + method_fixed, method_moving = fixed_image, moving_image + + registrar: RegisterImagesBase + if method_name == "greedy_icon_finetuned_dmap": + # Both stages are configured exactly as the standalone "greedy_dmap" + # and "icon_finetuned_dmap" rows above, so this row differs from + # "icon_finetuned_dmap" only by the Greedy transform ICON starts + # from. + chain = RegisterImagesGreedyICON(log_level=log_level) + chain.greedy.set_transform_type("Deformable") + chain.greedy.set_metric("CC") + if number_of_iterations_greedy is not None: + chain.greedy.set_number_of_iterations(number_of_iterations_greedy) + chain.icon.set_number_of_iterations(number_of_iterations_icon) + chain.icon.set_mass_preservation(False) + chain.icon.set_weights_path(str(method_weights)) + registrar = chain + elif method_name.startswith("greedy"): + registrar = RegisterImagesGreedy(log_level=log_level) + registrar.set_transform_type("Deformable") + # CC is what RegisterModelsDistanceMaps uses on distance maps. + registrar.set_metric("CC") + if number_of_iterations_greedy is not None: + registrar.set_number_of_iterations(number_of_iterations_greedy) + else: + registrar = RegisterImagesICON(log_level=log_level) + # None, not 0: icon_registration rejects 0 and takes None to mean + # "no test-time finetuning steps", so the comparison reflects what + # each set of weights predicts rather than per-pair optimization. + registrar.set_number_of_iterations(number_of_iterations_icon) + # Mass preservation models CT density; a distance map carries no + # mass, so it is enabled only on the CT reference row. + registrar.set_mass_preservation(method_input == "ct") + if method_weights is not None: + registrar.set_weights_path(str(method_weights)) + registrar.set_modality("ct") + registrar.set_fixed_image(method_fixed) + + start_time = time.perf_counter() + result = registrar.register(method_moving) + elapsed_s = time.perf_counter() - start_time + + registered_distance_maps[method_name] = transform_tools.transform_image( + moving_distance_map, result["forward_transform"], fixed_distance_map + ) + labelmaps[method_name] = transform_tools.transform_image( + moving_labelmap, + result["forward_transform"], + fixed_labelmap, + interpolation_method="nearest", + ) + rows.append( + { + "method": method_name, + "input": method_input, + "weights": str(method_weights) if method_weights else "-", + "registration_time_s": elapsed_s, + "loss": float(result["loss"]), + **landmark_metrics(landmark_errors(result["forward_transform"])), + **overlap_metrics(labelmaps[method_name]), + } + ) + + # Result saving + itk.imwrite( + fixed_distance_map, str(output_dir / "fixed_distance_map.mha"), compression=True + ) + # Difference against the fixed distance map, not the resampled result: + # residual structure is what distinguishes the methods, and it is invisible + # in the registered distance maps themselves. + fixed_arr = itk.GetArrayFromImage(fixed_distance_map).astype(np.float32) + for method_name, distance_map in registered_distance_maps.items(): + difference = itk.GetImageFromArray( + fixed_arr - itk.GetArrayFromImage(distance_map).astype(np.float32) + ) + difference.CopyInformation(fixed_distance_map) + itk.imwrite( + difference, + str(output_dir / f"difference_distance_map_{method_name}.mha"), + compression=True, + ) + for method_name, labelmap in labelmaps.items(): + itk.imwrite( + labelmap, + str(output_dir / f"labelmap_{method_name}.mha"), + compression=True, + ) + + summary_file = output_dir / "registration_summary.csv" + with summary_file.open("w", newline="", encoding="utf-8") as fh: + writer = csv.DictWriter(fh, fieldnames=list(rows[0].keys())) + writer.writeheader() + writer.writerows(rows) + + # Reporting + reporter.log_info( + "Case1Pack_T00 -> Case1Pack_T50, error at %d expert landmarks, mm", + len(fixed_landmarks), + ) + reporter.log_info( + " %-26s %7s %7s %7s %7s %9s", + "method", + "mean", + "std", + "p95", + "max", + "time_s", + ) + for row in rows: + elapsed = row["registration_time_s"] + reporter.log_info( + " %-26s %7.2f %7.2f %7.2f %7.2f %9s", + row["method"], + row["tre_mean"], + row["tre_std"], + row["tre_p95"], + row["tre_max"], + "-" if elapsed is None else f"{float(elapsed):.1f}", + ) + + reporter.log_info("Per-class Dice of the warped moving labelmap against the fixed") + reporter.log_info( + " %-26s %7s %7s %7s %7s %7s %7s %7s %12s", + "method", + "classes", + "mean", + "p5", + "median", + "p95", + "min", + "max", + "mislabeled", + ) + for row in rows: + reporter.log_info( + " %-26s %7d %7.4f %7.4f %7.4f %7.4f %7.4f %7.4f %12d", + row["method"], + row["n_classes"], + row["dice_mean"], + row["dice_p5"], + row["dice_median"], + row["dice_p95"], + row["dice_min"], + row["dice_max"], + row["mislabeled_voxels"], + ) + reporter.log_info("Wrote summary: %s", summary_file) + + # Testing + tt = TestTools( + class_name=class_name, + results_dir=output_dir, + baselines_dir=baselines_dir, + log_level=log_level, + ) + + screenshots: list[Path] = [ + tt.save_screenshot_image_slice( + fixed_distance_map, + "fixed_distance_map.png", + axis=0, + slice_fraction=0.5, + colormap="gray", + ) + ] + for method_name, distance_map in registered_distance_maps.items(): + screenshots.append( + tt.save_screenshot_image_slice( + distance_map, + f"registered_distance_map_{method_name}.png", + axis=0, + slice_fraction=0.5, + colormap="gray", + ) + ) + + tutorial_results = { + "weights_path": weights_path, + "registration_metrics": rows, + "labelmaps": labelmaps, + "summary_file": summary_file, + "registered_distance_maps": registered_distance_maps, + "screenshots": screenshots, + } diff --git a/tutorials/tutorial_02_lung_finetune_icon.py b/tutorials/tutorial_02_lung_finetune_icon.py index 0d04de6c..89b13e8d 100644 --- a/tutorials/tutorial_02_lung_finetune_icon.py +++ b/tutorials/tutorial_02_lung_finetune_icon.py @@ -4,11 +4,14 @@ Purpose ------- Finetune uniGradICON on every DIR-Lab 4D CT case except Case 1, then register -``Case1Pack_T00.mha`` (moving) to ``Case1Pack_T50.mha`` (fixed) three ways: +``Case1Pack_T00.mha`` (moving) to ``Case1Pack_T50.mha`` (fixed) four ways: ``RegisterImagesGreedy`` alone, deformable, with its default iteration -schedule, and ``RegisterImagesICON`` with the stock uniGradICON weights and -with the finetuned weights. Case 1 is never seen during finetuning, so it is a held-out -evaluation pair. +schedule; ``RegisterImagesICON`` with the stock uniGradICON weights and with +the finetuned weights; and ``RegisterImagesGreedyICON`` -- the same Greedy +stage initializing an ICON stage that carries the finetuned weights -- which +separates what the finetuned network adds on its own from what it adds on top +of a classical affine-plus-deformable initialization. Case 1 is never seen +during finetuning, so it is a held-out evaluation pair. Accuracy is measured two ways. The primary metric is target registration error: DIR-Lab ships 300 expert landmarks for the extreme phases (T00 and T50) @@ -59,6 +62,7 @@ PhysioTwin4DBase, RegisterImagesBase, RegisterImagesGreedy, + RegisterImagesGreedyICON, RegisterImagesICON, SegmentNVSegmentCTMRI, TestTools, @@ -79,6 +83,8 @@ class_name = "tutorial_02_lung_finetune_icon" output_dir = tutorials_dir / "output" / "tutorial_02_lung" + # Segmented labelmaps, cached so re-runs skip the segmentation. + labelmaps_dir = output_dir / "labelmaps" baselines_dir = repo_root / "tests" / "baselines" # The workflow writes its dataset JSON, YAML config, and checkpoint tree @@ -92,20 +98,18 @@ if test_mode: data_dir = repo_root / "data" / "test" / "DirLab-4DCT" number_of_iterations_greedy: Optional[list[int]] = [1, 0] + number_of_iterations_icon = None # [1] epochs = 1 else: data_dir = repo_root / "data" / "DirLab-4DCT" - number_of_iterations_greedy = [60, 30, 20] # Greedy defaults - # 90 training frames at batch_size 4 is 22 optimizer steps per epoch, so - # 100 epochs is ~2200 steps at a 5e-5 learning rate. Far fewer than - # that leaves the finetuned weights statistically indistinguishable - # from the stock weights they started from. + number_of_iterations_greedy = None # [60, 30, 20] + number_of_iterations_icon = None # [10] epochs = 100 log_level = logging.INFO reporter = PhysioTwin4DBase(class_name=class_name, log_level=log_level) - output_dir.mkdir(parents=True, exist_ok=True) + labelmaps_dir.mkdir(parents=True, exist_ok=True) # Held-out evaluation pair (Case 1 is excluded from finetuning). T00 and # T50 are the extreme inhale/exhale phases, the only pair DIR-Lab supplies @@ -185,9 +189,18 @@ weights_path = workflow.process() else: weights_path = ( - Path(__file__).resolve().parent - / "network_weights/icon_dirlab_4dct/icon_dirlab_4dct_model/checkpoints/network_weights_final.trch" + experiment_dir + / f"{finetune_name}_model" + / "checkpoints" + / "network_weights_final.trch" ) + # Checked here rather than at the first set_weights_path() call, which + # only happens after the greedy and stock-ICON rows have already run. + if not weights_path.exists(): + raise FileNotFoundError( + f"run_finetuning is False but no checkpoint at {weights_path}. " + "Set run_finetuning = True to finetune from scratch." + ) # Registration comparison fixed_image = itk.imread(str(fixed_file), pixel_type=itk.F) @@ -235,8 +248,24 @@ def landmark_errors(transform: itk.Transform) -> np.ndarray: # transform, so Dice reflects the transform rather than what the segmenter # does differently on each interpolated volume. segmenter = SegmentNVSegmentCTMRI(log_level=log_level) - fixed_labelmap = segmenter.segment(fixed_image)["labelmap"] - moving_labelmap = segmenter.segment(moving_image)["labelmap"] + + def segment_phase(image_file: Path, image: itk.Image) -> itk.Image: + """Segment one phase, caching the labelmap under ``labelmaps_dir``. + + An existing labelmap short-circuits the segmentation, which dominates + this tutorial's runtime outside of finetuning. + """ + labelmap_file = labelmaps_dir / f"{image_file.stem}_labelmap.mha" + if labelmap_file.exists(): + reporter.log_info("Reusing cached labelmap: %s", labelmap_file.name) + return itk.imread(str(labelmap_file)) + + labelmap = segmenter.segment(image)["labelmap"] + itk.imwrite(labelmap, str(labelmap_file), compression=True) + return labelmap + + fixed_labelmap = segment_phase(fixed_file, fixed_image) + moving_labelmap = segment_phase(moving_file, moving_image) fixed_labels = itk.array_from_image(fixed_labelmap) def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: @@ -299,6 +328,7 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: ("greedy", None), ("icon_stock", None), ("icon_finetuned", weights_path), + ("greedy_icon_finetuned", weights_path), ): registrar: RegisterImagesBase if method_name == "greedy": @@ -306,12 +336,24 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: registrar.set_transform_type("Deformable") if number_of_iterations_greedy is not None: registrar.set_number_of_iterations(number_of_iterations_greedy) + elif method_name == "greedy_icon_finetuned": + # Both stages are configured exactly as the standalone "greedy" and + # "icon_finetuned" rows above, so this row differs from + # "icon_finetuned" only by the Greedy transform ICON starts from. + chain = RegisterImagesGreedyICON(log_level=log_level) + chain.greedy.set_transform_type("Deformable") + if number_of_iterations_greedy is not None: + chain.greedy.set_number_of_iterations(number_of_iterations_greedy) + chain.icon.set_number_of_iterations(number_of_iterations_icon) + chain.icon.set_mass_preservation(True) # For non-contrast CT + chain.icon.set_weights_path(str(method_weights)) + registrar = chain else: registrar = RegisterImagesICON(log_level=log_level) # None, not 0: icon_registration rejects 0 and takes None to mean # "no test-time finetuning steps", so the comparison reflects what # each set of weights predicts rather than per-pair optimization. - registrar.set_number_of_iterations(None) + registrar.set_number_of_iterations(number_of_iterations_icon) registrar.set_mass_preservation(True) # For non-contrast CT if method_weights is not None: registrar.set_weights_path(str(method_weights)) @@ -346,10 +388,18 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: itk.imwrite( fixed_labelmap, str(output_dir / "fixed_labelmap.mha"), compression=True ) + # Difference against the fixed image, not the resampled result: residual + # structure is what distinguishes the methods, and it is invisible in the + # registered images themselves. + fixed_arr = itk.GetArrayFromImage(fixed_image).astype(np.float32) for method_name, image in registered_images.items(): + difference = itk.GetImageFromArray( + fixed_arr - itk.GetArrayFromImage(image).astype(np.float32) + ) + difference.CopyInformation(fixed_image) itk.imwrite( - image, - str(output_dir / f"registered_{method_name}.mha"), + difference, + str(output_dir / f"difference_{method_name}.mha"), compression=True, ) for method_name, labelmap in labelmaps.items(): @@ -371,12 +421,12 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: len(fixed_landmarks), ) reporter.log_info( - " %-13s %7s %7s %7s %7s %9s", "method", "mean", "std", "p95", "max", "time_s" + " %-21s %7s %7s %7s %7s %9s", "method", "mean", "std", "p95", "max", "time_s" ) for row in rows: elapsed = row["registration_time_s"] reporter.log_info( - " %-13s %7.2f %7.2f %7.2f %7.2f %9s", + " %-21s %7.2f %7.2f %7.2f %7.2f %9s", row["method"], row["tre_mean"], row["tre_std"], @@ -387,7 +437,7 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: reporter.log_info("Per-class Dice of the warped moving labelmap against the fixed") reporter.log_info( - " %-13s %7s %7s %7s %7s %7s %7s %7s %12s", + " %-21s %7s %7s %7s %7s %7s %7s %7s %12s", "method", "classes", "mean", @@ -400,7 +450,7 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: ) for row in rows: reporter.log_info( - " %-13s %7d %7.4f %7.4f %7.4f %7.4f %7.4f %7.4f %12d", + " %-21s %7d %7.4f %7.4f %7.4f %7.4f %7.4f %7.4f %12d", row["method"], row["n_classes"], row["dice_mean"], diff --git a/tutorials/tutorial_04_heart_ct_to_vtk.py b/tutorials/tutorial_04_heart_ct_to_vtk.py index 069f0d9d..00f5cec9 100644 --- a/tutorials/tutorial_04_heart_ct_to_vtk.py +++ b/tutorials/tutorial_04_heart_ct_to_vtk.py @@ -110,9 +110,18 @@ ) # Result saving + # + # Merging the per-structure surfaces, rather than the per-group ones, lets + # the combined file carry a per-cell SegmentationLabelIds array: structure + # identity survives the merge, so the file can still be split per structure + # downstream. Per-group surfaces are contoured from a merged binary mask + # and have no per-cell identity to record. + combined_input = ( + result["label_surfaces"] if save_label_surfaces else result["surfaces"] + ) surface_file = Path( ContourTools.save_combined_surfaces( - result["surfaces"], + combined_input, str(output_dir / "patient_surfaces.vtp"), ) ) diff --git a/tutorials/tutorial_04_lung_ct_to_vtk.py b/tutorials/tutorial_04_lung_ct_to_vtk.py index 5b3260c5..be1738c6 100644 --- a/tutorials/tutorial_04_lung_ct_to_vtk.py +++ b/tutorials/tutorial_04_lung_ct_to_vtk.py @@ -94,9 +94,18 @@ ) # Result saving + # + # Merging the per-structure surfaces, rather than the per-group ones, lets + # the combined file carry a per-cell SegmentationLabelIds array: structure + # identity survives the merge, so the file can still be split per structure + # downstream. Per-group surfaces are contoured from a merged binary mask + # and have no per-cell identity to record. + combined_input = ( + result["label_surfaces"] if save_label_surfaces else result["surfaces"] + ) surface_file = Path( ContourTools.save_combined_surfaces( - result["surfaces"], + combined_input, str(output_dir / "patient_surfaces.vtp"), ) ) diff --git a/tutorials/tutorial_05_heart_vtk_to_usd.py b/tutorials/tutorial_05_heart_vtk_to_usd.py index 4f5d0fc4..327d42ff 100644 --- a/tutorials/tutorial_05_heart_vtk_to_usd.py +++ b/tutorials/tutorial_05_heart_vtk_to_usd.py @@ -6,9 +6,16 @@ Convert the VTK surface output from Tutorial 4, or another VTK-compatible mesh, into a USD file with anatomy materials. +Tutorial 4 writes one VTP per anatomical structure, each annotated with its +structure name. Feeding those files in individually — rather than the single +combined surface — keeps that name attached through the conversion, so each +structure becomes its own named USD prim and gets its own material: bright +oxygenated red for the left chambers, darker deoxygenated red for the right, +red-brown myocardium, and so on, instead of one uniform heart material. + Data Required ------------- -Preferred input: ``tutorials/output/tutorial_04_heart/patient_surfaces.vtp`` +Preferred input: ``tutorials/output/tutorial_04_heart/patient_*.vtp`` """ # Imports @@ -43,8 +50,8 @@ project_name = "tutorial_04_heart" - # Preferred input: the combined surface saved by Tutorial 4. - vtk_file = tutorials_dir / "output" / "tutorial_04_heart" / "patient_surfaces.vtp" + # Preferred input: the per-structure surfaces saved by Tutorial 4. + input_dir = tutorials_dir / "output" / "tutorial_04_heart" log_level = logging.INFO @@ -52,16 +59,42 @@ output_dir.mkdir(parents=True, exist_ok=True) - mesh = pv.read(str(vtk_file)) + # Tutorial 4 writes both per-group and per-structure surfaces into one + # directory. A per-structure surface is the one carrying exactly one name + # in SegmentationLabelNames, and its filename ends with that name; the + # per-group surfaces list every structure in the group and are skipped, as + # their geometry would otherwise be exported twice. + meshes: list[pv.PolyData] = [] + structure_names: list[str] = [] + for vtk_file in sorted(input_dir.glob("patient_*.vtp")): + mesh = pv.read(str(vtk_file)) + label_names = mesh.field_data.get("SegmentationLabelNames") + if label_names is None or len(label_names) != 1: + continue + if not vtk_file.stem.endswith(str(label_names[0])): + continue + meshes.append(mesh) + structure_names.append(str(label_names[0])) + + if not meshes: + raise FileNotFoundError( + "No per-structure surfaces found. Checked:\n" + + f" - {input_dir}/patient_*.vtp\n" + + "Run tutorial_04_heart_ct_to_vtk.py with save_label_surfaces=True." + ) # Workflow initialization - + # + # static_merge=True treats the meshes as separate objects in one scene + # rather than as frames of a time series. Leaving anatomy_type unset lets + # each object's name select its material, and object names default to the + # structure names read from field_data above. workflow = WorkflowConvertVTKToUSD( - input_meshes=[mesh], + input_meshes=meshes, usd_project_name=project_name, output_directory=output_dir, appearance="anatomy", - anatomy_type="heart", + static_merge=True, separate_by_connectivity=True, log_level=log_level, ) @@ -84,4 +117,8 @@ ) ] - tutorial_results = {"usd_file": results["usd_file"], "screenshots": screenshots} + tutorial_results = { + "usd_file": results["usd_file"], + "structures": structure_names, + "screenshots": screenshots, + } diff --git a/tutorials/tutorial_06_heart_create_statistical_model.py b/tutorials/tutorial_06_heart_create_statistical_model.py index cf470bbc..b62a8499 100644 --- a/tutorials/tutorial_06_heart_create_statistical_model.py +++ b/tutorials/tutorial_06_heart_create_statistical_model.py @@ -48,10 +48,10 @@ test_mode = TestTools.running_as_test() if test_mode: data_dir = repo_root / "data" / "test" / "KCL-Heart-Model" - pca_components = 5 + number_of_pca_components = 5 else: data_dir = repo_root / "data" / "KCL-Heart-Model" - pca_components = 10 + number_of_pca_components = 10 log_level = logging.INFO @@ -87,7 +87,7 @@ workflow = WorkflowCreateStatisticalModel( sample_meshes=sample_meshes, reference_mesh=reference_mesh, - pca_number_of_components=pca_components, + number_of_pca_components=number_of_pca_components, log_level=log_level, ) @@ -127,7 +127,7 @@ components = pca_model.get("components", []) eigenvalues = pca_model.get("eigenvalues", []) mean_points = np.asarray(mean_surface.points) - mode_count = min(2, pca_components, len(components), len(eigenvalues)) + mode_count = min(2, number_of_pca_components, len(components), len(eigenvalues)) xvfb_started = False try: diff --git a/tutorials/tutorial_06_lung_create_statistical_model.py b/tutorials/tutorial_06_lung_create_statistical_model.py index 02fcf5df..8c475e6b 100644 --- a/tutorials/tutorial_06_lung_create_statistical_model.py +++ b/tutorials/tutorial_06_lung_create_statistical_model.py @@ -67,7 +67,9 @@ baselines_dir = repo_root / "tests" / "baselines" data_dir = repo_root / "data" / "DirLab-4DCT" - pca_components = 7 + + number_of_pca_components = 5 + # Atlas iterations used to build the reference surface; 1 is a single # template-biased pass. mean_surface_iterations = 3 @@ -126,7 +128,7 @@ workflow = WorkflowCreateStatisticalModel( sample_meshes=sample_surfaces, reference_mesh=reference_surface, - pca_number_of_components=pca_components, + number_of_pca_components=number_of_pca_components, log_level=log_level, ) @@ -165,7 +167,9 @@ components = pca_model.get("components", []) eigenvalues = pca_model.get("eigenvalues", []) mean_points = np.asarray(mean_surface.points) - mode_count = min(2, pca_components, len(components), len(eigenvalues)) + # PCA rank is capped by the sample count, so the model can hold fewer + # components than requested. + mode_count = min(number_of_pca_components, len(components), len(eigenvalues)) mode_surface_files: list[Path] = [] xvfb_started = False diff --git a/tutorials/tutorial_07_lung_fit_statistical_model_to_patient.py b/tutorials/tutorial_07_lung_fit_statistical_model_to_patient.py index fac26c14..20614dd5 100644 --- a/tutorials/tutorial_07_lung_fit_statistical_model_to_patient.py +++ b/tutorials/tutorial_07_lung_fit_statistical_model_to_patient.py @@ -14,6 +14,10 @@ Patient image: a routine clinical 3D chest CT, ``data/Chest-CT/Chest-CT.mha``, downloaded with ``physiotwin4d-download-data Chest-CT --directory data/Chest-CT`` +ICON weights: ``tutorial_02_lung_distancemap_finetune_icon.py`` output +(``network_weights/icon_dirlab_4dct_distancemap/ +icon_dirlab_4dct_distancemap_model/checkpoints/network_weights_final.trch``), +optional -- the stock uniGradICON weights are used when it is absent. """ # Imports @@ -57,8 +61,23 @@ pca_json = tutorial_06_dir / "pca_model.json" pca_mean_file = tutorial_06_dir / "pca_mean_surface.vtp" + number_of_pca_components = 5 + patient_image_file = repo_root / "data" / "Chest-CT" / "Chest-CT.mha" + # Distance-map weights finetuned on DIR-Lab by Tutorial 2; see + # WorkflowFinetuneICONRegistration.expected_weights_path(). The + # mask_dilation_mm set below must match the one that tutorial finetuned + # with, since it fixes the distance maps' saturation radius. + icon_weights_path = ( + tutorials_dir + / "network_weights" + / "icon_dirlab_4dct_distancemap" + / "icon_dirlab_4dct_distancemap_model" + / "checkpoints" + / "network_weights_final.trch" + ) + log_level = logging.INFO # The same segmenter and surface-extraction workflow used by Tutorial 6, so @@ -122,10 +141,26 @@ workflow.set_use_pca_registration( use_pca_registration=True, pca_model=pca_model, - pca_number_of_modes=6, + number_of_pca_components=number_of_pca_components, use_surface=False, ) + workflow.set_mask_dilation_mm(mask_dilation_mm=40) + + # The labelmap-to-labelmap stage registers distance maps, not intensities, + # so it uses the distance-map-finetuned weights when they exist; without + # them the tutorial still runs, on the stock uniGradICON weights. + if icon_weights_path.exists(): + workflow.set_labelmap_to_labelmap_icon_weights_path(str(icon_weights_path)) + else: + workflow.log_warning( + "Finetuned distance-map ICON weights not found at %s; fitting with " + "the stock uniGradICON weights. Run " + "tutorials/tutorial_02_lung_distancemap_finetune_icon.py to create " + "them.", + icon_weights_path, + ) + # Workflow execution workflow_results = workflow.process() @@ -152,6 +187,14 @@ str(output_dir / f"{project_name}_template_surface_registered.vtp") ) + registered_pca_surface = workflow.pca_template_model_surface + assert registered_pca_surface is not None, ( + "pca_template_model_surface must be set after process()" + ) + registered_pca_surface.save( + str(output_dir / f"{project_name}_pca_surface_registered.vtp") + ) + # Testing TestTools( class_name=project_name, diff --git a/tutorials/tutorial_08_lung_fit_model_to_4d_patients.py b/tutorials/tutorial_08_lung_fit_model_to_4d_patients.py index 176f4288..f40ec176 100644 --- a/tutorials/tutorial_08_lung_fit_model_to_4d_patients.py +++ b/tutorials/tutorial_08_lung_fit_model_to_4d_patients.py @@ -32,8 +32,11 @@ ``pca_mean_surface.vtp``) ICON weights: Tutorial 2 output (``network_weights/icon_dirlab_4dct/icon_dirlab_4dct_model/checkpoints/ -network_weights_final.trch``), optional — the stock uniGradICON weights are used -when it is absent. +network_weights_final.trch``) for the phase-to-phase CT registration, and +``network_weights/icon_dirlab_4dct_distancemap/ +icon_dirlab_4dct_distancemap_model/checkpoints/network_weights_final.trch`` for +the distance-map stage of the SSM fit. Both optional — the stock uniGradICON +weights are used when they are absent. Outputs (per case, under ``output/tutorial_08_lung//``) ------------------------------------------------------------ @@ -101,6 +104,21 @@ / "network_weights_final.trch" ) + # Distance-map weights finetuned on DIR-Lab by + # tutorial_02_lung_distancemap_finetune_icon.py, used by the + # labelmap-to-labelmap stage of the SSM fit. mask_dilation_mm below must + # match the value that tutorial finetuned with, since it fixes the distance + # maps' saturation radius. + icon_distancemap_weights_path = ( + tutorials_dir + / "network_weights" + / "icon_dirlab_4dct_distancemap" + / "icon_dirlab_4dct_distancemap_model" + / "checkpoints" + / "network_weights_final.trch" + ) + fit_mask_dilation_mm = 40.0 + # Phase the SSM is fitted to; Tutorial 6 builds the lung PCA model from the # T70 surfaces, so the same phase is used here as the fitting reference. reference_phase = "T70" @@ -137,6 +155,21 @@ icon_weights_path, ) + use_finetuned_distancemap_weights = icon_distancemap_weights_path.exists() + if use_finetuned_distancemap_weights: + logger.info( + "Fitting the SSM with finetuned distance-map ICON weights: %s", + icon_distancemap_weights_path, + ) + else: + logger.warning( + "Finetuned distance-map ICON weights not found at %s; fitting the SSM " + "with the stock uniGradICON weights. Run " + "tutorials/tutorial_02_lung_distancemap_finetune_icon.py to create " + "them.", + icon_distancemap_weights_path, + ) + reference_files = sorted(data_dir.glob(f"Case*_{reference_phase}.mha")) if not reference_files: raise FileNotFoundError( @@ -200,9 +233,14 @@ fit_workflow.set_use_pca_registration( use_pca_registration=True, pca_model=pca_model, - pca_number_of_modes=6, + number_of_pca_components=6, use_surface=False, ) + fit_workflow.set_mask_dilation_mm(mask_dilation_mm=fit_mask_dilation_mm) + if use_finetuned_distancemap_weights: + fit_workflow.set_labelmap_to_labelmap_icon_weights_path( + str(icon_distancemap_weights_path) + ) fit_result = fit_workflow.process() pca_coefficients_file = case_output_dir / f"{case_id}_ssm_pca_coefficients.json" diff --git a/tutorials/tutorial_09_lung_train_physicsnemo_mgn.py b/tutorials/tutorial_09_lung_train_physicsnemo_mgn.py index b5bc566d..24a14252 100644 --- a/tutorials/tutorial_09_lung_train_physicsnemo_mgn.py +++ b/tutorials/tutorial_09_lung_train_physicsnemo_mgn.py @@ -15,8 +15,10 @@ filenames and written explicitly into the manifest (the workflow never parses filenames). -2. Split the cases into train / validation / held-out test and train the - MeshGraphNet (``WorkflowTrainPhysicsNeMo`` driving ``TrainPhysicsNeMoMGN``). +2. Split the cases into train and held-out test — plus an optional validation + set, empty by default, which is what makes the intermittent validation RMSE + read ``n/a`` — and train the MeshGraphNet (``WorkflowTrainPhysicsNeMo`` + driving ``TrainPhysicsNeMoMGN``). 3. Evaluate the held-out test cases against their ground-truth phases with :class:`physiotwin4d.WorkflowInferPhysicsNeMo` wrapped in @@ -33,6 +35,15 @@ Edge features (per edge): [rel_x, rel_y, rel_z, distance] (from the mean shape) Output (per vertex): [dx, dy, dz] (displacement in mm) +Runtime +------- +Measured on the full 10-case DIR-Lab set with the Tutorial 6 lung template +(179k points, 1.07M mesh-graph edges): one training step of ``batch_size`` 4 +takes ~430 ms and peaks near 43 GiB of GPU memory, giving ~9 s per epoch and +roughly 4 hours for the 1500 epochs below. Lower ``batch_size``, or call +``training_method.set_num_processor_checkpoint_segments(...)`` to trade compute +for memory, on a smaller card. + Extra Install Required ---------------------- PhysicsNeMo and PyTorch Geometric must be installed:: @@ -111,17 +122,29 @@ def _write_target_mesh( return target_path -def _write_case_manifest(case_dir: Path, manifests_dir: Path) -> Optional[Path]: +def _write_case_manifest( + case_dir: Path, manifests_dir: Path, logger: logging.Logger +) -> Optional[Path]: """Write a per-case manifest JSON; return its path (or None if incomplete). A case needs a reference SSM surface, a PCA coefficient file, and at least - two respiratory-phase surfaces. + two respiratory-phase surfaces. A case that is missing any of them is skipped + with the reason logged, so a half-finished Tutorial 8 run is distinguishable + from one that never ran. """ case_id = case_dir.name ref_file = case_dir / f"{case_id}_ssm_surface.vtp" pca_file = case_dir / f"{case_id}_ssm_pca_coefficients.json" phase_files = sorted(case_dir.glob(f"{case_id}_T??_ssm_surface.vtp")) - if not ref_file.exists() or not pca_file.exists() or len(phase_files) < 2: + missing = [] + if not ref_file.exists(): + missing.append(f"reference surface {ref_file.name}") + if not pca_file.exists(): + missing.append(f"PCA coefficients {pca_file.name}") + if len(phase_files) < 2: + missing.append(f"at least 2 phase surfaces (found {len(phase_files)})") + if missing: + logger.warning("Skipping %s: missing %s", case_id, "; ".join(missing)) return None manifests_dir.mkdir(parents=True, exist_ok=True) @@ -178,7 +201,9 @@ def _write_case_manifest(case_dir: Path, manifests_dir: Path) -> Optional[Path]: num_layers = 2 # MLP layers inside each encoder / processor / decoder block # Explicit held-out splits; every other discovered case is used for training. - # Case1 is also the case held out of the Tutorial 2 ICON finetuning. + # Case1 is also the case held out of the Tutorial 2 ICON finetuning. Adding a + # case to val_cases spends it on the intermittent validation RMSE instead of + # training; empty means that RMSE is reported as "n/a". test_cases = ["Case1Pack"] val_cases: list[str] = [] log_level = logging.INFO @@ -198,17 +223,20 @@ def _write_case_manifest(case_dir: Path, manifests_dir: Path) -> Optional[Path]: "Run tutorials/tutorial_06_lung_create_statistical_model.py first." ) - # Step 1: build one manifest per valid case and partition into splits + # Step 1: build one manifest per valid case and partition into splits. + # DIR-Lab names case 8 "Case8Deploy" while every other case is "Case*Pack", + # so match on "Case*" to avoid silently dropping it. manifests: dict[str, Path] = {} - for case_dir in sorted(data_dir.glob("Case*Pack")): - manifest_path = _write_case_manifest(case_dir, manifests_dir) + for case_dir in sorted(p for p in data_dir.glob("Case*") if p.is_dir()): + manifest_path = _write_case_manifest(case_dir, manifests_dir, logger) if manifest_path is not None: manifests[case_dir.name] = manifest_path if len(manifests) < 3: raise RuntimeError( f"Found only {len(manifests)} valid case(s) under {data_dir}; need at " - "least 3 for a train / val / test split. Run " + "least 3 to hold one out and still train on a population. See the " + "skip reasons logged above, and run " "tutorials/tutorial_08_lung_fit_model_to_4d_patients.py first." )