|
| 1 | +# TPTBox – Registration subsystem notes |
| 2 | + |
| 3 | +This file complements `CLAUDE.md` with a targeted overview of the |
| 4 | +`TPTBox/registration/` sub-package. It is meant as a working memory for anyone |
| 5 | +extending or debugging the registration code paths, plus a running log of the |
| 6 | +changes that were made in the "Point-Reg" refactor branch. |
| 7 | + |
| 8 | +Update this file *incrementally* whenever you touch the registration code – add |
| 9 | +sections when you introduce new classes, and update the status table when you |
| 10 | +add or finish tasks. |
| 11 | + |
| 12 | +## Layout |
| 13 | + |
| 14 | +``` |
| 15 | +TPTBox/registration/ |
| 16 | +├── __init__.py # Aggregates public API. Optional imports guarded. |
| 17 | +├── script_ax2sag.py # CLI helper (unchanged). |
| 18 | +├── _ridged_points/ |
| 19 | +│ ├── point_registration.py # SITK closed-form rigid (VersorRigid3D) landmark fit. |
| 20 | +│ └── deepali_point_registration.py # NEW: DeepALI equivalent of the above. |
| 21 | +├── _ridged_intensity/ |
| 22 | +│ └── affine_deepali.py # Rigid intensity-based registration used by Rigid_Elements. |
| 23 | +├── _deepali/ |
| 24 | +│ ├── deepali_model.py # General_Registration wrapper around DeepaliPairwiseImageTrainer. |
| 25 | +│ ├── deepali_trainer.py # Multi-resolution pyramid training loop. |
| 26 | +│ └── spine_rigid_elements_reg.py # Per-vertebra rigid registration + weighted blending. |
| 27 | +└── _deformable/ |
| 28 | + ├── deformable_reg.py # Wraps General_Registration for BSpline / SVFFD. |
| 29 | + └── multilabel_segmentation.py # Template_Registration + Template_Registration2 (NEW). |
| 30 | +``` |
| 31 | + |
| 32 | +## Public entry points |
| 33 | + |
| 34 | +| Class | File | Purpose | |
| 35 | +| ------------------------------ | ----------------------------------------------------------------- | -------------------------------------------------------------------------------- | |
| 36 | +| `Point_Registration` | `_ridged_points/point_registration.py` | SITK `VersorRigid3D` landmark-based rigid registration. Serialisable. | |
| 37 | +| `Deepali_Point_Registration` | `_ridged_points/deepali_point_registration.py` **(NEW)** | Kabsch/Horn SVD fit on paired POI landmarks, wrapped as a DeepALI `HomogeneousTransform`. | |
| 38 | +| `General_Registration` | `_deepali/deepali_model.py` | Generic DeepALI pairwise image registration (rigid / affine / SVFFD / …). | |
| 39 | +| `Deformable_Registration` | `_deformable/deformable_reg.py` | Thin wrapper enforcing a non-rigid transform on `General_Registration`. | |
| 40 | +| `Template_Registration` | `_deformable/multilabel_segmentation.py` | Two-stage rigid-then-deformable atlas → target alignment (POI-based rigid). | |
| 41 | +| `Template_Registration2` | `_deformable/multilabel_segmentation.py` **(NEW)** | Same idea, but accepts a `Deepali_Point_Registration` as pre-registration to skip the SITK resample. | |
| 42 | +| `Rigid_Elements_Registration` | `_deepali/spine_rigid_elements_reg.py` | Per-vertebra rigid registration + inverse-distance blending field. | |
| 43 | + |
| 44 | +## Key invariants |
| 45 | + |
| 46 | +* **Coordinate conventions.** POIs and NIIs store voxel coords. `local_to_global(x, itk=True)` converts to LPS (ITK / DeepALI) world coords. `local_to_global(x)` (default) yields RAS (NIfTI) world coords. Never mix conventions. |
| 47 | +* **Rigid transform direction.** For resampling (SITK `Resample` **and** DeepALI `TransformImage`) the *forward* direction of the transform is `fixed → moving`. Landmark fits usually estimate `moving → fixed`; invert once and stick with fixed→moving thereafter. |
| 48 | +* **DeepALI transform tensor space.** `HomogeneousTransform.tensor()` returns a `(N, D, D+1)` matrix expressed in **target-grid cube coordinates** (`Axes.CUBE_CORNERS` if `align_corners=True`). When the fit is done in LPS world coords, convert via `M = A^-1 @ W @ A` with `A = target.transform(CUBE_CORNERS, WORLD)`. The moving-grid conversion is done by `SampleImage` at sampling time and must not be baked in – see `_build_deepali_transform` in `deepali_point_registration.py`. |
| 49 | +* **`SampleImage(target, source)` vs `TransformImage(target, source)`.** The existing `_warp_image` in `deepali_model.py` uses `source=target_grid`, which silently *requires* the moving image to already live on the fixed grid. This is where the "same-space" assumption of `General_Registration` comes from. |
| 50 | + |
| 51 | +## In-flight changes ("Point-Reg" branch) |
| 52 | + |
| 53 | +Status legend: 🟡 in progress · ✅ done · ⬜ pending |
| 54 | + |
| 55 | +| # | Task | Status | |
| 56 | +| - | ---- | ------ | |
| 57 | +| 1 | New `Deepali_Point_Registration` (closed-form rigid via DeepALI) | ✅ | |
| 58 | +| 2 | `General_Registration` accepts fixed/moving on different grids (new `same_space=True` flag) | ✅ | |
| 59 | +| 3 | `General_Registration` accepts `POI` / `POI_Global` landmark sets and matches shared IDs automatically | ✅ | |
| 60 | +| 4 | New `Template_Registration2` that consumes a `Deepali_Point_Registration` as pre-registration | ✅ | |
| 61 | +| 5 | Unit tests + speed / memory sanity checks | ✅ | |
| 62 | +| 6 | This file + `CLAUDE.md` cross-reference | ✅ | |
| 63 | +| 7 | Speed benchmark Deepali vs SimpleITK on CPU (recorded below) | ✅ | |
| 64 | + |
| 65 | +## Testing / running |
| 66 | + |
| 67 | +* Conda env: `/home/robert/anaconda3/envs/py3.12/bin/python` – DeepALI (`hf-deepali`) is installed there. |
| 68 | +* Sample data: |
| 69 | + * `TPTBox/tests/sample_ct/` and `TPTBox/tests/sample_mri/` – tiny NIfTIs + segmentations, checked into the repo. |
| 70 | + * `tutorials/tutorial_data_processing/` – DICOM + PixelPandemonium MR pair, downloaded by the tutorial. |
| 71 | +* Existing unit tests live in `unit_tests/`. New registration tests should follow the same pattern (no GPU-only paths in default tests; guard CUDA imports). |
| 72 | +* `test_deformable_stage_improves_over_rigid_only` needs `elasticdeform`. **Known issue: the PyPI wheel is built against NumPy 1.x and fails at import under NumPy 2.x.** Install the source tarball instead so it recompiles locally: |
| 73 | + |
| 74 | + ``` |
| 75 | + pip install https://github.com/gvtulder/elasticdeform/archive/refs/tags/v0.5.1.tar.gz |
| 76 | + ``` |
| 77 | + |
| 78 | + Confirmed working with NumPy 2.4.1 / SciPy 1.17.0. If `elasticdeform` is unavailable the test is skipped cleanly (`_HAS_ELASTIC = False`). |
| 79 | + |
| 80 | +## Benchmark: Deepali vs SimpleITK on CPU |
| 81 | + |
| 82 | +Two benchmarks were run in `py3.12`: |
| 83 | + |
| 84 | +**Tiny CT (73×47×73, 3 landmarks):** |
| 85 | + |
| 86 | +| step | SimpleITK | Deepali | notes | |
| 87 | +| ----- | --------- | --------- | ----- | |
| 88 | +| fit | ~7.3 ms | ~1.4 ms | Kabsch SVD is trivially cheap | |
| 89 | +| warp | ~9.6 ms | ~19.7 ms | SITK BSplineResampler is fast on this size | |
| 90 | +| accuracy (mean-abs voxel err on identity round-trip)| **54.2** HU | **0.003** HU | SITK BSpline shows heavy ringing on CT | |
| 91 | +| peak mem | – | 0.5 MiB (`tracemalloc`) | | |
| 92 | + |
| 93 | +**Tutorial MR volume (270×220×72, 6 landmarks):** |
| 94 | + |
| 95 | +| step | SimpleITK | Deepali | notes | |
| 96 | +| ----- | --------- | --------- | ----- | |
| 97 | +| fit | ~87 ms | ~1.4 ms | ~60× faster – closed-form SVD stays constant with landmark count | |
| 98 | +| warp | ~131 ms | ~243 ms | SITK still edges out on CPU for pure resampling | |
| 99 | +| accuracy | 6.53 | **0.0002** | Deepali linear sampler is drastically more accurate | |
| 100 | +| peak mem | 32.6 MiB | 32.6 MiB | Same order of magnitude | |
| 101 | + |
| 102 | +Short answer to the user's mid-run question: **Deepali is meaningfully more accurate and its fit is ~5–60× faster on CPU, but SITK is still ~1.5–2× faster on the actual image warp step on CPU**. Deepali becomes clearly faster once a GPU is used or once the warp is followed by more DeepALI work (no extra CPU→GPU copies). See `unit_tests/test_registration_deepali.py::TestSpeedAndMemory` – it prints a summary line every run. |
| 103 | + |
| 104 | +## Optional-dependency handling |
| 105 | + |
| 106 | +`hf-deepali` (and its prerequisite PyTorch) is an *optional* install. |
| 107 | +`TPTBox.registration/__init__.py` therefore imports each deepali-backed entry |
| 108 | +point inside its own `try/except ImportError`. On failure the name is replaced |
| 109 | +by a small class/function stub built by `_make_missing_deepali_stub` / |
| 110 | +`_make_missing_deepali_func` – instantiating or calling the stub raises |
| 111 | + |
| 112 | + ImportError: `<Name>` requires the optional dependency `hf-deepali` |
| 113 | + (which in turn requires PyTorch). Install both with: |
| 114 | + pip install torch hf-deepali |
| 115 | + |
| 116 | +so users get an actionable message instead of a bare `NameError`. The SITK |
| 117 | +`Point_Registration` path stays fully usable when deepali is absent. See the |
| 118 | +`TestOptionalDeepaliStubs` test for a regression guard. |
| 119 | + |
| 120 | +## Design notes / gotchas |
| 121 | + |
| 122 | +* The current `_warp_image` uses `TransformImage(target=target_grid, source=target_grid)`. Passing a source image that lives on the moving grid is only safe when moving grid == fixed grid. Task #2 addresses this properly by threading the moving grid through when `same_space=False`. |
| 123 | +* When adding DeepALI landmark loss (`LandmarkPointDistance`), the trainer expects target/source landmarks as `(N, M, D)` tensors in `Axes.CUBE_CORNERS` on the transform's grid. The wrapper in `General_Registration` should convert from POI voxel coords to that space. |
| 124 | +* `Template_Registration` mutates its inputs by resampling the atlas after each SITK point-reg attempt. Template_Registration2 avoids the resample by keeping the transform composable with the downstream `Deformable_Registration`. |
0 commit comments