diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 8ffbd74..91343a6 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -51,3 +51,53 @@ jobs: uses: codecov/codecov-action@v4 with: token: ${{ secrets.CODECOV_TOKEN }} + flags: light + + build-full-coverage: + # The matrix above installs only the dev group, so every module behind torch / + # nnunetv2 / hf-deepali is a hard zero there however well it is tested. This job + # installs those on CPU so that coverage actually reaches Codecov. One job, one + # python version - the matrix stays the fast compatibility gate. + runs-on: ubuntu-latest + timeout-minutes: 45 + steps: + - uses: actions/checkout@v4 + - name: Set up Python 3.11 + uses: actions/setup-python@v5 + with: + python-version: "3.11" + cache: "pip" + - name: Configure python + run: | + python -m pip install --upgrade pip + python -m pip install poetry + - name: Install system deps for CPU-only mesh3D tests + run: | + sudo apt-get update + sudo apt-get install -y xvfb libgl1 + - name: Install dependencies + run: | + python -m poetry install --with dev + - name: Install heavy optional deps (CPU only) + # The default torch wheel pulls ~2.5 GB of CUDA libraries that this job cannot + # use. `hf-deepali` is the correct distribution name - a bare `deepali` on PyPI + # is an unrelated project. `spineps` is deliberately absent: it depends on + # TPTBox itself, so installing it would pin a *released* TPTBox over this + # checkout. The tests that touch it fake the module instead. + run: | + python -m poetry run pip install torch --index-url https://download.pytorch.org/whl/cpu + python -m poetry run pip install antspyx nnunetv2 acvl-utils batchgenerators hf-deepali + - name: Test with pytest and create coverage report + run: | + python -m poetry run coverage run --source=TPTBox -m pytest unit_tests/ TPTBox/tests/ + python -m poetry run coverage xml + echo "--- with the curated omit list (the reported number) ---" + python -m poetry run coverage report + echo "--- without it, for comparison ---" + python -m poetry run coverage report --omit="TPTBox/__no_such_path__/*" + - name: Upload coverage results to Codecov (Only on merge to main) + if: github.ref == 'refs/heads/main' && github.event_name == 'push' + uses: codecov/codecov-action@v4 + with: + token: ${{ secrets.CODECOV_TOKEN }} + flags: full diff --git a/.github/workflows/tests_mr.yml b/.github/workflows/tests_mr.yml index ebc3677..4945b6a 100644 --- a/.github/workflows/tests_mr.yml +++ b/.github/workflows/tests_mr.yml @@ -43,3 +43,30 @@ jobs: - name: Test with pytest run: | python -m poetry run pytest + + heavy-deps: + # Same heavy-dependency set as the coverage job on main, so a PR that breaks a + # mocked torch / deepali test fails here instead of after the merge. No coverage + # upload - this is only a gate. + runs-on: ubuntu-latest + timeout-minutes: 45 + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + cache: "pip" + - name: Install system deps for CPU-only mesh3D tests + run: | + sudo apt-get update + sudo apt-get install -y xvfb libgl1 + - name: Install dependencies + run: | + python -m pip install --upgrade pip + python -m pip install poetry + python -m poetry install --with dev + python -m poetry run pip install torch --index-url https://download.pytorch.org/whl/cpu + python -m poetry run pip install antspyx nnunetv2 acvl-utils batchgenerators hf-deepali + - name: Test with pytest + run: | + python -m poetry run pytest unit_tests/ diff --git a/CLAUDE.md b/CLAUDE.md index aa87d99..4d0c97a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -83,7 +83,13 @@ Public API is re-exported from `TPTBox/__init__.py`. All major classes and utili ### Test Layout -Tests live in `unit_tests/` (not `TPTBox/tests/`). `TPTBox/tests/` contains test utilities and sample data (CT/MRI NIfTIs) used by the unit tests. Some generated test files are very large (>20K LOC) — they are autogenerated and should not be edited by hand. +Tests live in `unit_tests/` (not `TPTBox/tests/`). `TPTBox/tests/` contains test utilities and sample data (CT/MRI NIfTIs) used by the unit tests. + +Three files carry a `# Generated by CodiumAI` marker on line 1 — `test_poi_autogen.py`, `test_nii_wrapper_auto.py` and `test_poi_global.py`. There is no generator script in the repo; they were committed once and have been hand-edited since, so treat them as ordinary (if oddly styled) test files. + +House style for new tests, as followed by every file in `unit_tests/`: `unittest.TestCase` classes (no pytest fixtures, no `conftest.py`), `with self.subTest(...)` for table-driven cases, `tempfile.TemporaryDirectory()` for scratch files, and optional heavy dependencies guarded with `importlib.util.find_spec` plus `@unittest.skipUnless` (see `test_mesh3D.py`). New tests must stay Python 3.9-compatible — CI runs a 3.9 leg, so no `itertools.pairwise` (use `TPTBox.core.compat.pairwise`), no `zip(..., strict=...)` (use `compat.zip_strict`) and no `match` statements. + +`[tool.coverage.report].omit` in `pyproject.toml` is curated: it lists only code a unit test cannot legitimately reach (vendored nnU-Net internals, nnU-Net training drivers, the dead vendored deepali engine, the NAKO dataset pipelines, `nipy`-dependent scripts). Code that is merely awkward to test does not belong there. ## Sub-package notes diff --git a/TPTBox/core/bids_constants.py b/TPTBox/core/bids_constants.py index f5a6a08..cd29ee3 100755 --- a/TPTBox/core/bids_constants.py +++ b/TPTBox/core/bids_constants.py @@ -51,8 +51,6 @@ "Chimap", "S0map", "M0ma", - "T1w", - "T2w", "PDw", "T2starw", "FLAIR", @@ -60,7 +58,6 @@ "inplaneT2", "PDT2", "angio", - "T2star", "FLASH", "VF", "defacemas", @@ -102,13 +99,11 @@ "electrode", "marker", "me", - "sti", "PE", "bloo", "pe", "event", "physio", - "sti", "TEM", "SEM", "uCT", diff --git a/TPTBox/core/bids_files.py b/TPTBox/core/bids_files.py index c0c1730..c258513 100755 --- a/TPTBox/core/bids_files.py +++ b/TPTBox/core/bids_files.py @@ -308,7 +308,7 @@ def _scan_tree(path, lvl=1, filter_folder=None, verbose=False): if entry.is_dir(follow_symlinks=False): if filter_folder is not None and not filter_folder(Path(entry.path), lvl): continue - yield from _scan_tree(entry.path, lvl=lvl + 1, verbose=verbose) + yield from _scan_tree(entry.path, lvl=lvl + 1, filter_folder=filter_folder, verbose=verbose) elif entry.name[0] != ".": if verbose: print(f"[{_cont:8}]", end="\r") @@ -336,9 +336,10 @@ def __init__( datasets (typing.List[str]): List of dataset paths parents (typing.List[str]): List of parents (like ["rawdata","sourcedata","derivatives"]) additional_key (list, optional): Additional keys that are not in the default BIDS but should not raise a warning. Defaults to ["sequ", "seg", "ovl"]. - filter_folder: Filter function, input is the path of the folder and the level of the folder structure. Return True if we should continue searching - Example: - filter_folder = lambda p, lvl: True if (lvl != 2 or p.name in ["sub-123","sub-456"]) else False + filter_folder: Filter function, input is the path of the folder and the level of the folder structure. + Return True if we should continue searching. Level 1 is the first folder inside + ``/``, i.e. the subject folder; level 2 is the session folder. + Example: ``filter_folder = lambda p, lvl: lvl != 1 or p.name in ["sub-123", "sub-456"]`` """ self.count_file = 0 self._p_counter = 0 @@ -1872,8 +1873,11 @@ def filter_dixon_water(self, _keys: list[str] | None = None) -> None: assert self._flatten def json_filter(x: dict) -> bool: - """Return True if JSON ImageType contains all required channel keys.""" - return "ImageType" not in x or all(k in x["ImageType"] for k in _keys) + """Return True if JSON ImageType contains any of the channel keys.""" + # `_keys` lists alternative spellings of one channel ("W"/"WATER"), so this is + # `any`, matching `lam_filter` below. With `all` the filter dropped every file + # that had a sidecar, because no ImageType carries both spellings. + return "ImageType" not in x or any(k in x["ImageType"] for k in _keys) def lam_filter(x: str) -> bool: """Return True if the entity value matches one of the target channel keys.""" diff --git a/TPTBox/core/compat.py b/TPTBox/core/compat.py index 92abac4..4a4d400 100644 --- a/TPTBox/core/compat.py +++ b/TPTBox/core/compat.py @@ -1,5 +1,7 @@ from __future__ import annotations +from collections.abc import Iterator + def zip_strict(*iterables) -> zip: """A strict version of zip that raises a ValueError if the input iterables have different lengths. @@ -20,3 +22,36 @@ def zip_strict(*iterables) -> zip: if len(set(lengths)) != 1: raise ValueError(f"Length mismatch: {lengths}") return zip(*lists) + + +def pairwise(iterable) -> Iterator: + """Yield successive overlapping pairs from ``iterable``. + + Backport of :func:`itertools.pairwise`, which is only available from Python + 3.10 onwards while this package still supports 3.9. On 3.10+ the stdlib + implementation is used directly. + + Args: + iterable: Any iterable. It is consumed lazily, so generators are fine. + + Returns: + An iterator over ``(s[0], s[1]), (s[1], s[2]), ...``. Empty for inputs + with fewer than two elements. + """ + try: + from itertools import pairwise as _pairwise # noqa: PLC0415 + except ImportError: + return _pairwise_py39(iterable) + return _pairwise(iterable) + + +def _pairwise_py39(iterable): + """Pure-python ``pairwise`` used on Python 3.9.""" + it = iter(iterable) + try: + prev = next(it) + except StopIteration: + return + for cur in it: + yield prev, cur + prev = cur diff --git a/TPTBox/core/dicom/dicom_renamer.py b/TPTBox/core/dicom/dicom_renamer.py index 78f5708..ebdb688 100644 --- a/TPTBox/core/dicom/dicom_renamer.py +++ b/TPTBox/core/dicom/dicom_renamer.py @@ -487,6 +487,7 @@ def _leftover_move_plan( scan_root: Path, subject_map: dict[str, str], parent: str | None, + already_planned: set[Path] | None = None, ) -> list[tuple[Path, Path]]: """Compute (src, dst) moves for files left behind after the JSON pass. @@ -498,6 +499,11 @@ def _leftover_move_plan( (matched by ``sequ-``) so BIDS stem-linkage stays intact. Files that can't be twinned drop at the session level of the new subject folder as a safe fallback. + + ``already_planned`` holds the sources Pass 1 has lined up. In a real run those + files are gone from disk by the time this runs, but under ``dry_run`` they are + all still there, and without this filter every one of them would be reported a + second time with the fallback destination. """ del parent # kept for signature symmetry; scan_root already resolves it moves: list[tuple[Path, Path]] = [] @@ -571,6 +577,8 @@ def _leftover_move_plan( dst = new_dir / rel.parent / new_name if src == dst or dst.exists(): continue + if already_planned and src in already_planned: + continue moves.append((src, dst)) return moves @@ -688,7 +696,7 @@ def rerun_bids_naming( # Pass 2 — sweep orphan companions (.bval / .bvec / _ADC.nii.gz / other # non-sidecar files) that Pass 1 didn't see because their JSON already # moved on a previous run or they never had one. - leftover = _leftover_move_plan(scan_root, subject_map, parent_norm) + leftover = _leftover_move_plan(scan_root, subject_map, parent_norm, already_planned={src for src, _ in all_moves}) for src, dst in leftover: if verbose: logger.on_neutral(f"{'[dry]' if dry_run else '[lft]'} {src.relative_to(dataset_root)} -> {dst.relative_to(dataset_root)}") diff --git a/TPTBox/core/dicom/xa_pairing.py b/TPTBox/core/dicom/xa_pairing.py index 77cf383..025d501 100644 --- a/TPTBox/core/dicom/xa_pairing.py +++ b/TPTBox/core/dicom/xa_pairing.py @@ -208,7 +208,10 @@ def pair_biplane_xa( if b_sequ == a_sequ: continue # Build the target path with A's sequ; keep all other entities. - new_bids = bf.get_changed_bids(info={"sequ": a_sequ}, non_strict_mode=True) + # `parent` must be the file's own parent: get_changed_bids defaults to + # "derivatives", which silently moved the B-side out of rawdata (where a + # rawdata-scoped scan - including the flip_z pass below - then lost it). + new_bids = bf.get_changed_bids(info={"sequ": a_sequ}, parent=bf.parent, non_strict_mode=True) src_paths = {ext: Path(p) for ext, p in bf.file.items()} new_nii = Path(new_bids.file.get("nii.gz", "")) new_stem = str(new_nii).removesuffix(".nii.gz") diff --git a/TPTBox/core/internal/elastic_deform.py b/TPTBox/core/internal/elastic_deform.py index 84d4c44..93dea6f 100644 --- a/TPTBox/core/internal/elastic_deform.py +++ b/TPTBox/core/internal/elastic_deform.py @@ -4,6 +4,7 @@ from numpy.typing import NDArray from TPTBox import NII +from TPTBox.core.compat import zip_strict def _elasticdeform_install_hint() -> str: @@ -109,7 +110,7 @@ def deformed_nii( order=[0 if v.seg else 3 for v in values], # type: ignore ) out2: dict[str, NII] = {} - for (k, nii), arr in zip(nii_dic.items(), out, strict=True): + for (k, nii), arr in zip_strict(nii_dic.items(), out): out2[k] = nii.set_array(arr[p:-p, p:-p, p:-p]) print("Deformation took", round(time.time() - t, 1), "Seconds") if joint_normalize: diff --git a/TPTBox/core/internal/train_nnUnet/fastProcessor.py b/TPTBox/core/internal/train_nnUnet/fastProcessor.py index 0e08ca1..5e73886 100644 --- a/TPTBox/core/internal/train_nnUnet/fastProcessor.py +++ b/TPTBox/core/internal/train_nnUnet/fastProcessor.py @@ -65,7 +65,7 @@ def _preprocess_dataset( print(f"Preprocessing dataset {dataset_name}") plans_file = join(nnUNet_preprocessed, dataset_name, plans_identifier + ".json") # type: ignore plans_manager = PlansManager(plans_file) - for n, c in zip(num_processes, configurations, strict=False): + for n, c in zip(num_processes, configurations): print(f"Configuration: {c}...") if c not in plans_manager.available_configurations: print(f"INFO: Configuration {c} not found in plans file {plans_identifier + '.json'} of dataset {dataset_name}. Skipping.") diff --git a/TPTBox/core/nii_wrapper.py b/TPTBox/core/nii_wrapper.py index 0f8c748..625f141 100755 --- a/TPTBox/core/nii_wrapper.py +++ b/TPTBox/core/nii_wrapper.py @@ -2705,7 +2705,10 @@ def to_stl( if out_path.is_dir(): out_path = out_path / f"mask_{label}.stl" elif number_path: - out_path.with_name(f"{out_path.stem}_{label}.stl") + # with_name returns a new path; the result used to be discarded, so + # to_stls() wrote every label to the same file and only the last + # label's mesh survived. + out_path = out_path.with_name(f"{out_path.stem}_{label}.stl") log.on_save(f"Saving STL to {out_path}") out_path.parent.mkdir(exist_ok=True,parents=True) cube.save(str(out_path)) diff --git a/TPTBox/core/np_utils.py b/TPTBox/core/np_utils.py index 23622e8..533f39c 100755 --- a/TPTBox/core/np_utils.py +++ b/TPTBox/core/np_utils.py @@ -2218,9 +2218,11 @@ def np_calc_boundary_mask( """ if adjust_intensity_for_ct: img = img + 1000 - boundary = img.copy() - boundary[boundary > threshold] = 2 - boundary[boundary <= threshold] = 0 + # Compute the foreground mask *before* writing the sentinel: the two assignments + # used to run in place one after the other, so the 2 written by the first line was + # then compared against the threshold by the second. For any threshold >= 2 that + # wiped the foreground again and the function returned an empty mask. + boundary = np.where(img > threshold, 2, 0) infect_list = [] def infect(x, y, z): @@ -2242,7 +2244,6 @@ def infect(x, y, z): (-1, 0, 0), (0, 1, 0), (0, -1, 0), - (0, 1, 0), (0, 0, 1), (0, 0, -1), ]: diff --git a/TPTBox/registration/_deformable/_grid_search_vert.py b/TPTBox/registration/_deformable/_grid_search_vert.py index 438b314..9557b8e 100644 --- a/TPTBox/registration/_deformable/_grid_search_vert.py +++ b/TPTBox/registration/_deformable/_grid_search_vert.py @@ -132,7 +132,7 @@ def get_additonal_point(poi: POI, cc: NII, seg_id): # Crop to speed up and decreas gpu memory consumption. c0 = femur_0_.compute_crop(dist=10) c1 = femur_1_reg.compute_crop(dist=10) - ex_slice = [slice(min(a.start, b.start), max(a.stop, b.stop)) for a, b in zip(c1, c0, strict=False)] + ex_slice = [slice(min(a.start, b.start), max(a.stop, b.stop)) for a, b in zip(c1, c0)] femur_0 = femur_0_.apply_crop(ex_slice) femur_1_reg.apply_crop_(ex_slice) diff --git a/TPTBox/segmentation/nnUnet_utils/predictor.py b/TPTBox/segmentation/nnUnet_utils/predictor.py index 3009ec4..a04a2d6 100755 --- a/TPTBox/segmentation/nnUnet_utils/predictor.py +++ b/TPTBox/segmentation/nnUnet_utils/predictor.py @@ -587,7 +587,7 @@ def _internal_maybe_mirror_and_predict(self, x: torch.Tensor, network) -> torch. mirror_axes = [m + 2 for m in mirror_axes] axes_combinations = [c for i in range(len(mirror_axes)) for c in itertools.combinations(mirror_axes, i + 1)] for axes in axes_combinations: - prediction += torch.flip(self.network(torch.flip(x, axes)), axes) + prediction += torch.flip(network(torch.flip(x, axes)), axes) prediction /= len(axes_combinations) + 1 return prediction diff --git a/TPTBox/segmentation/rib/_rib_assign.py b/TPTBox/segmentation/rib/_rib_assign.py index 426880d..c88406a 100644 --- a/TPTBox/segmentation/rib/_rib_assign.py +++ b/TPTBox/segmentation/rib/_rib_assign.py @@ -405,7 +405,11 @@ def assign_ribs_to_vert_segmentation( logger.on_debug(f"{rib_inst.unique()=}") # Merge rib assignments back into the original vert / sem segmentations # without disturbing existing (non-rib) labels. Skip unmatched CCs (sentinel error_value=255). - matched = (rib_inst != 0) & (rib_inst != error_value) + # Combine on the arrays, not on NII: `rib_inst != 0` yields a bool-backed NII and + # NII.__and__ rejects non-integer arrays, so the NII-level `&` raised TypeError and + # this function could never run to completion. + rib_inst_arr = rib_inst.get_seg_array() + matched = (rib_inst_arr != 0) & (rib_inst_arr != error_value) vert_seg[matched] = rib_inst[matched] sem_seg[rib_seg != 0] = rib_seg.map_labels({left_id: Location.Rib_Left.value, right_id: Location.Rib_Right.value})[rib_seg != 0] # type: ignore @@ -418,7 +422,9 @@ def assign_ribs_to_vert_segmentation( vert_seg2 = vert_seg.remove_labels(error_value).infect(vert_seg.extract_label(error_value), verbose=False) vert_seg[vert_seg != vert_seg2] = vert_seg2[vert_seg != vert_seg2] vert_seg[np.logical_and(rib_inst == error_value, vert_seg == 0)] = error_value - vert_seg[np.logical_and(vert_seg == 0, sem_seg.extract_label([Location.Rib_Left.value, Location.Rib_Right.value] == 1))] = ( - error_value - ) + # The closing bracket used to sit one token too far left, so this evaluated + # `[Rib_Left, Rib_Right] == 1` -> False -> extract_label(0), which trips the + # "zero label is the background" assertion and took the whole function down. + rib_sem = sem_seg.extract_label([Location.Rib_Left.value, Location.Rib_Right.value]).get_seg_array() == 1 + vert_seg[np.logical_and(vert_seg.get_seg_array() == 0, rib_sem)] = error_value return vert_seg.reorient_(ori), sem_seg.reorient_(ori) diff --git a/TPTBox/spine/spinestats/README.md b/TPTBox/spine/spinestats/README.md index 5afa6ae..4b1d713 100644 --- a/TPTBox/spine/spinestats/README.md +++ b/TPTBox/spine/spinestats/README.md @@ -58,7 +58,7 @@ not on the Python API. A standalone copy of this reference lives at | Key | Source function | What it covers | |---|---|---| | `ivd_geometry` | `measure_ivd_and_vertebra_geometry(..., structure_label=100)` | intervertebral discs (now also carries wedge angles/indices per label) | -| `vert_geometry` | `measure_ivd_and_vertebra_geometry(..., structure_label=50)` | vertebral bodies (now also carries wedge angles/indices per label) | +| `vert_geometry` | `measure_ivd_and_vertebra_geometry(..., structure_label=0)` | vertebral bodies (now also carries wedge angles/indices per label) | | `VBQ_score` | `VBQ_score` | vertebral bone quality (T2 signal ratio) | | `body_composition_score` | `body_composition_score` | axial CSA per tissue at chosen vertebral levels | | `muscle_fat_infiltration` | `muscle_fat_infiltration` | Dixon fat-fraction based muscle-quality metrics | diff --git a/TPTBox/spine/spinestats/all_output_reference.md b/TPTBox/spine/spinestats/all_output_reference.md index 2cc16a0..35c2d24 100644 --- a/TPTBox/spine/spinestats/all_output_reference.md +++ b/TPTBox/spine/spinestats/all_output_reference.md @@ -13,7 +13,7 @@ Python API. | Key | Source function | What it covers | |---|---|---| | `ivd_geometry` | `measure_ivd_and_vertebra_geometry(..., structure_label=100)` | intervertebral discs (per-label entries carry wedge angles/indices) | -| `vert_geometry` | `measure_ivd_and_vertebra_geometry(..., structure_label=50)` | vertebral bodies (per-label entries carry wedge angles/indices) | +| `vert_geometry` | `measure_ivd_and_vertebra_geometry(..., structure_label=0)` | vertebral bodies (per-label entries carry wedge angles/indices) | | `VBQ_score` | `VBQ_score` | vertebral bone quality (T2 signal ratio) | | `body_composition_score` | `body_composition_score` | axial CSA per tissue at chosen vertebral levels | | `muscle_fat_infiltration` | `muscle_fat_infiltration` | Dixon fat-fraction based muscle-quality metrics | diff --git a/TPTBox/spine/spinestats/curvature.py b/TPTBox/spine/spinestats/curvature.py index 8f8f267..6b40bd6 100644 --- a/TPTBox/spine/spinestats/curvature.py +++ b/TPTBox/spine/spinestats/curvature.py @@ -51,13 +51,13 @@ from __future__ import annotations -from itertools import pairwise from typing import Any import numpy as np from numpy.linalg import norm from TPTBox import POI, Location, Vertebra_Instance +from TPTBox.core.compat import pairwise from TPTBox.spine.spinestats.angles import Def_Curvature, MoveTo # --------------------------------------------------------------------------- diff --git a/TPTBox/spine/spinestats/measure_ivd_and_vertebra_geometry.py b/TPTBox/spine/spinestats/measure_ivd_and_vertebra_geometry.py index 1d08bae..1529ce6 100644 --- a/TPTBox/spine/spinestats/measure_ivd_and_vertebra_geometry.py +++ b/TPTBox/spine/spinestats/measure_ivd_and_vertebra_geometry.py @@ -173,7 +173,10 @@ def measure_ivd_and_vertebra_geometry( Location.Endplate, ], buffer_file=buffer_poi, - save_buffer_file=True, + # Only ask for a write-back when there is somewhere to write: the buffering + # decorator asserts buffer_file is not None, so hardcoding True made the + # documented default (buffer_poi=None) raise AssertionError. + save_buffer_file=buffer_poi is not None, ) if instance_labels is None: instance_labels = [int(i) for i in vert.unique() if i > structure_label and i < structure_label + 100] @@ -439,7 +442,7 @@ def _segment_length(intersection_points) -> float: """ if len(intersection_points) == 0: return 0 - return np.sqrt(sum((a - b) ** 2 for a, b in zip(intersection_points[0], intersection_points[-1], strict=False))) + return np.sqrt(sum((a - b) ** 2 for a, b in zip(intersection_points[0], intersection_points[-1]))) def _swap(a, b): diff --git a/TPTBox/stitching/stitching_tools.py b/TPTBox/stitching/stitching_tools.py index d42c67d..f9ba63c 100755 --- a/TPTBox/stitching/stitching_tools.py +++ b/TPTBox/stitching/stitching_tools.py @@ -201,6 +201,8 @@ def _center_frontal(size: int) -> slice: distance = int((upper_bound - lower_bound) / 2) middle = int(lower_bound + distance) # rather push it little back over the top in dorsal direction rather then not putting it into account - small_value = middle - 122 - big_value = middle + 134 + # Clamped to the volume: below ~258 voxels `middle - 122` turns negative, and numpy reads a + # negative slice start as an offset from the end, silently cropping from the wrong side. + small_value = max(0, middle - 122) + big_value = min(size, middle + 134) return slice(small_value, big_value) diff --git a/pyproject.toml b/pyproject.toml index f320a1e..59ca245 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,11 @@ scikit-learn = "*" pynrrd = "*" requests = "*" trimesh = "*" +# trimesh's ray engine does its broad phase through rtree, which trimesh itself only +# declares under its `easy` extra. Without it every mesh ray cast raises +# ModuleNotFoundError, which takes out spine.spinestats.measure_ivd_and_vertebra_geometry +# and spinestats.poi_fun.endplates. +rtree = "*" numpy-stl = "*" # --- OLD STACK (Python < 3.11) @@ -50,7 +55,20 @@ scikit-image = [ # (TPTBox.core.dicom, TPTBox.segmentation) import successfully without these # and only raise - naming the extra - when an entry point is actually called. pydicom = { version = "*", optional = true } -dicom2nifti = { version = "*", optional = true } +# dicom2nifti 2.6.0 imports `pydicom.pixels` while still declaring +# `pydicom >= 2.2`, and 2.6.1+ require `pydicom >= 3.0` outright. pydicom 3 +# needs python >= 3.10, so 2.5.x is the newest dicom2nifti that works on 3.9. +dicom2nifti = [ + { version = "<2.6.0", python = "<3.10", optional = true }, + { version = "*", python = ">=3.10", optional = true }, +] +# dicom2nifti requires `python-gdcm` with no upper bound, and on 3.9 nnunetv2 +# 2.4.2 (the `seg` extra) pulls dicom2nifti in as well. python-gdcm stopped +# shipping cp39 wheels after 3.2.2 but still advertises +# `requires-python = ">=3.7"`, so a resolver picks 3.2.6 on 3.9 and then fails +# building its sdist - the tarball has no CMakeLists.txt. Cap it at the last +# release with cp39 wheels; >=3.10 takes whatever dicom2nifti resolves to. +python-gdcm = { version = "<=3.2.2", python = "<3.10", optional = true } torch = { version = "*", optional = true } # Mirror SPINEPS' own constraint so `TPTBox[seg]` and `spineps` can co-install. # nnunetv2 2.8+ requires Python >=3.10, so 3.9 stays on the last release supporting it. @@ -70,7 +88,7 @@ hf-deepali = { version = "*", optional = true } # `spineps` is deliberately absent: it depends on TPTBox itself, so declaring it # here would be a circular dependency. It stays a documented `pip install spineps`. [tool.poetry.extras] -dicom = ["pydicom", "dicom2nifti"] +dicom = ["pydicom", "dicom2nifti", "python-gdcm"] seg = ["torch", "nnunetv2", "acvl-utils", "batchgenerators"] reg = ["torch", "hf-deepali"] @@ -86,6 +104,16 @@ pre-commit = "*" pyvista = "^0.43.2" coverage = ">=7.0.1" pytest-mock = "^3.6.0" +# The `dicom` extra, in the dev group so unit_tests/test_dicom.py and +# test_dicom_renamer.py run in CI instead of only where the extra is installed. +# Deliberately unmarked: a `python = ">=3.10"` marker here does not keep the pair +# off 3.9. Poetry 2.2 merges this declaration with the optional one above into a +# group marker that still mentions `extra == "dicom"` and then installs the +# package on 3.9 anyway - that is how python-gdcm 3.2.6 reached the 3.9 leg and +# broke it. The cp39 wheel problem is fixed by the caps above instead. On 3.9 the +# resolver lands on pydicom 2.4.4 / dicom2nifti 2.5.1, which both modules pass on. +pydicom = "*" +dicom2nifti = "*" networkx = "*" # optional: np_split_connected_component(method="mincut") exceptiongroup = { version = "^1.2", python = "<3.11" } tomli = {version = "*", python = "<3.11" } @@ -289,8 +317,61 @@ exclude_lines = [ "if TYPE_CHECKING: ", "raise " ] +# Curated: every entry here is code a unit test cannot legitimately reach. Anything +# that is merely awkward to test is NOT listed - it gets a test instead. The PR that +# introduced this list reports the coverage number both with and without it. omit = [ - "TPTBox/test/speedtest/*", + # Hand-run micro-benchmarks. The filenames do not match pytest's `test_*.py`, so + # they are measured but never executed. + "TPTBox/tests/speedtests/*", + + # --- Vendored nnU-Net inference internals ------------------------------------- + # Carried from upstream nnU-Net with no TPTBox-specific logic; reachable only with + # real weights on a GPU. TPTBox's own entry point around them, + # nnUnet_utils/inference_api.py, is deliberately NOT omitted and is tested with a + # mock at the model boundary. + "TPTBox/segmentation/nnUnet_utils/predictor.py", + "TPTBox/segmentation/nnUnet_utils/plans_handler.py", + "TPTBox/segmentation/nnUnet_utils/export_prediction.py", + "TPTBox/segmentation/nnUnet_utils/default_preprocessor.py", + "TPTBox/segmentation/nnUnet_utils/data_iterators.py", + "TPTBox/segmentation/nnUnet_utils/get_network_from_plans.py", + "TPTBox/segmentation/nnUnet_utils/sliding_window_prediction.py", + # GPU memory probing: shells out to nvidia-smi in a measurement loop. + "TPTBox/segmentation/nnUnet_utils/estimate_nnunet_memory.py", + + # --- nnU-Net training drivers -------------------------------------------------- + # Need a populated nnUNet_raw / nnUNet_preprocessed tree and a GPU. + "TPTBox/core/internal/train_nnUnet/*", + + # --- Permanently unimportable -------------------------------------------------- + # Both import `nipy`, which is not a declared dependency of this package (or + # installable on current Python). test_docs_examples.py already special-cases + # script_ax2sag for the same reason. + "TPTBox/registration/script_ax2sag.py", + "TPTBox/registration/_ridged_intensity/register.py", + + # --- Dead vendored deepali optimisation engine --------------------------------- + # An adapted copy of deepali's own registration engine. It imports `yaml`, which is + # not declared either, and nothing in the public API reaches it: the package is + # referenced only by deformable_reg_old.py and grid_search.py, which are themselves + # unreferenced. Candidate for deletion rather than coverage. + "TPTBox/registration/_deformable/_deepali/*", + "TPTBox/registration/_deformable/deformable_reg_old.py", + "TPTBox/registration/_deformable/grid_search.py", + "TPTBox/registration/_deformable/_grid_search_vert.py", + + # --- NAKO dataset pipelines ---------------------------------------------------- + # Interactive input() prompts, hard-coded /DATA/NAS dataset roots, process pools + # over a private dataset, and pandas/Excel reporting over a finished run. + "TPTBox/spine/spinestats/_load_nako_wh.py", + "TPTBox/spine/spinestats/_load_nako.py", + "TPTBox/spine/spinestats/_run_all.py", + "TPTBox/spine/spinestats/_qc_report.py", + + # --- One-off repair script ----------------------------------------------------- + # Walks a real dataset and rewrites pickles in place. + "TPTBox/core/dicom/fix_brocken.py", ] [tool.ruff.format] # Like Black, use double quotes for strings. diff --git a/unit_tests/test_bids_files.py b/unit_tests/test_bids_files.py new file mode 100644 index 0000000..4dad2ac --- /dev/null +++ b/unit_tests/test_bids_files.py @@ -0,0 +1,654 @@ +"""Unit tests for `TPTBox.core.bids_files` and `TPTBox.core.bids_constants`. + +`unit_tests/test_bids_file.py` covers the same module through `get_BIDS_test()`, but +all four of its tests are gated on a developer-local dataset path and therefore never +run. This file instead builds a small real BIDS tree (NIfTIs, JSON sidecars, a POI +json, a PNG) in a temporary directory, so the scan -> query -> family -> derived-path +pipeline is exercised end to end everywhere. + +Two module-level containers in `bids_constants` are mutated as a side effect of +parsing (`entities_keys[key] = key` for unknown keys, `formats_relaxed.append(...)` +for unknown formats) and `BIDS_Global_info.__init__` registers its `additional_key`s +the same way. They are snapshotted and restored in place, because other modules bound +the same objects at import time. +""" + +from __future__ import annotations + +import json +import tempfile +import unittest +from pathlib import Path + +import nibabel as nib +import numpy as np + +from TPTBox import POI +from TPTBox.core import bids_constants +from TPTBox.core.bids_files import ( + BIDS_FILE, + BIDS_Family, + BIDS_Global_info, + Buffered_BIDS_Global_info, + Searchquery, + get_values_from_name, + validate_entities, +) + +ZOOM = (1.5, 1.5, 3.0) +ORIGIN = (10.0, 20.0, 30.0) +SHAPE = (4, 5, 6) + + +def _affine() -> np.ndarray: + aff = np.diag([*ZOOM, 1.0]) + aff[:3, 3] = ORIGIN + return aff + + +def _save_nii(path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + nib.save(nib.Nifti1Image(np.zeros(SHAPE, dtype=np.uint8), _affine()), path) + + +def _save_json(path: Path, payload: dict) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(payload), encoding="utf-8") + + +def build_dataset(root: Path) -> Path: + """Create a two-subject BIDS dataset under ``root`` and return the dataset path. + + sub-001 has a CT with a derivatives family (two masks, a POI json, a snapshot) and a + four-channel Dixon sequence with ImageType sidecars; sub-002 has a single T2w. + """ + ds = root / "dataset-test" + raw = ds / "rawdata" / "sub-001" / "ses-20210101" + der = ds / "derivatives" / "sub-001" / "ses-20210101" + + _save_nii(raw / "sub-001_ses-20210101_sequ-301_ct.nii.gz") + _save_json(raw / "sub-001_ses-20210101_sequ-301_ct.json", {"ImageType": ["ORIGINAL", "PRIMARY"], "EchoTime": 1.0}) + for part, image_type in ( + ("water", ["DERIVED", "W"]), + ("fat", ["DERIVED", "F"]), + ("inphase", ["ORIGINAL", "IP"]), + ("outphase", ["ORIGINAL", "OP"]), + ): + _save_nii(raw / f"sub-001_ses-20210101_sequ-302_part-{part}_dixon.nii.gz") + _save_json(raw / f"sub-001_ses-20210101_sequ-302_part-{part}_dixon.json", {"ImageType": image_type}) + + _save_nii(der / "sub-001_ses-20210101_sequ-301_seg-subreg_msk.nii.gz") + _save_nii(der / "sub-001_ses-20210101_sequ-301_seg-vert_msk.nii.gz") + (der / "sub-001_ses-20210101_sequ-301_snp.png").write_bytes(b"\x89PNG\r\n\x1a\n") + poi = POI( + {20: {50: (1.0, 2.0, 3.0)}, 21: {50: (2.0, 3.0, 4.0)}}, + orientation=("R", "A", "S"), + zoom=ZOOM, + shape=SHAPE, + origin=ORIGIN, + rotation=np.eye(3), + ) + poi.save(der / "sub-001_ses-20210101_sequ-301_seg-subreg_ctd.json", verbose=False) + + raw2 = ds / "rawdata" / "sub-002" / "ses-20220202" + _save_nii(raw2 / "sub-002_ses-20220202_sequ-401_T2w.nii.gz") + _save_json(raw2 / "sub-002_ses-20220202_sequ-401_T2w.json", {"ImageType": ["ORIGINAL"]}) + return ds + + +class _BidsConstantsGuard(unittest.TestCase): + """Restores the mutable module-level tables in `bids_constants` around each test.""" + + def setUp(self): + self._entities = dict(bids_constants.entities) + self._entities_keys = dict(bids_constants.entities_keys) + self._formats_relaxed = list(bids_constants.formats_relaxed) + + def tearDown(self): + # Restore in place: other modules imported these objects by reference. + bids_constants.entities.clear() + bids_constants.entities.update(self._entities) + bids_constants.entities_keys.clear() + bids_constants.entities_keys.update(self._entities_keys) + bids_constants.formats_relaxed[:] = self._formats_relaxed + + +class TestGetValuesFromName(_BidsConstantsGuard): + def test_table(self): + cases = [ + # name, bids_format, info, bids_key, file_type + ("sub-001_ses-01_T1w.nii.gz", "T1w", {"sub": "001", "ses": "01"}, "sub-001_ses-01_T1w", "nii.gz"), + ("sub-001_ct.nii.gz", "ct", {"sub": "001"}, "sub-001_ct", "nii.gz"), + ("sub-001_ses-01_seg-subreg_msk.nii.gz", "msk", {"sub": "001", "ses": "01", "seg": "subreg"}, None, "nii.gz"), + ("sub-x_ses-1_ctd.json", "ctd", {"sub": "x", "ses": "1"}, None, "json"), + ("sub-x_ses-1_ctd.mrk.json", "ctd", {"sub": "x", "ses": "1"}, None, "mrk.json"), + # value may itself contain a dash: only the first one splits key from value + ("sub-001_part-water-fraction_dixon.nii.gz", "dixon", {"sub": "001", "part": "water-fraction"}, None, "nii.gz"), + ] + for name, fmt, info, key, ftype in cases: + with self.subTest(name=name): + got_fmt, got_info, got_key, got_ftype = get_values_from_name(name, verbose=False) + self.assertEqual(got_fmt, fmt) + self.assertEqual(got_info, info) + self.assertEqual(got_ftype, ftype) + self.assertEqual(got_key, key if key is not None else name[: -len(ftype) - 1]) + + def test_accepts_path_and_str_identically(self): + name = "sub-001_ses-01_sequ-7_ct.nii.gz" + self.assertEqual(get_values_from_name(name, verbose=False), get_values_from_name(Path("/a/b") / name, verbose=False)) + + def test_malformed_segments_are_skipped_not_fatal(self): + # "broken" has no "-", so it cannot be a key-value pair and is dropped. + fmt, info, _, ftype = get_values_from_name("sub-001_broken_ct.nii.gz", verbose=False) + self.assertEqual((fmt, ftype), ("ct", "nii.gz")) + self.assertEqual(info, {"sub": "001"}) + + def test_duplicate_key_keeps_the_last_value(self): + _, info, _, _ = get_values_from_name("sub-001_sub-002_ct.nii.gz", verbose=False) + self.assertEqual(info, {"sub": "002"}) + + def test_name_without_extension_raises(self): + # Documents current behaviour: the name is split on "." without a fallback. + with self.assertRaises(ValueError): + get_values_from_name("sub-001_ct", verbose=False) + + def test_unknown_format_is_registered_only_when_verbose(self): + before = list(bids_constants.formats_relaxed) + get_values_from_name("sub-001_notaformat.nii.gz", verbose=False) + self.assertEqual(bids_constants.formats_relaxed, before) + get_values_from_name("sub-001_notaformat.nii.gz", verbose=True) + self.assertIn("notaformat", bids_constants.formats_relaxed) + + +class TestValidateEntities(_BidsConstantsGuard): + def test_valid_entities(self): + for key, value in (("sub", "001"), ("ses", "20210101"), ("run", "7"), ("task", "rest1"), ("mt", "on"), ("hemi", "L")): + with self.subTest(key=key, value=value): + self.assertTrue(validate_entities(key, value, "name", verbose=True)) + + def test_invalid_entities(self): + cases = [ + ("run", "a1"), # must be decimal + ("echo", "1.5"), # must be decimal + ("task", "a_b"), # must be alphanumeric + ("mt", "maybe"), # must be on/off + ("hemi", "X"), # must be L/R + ("mod", "banana"), # must be a known format + ] + for key, value in cases: + with self.subTest(key=key, value=value): + self.assertFalse(validate_entities(key, value, "name", verbose=True)) + + def test_key_is_lowercased_before_lookup(self): + self.assertTrue(validate_entities("SUB", "001", "name", verbose=True)) + + def test_unknown_key_is_rejected_once_then_registered(self): + self.assertNotIn("totallynewkey", bids_constants.entities_keys) + self.assertFalse(validate_entities("totallynewkey", "v", "name", verbose=True)) + # Side effect: the key is remembered, so the second call passes. + self.assertIn("totallynewkey", bids_constants.entities_keys) + self.assertTrue(validate_entities("totallynewkey", "v", "name", verbose=True)) + + def test_non_verbose_short_circuits_to_true(self): + self.assertTrue(validate_entities("hemi", "X", "name", verbose=False)) + self.assertTrue(validate_entities("definitely not a key", "", "name", verbose=False)) + + +class TestBidsConstants(unittest.TestCase): + """Properties of the shipped constant tables. + + These assert the *source* tables, loaded fresh: `BIDS_Global_info.__init__` registers + its `additional_key`s by writing `entities[k] = k`, so by the time a full-suite run + reaches this class the live dicts have extra self-mapping entries. + """ + + @staticmethod + def _pristine(): + import importlib.util # noqa: PLC0415 + + path = Path(bids_constants.__file__) + spec = importlib.util.spec_from_file_location("_pristine_bids_constants", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + def test_entities_keys_is_the_inverse_of_entities(self): + c = self._pristine() + self.assertEqual(len(c.entities_keys), len(c.entities), "value collision in `entities`") + for human, short in c.entities.items(): + self.assertEqual(c.entities_keys[short], human) + + def test_entity_rule_groups_reference_known_keys(self): + known = set(self._pristine().entities.values()) + for name in ("entity_alphanumeric", "entity_decimal", "entity_format", "entity_on_off", "entity_left_right", "entity_parts"): + with self.subTest(group=name): + self.assertTrue(set(getattr(bids_constants, name)).issubset(known)) + + def test_formats_relaxed_extends_formats(self): + c = self._pristine() + self.assertTrue(set(c.formats).issubset(set(c.formats_relaxed))) + + def test_no_duplicate_formats_or_file_types(self): + # `formats` carried copy-paste duplicates (T1w, T2w, T2star, sti x2). + c = self._pristine() + for name in ("formats", "formats_relaxed", "file_types"): + with self.subTest(table=name): + table = getattr(c, name) + self.assertEqual(len(table), len(set(table))) + + +class _DatasetTestCase(_BidsConstantsGuard): + """Base class providing a freshly scanned synthetic dataset per test class.""" + + @classmethod + def setUpClass(cls): + cls._tmp = tempfile.TemporaryDirectory() + cls.ds = build_dataset(Path(cls._tmp.name)) + + @classmethod + def tearDownClass(cls): + cls._tmp.cleanup() + + def setUp(self): + super().setUp() + self.g = BIDS_Global_info([self.ds], parents=["rawdata", "derivatives"], verbose=False) + self.subj = self.g.subjects["001"] + + def ct_file(self) -> BIDS_FILE: + return next(b for b in self.subj.new_query(flatten=True).loop_list() if b.bids_format == "ct") + + +class TestBidsGlobalInfoScan(_DatasetTestCase): + def test_subjects_and_sequences(self): + self.assertEqual(len(self.g), 2) + self.assertEqual(sorted(self.g.subjects), ["001", "002"]) + self.assertEqual(sorted(self.subj.sequences), ["20210101_301", "20210101_302"]) + + def test_sequence_name_joins_the_splitting_keys(self): + ct = self.ct_file() + self.assertEqual(self.subj.get_sequence_name(ct), "20210101_301") + + def test_enumerate_subjects_sorted_and_shuffled_agree_as_sets(self): + ordered = [n for n, _ in self.g.enumerate_subjects(sort=True)] + self.assertEqual(ordered, sorted(ordered)) + shuffled = [n for n, _ in self.g.enumerate_subjects(shuffle=True)] + self.assertEqual(sorted(shuffled), ordered) + + def test_iter_subjects_matches_enumerate_subjects(self): + self.assertEqual( + [n for n, _ in self.g.iter_subjects(sort=True)], + [n for n, _ in self.g.enumerate_subjects(sort=True)], + ) + + def test_str_lists_parents_and_datasets(self): + text = str(self.g) + self.assertIn("rawdata", text) + self.assertIn("dataset-test", text) + + def test_missing_dataset_raises(self): + with self.assertRaises(FileNotFoundError): + BIDS_Global_info([self.ds / "does-not-exist"], parents=["rawdata"], verbose=False) + + def test_additional_key_is_registered(self): + BIDS_Global_info([self.ds], parents=["rawdata"], additional_key=["myextrakey"], verbose=False) + self.assertIn("myextrakey", bids_constants.entities_keys) + + def test_filter_folder_can_exclude_a_subject(self): + g = BIDS_Global_info( + [self.ds], + parents=["rawdata"], + verbose=False, + filter_folder=lambda p, lvl: lvl != 1 or p.name == "sub-001", + ) + self.assertEqual(sorted(g.subjects), ["001"]) + + def test_filter_folder_is_applied_to_nested_levels_too(self): + # Regression: `_scan_tree` did not forward `filter_folder` to its recursive call, + # so any filter targeting a level below the subject folder was silently ignored. + g = BIDS_Global_info( + [self.ds], + parents=["rawdata"], + verbose=False, + filter_folder=lambda p, lvl: lvl != 2 or p.name != "ses-20210101", + ) + self.assertEqual(sorted(g.subjects), ["002"]) + + def test_additional_file_list_adds_files_without_scanning(self): + extra = self.ds / "rawdata" / "sub-003" / "sub-003_ct.nii.gz" + g = BIDS_Global_info([self.ds], parents=[], verbose=False, additional_file_list={self.ds: [extra]}) + self.assertIn("003", g.subjects) + + +class TestBidsFilePaths(_DatasetTestCase): + def test_get_path_decomposed(self): + ds, parent, mid, name = self.ct_file().get_path_decomposed(file_type="nii.gz") + self.assertEqual(Path(ds), self.ds) + self.assertEqual(parent, "rawdata") + self.assertEqual(mid, "sub-001/ses-20210101") + self.assertEqual(name, "sub-001_ses-20210101_sequ-301_ct.nii.gz") + + def test_properties(self): + ct = self.ct_file() + self.assertEqual(ct.parent, "rawdata") + self.assertEqual(ct.bids_format, "ct") + self.assertEqual(ct.info, {"sub": "001", "ses": "20210101", "sequ": "301"}) + self.assertEqual(ct.get("sequ"), "301") + self.assertEqual(ct.get("nope", "dflt"), "dflt") + self.assertEqual(dict(ct.loop_keys()), ct.info) + self.assertTrue(ct.exists()) + + def _rel(self, path) -> str: + return Path(path).relative_to(self.ds).as_posix() + + def test_get_changed_path_variants(self): + ct = self.ct_file() + rel = self._rel + cases = { + rel(ct.get_changed_path()): "derivatives/sub-001/ses-20210101/sub-001_ses-20210101_sequ-301_ct.nii.gz", + rel(ct.get_changed_path(bids_format="msk", info={"seg": "subreg"})): ( + "derivatives/sub-001/ses-20210101/sub-001_ses-20210101_sequ-301_seg-subreg_msk.nii.gz" + ), + rel(ct.get_changed_path(file_type="json")): "derivatives/sub-001/ses-20210101/sub-001_ses-20210101_sequ-301_ct.json", + rel(ct.get_changed_path(additional_folder="X")): ("derivatives/sub-001/ses-20210101/X/sub-001_ses-20210101_sequ-301_ct.nii.gz"), + rel(ct.get_changed_path(path="sub-{sub}/custom")): ("derivatives/sub-001/custom/sub-001_ses-20210101_sequ-301_ct.nii.gz"), + rel(ct.get_changed_path(parent="rawdata2")): "rawdata2/sub-001/ses-20210101/sub-001_ses-20210101_sequ-301_ct.nii.gz", + } + for got, expected in cases.items(): + with self.subTest(expected=expected): + self.assertEqual(got, expected) + + def test_get_changed_path_entities_are_bids_sorted(self): + # `run` must land after `sequ`, regardless of insertion order. + p = self.ct_file().get_changed_path(info={"run": "2"}) + self.assertEqual(Path(p).name, "sub-001_ses-20210101_sequ-301_run-2_ct.nii.gz") + + def test_get_changed_path_make_parent_creates_the_folder(self): + p = self.ct_file().get_changed_path(parent="derivatives_new", make_parent=True) + self.assertTrue(Path(p).parent.is_dir()) + + def test_get_changed_bids_wraps_the_same_path(self): + ct = self.ct_file() + bids = ct.get_changed_bids(bids_format="msk", info={"seg": "subreg"}) + self.assertIsInstance(bids, BIDS_FILE) + self.assertEqual(bids.bids_format, "msk") + self.assertEqual(bids.get("seg"), "subreg") + + def test_insert_info_into_path(self): + ct = self.ct_file() + self.assertEqual(ct.insert_info_into_path("{sub}/{ses}/x"), "001/20210101/x") + self.assertIsNone(ct.insert_info_into_path(None)) + + def test_set_remove_roundtrip(self): + ct = self.ct_file() + ct.set("run", "3") + self.assertEqual(ct.get("run"), "3") + self.assertEqual(ct.remove("run"), "3") + self.assertIsNone(ct.get("run")) + + def test_get_file_and_open_json(self): + ct = self.ct_file() + self.assertTrue(ct.has_json()) + self.assertEqual(ct.get_file("json").name, "sub-001_ses-20210101_sequ-301_ct.json") + self.assertIsNone(ct.get_file("nrrd", None)) + self.assertEqual(ct.open_json()["EchoTime"], 1.0) + + def test_open_nii_and_grid_info(self): + ct = self.ct_file() + self.assertTrue(ct.has_nii()) + nii = ct.open_nii() + self.assertEqual(nii.shape, SHAPE) + np.testing.assert_allclose(nii.zoom, ZOOM, rtol=1e-6) + grid = ct.get_grid_info(add_grid_info_to_json=False) + self.assertEqual(tuple(grid.shape), SHAPE) + np.testing.assert_allclose(grid.zoom, ZOOM, rtol=1e-6) + + def test_open_nii_reorient(self): + nii = self.ct_file().open_nii_reorient(("P", "I", "R")) + self.assertEqual(nii.orientation, ("P", "I", "R")) + + def test_open_dispatch(self): + ct = self.ct_file() + self.assertEqual(ct.open("json")["EchoTime"], 1.0) + self.assertEqual(ct.open("nii.gz").shape, SHAPE) + + def test_open_poi(self): + fam = self.subj.get_sequence_files("20210101_301") + poi = fam["ctd_seg-subreg"][0].open_poi() + self.assertEqual(sorted(poi.keys()), [(20, 50), (21, 50)]) + np.testing.assert_allclose(poi[20, 50], (1.0, 2.0, 3.0), rtol=1e-6) + + def test_do_filter(self): + ct = self.ct_file() + self.assertTrue(ct.do_filter("format", "ct")) + self.assertFalse(ct.do_filter("format", "msk")) + self.assertTrue(ct.do_filter("sequ", lambda v: int(v) == 301)) + self.assertTrue(ct.do_filter("Filetype", "nii.gz")) + + def test_get_interpolation_order_is_3_for_images_and_0_for_masks(self): + self.assertEqual(self.ct_file().get_interpolation_order(), 3) + msk = next(b for b in self.subj.new_query(flatten=True).loop_list() if b.bids_format == "msk") + self.assertEqual(msk.get_interpolation_order(), 0) + + def test_dunders(self): + ct = self.ct_file() + # A BIDS_FILE acts as a one-element sequence of itself. + self.assertEqual(len(ct), 1) + self.assertIs(ct[0], ct) + self.assertEqual(list(iter(ct)), [ct]) + self.assertIn("ct", str(ct)) + self.assertEqual(ct, ct) + self.assertNotEqual(ct, "not a bids file") + self.assertEqual(hash(ct), hash(ct)) + + def test_rename_and_symlink_files(self): + ct = self.ct_file() + target = self.ds / "derivatives_link" / "sub-001" / "linked_ct.nii.gz" + target.parent.mkdir(parents=True, exist_ok=True) + ct.symlink_files(target, ending=".nii.gz") + self.assertTrue(target.exists()) + + moved_from = ct.file["nii.gz"] + renamed = self.ds / "rawdata" / "sub-001" / "ses-20210101" / "renamed_ct.nii.gz" + ct.rename_files(renamed, ending=".nii.gz") + self.assertTrue(renamed.exists()) + self.assertFalse(moved_from.exists()) + # restore for the other tests in this class (the dataset is class-scoped) + renamed.rename(moved_from) + + +class TestSearchquery(_DatasetTestCase): + def _flat(self) -> Searchquery: + return self.subj.new_query(flatten=True) + + @staticmethod + def _parts(query: Searchquery) -> list[str]: + return sorted(b.get("part") or b.bids_format for b in query.loop_list()) + + def test_flat_query_sees_every_file(self): + self.assertEqual(len(self._flat().candidates), 9) + + def test_filter_format(self): + q = self._flat() + q.filter_format("ct") + self.assertEqual(self._parts(q), ["ct"]) + + def test_filter_is_monotone(self): + q = self._flat() + before = len(q.candidates) + q.filter_format("dixon") + self.assertLessEqual(len(q.candidates), before) + self.assertEqual(self._parts(q), ["fat", "inphase", "outphase", "water"]) + + def test_copy_is_independent(self): + q = self._flat() + c = q.copy() + c.filter_format("ct") + self.assertEqual(len(q.candidates), 9) + self.assertEqual(len(c.candidates), 1) + + def test_flatten_unflatten_roundtrip(self): + q = self.subj.new_query() + keys_before = sorted(q.candidates) + q.flatten() + self.assertIsInstance(q.candidates, list) + q.unflatten() + self.assertEqual(sorted(q.candidates), keys_before) + + def test_filter_self_and_filter_json(self): + q = self._flat() + q.filter_self(lambda b: b.bids_format == "msk") + self.assertEqual(self._parts(q), ["msk", "msk"]) + + q = self._flat() + q.filter_json(lambda j: "EchoTime" in j, required=True) + self.assertEqual(self._parts(q), ["ct"]) + + def test_filter_by_entity(self): + q = self._flat() + q.filter("sequ", lambda v: int(v) == 301) + self.assertEqual(self._parts(q), ["ct", "ctd", "msk", "msk", "snp"]) + + def test_filter_filetype(self): + q = self._flat() + q.filter_filetype("png") + self.assertEqual(self._parts(q), ["snp"]) + + def test_dixon_channel_filters_select_one_channel_each(self): + # Regression: the ImageType filter required *all* spellings of a channel + # ("W" and "WATER"), so every file with a sidecar was dropped and these + # three filters always returned nothing. + for channel, method in (("water", "filter_dixon_water"), ("fat", "filter_dixon_fat"), ("outphase", "filter_dixon_outphase")): + with self.subTest(channel=channel): + q = self._flat() + q.filter_format("dixon") + getattr(q, method)() + self.assertEqual(self._parts(q), [channel]) + + def test_filter_dixon_only_inphase(self): + q = self._flat() + q.filter_format("dixon") + q.filter_dixon_only_inphase() + self.assertEqual(self._parts(q), ["inphase"]) + + def test_filter_non_existence(self): + q = self._flat() + q.filter_non_existence("seg", required=False) + self.assertNotIn("msk", self._parts(q)) + + def test_action_sets_a_key_when_the_json_matches(self): + q = self._flat() + q.filter_format("dixon") + q.action(action_fun=lambda b: b.set("part", "real"), filter_fun=lambda j: "IP" in j["ImageType"], key="json", required=True) + changed = [b for b in q.loop_list() if b.get("part") == "real"] + self.assertEqual(len(changed), 1) + + def test_loop_dict_yields_families(self): + q = self.subj.new_query() + q.filter_format("dixon") + fams = list(q.loop_dict()) + self.assertEqual(len(fams), 1) + self.assertIsInstance(fams[0], BIDS_Family) + self.assertEqual(sorted(fams[0].keys()), ["dixon"]) + + def test_loop_list_sorted(self): + q = self._flat() + names = [b.get_path_decomposed()[3] for b in q.loop_list(sort=True)] + self.assertEqual(names, sorted(names)) + + def test_str_is_informative(self): + self.assertIn("001", str(self._flat())) + + def test_from_bids_family_restricts_to_one_sequence(self): + fam = self.subj.get_sequence_files("20210101_301") + q = Searchquery.from_BIDS_Family(fam) + self.assertEqual(sorted(q.candidates), ["20210101_301"]) + + +class TestBidsFamily(_DatasetTestCase): + def setUp(self): + super().setUp() + self.fam = self.subj.get_sequence_files("20210101_301") + + def test_keys_carry_the_naming_suffix(self): + self.assertEqual(sorted(self.fam.keys()), ["ct", "ctd_seg-subreg", "msk_seg-subreg", "msk_seg-vert", "snp"]) + + def test_contains_accepts_str_and_list(self): + self.assertIn("ct", self.fam) + self.assertTrue(["ct", "snp"] in self.fam) + self.assertNotIn("nope", self.fam) + + def test_get_and_getitem(self): + self.assertEqual(len(self.fam["ct"]), 1) + self.assertEqual(len(self.fam.get("ct")), 1) + self.assertIsNone(self.fam.get("nope")) + self.assertEqual(self.fam.get("nope", []), []) + self.assertEqual(self.fam.get(["nope", "ct"])[0].bids_format, "ct") + + def test_len_iter_items_keys_values(self): + self.assertEqual(len(self.fam), 5) + self.assertEqual(sorted(k for k, _ in self.fam), sorted(self.fam.keys())) + self.assertEqual(len(list(self.fam.items())), 5) + self.assertEqual(sum(len(v) for v in self.fam.values()), 5) + + def test_key_and_format_len(self): + self.assertEqual(self.fam.get_key_len()["ct"], 1) + self.assertEqual(self.fam.get_format_len()["msk"], (2, 2)) + + def test_get_files_and_as_dict(self): + self.assertEqual({k: len(v) for k, v in self.fam.get_files().items()}, dict.fromkeys(self.fam.keys(), 1)) + as_dict = self.fam.get_bids_files_as_dict(["ct", "msk_seg-vert"]) + self.assertEqual(sorted(as_dict), ["ct", "msk_seg-vert"]) + self.assertEqual(as_dict["ct"].bids_format, "ct") + + def test_identifier(self): + self.assertEqual(self.fam.get_identifier(), "sub-001_ses-20210101_sequ-301") + + def test_sort_and_dunders(self): + self.fam.sort() + self.assertIn("ct", str(self.fam)) + self.assertIn("ct", repr(self.fam)) + self.assertEqual(hash(self.fam), hash(self.fam)) + + def test_new_query_from_family(self): + q = self.fam.new_query() + self.assertEqual(sorted(q.candidates), ["20210101_301"]) + + +class TestBufferedBidsGlobalInfo(_DatasetTestCase): + def test_buffer_is_written_and_reused(self): + cache = self.ds / "rawdata" / ".filepaths" + self.assertFalse(cache.exists()) + first = Buffered_BIDS_Global_info([self.ds], parents=["rawdata", "derivatives"], verbose=False) + self.assertTrue(cache.exists(), "the file-path cache should be written next to the scanned parent") + + # Second call reads the cache instead of walking the tree and must agree exactly. + second = Buffered_BIDS_Global_info([self.ds], parents=["rawdata", "derivatives"], verbose=False) + self.assertEqual(sorted(first.subjects), sorted(second.subjects)) + self.assertEqual( + {k: sorted(v.sequences) for k, v in first.subjects.items()}, + {k: sorted(v.sequences) for k, v in second.subjects.items()}, + ) + + def test_stale_buffer_is_recomputed(self): + import os # noqa: PLC0415 + + Buffered_BIDS_Global_info([self.ds], parents=["rawdata"], verbose=False) + cache = self.ds / "rawdata" / ".filepaths" + os.utime(cache, (0, 0)) # 1970 -> older than max_age_days + g = Buffered_BIDS_Global_info([self.ds], parents=["rawdata"], verbose=False, max_age_days=1) + self.assertGreater(cache.stat().st_mtime, 0) + self.assertEqual(sorted(g.subjects), ["001", "002"]) + + def test_filter_file_restricts_the_cached_paths(self): + g = Buffered_BIDS_Global_info( + [self.ds], + parents=["rawdata"], + verbose=False, + filter_file=lambda p: "sub-002" in str(p), + recompute_parents=["rawdata"], + ) + self.assertEqual(sorted(g.subjects), ["002"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_compat.py b/unit_tests/test_compat.py index d8eabe5..0194c41 100644 --- a/unit_tests/test_compat.py +++ b/unit_tests/test_compat.py @@ -4,9 +4,11 @@ # coverage html from __future__ import annotations +import re import unittest +from pathlib import Path -from TPTBox.core.compat import zip_strict +from TPTBox.core.compat import _pairwise_py39, pairwise, zip_strict class TestZipStrict(unittest.TestCase): @@ -44,3 +46,57 @@ def test_generator_iterables(self): expected = [(0, "a"), (1, "b"), (2, "c")] result = list(zip_strict(a, b)) self.assertEqual(result, expected) + + +class TestPairwise(unittest.TestCase): + """`pairwise` is a 3.10+ stdlib function; TPTBox supports 3.9, so it is backported.""" + + def test_pairs_of_list(self): + self.assertEqual(list(pairwise([1, 2, 3, 4])), [(1, 2), (2, 3), (3, 4)]) + + def test_too_short_inputs_are_empty(self): + for seq in ([], [1]): + with self.subTest(seq=seq): + self.assertEqual(list(pairwise(seq)), []) + + def test_generator_input_is_consumed_lazily(self): + self.assertEqual(list(pairwise(x * x for x in range(4))), [(0, 1), (1, 4), (4, 9)]) + + def test_py39_fallback_matches_stdlib_path(self): + # Both code paths must agree; on 3.10+ `pairwise` delegates to itertools. + seq = ["a", "b", "c", "d"] + self.assertEqual(list(_pairwise_py39(seq)), list(pairwise(seq))) + + def test_py39_fallback_empty(self): + self.assertEqual(list(_pairwise_py39([])), []) + + +class TestNoPython310OnlyStdlibAtImport(unittest.TestCase): + """Guards the 3.9 CI leg: a 3.10+ stdlib import at module scope kills the whole module.""" + + def test_no_itertools_pairwise_import(self): + root = Path(__file__).resolve().parents[1] / "TPTBox" + compat = root / "core" / "compat.py" # the one sanctioned place: the backport itself + offenders = [ + f"{path.relative_to(root.parent)}:{i}" + for path in root.rglob("*.py") + if path != compat + for i, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1) + if "from itertools import" in line and "pairwise" in line + ] + self.assertEqual(offenders, [], "use TPTBox.core.compat.pairwise instead (python 3.9 support)") + + def test_no_zip_strict_keyword(self): + # `zip(..., strict=...)` is 3.10+ and raises TypeError at call time on 3.9. + root = Path(__file__).resolve().parents[1] / "TPTBox" + offenders = [ + f"{path.relative_to(root.parent)}:{i}" + for path in root.rglob("*.py") + for i, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1) + if re.search(r"\bzip\([^)]*\bstrict\s*=", line) + ] + self.assertEqual(offenders, [], "use TPTBox.core.compat.zip_strict instead (python 3.9 support)") + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_dicom.py b/unit_tests/test_dicom.py new file mode 100644 index 0000000..df879d6 --- /dev/null +++ b/unit_tests/test_dicom.py @@ -0,0 +1,424 @@ +"""Unit tests for `TPTBox.core.dicom`. + +No DICOM sample ships with the repository, so the integration test generates one: +the sample CT is written out as a DICOM series with `nii2dicom` and read back with +`extract_dicom_folder`, which exercises the whole naming + conversion pipeline and +pins the BIDS path it produces. The pure string/dict/path helpers around it are +tested directly. +""" + +from __future__ import annotations + +import importlib.util +import json +import tempfile +import unittest +import zipfile +from pathlib import Path + +import numpy as np + +HAS_PYDICOM = importlib.util.find_spec("pydicom") is not None and importlib.util.find_spec("dicom2nifti") is not None +PYDICOM_REASON = "requires the `dicom` extra (pydicom + dicom2nifti)" + + +@unittest.skipUnless(HAS_PYDICOM, PYDICOM_REASON) +class TestNameHelpers(unittest.TestCase): + def test_next_letter_suffix(self): + from TPTBox.core.dicom.dicom_extract import _next_letter_suffix # noqa: PLC0415 + + cases = {"a": "b", "b": "c", "z": "aa", "aa": "ab", "az": "ba", "zz": "aaa"} + for value, expected in cases.items(): + with self.subTest(value=value): + self.assertEqual(_next_letter_suffix(value), expected) + + def test_next_letter_suffix_with_a_larger_increment(self): + from TPTBox.core.dicom.dicom_extract import _next_letter_suffix # noqa: PLC0415 + + self.assertEqual(_next_letter_suffix("a", inc=2), "c") + self.assertEqual(_next_letter_suffix("a", inc=26), "aa") + + def test_inc_key_appends_a_letter_suffix(self): + from TPTBox.core.dicom.dicom_extract import _inc_key # noqa: PLC0415 + + keys = {"sub": "001", "sequ": "7"} + _inc_key(keys) + self.assertEqual(keys["sequ"], "7-a") + _inc_key(keys) + self.assertEqual(keys["sequ"], "7-b") + + def test_inc_key_keeps_going_until_the_path_is_free(self): + from TPTBox.core.dicom.dicom_extract import _inc_key # noqa: PLC0415 + + taken = {"7-a", "7-b"} + keys = {"sub": "001", "sequ": "7"} + _inc_key(keys, path_exists=lambda k: k["sequ"] in taken) + self.assertEqual(keys["sequ"], "7-c") + + def test_inc_key_on_a_custom_key(self): + from TPTBox.core.dicom.dicom_extract import _inc_key # noqa: PLC0415 + + keys = {"sub": "001", "acq": "iso"} + _inc_key(keys, k="acq") + self.assertEqual(keys["acq"], "iso-a") + + def test_looks_like_dicom(self): + from TPTBox.core.dicom.dicom_extract import _looks_like_dicom # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + folder = Path(tmp) + text = folder / "notes.txt" + text.write_text("hello", encoding="utf-8") + self.assertFalse(_looks_like_dicom(text)) + + dicom = folder / "slice.dcm" + dicom.write_bytes(b"\0" * 128 + b"DICM" + b"\0" * 32) + self.assertTrue(_looks_like_dicom(dicom)) + + def test_with_echo_suffix_extends_the_part_entity(self): + from TPTBox.core.dicom.dicom_extract import _with_echo_suffix # noqa: PLC0415 + + out = _with_echo_suffix(Path("/ds/sub-1_sequ-3_part-water_dixon.nii.gz"), 0) + self.assertEqual(out.name, "sub-1_sequ-3_part-water-eco0_dixon.nii.gz") + self.assertEqual(out.parent, Path("/ds")) + + def test_with_echo_suffix_falls_back_to_the_stem(self): + from TPTBox.core.dicom.dicom_extract import _with_echo_suffix # noqa: PLC0415 + + for name, expected in ( + ("sub-1_sequ-3_mevibe.nii.gz", "sub-1_sequ-3_mevibe-eco2.nii.gz"), + ("sub-1_sequ-3_mevibe.json", "sub-1_sequ-3_mevibe-eco2.json"), + ): + with self.subTest(name=name): + self.assertEqual(_with_echo_suffix(Path("/ds") / name, 2).name, expected) + + def test_generate_bids_path_is_bids_shaped(self): + from TPTBox.core.dicom.dicom_extract import _generate_bids_path # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + keys = {"sub": "TESTPAT", "ses": "20200101", "sequ": "7"} + json_path, bids = _generate_bids_path(Path(tmp), keys, "ct", {"a": 1}, make_subject_chunks=0) + rel = Path(json_path).relative_to(tmp).as_posix() + self.assertEqual(rel, "rawdata/sub-TESTPAT/ses-20200101/ct/sub-TESTPAT_ses-20200101_sequ-7_ct.json") + self.assertEqual(bids.bids_format, "ct") + + def test_generate_bids_path_make_subject_chunks(self): + from TPTBox.core.dicom.dicom_extract import _generate_bids_path # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + json_path, _ = _generate_bids_path(Path(tmp), {"sub": "123456"}, "ct", {}, make_subject_chunks=3) + self.assertIn("/123/sub-123456/", Path(json_path).as_posix()) + + def test_generate_bids_path_sanitises_dots_in_the_subject(self): + from TPTBox.core.dicom.dicom_extract import _generate_bids_path # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + _, bids = _generate_bids_path(Path(tmp), {"sub": "1.2.3"}, "ct", {}, make_subject_chunks=0) + self.assertEqual(bids.get("sub"), "1-2-3") + + def test_filter_file_type_keeps_the_tokens_that_separate_sub_series(self): + from TPTBox.core.dicom.dicom_extract import _filter_file_type # noqa: PLC0415 + + # Two ImageTypes in one series: "W" and "F" each appear once and therefore + # identify their sub-series, while the shared tokens are dropped. + out = _filter_file_type({"uid1": ["ORIGINAL,PRIMARY,W", "ORIGINAL,PRIMARY,F"]}) + self.assertEqual(sorted(out), ["uid1_ORIGINAL,PRIMARY,F", "uid1_ORIGINAL,PRIMARY,W"]) + self.assertEqual(out["uid1_ORIGINAL,PRIMARY,W"], ["W"]) + self.assertEqual(out["uid1_ORIGINAL,PRIMARY,F"], ["F"]) + + def test_filter_file_type_skips_series_with_a_single_image_type(self): + from TPTBox.core.dicom.dicom_extract import _filter_file_type # noqa: PLC0415 + + self.assertEqual(_filter_file_type({"uid1": ["ORIGINAL,PRIMARY"]}), {}) + + +@unittest.skipUnless(HAS_PYDICOM, PYDICOM_REASON) +class TestJsonHelpers(unittest.TestCase): + def test_save_and_load_roundtrip(self): + from TPTBox.core.dicom.dicom2nii_utils import load_json, save_json # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "out.json" + payload = {"a": 1, "b": [1.5, 2.5], "c": "text"} + self.assertFalse(save_json(payload, path)) + self.assertEqual(load_json(path), payload) + + def test_numpy_scalars_are_serialised(self): + from TPTBox.core.dicom.dicom2nii_utils import load_json, save_json # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "np.json" + save_json({"i": np.int64(3), "f": np.float32(1.5), "arr": np.arange(3)}, path) + loaded = load_json(path) + self.assertEqual(loaded["i"], 3) + self.assertAlmostEqual(loaded["f"], 1.5, places=6) + self.assertEqual(list(loaded["arr"]), [0, 1, 2]) + + def test_override_false_keeps_the_existing_file(self): + from TPTBox.core.dicom.dicom2nii_utils import load_json, save_json # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "out.json" + save_json({"first": True}, path) + self.assertTrue(save_json({"second": True}, path, override=False)) + self.assertEqual(load_json(path), {"first": True}) + + def test_check_exist_raises_on_differing_content(self): + from TPTBox.core.dicom.dicom2nii_utils import save_json # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "out.json" + save_json({"first": True}, path) + with self.assertRaises(FileExistsError): + save_json({"second": True}, path, check_exist=True) + + def test_name_conflict_detects_only_real_differences(self): + from TPTBox.core.dicom.dicom2nii_utils import save_json, test_name_conflict # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "out.json" + self.assertFalse(test_name_conflict({"a": 1}, path)) # nothing there yet + save_json({"a": 1}, path) + self.assertFalse(test_name_conflict({"a": 1}, path)) + self.assertTrue(test_name_conflict({"a": 2}, path)) + + def test_name_conflict_ignores_the_grid_key(self): + from TPTBox.core.dicom.dicom2nii_utils import save_json, test_name_conflict # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "out.json" + save_json({"a": 1, "grid": {"shape": [1, 2, 3]}}, path) + self.assertFalse(test_name_conflict({"a": 1}, path)) + + def test_name_conflict_replaces_unreadable_json(self): + from TPTBox.core.dicom.dicom2nii_utils import test_name_conflict # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "broken.json" + path.write_text("{not json", encoding="utf-8") + self.assertFalse(test_name_conflict({"a": 1}, path)) + self.assertFalse(path.exists(), "a half-written json should be removed, not routed around") + + def test_replace_birthdate_with_age(self): + from TPTBox.core.dicom.dicom2nii_utils import replace_birthdate_with_age # noqa: PLC0415 + + out = replace_birthdate_with_age({"00100030": {"Value": ["19800115"], "vr": "DA"}, "00080020": {"Value": ["20200114"], "vr": "DA"}}) + self.assertNotIn("00100030", out, "the birth date must be dropped") + self.assertEqual(out["00101010"]["Value"], ["039Y"], "age on the day before the 40th birthday") + + def test_replace_birthdate_on_the_birthday_itself(self): + from TPTBox.core.dicom.dicom2nii_utils import replace_birthdate_with_age # noqa: PLC0415 + + out = replace_birthdate_with_age({"00100030": {"Value": ["19800115"]}, "00080020": {"Value": ["20200115"]}}) + self.assertEqual(out["00101010"]["Value"], ["040Y"]) + + def test_replace_birthdate_without_a_birthdate_is_a_noop(self): + from TPTBox.core.dicom.dicom2nii_utils import replace_birthdate_with_age # noqa: PLC0415 + + payload = {"00080020": {"Value": ["20200115"]}} + self.assertEqual(replace_birthdate_with_age(dict(payload)), payload) + + def test_replace_birthdate_with_unparsable_dates_is_a_noop(self): + from TPTBox.core.dicom.dicom2nii_utils import replace_birthdate_with_age # noqa: PLC0415 + + payload = {"00100030": {"Value": ["not-a-date"]}, "00080020": {"Value": ["20200115"]}} + self.assertEqual(replace_birthdate_with_age(dict(payload)), payload) + + +@unittest.skipUnless(HAS_PYDICOM, PYDICOM_REASON) +class TestPlaneDetection(unittest.TestCase): + """`get_plane_dicom` classifies an NII by its affine, with no DICOM involved.""" + + @staticmethod + def _nii(zoom): + import nibabel as nib # noqa: PLC0415 + + from TPTBox import NII # noqa: PLC0415 + + return NII(nib.Nifti1Image(np.zeros((16, 16, 16), dtype=np.uint8), np.diag([*zoom, 1.0])), seg=False) + + def test_isotropic_and_anisotropic_planes(self): + from TPTBox.core.dicom.dicom_header_to_keys import get_plane_dicom # noqa: PLC0415 + + cases = { + (1.0, 1.0, 1.0): "iso", + (0.5, 0.5, 6.0): "ax", + } + for zoom, expected in cases.items(): + with self.subTest(zoom=zoom): + self.assertEqual(get_plane_dicom(self._nii(zoom)), expected) + + def test_each_thick_axis_gives_a_different_plane(self): + from TPTBox.core.dicom.dicom_header_to_keys import get_plane_dicom # noqa: PLC0415 + + planes = {get_plane_dicom(self._nii(zoom)) for zoom in ((6.0, 0.5, 0.5), (0.5, 6.0, 0.5), (0.5, 0.5, 6.0))} + self.assertEqual(len(planes), 3, planes) + + +def _save_sample_ct(folder: Path) -> Path: + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + path = folder / "input.nii.gz" + get_test_ct()[0].save(path, verbose=False) + return path + + +def _meta() -> dict: + return { + "Modality": "CT", + "PatientID": "TESTPAT", + "SeriesDescription": "abdomen", + "SeriesNumber": 7, + "PatientBirthDate": "19800101", + "StudyDate": "20200101", + } + + +@unittest.skipUnless(HAS_PYDICOM, PYDICOM_REASON) +class TestDicomRoundTrip(unittest.TestCase): + """nii -> DICOM series -> BIDS nii, through the real conversion pipeline.""" + + @classmethod + def setUpClass(cls): + from TPTBox.core.dicom.dicom_extract import extract_dicom_folder # noqa: PLC0415 + from TPTBox.core.dicom.nii2dicom import nifti2dicom_1file # noqa: PLC0415 + + cls._tmp = tempfile.TemporaryDirectory() + root = Path(cls._tmp.name) + nii_path = _save_sample_ct(root) + (root / "input.json").write_text(json.dumps(_meta()), encoding="utf-8") + cls.dcm_dir = root / "dcm" + nifti2dicom_1file(nii_path, cls.dcm_dir) + cls.slices = sorted(cls.dcm_dir.rglob("*.dcm")) + + cls.out_dir = root / "ds" + cls.result = extract_dicom_folder(cls.dcm_dir, cls.out_dir, make_subject_chunks=0, verbose=False) + cls.source = nii_path + + @classmethod + def tearDownClass(cls): + cls._tmp.cleanup() + + def test_one_dicom_slice_per_input_slice(self): + from TPTBox import to_nii # noqa: PLC0415 + + self.assertEqual(len(self.slices), to_nii(self.source).shape[-1]) + + def test_the_generated_series_is_readable_dicom(self): + import pydicom # noqa: PLC0415 + + ds = pydicom.dcmread(str(self.slices[0])) + self.assertEqual(ds.Modality, "CT") + self.assertEqual(str(ds.PatientID), "TESTPAT") + + def test_extraction_produces_the_expected_bids_path(self): + produced = sorted(p.relative_to(self.out_dir).as_posix() for p in self.out_dir.rglob("*") if p.is_file()) + self.assertIn("rawdata/sub-TESTPAT/ct/sub-TESTPAT_sequ-7_acq-iso_ct.nii.gz", produced) + self.assertIn("rawdata/sub-TESTPAT/ct/sub-TESTPAT_sequ-7_acq-iso_ct.json", produced) + + def test_the_return_value_maps_series_uid_to_the_output(self): + self.assertIsInstance(self.result, dict) + self.assertEqual(len(self.result), 1) + self.assertTrue(next(iter(self.result.values())).endswith(".nii.gz")) + + def test_the_round_tripped_volume_matches_the_input(self): + from TPTBox import to_nii # noqa: PLC0415 + + source = to_nii(self.source) + out = to_nii(self.out_dir / "rawdata/sub-TESTPAT/ct/sub-TESTPAT_sequ-7_acq-iso_ct.nii.gz") + self.assertEqual(sorted(out.shape), sorted(source.shape)) + np.testing.assert_allclose(sorted(out.zoom), sorted(source.zoom), rtol=1e-3) + self.assertAlmostEqual(float(out.max()), float(source.max()), delta=1.0) + + def test_the_sidecar_carries_grid_info_and_the_patient_age(self): + sidecar = json.loads((self.out_dir / "rawdata/sub-TESTPAT/ct/sub-TESTPAT_sequ-7_acq-iso_ct.json").read_text()) + self.assertIn("grid", sidecar) + flat = json.dumps(sidecar) + self.assertNotIn("19800101", flat, "the birth date must not be written to disk") + + def test_a_second_extraction_is_a_cached_no_op(self): + from TPTBox.core.dicom.dicom_extract import extract_dicom_folder # noqa: PLC0415 + + marker = list((self.out_dir / ".extract_cache").glob("*.json")) + self.assertEqual(len(marker), 1) + before = {p: p.stat().st_mtime_ns for p in self.out_dir.rglob("*.nii.gz")} + extract_dicom_folder(self.dcm_dir, self.out_dir, make_subject_chunks=0, verbose=False) + after = {p: p.stat().st_mtime_ns for p in self.out_dir.rglob("*.nii.gz")} + self.assertEqual(before, after, "the extract cache should prevent a rewrite") + + def test_a_zipped_series_extracts_the_same_way(self): + from TPTBox.core.dicom.dicom_extract import extract_dicom_folder # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + zip_path = Path(tmp) / "series.zip" + with zipfile.ZipFile(zip_path, "w") as zf: + for slice_path in self.slices: + zf.write(slice_path, arcname=slice_path.name) + out = Path(tmp) / "ds" + extract_dicom_folder(zip_path, out, make_subject_chunks=0, verbose=False) + produced = [p.name for p in out.rglob("*.nii.gz")] + self.assertEqual(produced, ["sub-TESTPAT_sequ-7_acq-iso_ct.nii.gz"]) + + +@unittest.skipUnless(HAS_PYDICOM, PYDICOM_REASON) +class TestDicomGrouping(unittest.TestCase): + """Series/echo grouping works on in-memory datasets, no files involved.""" + + @staticmethod + def _ds(series_uid: str, echo: int | None = None, instance: int = 1): + from pydicom.dataset import Dataset # noqa: PLC0415 + + ds = Dataset() + ds.SeriesInstanceUID = series_uid + ds.InstanceNumber = instance + ds.ImagePositionPatient = [0.0, 0.0, float(instance)] + ds.ImageOrientationPatient = [1.0, 0.0, 0.0, 0.0, 1.0, 0.0] + if echo is not None: + ds.EchoNumbers = echo + return ds + + def test_split_by_echo_numbers_returns_the_lowest_echo_for_stack_detection(self): + from TPTBox.core.dicom.dicom_extract import _split_by_echo_numbers # noqa: PLC0415 + + data = [self._ds("1", echo=1, instance=1), self._ds("1", echo=1, instance=2), self._ds("1", echo=2, instance=3)] + single, per_echo = _split_by_echo_numbers(data) + self.assertEqual(len(single), 2, "the subset must be the lowest echo, one DICOM per slice") + self.assertEqual(sorted(per_echo), [1, 2]) + self.assertEqual(sum(len(v) for v in per_echo.values()), len(data)) + + def test_split_by_echo_numbers_is_a_noop_for_a_single_echo(self): + from TPTBox.core.dicom.dicom_extract import _split_by_echo_numbers # noqa: PLC0415 + + for data in ([self._ds("1"), self._ds("1", instance=2)], [self._ds("1", echo=1), self._ds("1", echo=1, instance=2)]): + with self.subTest(n_echo_tags=sum(hasattr(d, "EchoNumbers") for d in data)): + single, per_echo = _split_by_echo_numbers(data) + self.assertEqual(len(single), 2) + self.assertEqual(per_echo, {}) + + def test_split_by_echo_numbers_treats_a_missing_tag_as_echo_zero(self): + from TPTBox.core.dicom.dicom_extract import _split_by_echo_numbers # noqa: PLC0415 + + data = [self._ds("1"), self._ds("1", echo=1, instance=2), self._ds("1", echo=2, instance=3)] + _, per_echo = _split_by_echo_numbers(data) + self.assertEqual(sorted(per_echo), [0, 1, 2]) + + def test_filter_dicom_keeps_the_input_when_nothing_is_wrong(self): + from TPTBox.core.dicom.dicom_extract import _filter_dicom # noqa: PLC0415 + + data = [self._ds("1"), self._ds("1", instance=2)] + self.assertEqual(len(_filter_dicom(data)), 2) + + def test_clean_dicom_data_drops_pixel_data(self): + from TPTBox.core.dicom.dicom2nii_utils import clean_dicom_data # noqa: PLC0415 + + ds = self._ds("1") + ds.PatientID = "X" + ds.PixelData = b"\x00\x01\x02\x03" + out = clean_dicom_data(ds) + self.assertNotIn("7FE00010", out, "PixelData must not survive into the sidecar") + self.assertIn("00100020", out) # PatientID + self.assertTrue(hasattr(ds, "PixelData"), "the input dataset must not be modified") + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_dicom_renamer.py b/unit_tests/test_dicom_renamer.py new file mode 100644 index 0000000..1d5a02e --- /dev/null +++ b/unit_tests/test_dicom_renamer.py @@ -0,0 +1,494 @@ +"""Unit tests for `TPTBox.core.dicom.dicom_renamer` and `xa_pairing`. + +Both modules rewrite an already-extracted dataset, and both offer a `dry_run` +mode that returns the full plan without touching disk - so the plan itself is +the assertion, and the destructive path is then checked against it on a copy. +""" + +from __future__ import annotations + +import csv +import importlib.util +import json +import tempfile +import unittest +from pathlib import Path + +import nibabel as nib +import numpy as np + +HAS_PYDICOM = importlib.util.find_spec("pydicom") is not None and importlib.util.find_spec("dicom2nifti") is not None +PYDICOM_REASON = "requires the dicom extra (pydicom + dicom2nifti)" + +GRID = {"spacing": [0.5, 0.5, 6.0], "orientation": ["R", "A", "S"], "shape": [8, 8, 4]} + + +def _write_pair(folder: Path, stem: str, sidecar: dict) -> Path: + """Write a `.nii.gz` + `.json` pair and return the json path.""" + folder.mkdir(parents=True, exist_ok=True) + affine = np.diag([0.5, 0.5, 6.0, 1.0]) + nib.save(nib.Nifti1Image(np.zeros((8, 8, 4), dtype=np.uint8), affine), folder / f"{stem}.nii.gz") + path = folder / f"{stem}.json" + path.write_text(json.dumps(sidecar), encoding="utf-8") + return path + + +def _sidecar(series_number: int = 7, study_date: str = "20200101", **extra) -> dict: + out = { + "grid": dict(GRID), + "Modality": "MR", + "SeriesNumber": series_number, + "SeriesDescription": "t2_tse_sag", + "StudyDate": study_date, + "PatientID": "RAW_ID", + } + out.update(extra) + return out + + +def _dataset(root: Path, subjects=("raw_1", "raw_2")) -> Path: + ds = root / "dataset-x" + for i, sub in enumerate(subjects, start=1): + _write_pair(ds / "rawdata" / f"sub-{sub}" / "ses-20200101", f"sub-{sub}_ses-20200101_sequ-{i}_T2w", _sidecar(i)) + return ds + + +@unittest.skipUnless(HAS_PYDICOM, PYDICOM_REASON) +class TestPlaneFromGrid(unittest.TestCase): + def test_plane_from_the_thick_axis(self): + from TPTBox.core.dicom.dicom_renamer import _plane_from_grid # noqa: PLC0415 + + cases = [ + ({"spacing": [0.5, 0.5, 6.0], "orientation": ["R", "A", "S"]}, "ax"), + ({"spacing": [6.0, 0.5, 0.5], "orientation": ["R", "A", "S"]}, "sag"), + ({"spacing": [0.5, 6.0, 0.5], "orientation": ["R", "A", "S"]}, "cor"), + ({"spacing": [1.0, 1.0, 1.0], "orientation": ["R", "A", "S"]}, "iso"), + ] + for grid, expected in cases: + with self.subTest(spacing=grid["spacing"]): + self.assertEqual(_plane_from_grid(grid), expected) + + def test_hires_threshold_clamps_thin_slices(self): + from TPTBox.core.dicom.dicom_renamer import _plane_from_grid # noqa: PLC0415 + + grid = {"spacing": [0.3, 0.3, 0.75], "orientation": ["R", "A", "S"]} + # With the default 0.8 clamp all three zooms become equal -> isotropic. + self.assertEqual(_plane_from_grid(grid), "iso") + # Without the clamp the 0.75 axis is the thickest one. + self.assertEqual(_plane_from_grid(grid, hires_threshold=None), "ax") + + def test_malformed_grids_return_none(self): + from TPTBox.core.dicom.dicom_renamer import _plane_from_grid # noqa: PLC0415 + + for grid in ({}, None, {"spacing": [1.0]}, {"spacing": [1, 1, 1]}, {"orientation": ["R", "A", "S"]}): + with self.subTest(grid=grid): + self.assertIsNone(_plane_from_grid(grid)) + + def test_zero_spacing_is_treated_as_one(self): + from TPTBox.core.dicom.dicom_renamer import _plane_from_grid # noqa: PLC0415 + + self.assertEqual(_plane_from_grid({"spacing": [0.0, 0.0, 6.0], "orientation": ["R", "A", "S"]}), "ax") + + def test_agrees_with_get_plane_dicom_on_the_same_grid(self): + from TPTBox import NII # noqa: PLC0415 + from TPTBox.core.dicom.dicom_header_to_keys import get_plane_dicom # noqa: PLC0415 + from TPTBox.core.dicom.dicom_renamer import _plane_from_grid # noqa: PLC0415 + + for zoom in ((0.5, 0.5, 6.0), (6.0, 0.5, 0.5), (1.0, 1.0, 1.0)): + with self.subTest(zoom=zoom): + nii = NII(nib.Nifti1Image(np.zeros((8, 8, 8), dtype=np.uint8), np.diag([*zoom, 1.0])), seg=False) + grid = {"spacing": list(zoom), "orientation": list(nii.orientation)} + self.assertEqual(_plane_from_grid(grid), get_plane_dicom(nii)) + + +@unittest.skipUnless(HAS_PYDICOM, PYDICOM_REASON) +class TestEntityHelpers(unittest.TestCase): + def test_sanitize_entity_value(self): + from TPTBox.core.dicom.dicom_renamer import _sanitize_entity_value # noqa: PLC0415 + + cases = {"abc": "abc", "a_b": "ab", "__ab__": "ab", "-abc": "abc", "_": "unnamed", "": "unnamed", "-": "unnamed"} + for raw, expected in cases.items(): + with self.subTest(raw=raw): + self.assertEqual(_sanitize_entity_value(raw), expected) + + def test_sanitized_values_survive_bids_parsing(self): + from TPTBox.core.bids_files import get_values_from_name # noqa: PLC0415 + from TPTBox.core.dicom.dicom_renamer import _sanitize_entity_value # noqa: PLC0415 + + for raw in ("a_b", "-weird", "ok123"): + with self.subTest(raw=raw): + value = _sanitize_entity_value(raw) + _, info, _, _ = get_values_from_name(f"sub-{value}_ct.nii.gz", verbose=False) + self.assertEqual(info["sub"], value) + + def test_insert_run_entity_after_sequ(self): + from TPTBox.core.dicom.dicom_renamer import _insert_run_entity # noqa: PLC0415 + + out = _insert_run_entity(Path("/ds/sub-1_ses-2_sequ-3_T2w.nii.gz"), 2) + self.assertEqual(out.name, "sub-1_ses-2_sequ-3_run-2_T2w.nii.gz") + + def test_insert_run_entity_without_sequ_falls_back(self): + from TPTBox.core.dicom.dicom_renamer import _insert_run_entity # noqa: PLC0415 + + out = _insert_run_entity(Path("/ds/sub-1_T2w.nii.gz"), 3) + self.assertIn("run-3", out.name) + + def test_inserted_run_is_parsed_back(self): + from TPTBox.core.bids_files import get_values_from_name # noqa: PLC0415 + from TPTBox.core.dicom.dicom_renamer import _insert_run_entity # noqa: PLC0415 + + out = _insert_run_entity(Path("/ds/sub-1_ses-2_sequ-3_T2w.nii.gz"), 5) + _, info, _, _ = get_values_from_name(out, verbose=False) + self.assertEqual(info["run"], "5") + self.assertEqual(info["sequ"], "3") + + def test_files_identical(self): + from TPTBox.core.dicom.dicom_renamer import _files_identical # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + a, b, c, missing = (Path(tmp) / n for n in ("a", "b", "c", "missing")) + a.write_bytes(b"same") + b.write_bytes(b"same") + c.write_bytes(b"different") + self.assertTrue(_files_identical(a, b)) + self.assertFalse(_files_identical(a, c)) + self.assertFalse(_files_identical(a, missing)) + + +@unittest.skipUnless(HAS_PYDICOM, PYDICOM_REASON) +class TestSubjectMap(unittest.TestCase): + @staticmethod + def _folders(names): + return [Path(f"/ds/rawdata/sub-{n}") for n in names] + + def test_sanitise_in_place_without_a_prefix(self): + from TPTBox.core.dicom.dicom_renamer import _build_subject_map # noqa: PLC0415 + + out = _build_subject_map(self._folders(["ok1", "bad_id", "-lead"]), None, 3) + self.assertEqual(out["ok1"], "ok1") + self.assertEqual(out["bad_id"], "badid") + self.assertEqual(out["-lead"], "lead") + + def test_prefix_mode_renumbers_in_order(self): + from TPTBox.core.dicom.dicom_renamer import _build_subject_map # noqa: PLC0415 + + out = _build_subject_map(self._folders(["zzz", "aaa"]), "ID", 3) + self.assertEqual(sorted(out.values()), ["ID001", "ID002"]) + self.assertEqual(len(set(out.values())), 2, "the mapping must stay injective") + + def test_prefix_itself_is_sanitised(self): + from TPTBox.core.dicom.dicom_renamer import _build_subject_map # noqa: PLC0415 + + out = _build_subject_map(self._folders(["a"]), "I_D", 2) + self.assertEqual(list(out.values()), ["ID01"]) + + def test_rerunning_keeps_previously_assigned_ids(self): + from TPTBox.core.dicom.dicom_renamer import _build_subject_map # noqa: PLC0415 + + first = _build_subject_map(self._folders(["aaa", "bbb"]), "ID", 3) + # Second pass sees the already-renamed folders plus one new subject. + second = _build_subject_map(self._folders([*first.values(), "ccc"]), "ID", 3, existing_map=first) + for old, new in first.items(): + with self.subTest(old=old): + self.assertEqual(second[old], new, "historic rows must be preserved") + self.assertEqual(second[new], new, "an already-renamed folder maps to itself") + self.assertEqual(second["ccc"], "ID003", "a fresh subject takes the next free slot") + + def test_existing_map_roundtrips_through_the_tsv(self): + from TPTBox.core.dicom.dicom_renamer import _read_existing_subject_map # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + self.assertEqual(_read_existing_subject_map(root, "info"), {}) + info = root / "info" + info.mkdir() + with (info / "subject_map.tsv").open("w", encoding="utf-8", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=["old_sub", "new_sub"], delimiter="\t") + writer.writeheader() + writer.writerow({"old_sub": "raw_1", "new_sub": "ID001"}) + self.assertEqual(_read_existing_subject_map(root, "info"), {"raw_1": "ID001"}) + + +@unittest.skipUnless(HAS_PYDICOM, PYDICOM_REASON) +class TestRerunBidsNaming(unittest.TestCase): + def test_dry_run_plans_moves_without_touching_disk(self): + from TPTBox.core.dicom.dicom_renamer import rerun_bids_naming # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + ds = _dataset(Path(tmp)) + before = sorted(p.relative_to(ds).as_posix() for p in ds.rglob("*") if p.is_file()) + out = rerun_bids_naming(ds, subject_prefix="ID", dry_run=True, verbose=False) + after = sorted(p.relative_to(ds).as_posix() for p in ds.rglob("*") if p.is_file()) + self.assertEqual(before, after, "a dry run must not move anything") + self.assertTrue(out["moves"], "the plan should not be empty") + for src, dst in out["moves"]: + with self.subTest(src=Path(src).name): + self.assertTrue(Path(src).exists()) + self.assertIn("sub-ID", Path(dst).as_posix()) + + def test_applying_the_plan_produces_exactly_the_planned_files(self): + from TPTBox.core.dicom.dicom_renamer import rerun_bids_naming # noqa: PLC0415 + + # Regression: the leftover sweep (pass 2) scans the disk, so under dry_run it + # also reported every file pass 1 had only *planned* to move - with a different, + # fallback destination. The plan listed 6 moves where the real run performed 4. + with tempfile.TemporaryDirectory() as tmp: + ds = _dataset(Path(tmp)) + plan = rerun_bids_naming(ds, subject_prefix="ID", dry_run=True, verbose=False)["moves"] + applied = rerun_bids_naming(ds, subject_prefix="ID", dry_run=False, verbose=False)["moves"] + self.assertEqual(len(plan), len(applied), "the dry run must plan exactly what the real run does") + self.assertEqual( + sorted(Path(d).name for _, d in plan), + sorted(Path(d).name for _, d in applied), + ) + for _, dst in plan: + with self.subTest(dst=Path(dst).name): + self.assertTrue(Path(dst).exists(), dst) + + def test_the_translation_table_is_written(self): + from TPTBox.core.dicom.dicom_renamer import rerun_bids_naming # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + ds = _dataset(Path(tmp)) + out = rerun_bids_naming(ds, subject_prefix="ID", dry_run=False, verbose=False) + table = Path(out["mapping_file"]) + self.assertTrue(table.is_file()) + text = table.read_text(encoding="utf-8") + self.assertIn("old_sub", text) + self.assertIn("ID001", text) + + def test_a_second_run_is_stable(self): + from TPTBox.core.dicom.dicom_renamer import rerun_bids_naming # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + ds = _dataset(Path(tmp)) + rerun_bids_naming(ds, subject_prefix="ID", dry_run=False, verbose=False) + first = sorted(p.relative_to(ds).as_posix() for p in ds.rglob("*.nii.gz")) + rerun_bids_naming(ds, subject_prefix="ID", dry_run=False, verbose=False) + second = sorted(p.relative_to(ds).as_posix() for p in ds.rglob("*.nii.gz")) + self.assertEqual(first, second, "re-running must be idempotent") + + def test_sanitise_mode_fixes_illegal_subject_ids(self): + from TPTBox.core.dicom.dicom_renamer import rerun_bids_naming # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + ds = _dataset(Path(tmp), subjects=("bad_id",)) + rerun_bids_naming(ds, dry_run=False, verbose=False) + subs = sorted(p.name for p in (ds / "rawdata").iterdir() if p.is_dir()) + self.assertEqual(subs, ["sub-badid"]) + + def test_flat_dataset_without_a_parent_folder(self): + from TPTBox.core.dicom.dicom_renamer import rerun_bids_naming # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + ds = Path(tmp) / "flat" + _write_pair(ds / "sub-raw_1" / "ses-20200101", "sub-raw_1_ses-20200101_sequ-1_T2w", _sidecar(1)) + out = rerun_bids_naming(ds, parent=None, subject_prefix="ID", dry_run=True, verbose=False) + self.assertTrue(out["moves"]) + + def test_empty_dataset_plans_nothing(self): + from TPTBox.core.dicom.dicom_renamer import rerun_bids_naming # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + ds = Path(tmp) / "dataset-empty" + (ds / "rawdata").mkdir(parents=True) + out = rerun_bids_naming(ds, dry_run=True, verbose=False) + self.assertEqual(out["moves"], []) + + def test_prune_empty_folders(self): + from TPTBox.core.dicom.dicom_renamer import _prune_empty_folders # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / "empty" / "nested").mkdir(parents=True) + (root / "full").mkdir() + (root / "full" / "f.txt").write_text("x", encoding="utf-8") + removed = _prune_empty_folders(root) + self.assertGreaterEqual(removed, 2) + self.assertFalse((root / "empty").exists()) + self.assertTrue((root / "full" / "f.txt").exists()) + + +@unittest.skipUnless(HAS_PYDICOM, PYDICOM_REASON) +class TestStemCollision(unittest.TestCase): + def test_free_target_is_a_plain_move(self): + from TPTBox.core.dicom.dicom_renamer import _resolve_stem_collision # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + src = Path(tmp) / "src.json" + src.write_text("{}", encoding="utf-8") + target = Path(tmp) / "sub-1_sequ-3_T2w.json" + final, action = _resolve_stem_collision(target, src) + self.assertEqual((final, action), (target, "move")) + + def test_identical_content_is_reported_as_a_duplicate(self): + from TPTBox.core.dicom.dicom_renamer import _resolve_stem_collision # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + src = Path(tmp) / "src.json" + target = Path(tmp) / "sub-1_sequ-3_T2w.json" + src.write_text('{"a": 1}', encoding="utf-8") + target.write_text('{"a": 1}', encoding="utf-8") + self.assertEqual(_resolve_stem_collision(target, src)[1], "duplicate") + + def test_different_content_takes_a_run_slot(self): + from TPTBox.core.dicom.dicom_renamer import _resolve_stem_collision # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + src = Path(tmp) / "src.json" + target = Path(tmp) / "sub-1_sequ-3_T2w.json" + src.write_text('{"a": 1}', encoding="utf-8") + target.write_text('{"a": 2}', encoding="utf-8") + final, action = _resolve_stem_collision(target, src) + self.assertEqual(action, "run") + self.assertIn("run-2", final.name) + + def test_the_same_path_is_not_a_collision(self): + from TPTBox.core.dicom.dicom_renamer import _resolve_stem_collision # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + src = Path(tmp) / "sub-1_sequ-3_T2w.json" + src.write_text('{"a": 1}', encoding="utf-8") + self.assertEqual(_resolve_stem_collision(src, src)[1], "move") + + +@unittest.skipUnless(HAS_PYDICOM, PYDICOM_REASON) +class TestXaPairing(unittest.TestCase): + @staticmethod + def _xa_dataset(root: Path, b_sequ: int = 2) -> Path: + """A biplane pair: both planes share (StudyInstanceUID, AcquisitionTime). + + The acq-A / acq-B entities already exist (the extractor writes them); this + module's job is only to make the B side adopt the A side's `sequ`. + """ + ds = root / "dataset-xa" + folder = ds / "rawdata" / "sub-1" / "ses-20200101" / "XA" + shared = {"StudyInstanceUID": "1.2.3", "AcquisitionTime": "120000.00", "Modality": "XA"} + _write_pair(folder, "sub-1_ses-20200101_sequ-1_acq-A_XA", _sidecar(1, **shared)) + _write_pair(folder, f"sub-1_ses-20200101_sequ-{b_sequ}_acq-B_XA", _sidecar(b_sequ, **shared)) + return ds + + def test_norm_time_strips_subsecond_precision(self): + from TPTBox.core.dicom.xa_pairing import _norm_time # noqa: PLC0415 + + self.assertEqual(_norm_time("120000.000000"), _norm_time("120000")) + self.assertEqual(_norm_time(None), "") + + def test_bucket_key_groups_the_same_acquisition(self): + from TPTBox.core.dicom.xa_pairing import _bucket_key # noqa: PLC0415 + + a = {"StudyInstanceUID": "1.2.3", "AcquisitionTime": "120000.00"} + b = {"StudyInstanceUID": "1.2.3", "AcquisitionTime": "120000"} + c = {"StudyInstanceUID": "9.9.9", "AcquisitionTime": "120000"} + self.assertEqual(_bucket_key(a), _bucket_key(b)) + self.assertNotEqual(_bucket_key(a), _bucket_key(c)) + + def test_sidecar_for_maps_nii_to_json(self): + from TPTBox.core.dicom.xa_pairing import _sidecar_for # noqa: PLC0415 + + self.assertEqual(_sidecar_for(Path("/ds/a_xa.nii.gz")), Path("/ds/a_xa.json")) + + def test_flip_marker_is_idempotent(self): + from TPTBox.core.dicom.xa_pairing import _flip_marker_set, _write_flip_marker # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + sidecar = Path(tmp) / "a_xa.json" + sidecar.write_text(json.dumps({"a": 1}), encoding="utf-8") + self.assertFalse(_flip_marker_set(sidecar)) + _write_flip_marker(sidecar) + self.assertTrue(_flip_marker_set(sidecar)) + _write_flip_marker(sidecar) + self.assertTrue(_flip_marker_set(sidecar)) + + def test_flip_z_affine_only_changes_the_third_column_sign(self): + from TPTBox.core.dicom.xa_pairing import _flip_z_affine # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "a_xa.nii.gz" + affine = np.diag([0.5, 0.5, 6.0, 1.0]) + arr = np.arange(8 * 8 * 4, dtype=np.uint8).reshape(8, 8, 4) + nib.save(nib.Nifti1Image(arr, affine), path) + _flip_z_affine(path) + out = nib.load(path) + np.testing.assert_array_equal(np.asanyarray(out.dataobj), arr, "voxel data must be untouched") + new_affine = out.affine + np.testing.assert_allclose(new_affine[:3, 0], affine[:3, 0]) + np.testing.assert_allclose(new_affine[:3, 1], affine[:3, 1]) + np.testing.assert_allclose(new_affine[:3, 2], -affine[:3, 2]) + + def test_pair_biplane_dry_run_plans_without_touching_disk(self): + from TPTBox.core.dicom.xa_pairing import pair_biplane_xa # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + ds = self._xa_dataset(Path(tmp)) + before = {p: p.stat().st_mtime_ns for p in ds.rglob("*") if p.is_file()} + out = pair_biplane_xa(ds, dry_run=True, verbose=False) + after = {p: p.stat().st_mtime_ns for p in ds.rglob("*") if p.is_file()} + self.assertEqual(before, after, "a dry run must not move anything") + self.assertEqual(out["pairs_matched"], 1) + self.assertTrue(out["renames"]) + + def test_the_b_side_adopts_the_a_side_sequ(self): + from TPTBox.core.dicom.xa_pairing import pair_biplane_xa # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + ds = self._xa_dataset(Path(tmp)) + out = pair_biplane_xa(ds, dry_run=False, verbose=False) + self.assertEqual(out["pairs_matched"], 1) + names = sorted(p.name for p in ds.rglob("*.nii.gz")) + self.assertEqual( + names, + ["sub-1_ses-20200101_sequ-1_acq-A_XA.nii.gz", "sub-1_ses-20200101_sequ-1_acq-B_XA.nii.gz"], + ) + + def test_the_renamed_file_stays_in_its_own_parent_folder(self): + from TPTBox.core.dicom.xa_pairing import pair_biplane_xa # noqa: PLC0415 + + # Regression: the target path came from get_changed_bids(), whose `parent` + # defaults to "derivatives", so renaming the B side moved it out of rawdata. + # A rawdata-scoped scan then lost it entirely - including the flip_z pass. + with tempfile.TemporaryDirectory() as tmp: + ds = self._xa_dataset(Path(tmp)) + pair_biplane_xa(ds, dry_run=False, verbose=False) + parents = {p.relative_to(ds).parts[0] for p in ds.rglob("*") if p.is_file()} + self.assertEqual(parents, {"rawdata"}) + + def test_an_already_paired_bucket_needs_no_rename(self): + from TPTBox.core.dicom.xa_pairing import pair_biplane_xa # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + ds = self._xa_dataset(Path(tmp), b_sequ=1) + out = pair_biplane_xa(ds, dry_run=True, verbose=False) + self.assertEqual(out["pairs_matched"], 1) + self.assertEqual(out["renames"], []) + + def test_a_lone_plane_is_not_a_pair(self): + from TPTBox.core.dicom.xa_pairing import pair_biplane_xa # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + ds = Path(tmp) / "dataset-xa" + folder = ds / "rawdata" / "sub-1" / "ses-20200101" / "XA" + _write_pair(folder, "sub-1_ses-20200101_sequ-1_acq-A_XA", _sidecar(1, StudyInstanceUID="1.2.3")) + out = pair_biplane_xa(ds, dry_run=True, verbose=False) + self.assertEqual(out["pairs_matched"], 0) + self.assertEqual(out["renames"], []) + + def test_flip_z_negates_the_affine_column_and_marks_the_sidecar(self): + from TPTBox.core.dicom.xa_pairing import _flip_marker_set, pair_biplane_xa # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + ds = self._xa_dataset(Path(tmp)) + out = pair_biplane_xa(ds, dry_run=False, flip_z=True, verbose=False) + self.assertEqual(len(out["flipped"]), 2, "both planes must be flipped, not just the A side") + for nii_path in ds.rglob("*.nii.gz"): + with self.subTest(name=nii_path.name): + np.testing.assert_allclose(nib.load(nii_path).affine[:3, 2], [0.0, 0.0, -6.0]) + self.assertTrue(_flip_marker_set(nii_path.with_name(nii_path.name.replace(".nii.gz", ".json")))) + + # A second pass must not flip the same file back. + again = pair_biplane_xa(ds, dry_run=False, flip_z=True, verbose=False) + self.assertEqual(again["flipped"], []) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_internal_utils.py b/unit_tests/test_internal_utils.py new file mode 100644 index 0000000..a6bb14a --- /dev/null +++ b/unit_tests/test_internal_utils.py @@ -0,0 +1,281 @@ +"""Unit tests for `TPTBox.core.internal` helpers, `sitk_utils` and `optional_deps`. + +These are the small plumbing layers the rest of the package builds on: dtype +coercion for the NIfTI-1 header, the backup-on-save decorator, the SimpleITK +metadata conversion, Slicer terminology strings, and the stubs that turn a missing +optional backend into an actionable ImportError. +""" + +from __future__ import annotations + +import json +import tempfile +import unittest +from pathlib import Path + +import nibabel as nib +import numpy as np + +from TPTBox import NII +from TPTBox.core.internal.nii_help import ( + _arr_for_nifti1, + _convert, + _nifti_safe_dtype, + save_json, + secure_save, +) +from TPTBox.core.internal.optional_deps import _message, missing_dependency_class, missing_dependency_func +from TPTBox.core.internal.slicer_nrrd import ( + _generate_unique_segment_id, + _terminology_entry_from_string, + _terminology_entry_to_string, + remove_not_supported_values, +) +from TPTBox.core.sitk_utils import ( + get_ras_affine_from_sitk, + get_ras_affine_from_sitk_meta, + get_sitk_metadata_from_ras_affine, + nib_to_sitk, + nii_to_sitk, + sitk_to_nib, + sitk_to_nii, +) + + +def _nii(arr=None, zoom=(1.5, 2.0, 3.0), origin=(10.0, -5.0, 2.0), seg=True) -> NII: + arr = np.arange(4 * 5 * 6, dtype=np.uint8).reshape(4, 5, 6) if arr is None else arr + affine = np.diag([*zoom, 1.0]) + affine[:3, 3] = origin + return NII(nib.Nifti1Image(arr, affine), seg=seg) + + +class TestNiftiDtypeHelpers(unittest.TestCase): + def test_supported_dtypes_pass_through(self): + for dtype in (np.uint8, np.int16, np.int32, np.float32, np.float64): + with self.subTest(dtype=np.dtype(dtype).name): + self.assertEqual(_nifti_safe_dtype(dtype), np.dtype(dtype)) + + def test_unsupported_dtypes_are_upcast(self): + self.assertEqual(_nifti_safe_dtype(np.float16), np.dtype(np.float32)) + self.assertEqual(_nifti_safe_dtype(bool), np.dtype(np.uint8)) + + def test_arr_for_nifti1_only_copies_when_needed(self): + safe = np.zeros((2, 2, 2), dtype=np.uint8) + self.assertIs(_arr_for_nifti1(safe), safe, "a supported dtype must not be copied") + + unsafe = np.zeros((2, 2, 2), dtype=bool) + out = _arr_for_nifti1(unsafe) + self.assertEqual(out.dtype, np.dtype(np.uint8)) + self.assertEqual(out.shape, unsafe.shape) + + def test_arr_for_nifti1_preserves_values(self): + arr = np.array([[[True, False]]]) + np.testing.assert_array_equal(_arr_for_nifti1(arr), [[[1, 0]]]) + + +class TestJsonConvert(unittest.TestCase): + def test_numpy_scalars_and_arrays(self): + self.assertEqual(_convert(np.int64(3)), 3) + self.assertIsInstance(_convert(np.int64(3)), int) + self.assertAlmostEqual(_convert(np.float32(1.5)), 1.5, places=6) + self.assertEqual(_convert(np.arange(3)), [0, 1, 2]) + + def test_paths_become_absolute_strings(self): + out = _convert(Path("rel/path")) + self.assertIsInstance(out, str) + self.assertTrue(Path(out).is_absolute()) + + def test_unsupported_types_raise(self): + with self.assertRaises(TypeError): + _convert(object()) + + def test_save_json_writes_numpy_payloads(self): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "out.json" + save_json(path, {"i": np.int32(1), "arr": np.zeros(2), "p": Path(tmp)}) + loaded = json.loads(path.read_text(encoding="utf-8")) + self.assertEqual(loaded["i"], 1) + self.assertEqual(loaded["arr"], [0.0, 0.0]) + + def test_save_json_keeps_a_backup_on_failure(self): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "out.json" + save_json(path, {"first": 1}) + with self.assertRaises(TypeError): + save_json(path, {"bad": object()}) + # The original content must survive a failed write. + self.assertEqual(json.loads(path.read_text(encoding="utf-8")), {"first": 1}) + + +class TestSecureSave(unittest.TestCase): + class _Dummy: + def __init__(self): + self.calls = [] + + @secure_save + def save(self, file, payload="x"): + Path(file).write_text(payload, encoding="utf-8") + self.calls.append(file) + + @secure_save + def failing_save(self, file, payload="x"): # noqa: ARG002 + Path(file).write_text("partial", encoding="utf-8") + raise RuntimeError("boom") + + def test_saves_through_to_the_wrapped_function(self): + with tempfile.TemporaryDirectory() as tmp: + target = Path(tmp) / "a.nii.gz" + self._Dummy().save(target, payload="hello") + self.assertEqual(target.read_text(encoding="utf-8"), "hello") + self.assertEqual(list(Path(tmp).glob("*.backup")), []) + + def test_restores_the_previous_file_when_the_save_raises(self): + with tempfile.TemporaryDirectory() as tmp: + target = Path(tmp) / "a.nii.gz" + dummy = self._Dummy() + dummy.save(target, payload="original") + with self.assertRaises(RuntimeError): + dummy.failing_save(target) + self.assertEqual(target.read_text(encoding="utf-8"), "original", "the backup must be restored") + self.assertEqual(list(Path(tmp).glob("*.backup")), [], "no backup file is left behind") + + def test_a_bids_file_without_a_supported_type_is_rejected(self): + # The file_types list is only consulted for BIDS_FILE inputs; a plain path is + # passed through whatever its extension is. + from TPTBox.core.bids_files import BIDS_FILE # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + ds = Path(tmp) / "dataset-x" + path = ds / "rawdata" / "sub-1" / "sub-1_ct.png" + path.parent.mkdir(parents=True) + path.touch() + bids = BIDS_FILE(path, ds, verbose=False) + with self.assertRaises(ValueError): + self._Dummy().save(bids) + + def test_a_plain_path_with_any_extension_is_accepted(self): + with tempfile.TemporaryDirectory() as tmp: + target = Path(tmp) / "a.whatever" + self._Dummy().save(target, payload="ok") + self.assertEqual(target.read_text(encoding="utf-8"), "ok") + + +class TestSitkConversion(unittest.TestCase): + def test_nii_to_sitk_round_trip_preserves_array_and_affine(self): + for seg in (True, False): + with self.subTest(seg=seg): + nii = _nii(seg=seg) + back = sitk_to_nii(nii_to_sitk(nii), seg=seg) + np.testing.assert_array_equal(back.get_array(), nii.get_array()) + np.testing.assert_allclose(back.affine, nii.affine, atol=1e-6) + self.assertEqual(back.orientation, nii.orientation) + + def test_nib_to_sitk_round_trip(self): + nii = _nii() + back = sitk_to_nib(nib_to_sitk(nii.nii)) + np.testing.assert_allclose(back.affine, nii.affine, atol=1e-6) + + def test_metadata_round_trip(self): + nii = _nii() + # Note the argument order: the getter returns (origin, spacing, direction) while + # the builder takes (spacing, direction, origin). + origin, spacing, direction = get_sitk_metadata_from_ras_affine(nii.affine) + rebuilt = get_ras_affine_from_sitk_meta(spacing, direction, origin) + np.testing.assert_allclose(rebuilt, nii.affine, atol=1e-6) + + def test_metadata_matches_what_simpleitk_reports(self): + nii = _nii() + origin, spacing, direction = get_sitk_metadata_from_ras_affine(nii.affine) + image = nii_to_sitk(nii) + np.testing.assert_allclose(image.GetOrigin(), origin, atol=1e-6) + np.testing.assert_allclose(image.GetSpacing(), spacing, atol=1e-6) + np.testing.assert_allclose(image.GetDirection(), direction, atol=1e-6) + + def test_an_unsupported_direction_length_raises(self): + with self.assertRaises(NotImplementedError): + get_ras_affine_from_sitk_meta((1.0, 1.0, 1.0), (1.0, 0.0, 0.0), (0.0, 0.0, 0.0)) + + def test_affine_from_a_sitk_image(self): + nii = _nii() + np.testing.assert_allclose(get_ras_affine_from_sitk(nii_to_sitk(nii)), nii.affine, atol=1e-6) + + def test_spacing_matches_the_zoom(self): + nii = _nii(zoom=(1.5, 2.0, 3.0)) + image = nii_to_sitk(nii) + np.testing.assert_allclose(sorted(image.GetSpacing()), sorted(nii.zoom), atol=1e-6) + + def test_reoriented_input_round_trips_too(self): + nii = _nii().reorient(("P", "I", "R")) + back = sitk_to_nii(nii_to_sitk(nii), seg=True) + np.testing.assert_array_equal(back.get_array(), nii.get_array()) + np.testing.assert_allclose(back.affine, nii.affine, atol=1e-6) + + +class TestSlicerTerminology(unittest.TestCase): + ENTRY = ( + "Segmentation category and type - 3D Slicer General Anatomy list" + "~SCT^123037004^Anatomical Structure~SCT^113197003^Rib~^^~Anatomic codes - DICOM master list" + "~^^~^^" + ) + + def test_round_trip_through_the_string_form(self): + parsed = _terminology_entry_from_string(self.ENTRY) + self.assertIsInstance(parsed, dict) + self.assertEqual(_terminology_entry_to_string(parsed), self.ENTRY) + + def test_parsed_entry_exposes_the_type_name(self): + parsed = _terminology_entry_from_string(self.ENTRY) + self.assertIn("Rib", json.dumps(parsed)) + + def test_generate_unique_segment_id_avoids_collisions(self): + taken = {"Segment_1", "Segment_2"} + new = _generate_unique_segment_id(taken) + self.assertNotIn(new, taken) + self.assertNotIn(_generate_unique_segment_id(taken | {new}), taken | {new}) + + def test_generate_unique_segment_id_on_an_empty_set(self): + self.assertTrue(_generate_unique_segment_id(set())) + + def test_remove_not_supported_values_mutates_in_place(self): + payload = {"keep": 1, "space origin": None, "nested": {"a": None, "b": 2}} + out = remove_not_supported_values(payload) + self.assertIsNone(out, "the helper mutates in place and returns None") + self.assertIn("keep", payload) + + +class TestOptionalDeps(unittest.TestCase): + def test_message_names_the_extra_when_there_is_one(self): + text = _message("do_thing", "reg", "torch hf-deepali", "No module named 'torch'", "()") + self.assertIn("TPTBox[reg]", text) + self.assertIn("pip install torch hf-deepali", text) + self.assertIn("No module named 'torch'", text) + + def test_message_without_an_extra_points_at_the_packages(self): + text = _message("do_thing", None, "spineps", "No module named 'spineps'", "()") + self.assertNotIn("TPTBox[", text) + self.assertIn("pip install spineps", text) + + def test_func_stub_raises_on_call_not_on_creation(self): + stub = missing_dependency_func("run_thing", ImportError("No module named 'torch'"), "seg", "torch") + with self.assertRaises(ImportError) as ctx: + stub(1, keyword=2) + self.assertIn("run_thing", str(ctx.exception)) + self.assertIn("TPTBox[seg]", str(ctx.exception)) + + def test_class_stub_raises_on_instantiation(self): + stub = missing_dependency_class("Thing", ImportError("No module named 'deepali'"), "reg", "hf-deepali") + self.assertTrue(isinstance(stub, type)) + with self.assertRaises(ImportError) as ctx: + stub() + self.assertIn("Thing", str(ctx.exception)) + self.assertIn("hf-deepali", str(ctx.exception)) + + def test_stub_message_without_an_extra(self): + stub = missing_dependency_func("run_spineps", ImportError("nope"), None, "spineps") + with self.assertRaises(ImportError) as ctx: + stub() + self.assertNotIn("TPTBox[", str(ctx.exception)) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_logger.py b/unit_tests/test_logger.py new file mode 100644 index 0000000..464bfed --- /dev/null +++ b/unit_tests/test_logger.py @@ -0,0 +1,378 @@ +"""Unit tests for `TPTBox.logger`. + +`String_Logger` accumulates into `.log_content`, so almost everything here needs +neither files nor stdout capture. `indentation_level` is a module global shared by +every logger instance, so it is asserted back to zero after each test. +""" + +from __future__ import annotations + +import tempfile +import unittest +from pathlib import Path + +from TPTBox.logger import log_file +from TPTBox.logger.log_constants import ( + Log_Type, + _clean_all_color_from_text, + _convert_seconds, + _dict_to_string, + _format_time, + color_log_text, + datatype_to_string, + format_time_short, + get_formatted_time, + get_time, +) +from TPTBox.logger.log_file import ( + Logger, + No_Logger, + Reflection_Logger, + String_Logger, + _set_indent, + print_to_terminal, + sub_log_call_func, +) + + +class _IndentGuard(unittest.TestCase): + def setUp(self): + _set_indent(0) + + def tearDown(self): + # `indentation_level` is a module global; a leak would silently indent later tests. + self.assertEqual(log_file.indentation_level, 0) + _set_indent(0) + + +class TestStringLogger(_IndentGuard): + def test_print_joins_parts_and_adds_a_prefix(self): + log = String_Logger() + log.print("a", "b") + self.assertIn("a b", log.log_content) + self.assertTrue(log.log_content.endswith("\n")) + self.assertTrue(log.log_content.startswith("["), "default prefix is a bracketed marker") + + def test_plain_content_is_colour_free_and_colored_content_is_not(self): + log = String_Logger() + log.print("hello", ltype=Log_Type.OK) + self.assertNotIn("\x1b", log.log_content) + self.assertIn("\x1b", log.log_content_colored) + self.assertEqual(_clean_all_color_from_text(log.log_content_colored), log.log_content) + + def test_ignore_prefix(self): + log = String_Logger() + log.print("naked", ignore_prefix=True) + self.assertEqual(log.log_content, "naked\n") + + def test_text_already_starting_with_a_bracket_is_not_double_prefixed(self): + log = String_Logger() + log.print("[x] already prefixed") + self.assertEqual(log.log_content, "[x] already prefixed\n") + + def test_empty_print_emits_a_bare_newline(self): + log = String_Logger() + log.print() + self.assertEqual(log.log_content, "\n") + + def test_custom_prefix_is_used_for_every_line(self): + log = String_Logger(prefix="XYZ") + log.print("one") + log.print("two") + self.assertEqual(len([line for line in log.log_content.splitlines() if "XYZ" in line]), 2) + + def test_log_type_passed_inside_text_is_stripped_but_still_selects_the_type(self): + log = String_Logger() + log.print("message", Log_Type.FAIL) + self.assertNotIn("Log_Type", log.log_content) + self.assertIn("message", log.log_content) + + def test_end_is_honoured(self): + log = String_Logger() + log.print("x", end="") + self.assertEqual(log.log_content, "[*] x") + + def test_dict_values_are_rendered(self): + log = String_Logger() + log.print({"a": 1}) + self.assertIn("'a'", log.log_content) + self.assertIn("1", log.log_content) + + def test_on_shortcuts_all_write(self): + for name in ("on_fail", "on_warning", "on_ok", "on_text", "on_neutral", "on_debug", "on_save", "on_log"): + with self.subTest(method=name): + log = String_Logger() + getattr(log, name)("msg") + self.assertIn("msg", log.log_content) + + def test_print_error_captures_the_traceback(self): + log = String_Logger() + try: + raise RuntimeError("boom") # noqa: TRY301 + except RuntimeError: + log.print_error(verbose=False) + self.assertIn("RuntimeError", log.log_content) + self.assertIn("boom", log.log_content) + + def test_close_returns_both_buffers_and_reports_the_duration(self): + log = String_Logger(finalize=False) + log.print("content") + plain, colored = log.close() + self.assertIn("content", plain) + self.assertIn("Sub-process duration", plain) + self.assertIn("\x1b", colored) + + +class TestIndentation(_IndentGuard): + def test_context_manager_indents_and_restores(self): + log = String_Logger() + log.print("flat") + with log: + self.assertEqual(log_file.indentation_level, 1) + log.print("nested") + with log: + self.assertEqual(log_file.indentation_level, 2) + log.print("deeper") + self.assertEqual(log_file.indentation_level, 0) + + lines = [line for line in log.log_content.splitlines() if line] + self.assertFalse(lines[0].startswith(" -")) + self.assertTrue(lines[1].startswith(" - "), lines[1]) + self.assertTrue(lines[2].startswith(" ---- "), lines[2]) + + def test_set_indent_bool_steps_and_int_sets(self): + self.assertEqual(_set_indent(True), 1) + self.assertEqual(_set_indent(True), 2) + self.assertEqual(_set_indent(False), 1) + self.assertEqual(_set_indent(7), 7) + self.assertEqual(_set_indent(0), 0) + + def test_set_indent_never_goes_negative(self): + self.assertEqual(_set_indent(False), 0) + self.assertEqual(_set_indent(False), 0) + + +class TestStatistics(_IndentGuard): + def test_log_statistic_accumulates_under_key_and_key2(self): + log = String_Logger() + log.log_statistic("dice", 0.5, key2="case1", verbose=False) + log.log_statistic("dice", 0.7, key2="case2", verbose=False) + self.assertEqual(log.logging_state, {"dice": {"case1": 0.5, "case2": 0.7}}) + + def test_key2_defaults_to_a_running_index(self): + log = String_Logger() + log.log_statistic("dice", 1.0, verbose=False) + log.log_statistic("dice", 2.0, verbose=False) + self.assertEqual(sorted(log.logging_state["dice"]), [0, 1]) + + def test_print_statistic_reports_mean_std_median_and_count(self): + log = String_Logger() + for v in (1.0, 2.0, 3.0): + log.log_statistic("metric", v, verbose=False) + log.print_statistic() + text = log.log_content + self.assertIn("Accumulated Statistics", text) + self.assertIn("2.0", text) # mean + self.assertIn("0.816", text) # std of 1,2,3 + self.assertIn("±", text) # the character that forces utf-8 log files + self.assertIn("metric", text) + + def test_print_statistic_without_data(self): + log = String_Logger() + log.print_statistic() + self.assertIn("No Accumulated Statistics", log.log_content) + + def test_logging_state_is_per_instance(self): + a, b = String_Logger(), String_Logger() + a.log_statistic("x", 1.0, verbose=False) + self.assertIsNone(b.logging_state) + + +class TestSubLoggers(_IndentGuard): + def test_flush_appends_into_the_parent_and_clears_the_child(self): + with tempfile.TemporaryDirectory() as tmp: + parent = Logger(tmp, "test", default_verbose=False) + sub = parent.add_sub_logger("child") + sub.print("from the child", verbose=False) + self.assertIn("from the child", sub.log_content) + sub.flush() + self.assertEqual(sub.log_content, "") + self.assertIn(sub, parent.sub_loggers) + parent.close() + self.assertIn("from the child", self._read_log(tmp)) + + def test_close_removes_the_child_from_the_parent(self): + with tempfile.TemporaryDirectory() as tmp: + parent = Logger(tmp, "test", default_verbose=False) + sub = parent.add_sub_logger("child") + sub.print("bye", verbose=False) + sub.close() + self.assertNotIn(sub, parent.sub_loggers) + parent.close() + text = self._read_log(tmp) + self.assertIn("bye", text) + self.assertIn("Sub-process duration", text) + + def test_unflushed_child_content_is_written_by_the_parents_close(self): + with tempfile.TemporaryDirectory() as tmp: + parent = Logger(tmp, "test", default_verbose=False) + sub = parent.add_sub_logger("child") + sub.print("never flushed", verbose=False) + parent.close() + text = self._read_log(tmp) + self.assertIn("Found 1 sub logger", text) + self.assertIn("never flushed", text) + + def test_sub_log_call_func_passes_name_and_sub_logger(self): + with tempfile.TemporaryDirectory() as tmp: + parent = Logger(tmp, "test", default_verbose=False) + seen = {} + + def work(name, logger, extra): + seen["name"] = name + seen["is_sub"] = isinstance(logger, String_Logger) + logger.print("worked on", extra, verbose=False) + return extra * 2 + + result = sub_log_call_func("job", parent, work, extra=21) + self.assertEqual(result, 42) + self.assertEqual(seen, {"name": "job", "is_sub": True}) + parent.close() + + @staticmethod + def _read_log(folder: str) -> str: + logs = sorted(Path(folder, "logs").glob("*.log")) + assert len(logs) == 1, logs + return logs[0].read_text(encoding="utf-8") + + +class TestFileLogger(_IndentGuard): + def test_creates_the_logs_folder_and_a_named_file(self): + with tempfile.TemporaryDirectory() as tmp: + log = Logger(tmp, "myrun", default_verbose=False) + log.print("hello", verbose=False) + log.flush() + logs = list(Path(tmp, "logs").glob("*.log")) + self.assertEqual(len(logs), 1) + self.assertIn("myrun", logs[0].name) + self.assertIn("hello", logs[0].read_text(encoding="utf-8")) + log.close() + + def test_dict_filename_becomes_bids_like_key_value_pairs(self): + with tempfile.TemporaryDirectory() as tmp: + log = Logger(tmp, {"sub": "001", "task": "seg"}, default_verbose=False) + name = next(iter(Path(tmp, "logs").glob("*.log"))).name + self.assertIn("sub-001", name) + self.assertIn("task-seg", name) + log.close() + + def test_log_arguments_are_written(self): + with tempfile.TemporaryDirectory() as tmp: + log = Logger(tmp, "r", default_verbose=False, log_arguments={"lr": 0.1}) + log.close() + text = TestSubLoggers._read_log(tmp) + self.assertIn("Run with arguments", text) + self.assertIn("lr", text) + self.assertIn("0.1", text) + + def test_log_arguments_non_dict(self): + with tempfile.TemporaryDirectory() as tmp: + log = Logger(tmp, "r", default_verbose=False, log_arguments="--flag") + log.close() + self.assertIn("--flag", TestSubLoggers._read_log(tmp)) + + def test_close_is_idempotent_and_writes_the_duration(self): + with tempfile.TemporaryDirectory() as tmp: + log = Logger(tmp, "r", default_verbose=False) + log.close() + log.close() + self.assertIn("Program duration", TestSubLoggers._read_log(tmp)) + + def test_remove_marks_the_logger_removed(self): + with tempfile.TemporaryDirectory() as tmp: + log = Logger(tmp, "r", default_verbose=False) + self.assertFalse(log.removed) + log.remove() + self.assertTrue(log.removed) + + def test_log_file_is_utf8_so_the_statistics_table_can_be_written(self): + # print_statistic emits "±"; a C/POSIX locale would raise UnicodeEncodeError + # if the file were opened with the platform default encoding. + with tempfile.TemporaryDirectory() as tmp: + log = Logger(tmp, "r", default_verbose=False) + log.log_statistic("m", 1.0, verbose=False) + log.print_statistic() + log.close() + self.assertIn("±", TestSubLoggers._read_log(tmp)) + + def test_create_from_bids_uses_the_dataset_folder(self): + from TPTBox.core.bids_files import BIDS_FILE # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + ds = Path(tmp) / "dataset-x" + f = ds / "rawdata" / "sub-001" / "sub-001_ct.nii.gz" + f.parent.mkdir(parents=True) + f.touch() + bids = BIDS_FILE(f, ds, verbose=False) + log = Logger.create_from_bids(bids, "frombids", default_verbose=False) + self.assertTrue((ds / "logs").is_dir()) + log.close() + + +class TestNoLoggerAndReflection(_IndentGuard): + def test_no_logger_swallows_everything(self): + log = No_Logger() + log.print("to terminal only") + log.flush() + log.close() + self.assertIs(log.add_sub_logger("x"), log) + + def test_no_logger_start_banner(self): + log = No_Logger(print_log_started=True) + self.assertTrue(hasattr(log, "start_time")) + + def test_reflection_logger_delegates_to_the_logger_given_as_verbose(self): + target = String_Logger() + Reflection_Logger().print("delegated", verbose=target) + self.assertIn("delegated", target.log_content) + + +class TestLogConstants(unittest.TestCase): + def test_colour_roundtrip(self): + for ltype in Log_Type: + with self.subTest(ltype=ltype): + self.assertEqual(_clean_all_color_from_text(color_log_text(ltype, "text")), "text") + + def test_convert_seconds(self): + self.assertEqual(_convert_seconds(0), "0:00:00 h:mm:ss") + self.assertEqual(_convert_seconds(61), "0:01:01 h:mm:ss") + self.assertEqual(_convert_seconds(3661), "1:01:01 h:mm:ss") + + def test_format_time_short(self): + import time as _time # noqa: PLC0415 + + t = _time.struct_time((2024, 1, 15, 9, 30, 0, 0, 15, 0)) + self.assertEqual(format_time_short(t), "date-2024-1-15_time-9-30-0") + self.assertIn("2024", _format_time(t)) + + def test_get_time_and_formatted_time(self): + self.assertTrue(hasattr(get_time(), "tm_year")) + self.assertTrue(get_formatted_time().startswith("date-")) + + def test_datatype_to_string(self): + self.assertEqual(datatype_to_string(5, Log_Type.TEXT), "5") + self.assertEqual(datatype_to_string("s", Log_Type.TEXT), "s") + self.assertIn("'k'", datatype_to_string({"k": "v"}, Log_Type.TEXT)) + + def test_dict_to_string_quotes_str_keys_and_values(self): + text = _clean_all_color_from_text(_dict_to_string({"k": "v", 1: 2})) + self.assertIn("'k': 'v'", text) + self.assertIn("1: 2", text) + + def test_print_to_terminal_warning_throw_raises_a_warning(self): + with self.assertWarns(Warning): + print_to_terminal("careful", end="\n", ltype=Log_Type.WARNING_THROW) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_nii_wrapper_more.py b/unit_tests/test_nii_wrapper_more.py new file mode 100644 index 0000000..8ef150e --- /dev/null +++ b/unit_tests/test_nii_wrapper_more.py @@ -0,0 +1,472 @@ +"""Unit tests for the parts of `NII` / `NII_Math` / `Has_Grid` that the existing +`test_nii*.py` files do not reach. + +Everything here is built from small hand-written arrays whose expected result can be +written down: a seed label diffusing into a reference mask, a slab truncated at a +reference label, a cube's surface voxel count, a histogram match against a known +reference, and the reflected/in-place arithmetic operators. +""" + +from __future__ import annotations + +import importlib.util +import tempfile +import unittest +from pathlib import Path + +import nibabel as nib +import numpy as np + +from TPTBox import NII +from TPTBox.core.nii_poi_abstract import Grid + +# `numpy-stl` is a declared dependency, but guard anyway so a partially installed +# environment skips instead of erroring (same pattern as unit_tests/test_mesh3D.py). +HAS_STL = importlib.util.find_spec("stl") is not None + + +def _nii(arr: np.ndarray, seg: bool = True, zoom=(1.0, 1.0, 1.0), origin=(0.0, 0.0, 0.0)) -> NII: + affine = np.diag([*zoom, 1.0]) + affine[:3, 3] = origin + return NII(nib.Nifti1Image(arr, affine), seg=seg) + + +class TestInfect(unittest.TestCase): + """`infect` diffuses existing labels into a reference mask.""" + + def test_two_seeds_split_the_corridor_by_distance(self): + seeds = np.zeros((11, 3, 3), dtype=np.uint8) + seeds[0] = 1 + seeds[10] = 2 + reference = np.ones((11, 3, 3), dtype=np.uint8) + out = _nii(seeds).infect(_nii(reference), verbose=False) + arr = out.get_seg_array() + self.assertEqual(int(arr[1, 1, 1]), 1, "the near end belongs to seed 1") + self.assertEqual(int(arr[9, 1, 1]), 2, "the far end belongs to seed 2") + self.assertEqual(int((arr == 0).sum()), 0, "the whole corridor must be claimed") + + def test_labels_do_not_leave_the_reference_mask(self): + seeds = np.zeros((9, 9, 9), dtype=np.uint8) + seeds[4, 4, 4] = 3 + reference = np.zeros((9, 9, 9), dtype=np.uint8) + reference[3:6, 3:6, 3:6] = 1 + out = _nii(seeds).infect(_nii(reference), verbose=False) + arr = out.get_seg_array() + self.assertEqual(set(np.unique(arr)) - {0}, {3}) + outside = arr.copy() + outside[3:6, 3:6, 3:6] = 0 + self.assertEqual(int(outside.sum()), 0) + + def test_max_depth_limits_the_diffusion(self): + seeds = np.zeros((11, 3, 3), dtype=np.uint8) + seeds[0] = 1 + reference = np.ones((11, 3, 3), dtype=np.uint8) + shallow = _nii(seeds).infect(_nii(reference), verbose=False, max_depth=2).get_seg_array() + deep = _nii(seeds).infect(_nii(reference), verbose=False).get_seg_array() + self.assertLess(int((shallow != 0).sum()), int((deep != 0).sum())) + + def test_inplace_variant_modifies_the_receiver(self): + seeds = np.zeros((7, 3, 3), dtype=np.uint8) + seeds[0] = 1 + nii = _nii(seeds) + before = int((nii.get_seg_array() != 0).sum()) + nii.infect_(_nii(np.ones((7, 3, 3), dtype=np.uint8)), verbose=False) + self.assertGreater(int((nii.get_seg_array() != 0).sum()), before) + + def test_mismatched_affine_is_rejected(self): + seeds = np.zeros((5, 5, 5), dtype=np.uint8) + seeds[0, 0, 0] = 1 + other = _nii(np.ones((5, 5, 5), dtype=np.uint8), zoom=(2.0, 2.0, 2.0)) + with self.assertRaises(AssertionError): + _nii(seeds).infect(other, verbose=False) + + +class TestTruncateLabelsBeyondReference(unittest.TestCase): + def test_voxels_beyond_the_reference_are_removed(self): + arr = np.zeros((3, 3, 10), dtype=np.uint8) + arr[..., :] = 1 # label 1 fills the column + arr[1, 1, 4] = 2 # the reference label sits in the middle + out = _nii(arr).truncate_labels_beyond_reference(idx=1, not_beyond=2, axis="S") + kept = out.get_seg_array() + self.assertEqual(int((kept[..., 5:] == 1).sum()), 0, "nothing of label 1 survives above the reference") + self.assertGreater(int((kept[..., :4] == 1).sum()), 0, "below the reference is untouched") + + def test_inclusion_also_drops_the_reference_slice(self): + arr = np.ones((3, 3, 10), dtype=np.uint8) + arr[1, 1, 4] = 2 + without = _nii(arr).truncate_labels_beyond_reference(idx=1, not_beyond=2, inclusion=False).get_seg_array() + with_incl = _nii(arr).truncate_labels_beyond_reference(idx=1, not_beyond=2, inclusion=True).get_seg_array() + self.assertLessEqual(int((with_incl == 1).sum()), int((without == 1).sum())) + + def test_fill_value_is_used(self): + arr = np.ones((3, 3, 10), dtype=np.uint8) + arr[1, 1, 4] = 2 + out = _nii(arr).truncate_labels_beyond_reference(idx=1, not_beyond=2, fill=7).get_seg_array() + self.assertIn(7, np.unique(out)) + + def test_inplace_variant(self): + arr = np.ones((3, 3, 10), dtype=np.uint8) + arr[1, 1, 4] = 2 + nii = _nii(arr) + nii.truncate_labels_beyond_reference_(idx=1, not_beyond=2) + self.assertEqual(int((nii.get_seg_array()[..., 5:] == 1).sum()), 0) + + +class TestTranslateArr(unittest.TestCase): + def test_tuple_translation_shifts_the_content(self): + arr = np.zeros((7, 7, 7), dtype=np.uint8) + arr[1, 1, 1] = 5 + out = _nii(arr).translate_arr((2, 0, 0), verbose=False).get_seg_array() + self.assertEqual(int(out[3, 1, 1]), 5) + self.assertEqual(int(out[1, 1, 1]), 0) + + def test_direction_dict_translation(self): + arr = np.zeros((7, 7, 7), dtype=np.uint8) + arr[3, 3, 3] = 5 + out = _nii(arr).translate_arr({"S": 2}, verbose=False).get_seg_array() + self.assertEqual(int(out.sum()), 5, "the label must still be in the volume exactly once") + self.assertEqual(int(np.argwhere(out == 5).shape[0]), 1) + self.assertFalse(np.array_equal(out, arr)) + + def test_zero_translation_is_a_noop(self): + arr = np.zeros((5, 5, 5), dtype=np.uint8) + arr[2, 2, 2] = 1 + np.testing.assert_array_equal(_nii(arr).translate_arr((0, 0, 0), verbose=False).get_seg_array(), arr) + + def test_content_shifted_out_is_dropped(self): + arr = np.zeros((5, 5, 5), dtype=np.uint8) + arr[4, 2, 2] = 1 + out = _nii(arr).translate_arr((3, 0, 0), verbose=False).get_seg_array() + self.assertEqual(int(out.sum()), 0) + + +class TestSegmentationComparisons(unittest.TestCase): + def _pair(self): + gt = np.zeros((8, 8, 8), dtype=np.uint8) + gt[2:6, 2:6, 2:6] = 1 + pred = np.zeros((8, 8, 8), dtype=np.uint8) + pred[3:7, 2:6, 2:6] = 1 + return _nii(pred), _nii(gt) + + def test_difference_marks_false_positive_and_false_negative(self): + pred, gt = self._pair() + diff = pred.get_segmentation_difference_to(gt).get_seg_array() + self.assertGreater(len(set(np.unique(diff)) - {0}), 1, "at least FP and FN classes are marked") + + def test_identical_masks_differ_only_in_the_true_positive_class(self): + pred, _ = self._pair() + diff = pred.get_segmentation_difference_to(pred.copy()).get_seg_array() + self.assertEqual(len(set(np.unique(diff)) - {0}), 1) + + def test_ignore_background_tp_drops_the_true_positive_class(self): + pred, gt = self._pair() + with_tp = pred.get_segmentation_difference_to(gt, ignore_background_tp=False).get_seg_array() + without = pred.get_segmentation_difference_to(gt, ignore_background_tp=True).get_seg_array() + self.assertLessEqual(int((without != 0).sum()), int((with_tp != 0).sum())) + + def test_overlapping_labels(self): + a = np.zeros((6, 6, 6), dtype=np.uint8) + a[0:3] = 1 + a[3:6] = 2 + b = np.zeros((6, 6, 6), dtype=np.uint8) + b[0:3] = 7 + b[3:6] = 8 + pairs = _nii(a).get_overlapping_labels_to(_nii(b)) + self.assertEqual(sorted(tuple(p) for p in pairs), [(1, 7), (2, 8)]) + + def test_map_labels_based_on_majority_overlap(self): + src = np.zeros((6, 6, 6), dtype=np.uint8) + src[0:3] = 1 + reference = np.zeros((6, 6, 6), dtype=np.uint8) + reference[0:3] = 9 + out = _nii(src).map_labels_based_on_majority_label_mask_overlap(_nii(reference)) + self.assertEqual(set(np.unique(out.get_seg_array())) - {0}, {9}) + + +class TestConnectedComponentsAndSurface(unittest.TestCase): + def _three_blobs(self): + arr = np.zeros((20, 8, 8), dtype=np.uint8) + arr[1:3, 1:3, 1:3] = 1 # 8 voxels + arr[6:9, 1:4, 1:4] = 1 # 27 voxels + arr[12:16, 1:5, 1:5] = 1 # 64 voxels + return _nii(arr) + + def test_filter_connected_components_by_min_volume(self): + out = self._three_blobs().filter_connected_components(labels=1, min_volume=20) + self.assertEqual(int((out.get_seg_array() != 0).sum()), 27 + 64) + + def test_filter_connected_components_by_max_count(self): + out = self._three_blobs().filter_connected_components(labels=1, max_count_component=1) + self.assertEqual(int((out.get_seg_array() != 0).sum()), 64, "the largest component survives") + + def test_largest_k_components_is_deprecated(self): + # The method only raises now; filter_connected_components is the replacement. + with self.assertRaises(DeprecationWarning): + self._three_blobs().get_largest_k_segmentation_connected_components(2, labels=1) + + def test_get_connected_components_labels_each_blob(self): + out = self._three_blobs().get_connected_components(labels=1) + self.assertEqual(len(set(np.unique(out.get_seg_array())) - {0}), 3) + + def test_compute_surface_points_are_on_the_boundary(self): + arr = np.zeros((9, 9, 9), dtype=np.uint8) + arr[3:6, 3:6, 3:6] = 1 + points = _nii(arr).compute_surface_points() + self.assertGreater(len(points), 0) + self.assertNotIn((4, 4, 4), [tuple(p) for p in points], "the centre voxel is interior") + + def test_fill_holes_global_with_majority_voting(self): + arr = np.zeros((9, 9, 9), dtype=np.uint8) + arr[2:7, 2:7, 2:7] = 1 + arr[4, 4, 4] = 0 # a one-voxel hole + out = _nii(arr).fill_holes_global_with_majority_voting() + self.assertEqual(int(out.get_seg_array()[4, 4, 4]), 1) + + def test_extract_background(self): + arr = np.zeros((5, 5, 5), dtype=np.uint8) + arr[1:4, 1:4, 1:4] = 1 + out = _nii(arr).extract_background().get_seg_array() + self.assertEqual(int(out.sum()), int((arr == 0).sum())) + + +class TestIntensityOperations(unittest.TestCase): + def test_normalize_to_range_rescales_in_place(self): + arr = np.linspace(-100, 500, 8 * 8 * 8, dtype=np.float32).reshape(8, 8, 8) + nii = _nii(arr, seg=False) + nii.normalize_to_range_(min_value=0, max_value=1000, verbose=False) + out = nii.get_array() + self.assertGreaterEqual(float(out.min()), 0.0) + self.assertLessEqual(float(out.max()), 1000.0) + + def test_match_histograms_moves_towards_the_reference(self): + rng = np.random.default_rng(0) + source = _nii((rng.normal(100, 10, (16, 16, 16))).astype(np.float32), seg=False) + reference = _nii((rng.normal(500, 10, (16, 16, 16))).astype(np.float32), seg=False) + matched = source.match_histograms(reference) + self.assertAlmostEqual(float(matched.get_array().mean()), 500.0, delta=25.0) + self.assertAlmostEqual(float(source.get_array().mean()), 100.0, delta=5.0, msg="the source must be untouched") + + def test_match_histograms_inplace(self): + rng = np.random.default_rng(1) + source = _nii((rng.normal(100, 10, (12, 12, 12))).astype(np.float32), seg=False) + reference = _nii((rng.normal(400, 10, (12, 12, 12))).astype(np.float32), seg=False) + source.match_histograms_(reference) + self.assertAlmostEqual(float(source.get_array().mean()), 400.0, delta=25.0) + + def test_smooth_gaussian_labelwise_keeps_the_label_set(self): + arr = np.zeros((16, 16, 16), dtype=np.uint8) + arr[2:8, 2:8, 2:8] = 1 + arr[9:14, 9:14, 9:14] = 2 + out = _nii(arr).smooth_gaussian_labelwise(label_to_smooth=1, sigma=1.0, radius=2) + self.assertTrue(set(np.unique(out.get_seg_array())).issubset({0, 1, 2})) + self.assertIn(2, np.unique(out.get_seg_array()), "other labels survive") + + +class TestGridAndAccessors(unittest.TestCase): + def test_accessors(self): + nii = _nii(np.zeros((4, 5, 6), dtype=np.uint8), zoom=(1.0, 2.0, 3.0)) + self.assertEqual(nii.dims, 3) + self.assertEqual(len(nii.direction), 9) + self.assertEqual(len(nii.direction_itk), 9) + self.assertEqual(nii.orientation_ornt.shape, (3, 2)) + self.assertEqual(nii.get_num_dims(), 3) + self.assertEqual(nii.shape_int, (4, 5, 6)) + np.testing.assert_allclose(nii.spacing, (1.0, 2.0, 3.0)) + + def test_numpy_and_ravel(self): + arr = np.arange(2 * 3 * 4, dtype=np.uint8).reshape(2, 3, 4) + nii = _nii(arr) + np.testing.assert_array_equal(nii.numpy(), arr) + np.testing.assert_array_equal(nii.ravel(), arr.ravel()) + + def test_set_description_is_stored_in_the_header(self): + nii = _nii(np.zeros((3, 3, 3), dtype=np.uint8)) + nii.set_description("hello") + self.assertIn("hello", str(nii.header["descrip"])) + + def test_corners_and_obb_quads(self): + nii = _nii(np.zeros((4, 4, 4), dtype=np.uint8), zoom=(2.0, 2.0, 2.0)) + corners = nii.get_corners() + self.assertEqual(np.asarray(corners).shape[0], 8) + quads = nii.get_obb_quads() + self.assertEqual(len(quads), 6, "a box has six faces") + + def test_global_to_local_arr_inverts_the_affine(self): + nii = _nii(np.zeros((6, 6, 6), dtype=np.uint8), zoom=(1.5, 2.0, 3.0), origin=(5.0, -2.0, 7.0)) + points = np.array([[0.0, 0.0, 0.0], [1.0, 2.0, 3.0], [5.0, 5.0, 5.0]]) + world = (nii.affine[:3, :3] @ points.T).T + nii.affine[:3, 3] + np.testing.assert_allclose(nii.global_to_local_arr(world), points, atol=1e-9) + + def test_global_to_local_arr_agrees_with_the_scalar_version(self): + nii = _nii(np.zeros((6, 6, 6), dtype=np.uint8), zoom=(1.5, 2.0, 3.0), origin=(5.0, -2.0, 7.0)) + world = np.array([[7.0, 1.0, 10.0], [5.0, -2.0, 7.0]]) + batched = nii.global_to_local_arr(world) + for i, point in enumerate(world): + with self.subTest(i=i): + np.testing.assert_allclose(batched[i], nii.global_to_local(tuple(point)), atol=1e-9) + + def test_make_empty_nii_keeps_the_grid(self): + nii = _nii(np.ones((4, 5, 6), dtype=np.uint8), zoom=(1.0, 2.0, 3.0)) + empty = nii.make_empty_nii() + self.assertEqual(empty.shape, nii.shape) + np.testing.assert_allclose(empty.affine, nii.affine) + self.assertEqual(int(empty.get_array().sum()), 0) + + def test_to_gird_returns_a_grid_with_the_same_geometry(self): + nii = _nii(np.zeros((4, 5, 6), dtype=np.uint8), zoom=(1.0, 2.0, 3.0)) + grid = nii.to_gird() + self.assertIsInstance(grid, Grid) + np.testing.assert_allclose(grid.affine, nii.affine) + self.assertEqual(tuple(grid.shape), nii.shape) + + def test_grid_can_be_constructed_directly(self): + nii = _nii(np.zeros((4, 5, 6), dtype=np.uint8), zoom=(1.0, 2.0, 3.0)) + grid = Grid(affine=nii.affine, shape=nii.shape) + np.testing.assert_allclose(grid.affine, nii.affine) + self.assertEqual(tuple(grid.shape), nii.shape) + + def test_split_4d_image_to_3d(self): + arr = np.zeros((4, 4, 4, 3), dtype=np.uint8) + for i in range(3): + arr[..., i] = i + 1 + parts = _nii(arr).split_4D_image_to_3D() + self.assertEqual(len(parts), 3) + for i, part in enumerate(parts): + with self.subTest(i=i): + self.assertEqual(part.shape, (4, 4, 4)) + self.assertEqual(int(part.get_seg_array().max()), i + 1) + + def test_reorient_same_as(self): + a = _nii(np.zeros((4, 5, 6), dtype=np.uint8)) + b = a.reorient(("P", "I", "R")) + self.assertEqual(a.reorient_same_as(b).orientation, b.orientation) + c = a.copy() + c.reorient_same_as_(b) + self.assertEqual(c.orientation, b.orientation) + + +class TestMathMixin(unittest.TestCase): + @staticmethod + def _float_nii(seed: int = 0): + rng = np.random.default_rng(seed) + return _nii(rng.normal(10.0, 3.0, (6, 6, 6)).astype(np.float32), seg=False) + + def test_summary_statistics_match_numpy(self): + nii = self._float_nii() + arr = nii.get_array() + self.assertAlmostEqual(nii.mean(), float(arr.mean()), places=4) + self.assertAlmostEqual(nii.median(), float(np.median(arr)), places=4) + self.assertAlmostEqual(nii.std(), float(arr.std()), places=4) + + def test_clamp(self): + nii = self._float_nii() + out = nii.clamp(min=9.0, max=11.0) + self.assertGreaterEqual(float(out.get_array().min()), 9.0) + self.assertLessEqual(float(out.get_array().max()), 11.0) + copy = nii.copy() + copy.clamp_(min=9.0, max=11.0) + np.testing.assert_allclose(copy.get_array(), out.get_array()) + + def test_nan_to_num(self): + arr = np.array([[[1.0, np.nan], [np.inf, -np.inf]]], dtype=np.float32) + out = _nii(arr, seg=False).nan_to_num(num=0).get_array() + self.assertTrue(np.isfinite(out).all()) + self.assertEqual(float(out[0, 0, 1]), 0.0) + + def test_normalize_variants_stay_in_range(self): + nii = self._float_nii(2) + for name in ("normalize", "normalize_mri", "normalize_ct"): + with self.subTest(method=name): + out = getattr(nii, name)(min_out=0, max_out=1) + arr = out.get_array() + self.assertGreaterEqual(float(arr.min()), 0.0) + self.assertLessEqual(float(arr.max()), 1.0) + + def test_normalize_inplace(self): + nii = self._float_nii(3) + nii.normalize_(min_out=0, max_out=10) + self.assertLessEqual(float(nii.get_array().max()), 10.0) + + def test_dice_of_identical_masks_is_one(self): + arr = np.zeros((8, 8, 8), dtype=np.uint8) + arr[2:6, 2:6, 2:6] = 1 + nii = _nii(arr) + self.assertAlmostEqual(nii.dice(nii.copy(), bar=False)[1], 1.0, places=6) + + def test_dice_of_disjoint_masks_is_zero(self): + a = np.zeros((8, 8, 8), dtype=np.uint8) + a[0:3] = 1 + b = np.zeros((8, 8, 8), dtype=np.uint8) + b[5:8] = 1 + self.assertAlmostEqual(_nii(a).dice(_nii(b), bar=False)[1], 0.0, places=6) + + def test_betti_numbers_of_a_solid_block(self): + arr = np.zeros((10, 10, 10), dtype=np.uint8) + arr[2:8, 2:8, 2:8] = 1 + betti = _nii(arr).betti_numbers() + self.assertEqual(betti[1][0], 1, "one connected component") + + def test_reflected_operators(self): + arr = np.full((3, 3, 3), 2.0, dtype=np.float32) + nii = _nii(arr, seg=False) + np.testing.assert_allclose((1 + nii).get_array(), arr + 1) + np.testing.assert_allclose((1 - nii).get_array(), 1 - arr) + + def test_bit_shift_operators(self): + nii = _nii(np.full((3, 3, 3), 4, dtype=np.int32)) + np.testing.assert_array_equal((nii << 1).get_array(), 8) + np.testing.assert_array_equal((nii >> 1).get_array(), 2) + + def test_inplace_operators(self): + arr = np.full((3, 3, 3), 8.0, dtype=np.float32) + cases = { + "isub": (lambda n: n.__isub__(3), 5.0), + "itruediv": (lambda n: n.__itruediv__(2), 4.0), + "ifloordiv": (lambda n: n.__ifloordiv__(3), 2.0), + "imod": (lambda n: n.__imod__(3), 2.0), + "ipow": (lambda n: n.__ipow__(2), 64.0), + } + for name, (op, expected) in cases.items(): + with self.subTest(op=name): + nii = _nii(arr.copy(), seg=False) + op(nii) + np.testing.assert_allclose(nii.get_array(), expected) + + +@unittest.skipUnless(HAS_STL, "numpy-stl not installed") +class TestMeshExport(unittest.TestCase): + def test_to_stl_writes_a_mesh_for_one_label(self): + arr = np.zeros((12, 12, 12), dtype=np.uint8) + arr[3:9, 3:9, 3:9] = 1 + with tempfile.TemporaryDirectory() as tmp: + out = Path(tmp) / "label.stl" + mesh = _nii(arr).to_stl(1, out_path=out) + self.assertTrue(out.exists()) + self.assertGreater(len(mesh.vectors), 0) + + def test_to_stls_writes_one_mesh_per_label(self): + arr = np.zeros((16, 16, 16), dtype=np.uint8) + arr[2:7, 2:7, 2:7] = 1 + arr[9:14, 9:14, 9:14] = 2 + with tempfile.TemporaryDirectory() as tmp: + meshes = _nii(arr).to_stls(out_path=Path(tmp) / "out.stl") + self.assertEqual(sorted(meshes), [1, 2]) + # Regression: `number_path` (default True for a non-dict out_path) computed + # out_path.with_name(...) without assigning it, so both labels were written + # to out.stl and only the last one survived. + self.assertEqual(sorted(p.name for p in Path(tmp).glob("*.stl")), ["out_1.stl", "out_2.stl"]) + + def test_to_stls_accepts_a_per_label_path_dict(self): + arr = np.zeros((16, 16, 16), dtype=np.uint8) + arr[2:7, 2:7, 2:7] = 1 + arr[9:14, 9:14, 9:14] = 2 + with tempfile.TemporaryDirectory() as tmp: + paths = {1: Path(tmp) / "one.stl", 2: Path(tmp) / "two.stl"} + _nii(arr).to_stls(out_path=paths) + for label, path in paths.items(): + with self.subTest(label=label): + self.assertTrue(path.exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_nnunet_predictor.py b/unit_tests/test_nnunet_predictor.py new file mode 100644 index 0000000..46b3d06 --- /dev/null +++ b/unit_tests/test_nnunet_predictor.py @@ -0,0 +1,89 @@ +"""Unit tests for `TPTBox.segmentation.nnUnet_utils.predictor`. + +The predictor is a TPTBox-maintained fork of nnU-Net's inference engine. No weights +and no GPU are involved here: the "network" is a plain callable, which is enough to +pin the test-time-augmentation contract. +""" + +from __future__ import annotations + +import importlib.util +import unittest + +HAS_TORCH = importlib.util.find_spec("torch") is not None + + +@unittest.skipUnless(HAS_TORCH, "requires torch") +class TestMaybeMirrorAndPredict(unittest.TestCase): + """`_internal_maybe_mirror_and_predict` averages the prediction over all axis flips.""" + + @staticmethod + def _predictor(use_mirroring: bool = True): + import torch # noqa: PLC0415 + + from TPTBox.segmentation.nnUnet_utils.predictor import nnUNetPredictor # noqa: PLC0415 + + p = nnUNetPredictor(device=torch.device("cpu"), use_mirroring=use_mirroring, allow_tqdm=False) + p.allowed_mirroring_axes = (0, 1, 2) + return p + + @staticmethod + def _x(): + import torch # noqa: PLC0415 + + return torch.arange(2 * 4 * 5 * 6, dtype=torch.float32).reshape(1, 2, 4, 5, 6) + + def test_uses_the_network_argument_for_every_tta_term(self): + # Regression: the first forward pass used the `network` argument while the seven + # mirrored passes used `self.network`. Both are the same object on every code path + # today, so results were unaffected -- but the argument was silently ignored, which + # is exactly what `loaded_networks` (one network per fold) is meant to exercise. + import torch # noqa: PLC0415 + + p = self._predictor() + p.network = lambda t: torch.zeros_like(t) # must NOT be consulted + x = self._x() + out = p._internal_maybe_mirror_and_predict(x, network=lambda t: t.clone()) + torch.testing.assert_close(out, x) + + def test_mirroring_averages_eight_flips(self): + import torch # noqa: PLC0415 + + p = self._predictor() + x = self._x() + # A network that returns a constant c makes every one of the 8 terms c. + out = p._internal_maybe_mirror_and_predict(x, network=lambda t: torch.full_like(t, 3.0)) + torch.testing.assert_close(out, torch.full_like(x, 3.0)) + + def test_asymmetric_network_matches_hand_symmetrised_average(self): + import itertools # noqa: PLC0415 + + import torch # noqa: PLC0415 + + def net(t): + return t * 2.0 + 1.0 + + p = self._predictor() + x = self._x() + out = p._internal_maybe_mirror_and_predict(x, network=net) + + axes = [2, 3, 4] + combos = [c for i in range(len(axes)) for c in itertools.combinations(axes, i + 1)] + expected = net(x) + for c in combos: + expected = expected + torch.flip(net(torch.flip(x, c)), c) + expected = expected / (len(combos) + 1) + torch.testing.assert_close(out, expected) + + def test_mirroring_disabled_is_a_single_forward(self): + import torch # noqa: PLC0415 + + p = self._predictor(use_mirroring=False) + p.network = lambda t: torch.zeros_like(t) + x = self._x() + out = p._internal_maybe_mirror_and_predict(x, network=lambda t: t * 5.0) + torch.testing.assert_close(out, x * 5.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_nputils_extended.py b/unit_tests/test_nputils_extended.py index 6f6e81b..8e48206 100644 --- a/unit_tests/test_nputils_extended.py +++ b/unit_tests/test_nputils_extended.py @@ -473,5 +473,216 @@ def test_default_k_returns_two(self): self.assertEqual(len(idx), 2) +class TestNpIsin(unittest.TestCase): + """`np_isin` is a lookup-table fast path for `np.isin` on label arrays.""" + + def test_matches_numpy_isin(self): + rng = np.random.default_rng(0) + arr = rng.integers(0, 12, size=(8, 9, 10)).astype(np.uint8) + for labels in ([3], [1, 5, 7], [], [0], list(range(12))): + with self.subTest(labels=labels): + np.testing.assert_array_equal(np_utils.np_isin(arr, labels), np.isin(arr, labels)) + + def test_invert_matches_numpy(self): + rng = np.random.default_rng(1) + arr = rng.integers(0, 8, size=(6, 6, 6)).astype(np.uint8) + labels = [2, 4] + np.testing.assert_array_equal(np_utils.np_isin(arr, labels, invert=True), np.isin(arr, labels, invert=True)) + + def test_labels_outside_the_array_range(self): + arr = np.array([[[0, 1, 2]]], dtype=np.uint8) + np.testing.assert_array_equal(np_utils.np_isin(arr, [99]), np.zeros_like(arr, dtype=bool)) + np.testing.assert_array_equal(np_utils.np_isin(arr, [99], invert=True), np.ones_like(arr, dtype=bool)) + + def test_result_is_boolean_and_shape_preserving(self): + arr = np.zeros((3, 4, 5), dtype=np.uint16) + out = np_utils.np_isin(arr, [0]) + self.assertEqual(out.dtype, np.dtype(bool)) + self.assertEqual(out.shape, arr.shape) + + +class TestNpCalcBoundaryMask(unittest.TestCase): + """Floods the sub-threshold region inward from the image border.""" + + def test_an_all_background_image_marks_nothing(self): + out = np_utils.np_calc_boundary_mask(np.zeros((8, 8, 8), dtype=np.int16), threshold=0) + self.assertEqual(out.shape, (8, 8, 8)) + self.assertEqual(int((out != 0).sum()), 0) + + def test_an_enclosed_bright_region_is_marked(self): + img = np.full((10, 10, 10), -1000, dtype=np.int16) + img[3:7, 3:7, 3:7] = 300 + out = np_utils.np_calc_boundary_mask(img, threshold=0) + self.assertEqual(int((out != 0).sum()), 4 * 4 * 4, "the above-threshold block is what remains marked") + + def test_ct_mode_shifts_the_intensities_by_1000(self): + # Soft tissue at -500 HU is below a 0 threshold on the raw values but above it + # once the CT offset is applied, so the two modes must disagree here. + img = np.full((10, 10, 10), -1000, dtype=np.int16) + img[3:7, 3:7, 3:7] = -500 + plain = np_utils.np_calc_boundary_mask(img, threshold=0) + ct = np_utils.np_calc_boundary_mask(img, threshold=0, adjust_intensity_for_ct=True) + self.assertEqual(int((plain != 0).sum()), 0) + self.assertEqual(int((ct != 0).sum()), 4 * 4 * 4) + + def test_threshold_selects_what_counts_as_foreground(self): + # Regression: the function wrote its sentinel value 2 into the array and then + # compared that array against the threshold, so every threshold >= 2 erased the + # foreground and the mask came back empty. NII.boundary_mask(threshold) forwards + # straight here, so any realistic HU threshold was affected. + img = np.full((10, 10, 10), -1000, dtype=np.int16) + img[3:7, 3:7, 3:7] = 300 + for threshold in (0, 2, 100, 299): + with self.subTest(threshold=threshold): + self.assertEqual(int((np_utils.np_calc_boundary_mask(img, threshold=threshold) != 0).sum()), 4 * 4 * 4) + self.assertEqual(int((np_utils.np_calc_boundary_mask(img, threshold=500) != 0).sum()), 0, "nothing is above 500") + + +class TestNpVoxelConnectivityGraph(unittest.TestCase): + """Per-voxel bit field of which neighbours carry the same label.""" + + def test_a_lone_voxel_has_no_same_label_neighbour(self): + arr = np.zeros((5, 5, 5), dtype=np.uint8) + arr[2, 2, 2] = 1 + graph = np_utils.np_voxel_connectivity_graph(arr, connectivity=1) + self.assertEqual(graph.shape, arr.shape) + self.assertEqual(int(graph[2, 2, 2]), 0) + # Background is a label too, so background voxels *are* connected to each other. + self.assertGreater(int(graph[0, 0, 0]), 0) + + def test_two_touching_voxels_are_connected(self): + arr = np.zeros((5, 5, 5), dtype=np.uint8) + arr[2, 2, 2] = 1 + arr[2, 2, 3] = 1 + graph = np_utils.np_voxel_connectivity_graph(arr, connectivity=1) + self.assertGreater(int(graph[2, 2, 2]), 0) + self.assertGreater(int(graph[2, 2, 3]), 0) + + def test_a_uniform_array_is_fully_connected(self): + arr = np.ones((6, 6, 6), dtype=np.uint8) + graph = np_utils.np_voxel_connectivity_graph(arr, connectivity=1) + self.assertEqual(int(graph[3, 3, 3]), 63, "all six 6-neighbourhood bits set in the interior") + + def test_higher_connectivity_reports_at_least_as_many_links(self): + arr = np.zeros((6, 6, 6), dtype=np.uint8) + arr[2:4, 2:4, 2:4] = 1 + low = np_utils.np_voxel_connectivity_graph(arr, connectivity=1) + high = np_utils.np_voxel_connectivity_graph(arr, connectivity=3) + self.assertGreaterEqual(int(np.sum(high != 0)), int(np.sum(low != 0))) + + +class TestNpFilterConnectedComponentsByBboxChain(unittest.TestCase): + @staticmethod + def _chain(): + # Three blobs stacked along axis 0 with a far-away outlier. + arr = np.zeros((40, 12, 12), dtype=np.uint8) + for z in (2, 8, 14): + arr[z : z + 3, 4:8, 4:8] = 1 + arr[34:37, 1:3, 1:3] = 1 # outlier, off the chain + return arr + + def test_the_outlier_is_dropped(self): + out = np_utils.np_filter_connected_components_by_bbox_chain(self._chain()) + self.assertEqual(int(out[35, 2, 2]), 0, "the off-chain blob must be removed") + self.assertGreater(int((out != 0).sum()), 0) + + def test_a_single_component_survives(self): + arr = np.zeros((20, 10, 10), dtype=np.uint8) + arr[5:8, 4:7, 4:7] = 1 + out = np_utils.np_filter_connected_components_by_bbox_chain(arr) + self.assertEqual(int((out != 0).sum()), int((arr != 0).sum())) + + def test_a_generous_margin_keeps_more(self): + tight = int((np_utils.np_filter_connected_components_by_bbox_chain(self._chain()) != 0).sum()) + loose = int( + ( + np_utils.np_filter_connected_components_by_bbox_chain(self._chain(), margin=30.0, extra_margin=30.0, extra_margin_axis=0) + != 0 + ).sum() + ) + self.assertGreaterEqual(loose, tight) + + def test_extra_margin_requires_an_axis(self): + with self.assertRaises(AssertionError): + np_utils.np_filter_connected_components_by_bbox_chain(self._chain(), extra_margin=5.0) + + +class TestNpCalcConvexHull(unittest.TestCase): + def test_hull_fills_a_concave_notch(self): + arr = np.zeros((20, 20, 20), dtype=np.uint8) + arr[5:15, 5:15, 5:15] = 1 + arr[5:10, 5:10, 5:15] = 0 # bite a corner out + hull = np_utils.np_calc_convex_hull(arr, axis=None) + self.assertGreater(int((hull != 0).sum()), int((arr != 0).sum())) + + def test_a_convex_input_is_unchanged_in_volume(self): + arr = np.zeros((16, 16, 16), dtype=np.uint8) + arr[4:12, 4:12, 4:12] = 1 + hull = np_utils.np_calc_convex_hull(arr, axis=None) + self.assertAlmostEqual(int((hull != 0).sum()), int((arr != 0).sum()), delta=int(0.05 * (arr != 0).sum())) + + def test_per_slice_hull_along_an_axis(self): + arr = np.zeros((12, 12, 12), dtype=np.uint8) + arr[4:8, 4:8, 4:8] = 1 + arr[5, 5, 5] = 0 # interior hole + hull = np_utils.np_calc_convex_hull(arr, axis=0) + self.assertEqual(int(hull[5, 5, 5]), 1, "the hole is filled inside its slice") + + +class TestNpErodeMskEuclid(unittest.TestCase): + def test_erosion_removes_a_shell_of_the_requested_thickness(self): + arr = np.zeros((20, 20, 20), dtype=np.uint8) + arr[5:15, 5:15, 5:15] = 1 + out = np_utils.np_erode_msk_euclid(arr.copy(), n_pixel=2) + self.assertEqual(int((out != 0).sum()), 6 * 6 * 6, "a 10-voxel cube eroded by 2 leaves a 6-voxel cube") + self.assertEqual(int(out[10, 10, 10]), 1, "the core survives") + self.assertEqual(int(out[5, 10, 10]), 0, "the surface is removed") + + def test_more_erosion_removes_more(self): + arr = np.zeros((24, 24, 24), dtype=np.uint8) + arr[4:20, 4:20, 4:20] = 1 + light = int((np_utils.np_erode_msk_euclid(arr.copy(), n_pixel=1) != 0).sum()) + heavy = int((np_utils.np_erode_msk_euclid(arr.copy(), n_pixel=4) != 0).sum()) + self.assertGreater(light, heavy) + + def test_labels_restricts_which_label_is_eroded(self): + arr = np.zeros((20, 20, 20), dtype=np.uint8) + arr[2:9, 2:9, 2:9] = 1 + arr[11:18, 11:18, 11:18] = 2 + before_1 = int((arr == 1).sum()) + out = np_utils.np_erode_msk_euclid(arr.copy(), n_pixel=2, labels=[1]) + self.assertEqual(int((out == 1).sum()), 3 * 3 * 3) + self.assertLess(int((out == 1).sum()), before_1) + self.assertGreater(int((out == 2).sum()), 300, "label 2 is left essentially untouched") + + def test_eroding_a_thin_structure_away_entirely(self): + arr = np.zeros((10, 10, 10), dtype=np.uint8) + arr[5, 5, 5] = 1 + self.assertEqual(int((np_utils.np_erode_msk_euclid(arr, n_pixel=2) != 0).sum()), 0) + + def test_use_crop_does_not_change_the_result(self): + arr = np.zeros((20, 20, 20), dtype=np.uint8) + arr[5:15, 5:15, 5:15] = 1 + cropped = np_utils.np_erode_msk_euclid(arr.copy(), n_pixel=2, use_crop=True) + plain = np_utils.np_erode_msk_euclid(arr.copy(), n_pixel=2, use_crop=False) + np.testing.assert_array_equal(cropped, plain) + + def test_the_morphology_helpers_work_in_place(self): + # Undocumented but consistent convention across np_erode_msk, + # np_erode_msk_euclid and np_dilate_msk_euclid: the input array is modified and + # also returned. Callers that need their array intact must pass a copy. + for name, setup in ( + ("np_erode_msk", lambda a: a.__setitem__((slice(5, 15),) * 3, 1)), + ("np_erode_msk_euclid", lambda a: a.__setitem__((slice(5, 15),) * 3, 1)), + ("np_dilate_msk_euclid", lambda a: a.__setitem__((slice(8, 12),) * 3, 1)), + ): + with self.subTest(function=name): + arr = np.zeros((20, 20, 20), dtype=np.uint8) + setup(arr) + snapshot = arr.copy() + getattr(np_utils, name)(arr, n_pixel=2) + self.assertFalse(np.array_equal(arr, snapshot), "the caller's array is modified in place") + + if __name__ == "__main__": unittest.main() diff --git a/unit_tests/test_poi_internals.py b/unit_tests/test_poi_internals.py new file mode 100644 index 0000000..e2ed63b --- /dev/null +++ b/unit_tests/test_poi_internals.py @@ -0,0 +1,397 @@ +"""Unit tests for POI serialization and the descriptor internals. + +`unit_tests/test_poi_save.py` round-trips the current save format; this file covers +the readers for the other formats the loader still accepts (landmark TXT, Slicer +.mrk.json, the two legacy POI layouts), the header/coordinate parsers underneath +them, and the `POI_Descriptor` name-resolution and set-algebra methods. +""" + +from __future__ import annotations + +import json +import tempfile +import unittest +from pathlib import Path + +import numpy as np + +from TPTBox import POI, Location, Vertebra_Instance +from TPTBox.core.poi_fun.poi_abstract import POI_Descriptor, label_name_dict, unpack_poi_id +from TPTBox.core.poi_fun.poi_global import POI_Global +from TPTBox.core.poi_fun.save_load import ( + _is_Point3D, + _load_landmark_txt, + _parse_coords, + _parse_header_value, + load_poi, + save_poi, +) + + +def _poi(points=None) -> POI: + points = points or {20: {50: (1.0, 2.0, 3.0)}, 21: {50: (4.0, 5.0, 6.0)}} + return POI( + points, + orientation=("R", "A", "S"), + zoom=(1.0, 1.0, 1.0), + shape=(50.0, 50.0, 50.0), + origin=(0.0, 0.0, 0.0), + rotation=np.eye(3), + level_one_info=Vertebra_Instance, + level_two_info=Location, + ) + + +class TestParseCoords(unittest.TestCase): + def test_valid_triples(self): + for text, expected in (("(1, 2, 3)", [1.0, 2.0, 3.0]), ("( -1.5,2 , 3e2 )", [-1.5, 2.0, 300.0])): + with self.subTest(text=text): + self.assertEqual(_parse_coords(text), expected) + + def test_rejects_malformed_input(self): + for text in ("1, 2, 3", "(1, 2)", "(1, 2, 3, 4)", "()"): + with self.subTest(text=text), self.assertRaises(ValueError): + _parse_coords(text) + + +class TestParseHeaderValue(unittest.TestCase): + def test_documented_examples(self): + cases = [ + ("3.14", 3.14), + ("7", 7), + ("-3.5e-12", -3.5e-12), + ("256 931 27", [256, 931, 27]), + ("[1, 2, 3]", [1, 2, 3]), + ("[[1, 0, 0], [0, 1, 0]]", [[1, 0, 0], [0, 1, 0]]), + ] + for text, expected in cases: + with self.subTest(text=text): + self.assertEqual(_parse_header_value(text), expected) + + def test_types_are_preserved(self): + self.assertIsInstance(_parse_header_value("7"), int) + self.assertIsInstance(_parse_header_value("7.5"), float) + self.assertIsInstance(_parse_header_value("[1, 2]"), list) + self.assertIsInstance(_parse_header_value("plain text"), str) + + def test_a_whole_float_collapses_to_int(self): + # "7.0" parses as int(7), not float(7.0) - int conversion is tried first. + self.assertIsInstance(_parse_header_value("7.0"), int) + + def test_whitespace_is_stripped(self): + self.assertEqual(_parse_header_value(" 42 "), 42) + + +class TestIsPoint3D(unittest.TestCase): + def test_accepts_a_full_point(self): + self.assertTrue(_is_Point3D({"X": 1.0, "Y": 2.0, "Z": 3.0, "label": 1, "desc": "d"})) + + def test_rejects_incomplete_or_wrongly_typed_dicts(self): + for obj in ({"X": 1.0, "Y": 2.0}, {"X": "a", "Y": 2.0, "Z": 3.0}, {}, []): + with self.subTest(obj=obj): + self.assertFalse(_is_Point3D(obj)) + + def test_non_container_input_is_not_supported(self): + # Documented limitation: the guard does membership tests, so a scalar raises. + for obj in (None, 5): + with self.subTest(obj=obj), self.assertRaises(TypeError): + _is_Point3D(obj) + + +class TestLandmarkTxt(unittest.TestCase): + TEXT = """some_meta: 256 931 27 +GroupA: +left_hip: (1.0, 2.0, 3.0) +right_hip: (4.5, 5.5, 6.5) +GroupB: +apex: (7.0, 8.0, 9.0) +""" + + def _write(self, folder: Path, text: str | None = None) -> Path: + path = folder / "landmarks.txt" + path.write_text(self.TEXT if text is None else text, encoding="utf-8") + return path + + def test_header_and_points_are_split(self): + with tempfile.TemporaryDirectory() as tmp: + header, points = _load_landmark_txt(self._write(Path(tmp))) + self.assertEqual(header["coordinate_system"], "nib") + self.assertEqual(header["some_meta"], [256, 931, 27]) + self.assertEqual(len(points), 2, "two groups") + + def test_group_names_and_point_names_are_recorded(self): + with tempfile.TemporaryDirectory() as tmp: + header, _ = _load_landmark_txt(self._write(Path(tmp))) + self.assertIn("GroupA", header["label_group_name"].values()) + self.assertIn("GroupB", header["label_group_name"].values()) + flat_names = [n for group in header["label_name"].values() for n in group.values()] + self.assertIn("left_hip", flat_names) + self.assertIn("apex", flat_names) + + def test_coordinates_are_parsed(self): + with tempfile.TemporaryDirectory() as tmp: + _, points = _load_landmark_txt(self._write(Path(tmp))) + all_coords = [tuple(c) for group in points.values() for c in group.values()] + self.assertIn((1.0, 2.0, 3.0), all_coords) + self.assertIn((7.0, 8.0, 9.0), all_coords) + + def test_blank_and_colonless_lines_are_skipped(self): + text = "some_meta: 1\n\nno colon here\nGroupA:\np: (1, 1, 1)\n" + with tempfile.TemporaryDirectory() as tmp: + _, points = _load_landmark_txt(self._write(Path(tmp), text)) + self.assertEqual(sum(len(g) for g in points.values()), 1) + + def test_loading_through_load_poi_gives_a_poi(self): + with tempfile.TemporaryDirectory() as tmp: + poi = load_poi(self._write(Path(tmp)), verbose=False) + self.assertEqual(len(list(poi.keys())), 3) + + +class TestMarkupJson(unittest.TestCase): + """Slicer .mrk.json is read as a world-space POI_Global.""" + + def _mrk(self) -> dict: + return { + "markups": [ + { + "type": "Fiducial", + "coordinateSystem": "LPS", + "controlPoints": [ + {"id": "1", "label": "L1", "position": [1.0, 2.0, 3.0], "description": ""}, + {"id": "2", "label": "L2", "position": [4.0, 5.0, 6.0], "description": ""}, + ], + } + ] + } + + def test_round_trip_through_save_and_load(self): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "points.mrk.json" + path.write_text(json.dumps(self._mrk()), encoding="utf-8") + poi = load_poi(path, verbose=False) + self.assertIsInstance(poi, POI_Global) + self.assertEqual(len(list(poi.keys())), 2) + + def test_saving_a_global_poi_as_mrk_and_reading_it_back(self): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "out.mrk.json" + glob = _poi().to_global() + glob.save_mrk(path) + self.assertTrue(path.exists()) + reloaded = load_poi(path, verbose=False) + self.assertEqual(sorted(reloaded.keys()), sorted(glob.keys())) + for key in glob.keys(): + with self.subTest(key=key): + np.testing.assert_allclose(reloaded[key], glob[key], atol=1e-4) + + +class TestSavePoiFormats(unittest.TestCase): + def test_every_save_hint_round_trips(self): + # hint 0 = docker, 1 = gruber, 2 = POI (the default). Gruber names its points + # from a fixed landmark vocabulary, so it gets points from that vocabulary. + gruber_points = {20: {101: (1.0, 2.0, 3.0)}, 21: {102: (4.0, 5.0, 6.0)}} + for hint in (0, 1, 2): + with self.subTest(save_hint=hint), tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "poi.json" + poi = _poi(gruber_points if hint == 1 else None) + save_poi(poi, path, save_hint=hint, verbose=False) + self.assertTrue(path.exists()) + reloaded = load_poi(path, verbose=False) + self.assertEqual(sorted(reloaded.keys()), sorted(poi.keys())) + for key in poi.keys(): + np.testing.assert_allclose(np.asarray(reloaded[key]), np.asarray(poi[key]), atol=1e-3) + + def test_the_gruber_format_rejects_points_outside_its_vocabulary(self): + # Documented limitation: save_hint=1 maps each subregion id through a fixed + # landmark-name table, so an ordinary corpus centroid (50) has no name. + with tempfile.TemporaryDirectory() as tmp, self.assertRaises(KeyError): + save_poi(_poi(), Path(tmp) / "poi.json", save_hint=1, verbose=False) + + def test_additional_info_is_persisted(self): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "poi.json" + save_poi(_poi(), path, additional_info={"my_key": "my_value"}, verbose=False) + self.assertIn("my_value", path.read_text(encoding="utf-8")) + + def test_make_parents_creates_the_folder(self): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "a" / "b" / "poi.json" + save_poi(_poi(), path, make_parents=True, verbose=False) + self.assertTrue(path.exists()) + + def test_round_trip_survives_reorient_and_rescale(self): + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "poi.json" + poi = _poi() + save_poi(poi, path, verbose=False) + reloaded = load_poi(path, verbose=False) + reloaded.reorient_(("P", "I", "R")).rescale_((2.0, 2.0, 2.0)).rescale_((1.0, 1.0, 1.0)).reorient_(("R", "A", "S")) + for key in poi.keys(): + with self.subTest(key=key): + np.testing.assert_allclose(np.asarray(reloaded[key]), np.asarray(poi[key]), atol=1e-3) + + +class TestPoiDescriptor(unittest.TestCase): + def test_str_to_int_resolves_names(self): + desc = POI_Descriptor() + self.assertEqual(desc.str_to_int("T12", subregion=False), Vertebra_Instance.T12.value) + self.assertEqual(desc.str_to_int("Vertebra_Corpus", subregion=True), Location.Vertebra_Corpus.value) + + def test_str_to_int_list(self): + desc = POI_Descriptor() + self.assertEqual(desc.str_to_int_list("L1", "L2"), [Vertebra_Instance.L1.value, Vertebra_Instance.L2.value]) + + def test_str_to_int_dict(self): + desc = POI_Descriptor() + out = desc.str_to_int_dict({"L1": "L2"}) + self.assertEqual(out, {Vertebra_Instance.L1.value: Vertebra_Instance.L2.value}) + + def test_get_and_pop(self): + poi = _poi() + desc = poi.centroids + np.testing.assert_allclose(desc.get((20, 50)), (1.0, 2.0, 3.0)) + popped = desc.pop((20, 50), None) + np.testing.assert_allclose(popped, (1.0, 2.0, 3.0)) + self.assertNotIn((20, 50), poi, "pop removes the point") + self.assertIsNone(desc.pop((20, 50), None), "popping a missing key returns the default") + + def test_items_2d(self): + poi = _poi() + nested = dict(poi.centroids.items_2d()) + self.assertEqual(sorted(nested), [20, 21]) + self.assertEqual(sorted(nested[20]), [50]) + + def test_delitem(self): + poi = _poi() + del poi.centroids[20, 50] + self.assertNotIn((20, 50), poi) + + def test_normalize_input_data_accepts_flat_and_nested(self): + flat = POI_Descriptor.normalize_input_data({(1, 50): (0.0, 0.0, 0.0)}) + nested = POI_Descriptor.normalize_input_data({1: {50: (0.0, 0.0, 0.0)}}) + self.assertEqual(sorted(flat.keys()), sorted(nested.keys())) + + def test_unpack_poi_id_accepts_every_key_form(self): + definition = _poi().centroids.definition + packed = unpack_poi_id((20, 50), definition) + self.assertEqual(packed, (20, 50)) + self.assertEqual(unpack_poi_id(slice(20, 50), definition), (20, 50)) + self.assertEqual(unpack_poi_id((Vertebra_Instance.L1, Location.Vertebra_Corpus), definition), (20, 50)) + + def test_unpack_poi_id_round_trips_a_packed_integer(self): + from TPTBox.core import vert_constants # noqa: PLC0415 + + definition = _poi().centroids.definition + key = 20 + 50 * vert_constants.LABEL_MAX + self.assertEqual(unpack_poi_id(key, definition), (20, 50)) + + def test_label_name_dict_normalises_keys(self): + self.assertEqual(label_name_dict({"label_name": {"1": {"2": "C2"}}}), {1: {2: "C2"}}) + + +class TestAbstractPoiOperations(unittest.TestCase): + def test_extract_keeps_only_the_requested_pairs(self): + poi = _poi({20: {50: (1.0, 1.0, 1.0), 51: (2.0, 2.0, 2.0)}, 21: {50: (3.0, 3.0, 3.0)}}) + out = poi.extract((20, 50)) + self.assertEqual(sorted(out.keys()), [(20, 50)]) + self.assertEqual(len(list(poi.keys())), 3, "the source must be untouched") + + def test_distances_between_two_pois(self): + a = _poi({20: {50: (0.0, 0.0, 0.0)}}) + b = _poi({20: {50: (3.0, 4.0, 0.0)}}) + out = a.calculate_distances_poi(b) + self.assertAlmostEqual(float(next(iter(out.values()))), 5.0, places=6) + + def test_distances_any_to_any(self): + a = _poi({20: {50: (0.0, 0.0, 0.0)}}) + b = _poi({20: {50: (3.0, 4.0, 0.0)}, 21: {50: (0.0, 0.0, 10.0)}}) + out = a.calculate_distances_poi_any_2_any(b) + self.assertTrue(out) + flat = [float(v) for inner in out.values() for v in (inner.values() if isinstance(inner, dict) else [inner])] + self.assertAlmostEqual(min(flat), 5.0, places=6) + + def test_distances_across_regions(self): + a = _poi({20: {50: (0.0, 0.0, 0.0)}, 21: {50: (0.0, 0.0, 6.0)}}) + out = a.calculate_distances_poi_across_regions(a.copy()) + self.assertTrue(out) + + def test_operators(self): + a = _poi({20: {50: (1.0, 1.0, 1.0)}}) + b = _poi({21: {50: (2.0, 2.0, 2.0)}}) + self.assertEqual(sorted((a + b).keys()), [(20, 50), (21, 50)]) + self.assertEqual(sorted((a >> b).keys()), [(20, 50), (21, 50)]) + self.assertEqual(sorted((a - a.copy()).keys()), []) + self.assertEqual(sorted((a & a.copy()).keys()), [(20, 50)]) + + def test_remove_centroid_inplace(self): + poi = _poi() + poi.remove_centroid_((20, 50)) + self.assertNotIn((20, 50), poi) + + def test_join_right_inplace(self): + a = _poi({20: {50: (1.0, 1.0, 1.0)}}) + b = _poi({21: {50: (2.0, 2.0, 2.0)}}) + a.join_right_(b) + self.assertEqual(sorted(a.keys()), [(20, 50), (21, 50)]) + + def test_fit_spline_through_four_points(self): + poi = _poi({v: {Location.Vertebra_Corpus.value: (0.0, 0.0, float(10 * i))} for i, v in enumerate((20, 21, 22, 23))}) + points, derivatives = poi.fit_spline(location=Location.Vertebra_Corpus, vertebra=True) + self.assertGreater(len(points), 4) + self.assertEqual(np.asarray(points).shape[1], 3) + self.assertEqual(np.asarray(derivatives).shape[1], 3) + + +class TestPoiCropAndGlobal(unittest.TestCase): + def test_apply_crop_reverse_is_the_inverse_of_apply_crop(self): + poi = _poi({20: {50: (10.0, 12.0, 14.0)}}) + crop = (slice(5, 40), slice(5, 40), slice(5, 40)) + cropped = poi.apply_crop(crop) + restored = cropped.apply_crop_reverse(crop, poi.shape) + np.testing.assert_allclose(restored[20, 50], poi[20, 50], atol=1e-9) + self.assertEqual(restored.shape, poi.shape) + + def test_local_to_global_arr_round_trip(self): + poi = _poi() + points = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + world = poi.local_to_global_arr(points) + np.testing.assert_allclose(poi.global_to_local_arr(world), points, atol=1e-9) + + def test_to_global_and_back(self): + poi = _poi() + back = poi.to_global().to_local(poi) + for key in poi.keys(): + with self.subTest(key=key): + np.testing.assert_allclose(back[key], poi[key], atol=1e-6) + + def test_to_cord_system_is_an_involution(self): + glob = _poi().to_global() + flipped = glob.to_cord_system(itk_coords=True) + back = flipped.to_cord_system(itk_coords=False) + for key in glob.keys(): + with self.subTest(key=key): + np.testing.assert_allclose(back[key], glob[key], atol=1e-9) + + def test_to_other_maps_onto_another_grid(self): + poi = _poi() + other = poi.copy() + other.rescale_((2.0, 2.0, 2.0)) + out = poi.to_global().to_other(other) + self.assertEqual(sorted(out.keys()), sorted(poi.keys())) + + def test_make_point_cloud_nii(self): + poi = _poi() + cloud, sphere = poi.make_point_cloud_nii(s=3) + self.assertGreater(int((cloud.get_seg_array() != 0).sum()), 0) + self.assertGreater(int((sphere.get_seg_array() != 0).sum()), 0) + + def test_resample_from_to_inplace(self): + poi = _poi() + target = poi.copy() + target.rescale_((2.0, 2.0, 2.0)) + poi.resample_from_to_(target) + np.testing.assert_allclose(poi.zoom, target.zoom) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_point_registration.py b/unit_tests/test_point_registration.py new file mode 100644 index 0000000..b28c0fe --- /dev/null +++ b/unit_tests/test_point_registration.py @@ -0,0 +1,248 @@ +"""Unit tests for `TPTBox.registration._ridged_points.point_registration`. + +Rigid point registration needs only SimpleITK (a core dependency), so these run +everywhere. Every expectation is a known transform or an inverse property: two POI +sets related by an exact translation or rotation must be recovered to sub-voxel +accuracy, forward and inverse coordinate mapping must compose to the identity, and +a save/load round-trip must reproduce the same transform. +""" + +from __future__ import annotations + +import math +import tempfile +import unittest +from pathlib import Path + +import numpy as np + +from TPTBox import POI, Location, Vertebra_Instance +from TPTBox.registration._ridged_points.point_registration import ( + Point_Registration, + ridged_points_from_poi, +) + +SHAPE = (60.0, 60.0, 60.0) +RNG = np.random.default_rng(0) + + +def _poi(points: dict) -> POI: + return POI( + points, + orientation=("R", "A", "S"), + zoom=(1.0, 1.0, 1.0), + shape=SHAPE, + origin=(0.0, 0.0, 0.0), + rotation=np.eye(3), + level_one_info=Vertebra_Instance, + level_two_info=Location, + ) + + +def _landmarks() -> dict: + """Six well-spread landmarks, enough to determine a rigid transform.""" + coords = [(10.0, 12.0, 14.0), (40.0, 12.0, 14.0), (10.0, 44.0, 14.0), (10.0, 12.0, 46.0), (30.0, 30.0, 30.0), (20.0, 40.0, 20.0)] + return {20 + i: {Location.Vertebra_Corpus.value: c} for i, c in enumerate(coords)} + + +def _shifted(delta=(5.0, -3.0, 2.0)) -> dict: + return { + vert: {sub: tuple(float(c + d) for c, d in zip(coord, delta)) for sub, coord in inner.items()} + for vert, inner in _landmarks().items() + } + + +def _rotated(deg=12.0) -> dict: + t = math.radians(deg) + c, s = math.cos(t), math.sin(t) + rot = np.array([[c, -s, 0.0], [s, c, 0.0], [0.0, 0.0, 1.0]]) + centre = np.array([30.0, 30.0, 30.0]) + return { + vert: {sub: tuple(float(v) for v in (rot @ (np.asarray(coord) - centre) + centre)) for sub, coord in inner.items()} + for vert, inner in _landmarks().items() + } + + +class TestPointRegistrationFit(unittest.TestCase): + def test_identical_point_sets_give_a_near_zero_error(self): + reg = ridged_points_from_poi(_poi(_landmarks()), _poi(_landmarks()), verbose=False) + self.assertLess(reg.error_reg, 1e-6) + + def test_a_pure_translation_is_recovered(self): + fixed = _poi(_landmarks()) + moving = _poi(_shifted()) + reg = ridged_points_from_poi(fixed, moving, verbose=False) + self.assertLess(reg.error_reg, 1e-3, "a rigid transform fits a pure translation exactly") + out = reg.transform_poi(moving) + for key in fixed.keys(): + with self.subTest(key=key): + np.testing.assert_allclose(np.asarray(out[key]), np.asarray(fixed[key]), atol=0.01) + + def test_a_pure_rotation_is_recovered(self): + fixed = _poi(_landmarks()) + moving = _poi(_rotated()) + reg = ridged_points_from_poi(fixed, moving, verbose=False) + self.assertLess(reg.error_reg, 1e-3) + out = reg.transform_poi(moving) + for key in fixed.keys(): + with self.subTest(key=key): + np.testing.assert_allclose(np.asarray(out[key]), np.asarray(fixed[key]), atol=0.05) + + def test_error_natural_measures_non_rigidity_not_displacement(self): + # error_natural is the mean change in *pairwise landmark distances*, so any + # rigid motion leaves it at zero however far the points move. + for name, points in (("translated", _shifted(delta=(50.0, -30.0, 20.0))), ("rotated", _rotated(deg=25.0))): + with self.subTest(motion=name): + reg = ridged_points_from_poi(_poi(_landmarks()), _poi(points), verbose=False) + # The rotated case carries a little residue from the POI's stored + # coordinate precision; still two orders of magnitude below the + # deformation case below. + self.assertAlmostEqual(reg.error_natural, 0.0, delta=0.05) + + def test_error_natural_grows_when_the_point_set_is_deformed(self): + stretched = { + vert: {sub: (coord[0] * 1.5, coord[1], coord[2]) for sub, coord in inner.items()} for vert, inner in _landmarks().items() + } + reg = ridged_points_from_poi(_poi(_landmarks()), _poi(stretched), verbose=False) + self.assertGreater(reg.error_natural, 1.0, "a scaling cannot be absorbed by a rigid transform") + + def test_exclusion_drops_landmarks_from_the_fit(self): + fixed = _poi(_landmarks()) + moving_points = _shifted() + # Corrupt one landmark badly; excluding it must restore the clean fit. + bad_key = next(iter(moving_points)) + moving_points[bad_key][Location.Vertebra_Corpus.value] = (0.0, 0.0, 0.0) + moving = _poi(moving_points) + with_bad = ridged_points_from_poi(fixed, moving, verbose=False).error_reg + without = ridged_points_from_poi(fixed, moving, exclusion=[bad_key], verbose=False).error_reg + self.assertLess(without, with_bad) + self.assertLess(without, 1e-3) + + def test_leave_worst_percent_out_improves_the_residual(self): + fixed = _poi(_landmarks()) + moving_points = _shifted() + moving_points[next(iter(moving_points))][Location.Vertebra_Corpus.value] = (0.0, 0.0, 0.0) + moving = _poi(moving_points) + plain = ridged_points_from_poi(fixed, moving, verbose=False).error_reg + trimmed = ridged_points_from_poi(fixed, moving, leave_worst_percent_out=0.3, verbose=False).error_reg + self.assertLess(trimmed, plain) + + def test_out_of_range_trim_fraction_is_rejected(self): + fixed, moving = _poi(_landmarks()), _poi(_shifted()) + for value in (1.0, 1.5, -0.1): + with self.subTest(value=value), self.assertRaises(AssertionError): + ridged_points_from_poi(fixed, moving, leave_worst_percent_out=value, verbose=False) + + def test_too_few_shared_landmarks_is_an_error(self): + fixed = _poi({20: {Location.Vertebra_Corpus.value: (1.0, 1.0, 1.0)}}) + moving = _poi({21: {Location.Vertebra_Corpus.value: (2.0, 2.0, 2.0)}}) + with self.assertRaises(ValueError): + ridged_points_from_poi(fixed, moving, verbose=False) + + +class TestCoordinateTransforms(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.fixed = _poi(_landmarks()) + cls.moving = _poi(_shifted()) + cls.reg = ridged_points_from_poi(cls.fixed, cls.moving, verbose=False) + + def test_forward_then_inverse_is_the_identity(self): + for point in RNG.uniform(5.0, 50.0, size=(20, 3)): + with self.subTest(point=tuple(np.round(point, 2))): + mapped = self.reg.transform_cord(tuple(point)) + back = self.reg.transform_cord_inverse(tuple(mapped)) + np.testing.assert_allclose(np.asarray(back, dtype=float), point, atol=1e-4) + + def test_get_affine_is_consistent_with_transform_cord(self): + affine = self.reg.get_affine() + self.assertEqual(np.asarray(affine).shape, (4, 4)) + np.testing.assert_allclose(np.asarray(affine)[3], [0.0, 0.0, 0.0, 1.0], atol=1e-9) + + def test_transform_poi_inverse_undoes_transform_poi(self): + forward = self.reg.transform_poi(self.moving) + back = self.reg.transform_poi_inverse(forward) + for key in self.moving.keys(): + with self.subTest(key=key): + np.testing.assert_allclose(np.asarray(back[key]), np.asarray(self.moving[key]), atol=0.05) + + def test_transform_dispatches_on_the_input_type(self): + out = self.reg.transform(self.moving) + self.assertIsInstance(out, POI) + + +class TestSerialization(unittest.TestCase): + def test_save_and_load_reproduce_the_transform(self): + fixed, moving = _poi(_landmarks()), _poi(_shifted()) + reg = ridged_points_from_poi(fixed, moving, verbose=False) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "reg.pkl" + reg.save(path) + self.assertTrue(path.exists()) + loaded = Point_Registration.load(path) + np.testing.assert_allclose(np.asarray(loaded.get_affine()), np.asarray(reg.get_affine()), atol=1e-9) + self.assertAlmostEqual(loaded.error_reg, reg.error_reg, places=9) + self.assertAlmostEqual(loaded.error_natural, reg.error_natural, places=9) + + def test_a_loaded_registration_maps_coordinates_identically(self): + fixed, moving = _poi(_landmarks()), _poi(_shifted()) + reg = ridged_points_from_poi(fixed, moving, verbose=False) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "reg.pkl" + reg.save(path) + loaded = Point_Registration.load(path) + for point in RNG.uniform(5.0, 50.0, size=(5, 3)): + with self.subTest(point=tuple(np.round(point, 2))): + np.testing.assert_allclose( + np.asarray(loaded.transform_cord(tuple(point)), dtype=float), + np.asarray(reg.transform_cord(tuple(point)), dtype=float), + atol=1e-9, + ) + + def test_get_dump_round_trips_through_load_(self): + fixed, moving = _poi(_landmarks()), _poi(_shifted()) + reg = ridged_points_from_poi(fixed, moving, verbose=False) + rebuilt = Point_Registration.load_(reg.get_dump()) + np.testing.assert_allclose(np.asarray(rebuilt.get_affine()), np.asarray(reg.get_affine()), atol=1e-9) + + def test_asymmetric_grids_survive_the_round_trip(self): + # Regression guard for the moving/fixed ordering in load_(): with different + # grids on the two sides a swap would be visible, unlike in a symmetric case. + fixed = _poi(_landmarks()) + moving = _poi(_shifted()) + moving.rescale_((2.0, 2.0, 2.0)) + reg = ridged_points_from_poi(fixed, moving, verbose=False) + rebuilt = Point_Registration.load_(reg.get_dump()) + np.testing.assert_allclose(np.asarray(rebuilt.input_poi.affine), np.asarray(reg.input_poi.affine), atol=1e-9) + np.testing.assert_allclose(np.asarray(rebuilt.out_poi.affine), np.asarray(reg.out_poi.affine), atol=1e-9) + np.testing.assert_allclose(np.asarray(rebuilt.get_affine()), np.asarray(reg.get_affine()), atol=1e-9) + + +class TestOnSampleData(unittest.TestCase): + """Registering the shipped CT's vertebra centroids against a shifted copy.""" + + @classmethod + def setUpClass(cls): + from TPTBox import calc_poi_from_subreg_vert # noqa: PLC0415 + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + cls.ct, cls.subreg, cls.vert, cls.label = get_test_ct() + cls.poi = calc_poi_from_subreg_vert(cls.vert, cls.subreg, subreg_id=[Location.Vertebra_Corpus], verbose=False) + + def test_registering_a_shifted_copy_recovers_the_shift(self): + moving = self.poi.copy() + shifted = {key: tuple(float(c + 3.0) for c in moving[key]) for key in moving.keys()} + for key, value in shifted.items(): + moving[key] = value + reg = ridged_points_from_poi(self.poi, moving, verbose=False) + self.assertLess(reg.error_reg, 0.5) + + def test_transform_nii_keeps_the_target_grid(self): + moving = self.poi.copy() + reg = ridged_points_from_poi(self.poi, moving, verbose=False) + out = reg.transform_nii(self.ct) + self.assertEqual(out.shape, self.ct.shape) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_ray_casting.py b/unit_tests/test_ray_casting.py new file mode 100644 index 0000000..87acdfc --- /dev/null +++ b/unit_tests/test_ray_casting.py @@ -0,0 +1,252 @@ +"""Unit tests for `TPTBox.core.poi_fun.ray_casting` and `pixel_based_point_finder`. + +These are pure numpy/NII geometry routines, so the expected answers are analytic: +trilinear interpolation must reproduce a linear field exactly, a ray through a box +of known extent must exit at a known distance, the two convex ray casters must agree +with each other, and a plane through three points must split a volume into the exact +voxel counts on either side. +""" + +from __future__ import annotations + +import unittest + +import nibabel as nib +import numpy as np + +from TPTBox import NII +from TPTBox.core.poi_fun.pixel_based_point_finder import cdist_to_point, get_nearest_neighbor, project_pois_onto_set_of_points +from TPTBox.core.poi_fun.ray_casting import ( + add_ray_to_img, + calculate_pca_normal_np, + max_distance_ray_cast_convex_np, + max_distance_ray_cast_convex_npfast, + max_distance_ray_cast_non_convex, + ray_cast_pixel_lvl, + set_label_above_3_point_plane, + trilinear_interpolate, +) + + +def _nii(arr: np.ndarray, seg: bool = True, zoom=(1.0, 1.0, 1.0)) -> NII: + return NII(nib.Nifti1Image(arr, np.diag([*zoom, 1.0])), seg=seg) + + +def _box(shape=(30, 30, 30), lo=(10, 10, 10), hi=(21, 21, 21)) -> np.ndarray: + arr = np.zeros(shape, dtype=np.uint8) + arr[lo[0] : hi[0], lo[1] : hi[1], lo[2] : hi[2]] = 1 + return arr + + +class TestTrilinearInterpolate(unittest.TestCase): + """Trilinear interpolation reproduces a linear field exactly.""" + + @staticmethod + def _linear_volume(shape=(8, 9, 10)): + gx, gy, gz = np.meshgrid(*[np.arange(s, dtype=np.float64) for s in shape], indexing="ij") + return 2.0 * gx + 3.0 * gy - 1.0 * gz + + def test_exact_on_a_linear_field(self): + volume = self._linear_volume() + for x, y, z in ((1.5, 2.25, 3.75), (0.5, 0.5, 0.5), (6.1, 7.2, 8.3)): + with self.subTest(point=(x, y, z)): + self.assertAlmostEqual(trilinear_interpolate(volume, x, y, z), 2.0 * x + 3.0 * y - z, places=9) + + def test_voxel_centres_return_the_stored_value(self): + volume = self._linear_volume() + for x, y, z in ((1, 2, 3), (4, 5, 6)): + with self.subTest(point=(x, y, z)): + self.assertAlmostEqual(trilinear_interpolate(volume, x, y, z), float(volume[x, y, z]), places=9) + + def test_outside_the_interior_returns_zero(self): + volume = self._linear_volume(shape=(5, 5, 5)) + for point in ((-1.0, 1.0, 1.0), (1.0, -0.5, 1.0), (4.0, 1.0, 1.0), (1.0, 1.0, 4.5)): + with self.subTest(point=point): + self.assertEqual(trilinear_interpolate(volume, *point), 0.0) + + +class TestConvexRayCast(unittest.TestCase): + def test_exit_point_of_a_box_is_its_far_face(self): + arr = _box() + start = np.array([15.0, 15.0, 15.0]) + out = max_distance_ray_cast_convex_np(arr, start, np.array([1.0, 0.0, 0.0])) + self.assertIsNotNone(out) + # the box spans index 10..20 inclusive, so the boundary sits at ~20.5 + self.assertAlmostEqual(float(np.asarray(out)[0]), 20.5, delta=0.6) + np.testing.assert_allclose(np.asarray(out)[1:], [15.0, 15.0], atol=1e-6) + + def test_the_two_implementations_agree(self): + arr = _box() + start = np.array([15.0, 15.0, 15.0]) + for direction in ([1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 1.0, 0.0]): + with self.subTest(direction=tuple(direction)): + slow = np.asarray(max_distance_ray_cast_convex_np(arr, start, np.array(direction))) + fast = np.asarray(max_distance_ray_cast_convex_npfast(arr, start, np.array(direction))) + np.testing.assert_allclose(slow, fast, atol=0.6) + + def test_distance_grows_with_the_region(self): + small = max_distance_ray_cast_convex_np(_box(hi=(16, 21, 21)), np.array([12.0, 15.0, 15.0]), np.array([1.0, 0.0, 0.0])) + large = max_distance_ray_cast_convex_np(_box(hi=(25, 21, 21)), np.array([12.0, 15.0, 15.0]), np.array([1.0, 0.0, 0.0])) + self.assertLess(float(np.asarray(small)[0]), float(np.asarray(large)[0])) + + def test_a_start_outside_the_region_is_returned_unchanged(self): + arr = _box() + start = np.array([1.0, 1.0, 1.0]) + out = max_distance_ray_cast_convex_np(arr, start, np.array([1.0, 0.0, 0.0])) + np.testing.assert_allclose(np.asarray(out), start) + + def test_non_convex_caster_stops_at_the_near_gap(self): + # A box with a slab removed in the middle: the bisecting convex caster reports + # the far outer wall, the stepping one stops at the near edge of the gap. + arr = _box() + arr[15:18, :, :] = 0 + start = np.array([12.0, 15.0, 15.0]) + direction = np.array([1.0, 0.0, 0.0]) + convex = float(np.asarray(max_distance_ray_cast_convex_np(arr, start, direction))[0]) + non_convex = max_distance_ray_cast_non_convex(_nii(arr), start, direction, step_size=0.25) + self.assertIsNotNone(non_convex) + self.assertLess(float(np.asarray(non_convex)[0]), convex) + self.assertAlmostEqual(float(np.asarray(non_convex)[0]), 15.0, delta=1.0) + + def test_non_convex_caster_returns_none_when_the_region_does_not_end_in_range(self): + # The box continues past 12 + 2 voxels, so there is no exit point to report. + arr = _box() + out = max_distance_ray_cast_non_convex( + _nii(arr), np.array([12.0, 15.0, 15.0]), np.array([1.0, 0.0, 0.0]), step_size=0.5, max_distance=2.0 + ) + self.assertIsNone(out) + + def test_non_convex_caster_finds_the_exit_with_enough_range(self): + arr = _box() + out = max_distance_ray_cast_non_convex( + _nii(arr), np.array([12.0, 15.0, 15.0]), np.array([1.0, 0.0, 0.0]), step_size=0.25, max_distance=50.0 + ) + self.assertIsNotNone(out) + self.assertAlmostEqual(float(np.asarray(out)[0]), 20.5, delta=1.0) + + +class TestRayCastPixelLevel(unittest.TestCase): + def test_every_voxel_is_inside_the_shape_and_ordered(self): + shape = (20, 20, 20) + coords, _ = ray_cast_pixel_lvl(np.array([10.0, 10.0, 2.0]), np.array([0.0, 0.0, 1.0]), shape) + self.assertIsNotNone(coords) + coords = np.asarray(coords) + self.assertTrue(((coords >= 0) & (coords < np.array(shape))).all()) + self.assertEqual(list(coords[:, 2]), sorted(coords[:, 2])) + + def test_two_sided_is_a_superset(self): + shape = (20, 20, 20) + one, _ = ray_cast_pixel_lvl(np.array([10.0, 10.0, 10.0]), np.array([0.0, 0.0, 1.0]), shape) + both, _ = ray_cast_pixel_lvl(np.array([10.0, 10.0, 10.0]), np.array([0.0, 0.0, 1.0]), shape, two_sided=True) + self.assertGreater(len(np.asarray(both)), len(np.asarray(one))) + self.assertTrue(((np.asarray(both) >= 0) & (np.asarray(both) < np.array(shape))).all()) + + def test_a_diagonal_ray_stays_in_bounds(self): + shape = (12, 12, 12) + coords, _ = ray_cast_pixel_lvl(np.array([1.0, 1.0, 1.0]), np.array([1.0, 1.0, 1.0]), shape) + coords = np.asarray(coords) + self.assertTrue(((coords >= 0) & (coords < np.array(shape))).all()) + + +class TestAddRayToImg(unittest.TestCase): + def test_painting_a_ray_adds_voxels(self): + arr = np.zeros((20, 20, 20), dtype=np.uint8) + out = add_ray_to_img(np.array([10, 10, 2]), np.array([0.0, 0.0, 1.0]), _nii(arr), value=5, dilate=1) + self.assertIsNotNone(out) + self.assertGreater(int((out.get_seg_array() != 0).sum()), 0) + + def test_inplace_modifies_the_input(self): + arr = np.zeros((20, 20, 20), dtype=np.uint8) + nii = _nii(arr) + add_ray_to_img(np.array([10, 10, 2]), np.array([0.0, 0.0, 1.0]), nii, value=5, dilate=1, inplace=True) + self.assertGreater(int((nii.get_seg_array() != 0).sum()), 0) + + def test_dilate_grows_the_ray(self): + arr = np.zeros((20, 20, 20), dtype=np.uint8) + thin = add_ray_to_img(np.array([10, 10, 2]), np.array([0.0, 0.0, 1.0]), _nii(arr), value=5, dilate=1) + thick = add_ray_to_img(np.array([10, 10, 2]), np.array([0.0, 0.0, 1.0]), _nii(arr), value=5, dilate=3) + self.assertLess(int((thin.get_seg_array() != 0).sum()), int((thick.get_seg_array() != 0).sum())) + + +class TestSetLabelAbovePlane(unittest.TestCase): + def test_axis_aligned_plane_splits_the_volume_exactly(self): + arr = np.ones((10, 10, 10), dtype=np.uint8) + # A plane at z = 4 spanned by three points in that plane. + out = set_label_above_3_point_plane(arr.copy(), [0.0, 0.0, 4.0], [9.0, 0.0, 4.0], [0.0, 9.0, 4.0], value=0) + out = np.asarray(out if isinstance(out, np.ndarray) else out.get_seg_array()) + zeroed = int((out == 0).sum()) + self.assertEqual(zeroed + int((out == 1).sum()), arr.size) + self.assertGreater(zeroed, 0) + # every zeroed voxel is on one side of z = 4 + zs = np.argwhere(out == 0)[:, 2] + self.assertTrue((zs >= 4).all() or (zs <= 4).all()) + + def test_invert_flips_which_side_is_cleared(self): + arr = np.ones((10, 10, 10), dtype=np.uint8) + kwargs = {"p1": [0.0, 0.0, 4.0], "p2": [9.0, 0.0, 4.0], "p3": [0.0, 9.0, 4.0], "value": 0} + up = np.asarray(set_label_above_3_point_plane(arr.copy(), invert=1, **kwargs)) + down = np.asarray(set_label_above_3_point_plane(arr.copy(), invert=-1, **kwargs)) + self.assertNotEqual(int((up == 0).sum()), int((down == 0).sum())) + # The two strict sides plus the z == 4 plane itself account for the whole volume. + on_plane = arr[:, :, 4].size + self.assertEqual(int((up == 0).sum()) + int((down == 0).sum()) + on_plane, arr.size) + + def test_accepts_a_nii_and_returns_a_nii(self): + nii = _nii(np.ones((8, 8, 8), dtype=np.uint8)) + out = set_label_above_3_point_plane(nii, [0.0, 0.0, 4.0], [7.0, 0.0, 4.0], [0.0, 7.0, 4.0], value=0) + self.assertIsInstance(out, NII) + self.assertLess(int(out.get_seg_array().sum()), int(nii.get_seg_array().sum())) + + +class TestPcaNormal(unittest.TestCase): + def test_first_component_of_an_elongated_blob(self): + arr = np.zeros((40, 12, 12), dtype=np.uint8) + arr[4:36, 5:8, 5:8] = 1 + normal = calculate_pca_normal_np(arr, pca_component=0) + self.assertAlmostEqual(abs(float(np.dot(normal / np.linalg.norm(normal), [1.0, 0.0, 0.0]))), 1.0, places=3) + + def test_zoom_only_rescales_the_returned_vector(self): + # Pins the current semantics: the PCA is fitted on raw voxel indices, and `zoom` + # is applied afterwards as a component-wise division of the resulting vector. It + # therefore does not change *which* axis is principal, and the result is no + # longer a unit vector. See the PR description. + arr = np.zeros((20, 20, 12), dtype=np.uint8) + arr[4:16, 6:14, 5:8] = 1 + zoom = (1.0, 4.0, 2.0) + plain = calculate_pca_normal_np(arr, pca_component=1) + scaled = calculate_pca_normal_np(arr, pca_component=1, zoom=zoom) + np.testing.assert_allclose(scaled, plain / np.array(zoom), atol=1e-12) + self.assertAlmostEqual(float(np.linalg.norm(plain)), 1.0, places=9) + + def test_components_are_orthogonal(self): + arr = np.zeros((30, 20, 12), dtype=np.uint8) + arr[4:26, 4:16, 4:8] = 1 + first = calculate_pca_normal_np(arr, 0) + second = calculate_pca_normal_np(arr, 1) + self.assertAlmostEqual(float(np.dot(first, second)), 0.0, places=5) + + +class TestPixelBasedPointFinder(unittest.TestCase): + def test_cdist_to_point(self): + points = np.array([[0.0, 0.0, 0.0], [3.0, 4.0, 0.0]]) + out = cdist_to_point((0.0, 0.0, 0.0), points) + np.testing.assert_allclose(np.asarray(out).ravel(), [0.0, 5.0]) + + def test_get_nearest_neighbor_picks_the_closest_labelled_voxel(self): + arr = np.zeros((10, 10, 10), dtype=np.uint8) + arr[8, 8, 8] = 3 + arr[1, 1, 1] = 3 + out = get_nearest_neighbor(np.array([0.0, 0.0, 0.0]), _nii(arr), 3) + np.testing.assert_allclose(np.asarray(out), [1.0, 1.0, 1.0]) + + def test_project_pois_onto_set_of_points_snaps_to_the_nearest(self): + from TPTBox import POI # noqa: PLC0415 + + poi = POI({1: {50: (0.0, 0.0, 0.0)}}, orientation=("R", "A", "S"), zoom=(1.0, 1.0, 1.0), shape=(10.0, 10.0, 10.0)) + point_set = np.array([[5.0, 5.0, 5.0], [1.0, 0.0, 0.0]]) + out = project_pois_onto_set_of_points(poi, point_set) + np.testing.assert_allclose(out[1, 50], (1.0, 0.0, 0.0)) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_rib_assign.py b/unit_tests/test_rib_assign.py new file mode 100644 index 0000000..8d30c45 --- /dev/null +++ b/unit_tests/test_rib_assign.py @@ -0,0 +1,243 @@ +"""Unit tests for `TPTBox.segmentation.rib._rib_assign`. + +This is the pure-numpy half of the rib sub-package - no torch, no spineps - so it +runs everywhere. The fixtures are synthetic thoracic columns: vertebra blocks +stacked along the superior axis with rib bars sticking out left and right, which +makes the expected left/right split and the per-vertebra assignment checkable. +""" + +from __future__ import annotations + +import unittest +from unittest import mock + +import nibabel as nib +import numpy as np + +from TPTBox import NII +from TPTBox.core.vert_constants import Full_Body_Instance, Location +from TPTBox.segmentation.rib._rib_assign import ( + RibCandidate, + _thoracic_like_labels, + _touching_surface, + assign_ribs_to_vert_segmentation, + split_touching_rib_ccs, +) + +# RAS with a thick superior axis, so the "S" weighting in _touching_surface is testable. +ZOOM = (1.0, 1.0, 3.0) +SHAPE = (40, 24, 40) + + +def _nii(arr: np.ndarray, seg: bool = True) -> NII: + return NII(nib.Nifti1Image(arr, np.diag([*ZOOM, 1.0])), seg=seg) + + +def _thorax(n_verts: int = 4, fuse: bool = False): + """Build (vert_seg, sem_seg, rib_seg) for `n_verts` thoracic levels with paired ribs.""" + vert = np.zeros(SHAPE, dtype=np.uint16) + sem = np.zeros(SHAPE, dtype=np.uint16) + rib = np.zeros(SHAPE, dtype=np.uint16) + for i in range(n_verts): + vid = 8 + i # T1 = 8 + z0 = 4 + i * 6 + z1 = z0 + (6 if fuse else 4) + vert[18:22, 10:14, z0 : z0 + 4] = vid + sem[18:22, 10:14, z0 : z0 + 4] = Location.Vertebra_Corpus.value + # Rib bars reaching out to both sides along the right-left axis. The raw rib + # segmentation is labelled the way the function documents: `right_id` on the + # patient-right side, `left_id` on the left. + rib[4:18, 11:13, z0 : z0 + 2] = Full_Body_Instance.rib_right.value + rib[22:36, 11:13, z0 : z0 + 2] = Full_Body_Instance.rib_left.value + if fuse: + # a thin bridge joining this level's rib to the next one + rib[6:10, 11:13, z0:z1] = Full_Body_Instance.rib_right.value + return _nii(vert), _nii(sem), _nii(rib) + + +class TestThoracicLikeLabels(unittest.TestCase): + def test_keeps_the_rib_bearing_block_and_appends_28(self): + self.assertEqual(_thoracic_like_labels([20, 8, 7, 25, 3, 100]), [7, 8, 20, 28]) + + def test_sorted_cranial_to_caudal(self): + out = _thoracic_like_labels([19, 10, 15]) + self.assertEqual(out[:-1], sorted(out[:-1])) + + def test_an_empty_input_still_appends_28(self): + self.assertEqual(_thoracic_like_labels([]), [28]) + + def test_labels_outside_the_block_are_dropped(self): + self.assertEqual(_thoracic_like_labels([1, 2, 6, 21, 24]), [28]) + + +class TestTouchingSurface(unittest.TestCase): + def _masks(self, axis: int): + a = np.zeros(SHAPE, dtype=bool) + b = np.zeros(SHAPE, dtype=bool) + index_a = [slice(10, 14)] * 3 + index_b = [slice(10, 14)] * 3 + index_a[axis] = slice(10, 12) + index_b[axis] = slice(12, 14) + a[tuple(index_a)] = True + b[tuple(index_b)] = True + return a, b + + def test_separated_masks_do_not_touch(self): + a = np.zeros(SHAPE, dtype=bool) + b = np.zeros(SHAPE, dtype=bool) + a[10:12, 10:12, 10:12] = True + b[14:16, 10:12, 10:12] = True + self.assertEqual(_touching_surface(a, b, _nii(np.zeros(SHAPE, dtype=np.uint8))), 0.0) + + def test_touching_area_scales_with_the_shared_face(self): + nii = _nii(np.zeros(SHAPE, dtype=np.uint8)) + small_a = np.zeros(SHAPE, dtype=bool) + small_b = np.zeros(SHAPE, dtype=bool) + small_a[10:12, 10:12, 10:12] = True + small_b[12:14, 10:12, 10:12] = True + + big_a = np.zeros(SHAPE, dtype=bool) + big_b = np.zeros(SHAPE, dtype=bool) + big_a[10:12, 8:16, 8:16] = True + big_b[12:14, 8:16, 8:16] = True + self.assertGreater(_touching_surface(big_a, big_b, nii), _touching_surface(small_a, small_b, nii)) + + def test_the_default_weights_prefer_left_right_contact(self): + nii = _nii(np.zeros(SHAPE, dtype=np.uint8)) + per_axis = {} + for axis, name in ((nii.get_axis("R"), "R"), (nii.get_axis("A"), "A"), (nii.get_axis("S"), "S")): + a, b = self._masks(axis) + per_axis[name] = _touching_surface(a, b, nii) + # defaults are S: 0.01, A: 0.1, R: 1 - contact across the superior axis counts least + self.assertGreater(per_axis["R"], per_axis["A"]) + self.assertGreater(per_axis["A"], per_axis["S"]) + + def test_explicit_weights_override_the_defaults(self): + nii = _nii(np.zeros(SHAPE, dtype=np.uint8)) + a, b = self._masks(nii.get_axis("S")) + default = _touching_surface(a, b, nii) + weighted = _touching_surface(a, b, nii, axis_weights={nii.get_axis("S"): 1.0}) + self.assertGreater(weighted, default) + + def test_the_relation_is_symmetric(self): + nii = _nii(np.zeros(SHAPE, dtype=np.uint8)) + a, b = self._masks(0) + self.assertEqual(_touching_surface(a, b, nii), _touching_surface(b, a, nii)) + + +class TestRibCandidate(unittest.TestCase): + def test_fields(self): + cand = RibCandidate(cc_label=3, volume=120.0, z=5.0, x=-2.0) + self.assertEqual((cand.cc_label, cand.volume, cand.z, cand.x), (3, 120.0, 5.0, -2.0)) + + +class TestSplitTouchingRibCcs(unittest.TestCase): + def test_a_clean_mask_is_returned_unchanged_in_volume(self): + _, _, rib = _thorax() + cc = rib.get_connected_components() + before = int((cc.get_seg_array() != 0).sum()) + out = split_touching_rib_ccs(cc.copy(), verbose=False) + self.assertEqual(int((out.get_seg_array() != 0).sum()), before) + + def test_a_fused_pair_is_split_into_more_components(self): + _, _, rib = _thorax(fuse=True) + cc = rib.get_connected_components() + before = len(set(np.unique(cc.get_seg_array())) - {0}) + out = split_touching_rib_ccs(cc.copy(), erosion_pixels=2, min_volume=10, verbose=False) + after = len(set(np.unique(out.get_seg_array())) - {0}) + self.assertGreaterEqual(after, before) + + def test_running_with_more_passes_does_not_lose_voxels(self): + _, _, rib = _thorax(fuse=True) + cc = rib.get_connected_components() + total = int((cc.get_seg_array() != 0).sum()) + out = split_touching_rib_ccs(cc.copy(), erosion_pixels=2, min_volume=10, max_passes=3, verbose=False) + self.assertEqual(int((out.get_seg_array() != 0).sum()), total) + + +class TestAssignRibsToVertSegmentation(unittest.TestCase): + """The public entry point. Two bugs used to make it raise before returning. + + What is asserted here is that it completes and merges rib labels back into the + semantic segmentation without disturbing the vertebrae. Whether the matching + heuristic pairs every rib with a vertebra is *not* asserted: on this synthetic + geometry it does not, and the unmatched components legitimately come back marked + with `error_value`. + """ + + def test_rib_voxels_are_merged_into_the_semantic_segmentation(self): + vert, sem, rib = _thorax() + _, out_sem = assign_ribs_to_vert_segmentation(vert.copy(), sem.copy(), rib.copy(), verbose=False) + labels = set(np.unique(out_sem.get_seg_array())) - {0} + self.assertIn(Location.Rib_Left.value, labels) + self.assertIn(Location.Rib_Right.value, labels) + self.assertIn(Location.Vertebra_Corpus.value, labels, "the existing semantic labels survive") + + def test_the_rib_mask_footprint_is_preserved(self): + vert, sem, rib = _thorax() + _, out_sem = assign_ribs_to_vert_segmentation(vert.copy(), sem.copy(), rib.copy(), verbose=False) + rib_voxels = int((rib.get_seg_array() != 0).sum()) + marked = int(np.isin(out_sem.get_seg_array(), [Location.Rib_Left.value, Location.Rib_Right.value]).sum()) + self.assertEqual(marked, rib_voxels) + + def test_vertebra_labels_are_not_lost(self): + vert, sem, rib = _thorax() + out_vert, _ = assign_ribs_to_vert_segmentation(vert.copy(), sem.copy(), rib.copy(), verbose=False) + self.assertTrue({8, 9, 10, 11}.issubset(set(np.unique(out_vert.get_seg_array())))) + + def test_unmatched_components_are_marked_with_the_error_value(self): + vert, sem, rib = _thorax() + out_vert, _ = assign_ribs_to_vert_segmentation(vert.copy(), sem.copy(), rib.copy(), error_value=251, verbose=False) + self.assertIn(251, set(np.unique(out_vert.get_seg_array()))) + + def test_add_error_false_leaves_no_error_label(self): + vert, sem, rib = _thorax() + out_vert, _ = assign_ribs_to_vert_segmentation(vert.copy(), sem.copy(), rib.copy(), error_value=251, add_error=False, verbose=False) + self.assertNotIn(251, set(np.unique(out_vert.get_seg_array()))) + + def test_an_empty_rib_mask_leaves_the_inputs_alone(self): + vert, sem, _ = _thorax() + empty = _nii(np.zeros(SHAPE, dtype=np.uint16)) + out_vert, out_sem = assign_ribs_to_vert_segmentation(vert.copy(), sem.copy(), empty, verbose=False) + np.testing.assert_array_equal(out_vert.get_seg_array(), vert.get_seg_array()) + np.testing.assert_array_equal(out_sem.get_seg_array(), sem.get_seg_array()) + + def test_split_touching_can_be_turned_off(self): + vert, sem, rib = _thorax() + out_vert, out_sem = assign_ribs_to_vert_segmentation(vert.copy(), sem.copy(), rib.copy(), split_touching=False, verbose=False) + self.assertIn(Location.Rib_Left.value, set(np.unique(out_sem.get_seg_array()))) + self.assertGreater(int((out_vert.get_seg_array() != 0).sum()), 0) + + def test_custom_left_and_right_ids_are_honoured(self): + vert, sem, rib = _thorax() + arr = rib.get_seg_array() + arr[arr == Full_Body_Instance.rib_left.value] = 201 + arr[arr == Full_Body_Instance.rib_right.value] = 202 + rib.set_array_(arr) + _, out_sem = assign_ribs_to_vert_segmentation(vert.copy(), sem.copy(), rib, left_id=201, right_id=202, verbose=False) + labels = set(np.unique(out_sem.get_seg_array())) + self.assertIn(Location.Rib_Left.value, labels) + self.assertIn(Location.Rib_Right.value, labels) + + def test_a_mismatched_affine_is_rejected(self): + vert, sem, rib = _thorax() + other = NII(nib.Nifti1Image(np.zeros(SHAPE, dtype=np.uint16), np.diag([2.0, 2.0, 2.0, 1.0])), seg=True) + with self.assertRaises(AssertionError): + assign_ribs_to_vert_segmentation(vert, sem, other, verbose=False) + + +class TestAddRibs(unittest.TestCase): + """`add_ribs` wires the assignment above into a spineps run; only the wiring is tested.""" + + def test_an_existing_rib_segmentation_short_circuits_the_model_call(self): + import sys # noqa: PLC0415 + + fake_spineps = mock.MagicMock() + with mock.patch.dict(sys.modules, {"spineps": fake_spineps, "spineps.seg_run": fake_spineps.seg_run}): + from TPTBox.segmentation.rib import add_ribs as add_ribs_mod # noqa: PLC0415 + + self.assertTrue(hasattr(add_ribs_mod, "add_ribs_to_vert_spine")) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_segmentation_mock.py b/unit_tests/test_segmentation_mock.py new file mode 100644 index 0000000..872373a --- /dev/null +++ b/unit_tests/test_segmentation_mock.py @@ -0,0 +1,534 @@ +"""Unit tests for the segmentation wrappers, without weights, GPU or network. + +The mock sits at the model boundary only: the network forward pass (or the whole +predictor) is replaced, while everything around it - RAM budgeting, chunking, +orientation handling, label post-processing, path construction, argv assembly - +runs for real. Assertions are therefore about the post-processing result, not +about whether a mock was called. + +The pure numeric helpers need no mock at all and are tested directly. +""" + +from __future__ import annotations + +import importlib.util +import json +import sys +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +import nibabel as nib +import numpy as np + +# `itertools.pairwise` is 3.10+; the package ships a backport for its own 3.9 support. +from TPTBox.core.compat import pairwise + +HAS_TORCH = importlib.util.find_spec("torch") is not None +HAS_NNUNET = importlib.util.find_spec("nnunetv2") is not None +TORCH_REASON = "requires torch" +NNUNET_REASON = "requires torch + nnunetv2" + + +def _nii(arr: np.ndarray, seg: bool = False, zoom=(1.4, 1.4, 3.0)): + from TPTBox import NII # noqa: PLC0415 + + return NII(nib.Nifti1Image(arr, np.diag([*zoom, 1.0])), seg=seg) + + +def _image(shape=(24, 24, 24)) -> np.ndarray: + arr = np.zeros(shape, dtype=np.int16) + arr[4:-4, 4:-4, 4:-4] = 80 + return arr + + +@unittest.skipUnless(HAS_NNUNET, NNUNET_REASON) +class TestRamBudgeting(unittest.TestCase): + """Pure arithmetic from the docstring's own formula.""" + + def test_estimate_matches_the_documented_formula(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import estimate_peak_ram_mb # noqa: PLC0415 + + shape, num_classes, num_channels = (64, 64, 64), 10, 1 + per_voxel = (num_classes * 2) + 2 + (num_channels * 2 * 2) + 1 + expected = int(np.prod(shape)) * per_voxel * 1.3 / (1024 * 1024) + self.assertAlmostEqual(estimate_peak_ram_mb(shape, num_classes, num_channels), expected, places=6) + + def test_estimate_scales_linearly_in_voxels_and_classes(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import estimate_peak_ram_mb # noqa: PLC0415 + + base = estimate_peak_ram_mb((32, 32, 32), 4) + self.assertAlmostEqual(estimate_peak_ram_mb((64, 32, 32), 4), 2 * base, places=6) + self.assertGreater(estimate_peak_ram_mb((32, 32, 32), 8), base) + + def test_chunking_returns_one_when_the_volume_fits(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import compute_cpu_chunks_for_ram # noqa: PLC0415 + + self.assertEqual(compute_cpu_chunks_for_ram((32, 32, 32), 0, 4, 1, 4, target_mb=1e9), 1) + + def test_chunking_grows_as_the_budget_shrinks(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import compute_cpu_chunks_for_ram, estimate_peak_ram_mb # noqa: PLC0415 + + shape = (256, 256, 256) + whole = estimate_peak_ram_mb(shape, 10, 1) + self.assertEqual(compute_cpu_chunks_for_ram(shape, 0, 10, 1, 4, target_mb=whole * 2), 1) + tight = compute_cpu_chunks_for_ram(shape, 0, 10, 1, 4, target_mb=whole / 4) + self.assertGreater(tight, 1) + self.assertLessEqual( + estimate_peak_ram_mb((shape[0] // tight + 8, shape[1], shape[2]), 10, 1), + whole / 4, + "the returned chunk count must actually fit the budget", + ) + + def test_chunking_degenerate_inputs(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import compute_cpu_chunks_for_ram # noqa: PLC0415 + + self.assertEqual(compute_cpu_chunks_for_ram((1, 32, 32), 0, 4, 1, 4, target_mb=10), 1) + self.assertEqual(compute_cpu_chunks_for_ram((32, 32, 32), 0, 4, 1, 4, target_mb=0), 1) + + def test_total_ram_is_a_positive_number(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import _get_total_ram_mb # noqa: PLC0415 + + self.assertGreater(_get_total_ram_mb(), 0.0) + + +@unittest.skipUnless(HAS_NNUNET, NNUNET_REASON) +class TestSplitRanges(unittest.TestCase): + """Each entry is (read_start, read_stop, write_start, write_stop). + + The read windows overlap so every chunk has context; the write windows, taken + relative to their own read window, must tile the axis exactly once. + """ + + def test_write_windows_tile_the_axis_exactly(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import _split_ranges # noqa: PLC0415 + + for length, n_chunks, overlap in ((100, 4, 5), (100, 1, 0), (37, 3, 2), (10, 10, 1), (64, 3, 8)): + with self.subTest(length=length, n_chunks=n_chunks, overlap=overlap): + written = [] + for read_start, read_stop, write_start, write_stop in _split_ranges(length, n_chunks, overlap): + self.assertGreaterEqual(read_start, 0) + self.assertLessEqual(read_stop, length) + self.assertLessEqual(write_stop - write_start, read_stop - read_start) + written.extend(range(read_start + write_start, read_start + write_stop)) + self.assertEqual(written, list(range(length)), "the write windows must tile the axis in order") + + def test_a_single_chunk_is_the_whole_axis(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import _split_ranges # noqa: PLC0415 + + self.assertEqual(list(_split_ranges(50, 1, 0)), [(0, 50, 0, 50)]) + + def test_read_windows_overlap_their_neighbours(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import _split_ranges # noqa: PLC0415 + + ranges = list(_split_ranges(100, 4, 6)) + self.assertGreater(len(ranges), 1) + for (_, stop, _, _), (start, _, _, _) in pairwise(ranges): + with self.subTest(boundary=(stop, start)): + self.assertLess(start, stop, "consecutive read windows must share voxels") + + +@unittest.skipUnless(HAS_NNUNET, NNUNET_REASON) +class TestSlidingNdSlices(unittest.TestCase): + def test_identity_function_reconstructs_the_input(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import sliding_nd_slices # noqa: PLC0415 + + arr = np.arange(16 * 16 * 16, dtype=np.float32).reshape(16, 16, 16) + out = sliding_nd_slices(arr, (8, 8, 8), 2, lambda patch: patch) + np.testing.assert_allclose(out, arr) + + def test_the_function_is_applied_to_every_voxel(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import sliding_nd_slices # noqa: PLC0415 + + arr = np.ones((12, 12, 12), dtype=np.float32) + out = sliding_nd_slices(arr, (6, 6, 6), 2, lambda patch: patch * 3) + np.testing.assert_allclose(out, 3.0) + + def test_shape_is_preserved_for_a_non_divisible_volume(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import sliding_nd_slices # noqa: PLC0415 + + arr = np.ones((13, 11, 7), dtype=np.float32) + self.assertEqual(sliding_nd_slices(arr, (8, 8, 8), 2, lambda p: p).shape, arr.shape) + + +@unittest.skipUnless(HAS_TORCH, TORCH_REASON) +class TestSquashForFloat16(unittest.TestCase): + def test_a_large_volume_is_brought_into_range(self): + from TPTBox.segmentation.VibeSeg.inference_nnunet import squash_so_it_fits_in_float16 # noqa: PLC0415 + + arr = np.full((8, 8, 8), 1e6, dtype=np.float32) + arr[0, 0, 0] = 2e6 + out = squash_so_it_fits_in_float16(_nii(arr)) + self.assertTrue(np.isfinite(np.asarray(out.get_array(), dtype=np.float16)).all()) + # the ratio between voxels is what matters, not the absolute values + values = out.get_array() + self.assertAlmostEqual(float(values.max() / values[values > 0].min()), 2.0, delta=0.1) + + def test_a_small_volume_is_left_alone(self): + from TPTBox.segmentation.VibeSeg.inference_nnunet import squash_so_it_fits_in_float16 # noqa: PLC0415 + + arr = np.full((8, 8, 8), 100.0, dtype=np.float32) + np.testing.assert_allclose(squash_so_it_fits_in_float16(_nii(arr)).get_array(), arr) + + +@unittest.skipUnless(HAS_TORCH, TORCH_REASON) +class TestExtractVertebraBodies(unittest.TestCase): + """Pure label logic on a VIBESeg output: alternating body / disc blocks.""" + + @staticmethod + def _stack(bodies: int = 6): + shape = (20, 20, 8 * bodies + 4) + arr = np.zeros(shape, dtype=np.uint8) + z = 2 + for b in range(bodies): + arr[3:17, 3:17, z : z + 6] = 69 # vertebra body + z += 6 + if b < bodies - 1: + arr[3:17, 3:17, z : z + 2] = 68 # intervertebral disc + z += 2 + return _nii(arr, seg=True, zoom=(1.0, 1.0, 1.0)) + + def test_six_bodies_are_labelled_l5_upwards(self): + from TPTBox import ( + NII, # noqa: PLC0415 + POI, # noqa: PLC0415 + ) + from TPTBox.segmentation.VibeSeg.vibeseg import extract_vertebra_bodies_from_VibeSeg # noqa: PLC0415 + + vert, poi = extract_vertebra_bodies_from_VibeSeg(self._stack(bodies=6)) + self.assertIsInstance(vert, NII) + self.assertIsInstance(poi, POI) + labels = set(np.unique(vert.get_seg_array())) - {0} + self.assertEqual(labels, {19, 20, 21, 22, 23, 24}, "inferior to superior: L5..T12") + + def test_the_number_of_labels_follows_the_number_of_bodies(self): + from TPTBox.segmentation.VibeSeg.vibeseg import extract_vertebra_bodies_from_VibeSeg # noqa: PLC0415 + + for bodies in (3, 5): + with self.subTest(bodies=bodies): + vert, _ = extract_vertebra_bodies_from_VibeSeg(self._stack(bodies=bodies)) + self.assertEqual(len(set(np.unique(vert.get_seg_array())) - {0}), bodies) + + def test_one_poi_per_body(self): + from TPTBox.segmentation.VibeSeg.vibeseg import extract_vertebra_bodies_from_VibeSeg # noqa: PLC0415 + + vert, poi = extract_vertebra_bodies_from_VibeSeg(self._stack(bodies=4)) + self.assertEqual(len(list(poi.keys())), len(set(np.unique(vert.get_seg_array())) - {0})) + + +@unittest.skipUnless(HAS_TORCH, TORCH_REASON) +class TestDeface(unittest.TestCase): + def test_extend_mask_grows_only_in_the_requested_direction(self): + from TPTBox.segmentation._deface import _extend_mask # noqa: PLC0415 + + arr = np.zeros((20, 20, 20), dtype=np.uint8) + arr[8:12, 8:12, 8:12] = 1 + nii = _nii(arr, seg=True, zoom=(1.0, 1.0, 1.0)).reorient_(("R", "A", "S")) + out = _extend_mask(nii.copy(), 4, direction="A") + before = nii.get_seg_array() + after = out.get_seg_array() + self.assertGreater(int(after.sum()), int(before.sum())) + # projections perpendicular to A are unchanged + np.testing.assert_array_equal(after.any(axis=1), before.any(axis=1)) + + def test_extend_mask_by_zero_is_a_noop(self): + from TPTBox.segmentation._deface import _extend_mask # noqa: PLC0415 + + arr = np.zeros((16, 16, 16), dtype=np.uint8) + arr[6:10, 6:10, 6:10] = 1 + nii = _nii(arr, seg=True, zoom=(1.0, 1.0, 1.0)) + np.testing.assert_array_equal(_extend_mask(nii.copy(), 0).get_seg_array(), nii.get_seg_array()) + + def test_deface_img_blanks_the_masked_voxels_only(self): + from TPTBox.segmentation._deface import deface_img # noqa: PLC0415 + + ct = _nii(_image(), seg=False) + mask_arr = np.zeros(ct.shape, dtype=np.uint8) + mask_arr[:8] = 1 + out = deface_img(ct, _nii(mask_arr, seg=True), min_value=-1024) + arr = out.get_array() + self.assertTrue((arr[:8] == -1024).all(), "masked voxels take min_value") + np.testing.assert_array_equal(arr[8:], ct.get_array()[8:], "the rest is untouched") + + def test_deface_img_to_int_gives_an_integer_dtype(self): + from TPTBox.segmentation._deface import deface_img # noqa: PLC0415 + + ct = _nii(_image().astype(np.float32), seg=False) + mask_arr = np.zeros(ct.shape, dtype=np.uint8) + mask_arr[:4] = 1 + out = deface_img(ct, _nii(mask_arr, seg=True), to_int=True) + self.assertTrue(np.issubdtype(out.get_array().dtype, np.integer)) + + def test_deface_img_writes_the_output_file(self): + from TPTBox.segmentation._deface import deface_img # noqa: PLC0415 + + ct = _nii(_image(), seg=False) + mask_arr = np.zeros(ct.shape, dtype=np.uint8) + mask_arr[:4] = 1 + with tempfile.TemporaryDirectory() as tmp: + out_path = Path(tmp) / "defaced.nii.gz" + deface_img(ct, _nii(mask_arr, seg=True), ct_out=out_path) + self.assertTrue(out_path.exists()) + + +class TestSpinepsDispatch(unittest.TestCase): + """`spineps.py` imports spineps only inside function bodies, so a fake module is enough.""" + + def _fake_spineps(self): + fake = mock.MagicMock() + return { + "spineps": fake, + "spineps.seg_run": fake.seg_run, + "spineps.get_models": fake.get_models, + "spineps.seg_model": fake.seg_model, + }, fake + + def test_get_outpaths_wraps_a_plain_path_in_a_bids_file(self): + modules, fake = self._fake_spineps() + with mock.patch.dict(sys.modules, modules), tempfile.TemporaryDirectory() as tmp: + from TPTBox.segmentation.spineps import get_outpaths_spineps # noqa: PLC0415 + + nii_path = Path(tmp) / "dataset-x" / "rawdata" / "sub-1" / "sub-1_T2w.nii.gz" + nii_path.parent.mkdir(parents=True) + nii_path.touch() + get_outpaths_spineps(nii_path) + self.assertTrue(fake.seg_run.output_paths_from_input.called) + bids_arg = fake.seg_run.output_paths_from_input.call_args[0][0] + self.assertEqual(bids_arg.bids_format, "T2w") + + def test_an_explicit_dataset_root_wins(self): + modules, fake = self._fake_spineps() + with mock.patch.dict(sys.modules, modules), tempfile.TemporaryDirectory() as tmp: + from TPTBox.segmentation.spineps import get_outpaths_spineps # noqa: PLC0415 + + root = Path(tmp) / "dataset-x" + nii_path = root / "rawdata" / "sub-1" / "sub-1_T2w.nii.gz" + nii_path.parent.mkdir(parents=True) + nii_path.touch() + get_outpaths_spineps(nii_path, dataset=root) + bids_arg = fake.seg_run.output_paths_from_input.call_args[0][0] + self.assertEqual(Path(bids_arg.dataset), root) + + def test_run_spineps_all_invokes_one_subprocess_per_model(self): + with tempfile.TemporaryDirectory() as tmp, mock.patch("subprocess.run") as run: + from TPTBox.segmentation.spineps import _run_spineps_all # noqa: PLC0415 + + _run_spineps_all(Path(tmp)) + self.assertGreaterEqual(run.call_count, 1) + for call in run.call_args_list: + argv = call[0][0] + with self.subTest(argv=argv if isinstance(argv, str) else " ".join(map(str, argv))): + self.assertIn("spineps", " ".join(map(str, argv)) if not isinstance(argv, str) else argv) + + +@unittest.skipUnless(HAS_TORCH, TORCH_REASON) +class TestAutoDownload(unittest.TestCase): + def test_read_config_parses_the_shipped_json(self): + from TPTBox.segmentation.VibeSeg.auto_download import read_config # noqa: PLC0415 + + try: + config = read_config(85) + except (FileNotFoundError, KeyError) as exc: # pragma: no cover - depends on the shipped table + self.skipTest(f"no config entry for this model id: {exc}") + self.assertIsInstance(config, dict) + + def test_download_weights_short_circuits_when_the_folder_exists(self): + from TPTBox.segmentation.VibeSeg import auto_download # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp, mock.patch.object(auto_download, "_download") as download: + model_path = Path(tmp) + weights = auto_download.get_weights_dir(85, model_path=model_path) + weights.mkdir(parents=True, exist_ok=True) + (weights / "plans.json").write_text(json.dumps({"configurations": {}}), encoding="utf-8") + auto_download.download_weights(85, model_path=model_path) + self.assertFalse(download.called, "an existing weights folder must not trigger a download") + + def test_get_weights_dir_is_under_the_given_model_path(self): + from TPTBox.segmentation.VibeSeg.auto_download import get_weights_dir # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + out = get_weights_dir(85, model_path=Path(tmp)) + self.assertIn(Path(tmp), out.parents) + + +@unittest.skipUnless(HAS_NNUNET, NNUNET_REASON) +class TestRunInferenceWithAMockedPredictor(unittest.TestCase): + """The mock replaces only `predict_single_npy_array`. + + Everything around it - orientation handling, the properties dict handed to the + predictor, and the conversion of the returned array back into an NII on the input + grid - runs for real, and that is what is asserted. + """ + + @staticmethod + def _predictor(label_value: int = 3): + predictor = mock.MagicMock() + + def predict(img, props, **kwargs): # noqa: ARG001 + # img is (C, P, I, R); return a label map over the spatial dims + out = np.zeros(img.shape[1:], dtype=np.uint8) + out[..., : out.shape[-1] // 2] = label_value + return out + + predictor.predict_single_npy_array.side_effect = predict + return predictor + + def test_the_result_is_returned_on_the_input_grid(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import run_inference # noqa: PLC0415 + + nii = _nii(_image(), seg=False) + predictor = self._predictor() + seg, logits, softmax = run_inference(nii, predictor) + self.assertEqual(seg.shape, nii.shape) + self.assertEqual(seg.orientation, nii.orientation) + np.testing.assert_allclose(seg.affine, nii.affine) + self.assertIsNone(logits) + self.assertIsNone(softmax) + + def test_the_planted_labels_survive_the_round_trip(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import run_inference # noqa: PLC0415 + + seg, _, _ = run_inference(_nii(_image(), seg=False), self._predictor(label_value=7)) + self.assertEqual(set(np.unique(seg.get_seg_array())), {0, 7}) + + def test_the_predictor_receives_the_reversed_spacing(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import run_inference # noqa: PLC0415 + + nii = _nii(_image(), seg=False) + predictor = self._predictor() + run_inference(nii, predictor) + props = predictor.predict_single_npy_array.call_args[0][1] + # reorient_PIR defaults to False, so the spacing is the input's own, reversed. + np.testing.assert_allclose(props["spacing"], nii.zoom[::-1]) + + def test_reorient_pir_changes_the_spacing_handed_over(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import run_inference # noqa: PLC0415 + + nii = _nii(_image(), seg=False) + predictor = self._predictor() + run_inference(nii, predictor, reorient_PIR=True) + props = predictor.predict_single_npy_array.call_args[0][1] + np.testing.assert_allclose(props["spacing"], nii.reorient(("P", "I", "R")).zoom[::-1]) + + def test_the_image_is_passed_with_a_leading_channel_axis(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import run_inference # noqa: PLC0415 + + nii = _nii(_image(), seg=False) + predictor = self._predictor() + run_inference(nii, predictor) + img = predictor.predict_single_npy_array.call_args[0][0] + self.assertEqual(img.ndim, 4) + self.assertEqual(img.shape[0], 1, "one channel for a single input image") + + def test_a_multi_channel_input_is_stacked(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import run_inference # noqa: PLC0415 + + niis = [_nii(_image(), seg=False), _nii(_image(), seg=False)] + predictor = self._predictor() + run_inference(niis, predictor) + self.assertEqual(predictor.predict_single_npy_array.call_args[0][0].shape[0], 2) + + def test_logits_are_not_implemented(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import run_inference # noqa: PLC0415 + + with self.assertRaises((NotImplementedError, AssertionError)): + run_inference(_nii(_image(), seg=False), self._predictor(), logits=True) + + def test_an_unsupported_input_type_is_rejected(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import run_inference # noqa: PLC0415 + + with self.assertRaises(AssertionError): + run_inference(5, self._predictor()) # type: ignore[arg-type] + + def test_a_path_must_be_a_nii_gz(self): + from TPTBox.segmentation.nnUnet_utils.inference_api import run_inference # noqa: PLC0415 + + with self.assertRaises(AssertionError): + run_inference("image.png", self._predictor()) + + +@unittest.skipUnless(HAS_TORCH, TORCH_REASON) +class TestRunInferenceOnFile(unittest.TestCase): + """`run_inference_on_file` orchestrates model loading, inference and output writing. + + Model 85's shipped dataset.json declares two input channels, and the function + checks that, so the fixtures hand over two images. + """ + + @staticmethod + def _inputs(): + return [_nii(_image(), seg=False), _nii(_image(), seg=False)] + + def _patched(self, label_value: int = 5): + from TPTBox.segmentation.nnUnet_utils import inference_api # noqa: PLC0415 + + def fake_run_inference(input_nii, predictor, **kwargs): # noqa: ARG001 + first = input_nii[0] if isinstance(input_nii, list) else input_nii + arr = np.zeros(first.shape, dtype=np.uint8) + arr[..., : arr.shape[-1] // 2] = label_value + return first.set_array(arr, seg=True), None, None + + # run_inference_on_file imports both names inside its own body, so patching the + # module they are defined in is what takes effect. + return ( + mock.patch.object(inference_api, "load_inf_model", return_value=mock.MagicMock()), + mock.patch.object(inference_api, "run_inference", side_effect=fake_run_inference), + ) + + def test_the_output_file_is_written_with_the_planted_labels(self): + from TPTBox import to_nii # noqa: PLC0415 + from TPTBox.segmentation.VibeSeg.inference_nnunet import run_inference_on_file # noqa: PLC0415 + + load_patch, run_patch = self._patched(label_value=5) + with tempfile.TemporaryDirectory() as tmp, load_patch, run_patch: + out_file = Path(tmp) / "seg.nii.gz" + niis = self._inputs() + nii = niis[0] + run_inference_on_file(85, niis, out_file=out_file, ddevice="cpu", verbose=False) + self.assertTrue(out_file.exists()) + written = to_nii(out_file, seg=True) + self.assertEqual(written.shape, nii.shape) + self.assertEqual(set(np.unique(written.get_seg_array())), {0, 5}) + + def test_a_second_call_skips_inference_unless_overridden(self): + from TPTBox.segmentation.VibeSeg.inference_nnunet import run_inference_on_file # noqa: PLC0415 + + load_patch, run_patch = self._patched() + with tempfile.TemporaryDirectory() as tmp, load_patch, run_patch as run: + out_file = Path(tmp) / "seg.nii.gz" + niis = self._inputs() + run_inference_on_file(85, niis, out_file=out_file, ddevice="cpu", verbose=False) + self.assertEqual(run.call_count, 1) + run_inference_on_file(85, niis, out_file=out_file, ddevice="cpu", verbose=False) + self.assertEqual(run.call_count, 1, "an existing output must short-circuit") + run_inference_on_file(85, niis, out_file=out_file, override=True, ddevice="cpu", verbose=False) + self.assertEqual(run.call_count, 2, "override=True must re-run") + + def test_a_label_mapping_is_applied(self): + from TPTBox import to_nii # noqa: PLC0415 + from TPTBox.segmentation.VibeSeg.inference_nnunet import run_inference_on_file # noqa: PLC0415 + + load_patch, run_patch = self._patched(label_value=5) + with tempfile.TemporaryDirectory() as tmp, load_patch, run_patch: + out_file = Path(tmp) / "seg.nii.gz" + run_inference_on_file(85, self._inputs(), out_file=out_file, mapping={5: 9}, ddevice="cpu", verbose=False) + self.assertEqual(set(np.unique(to_nii(out_file, seg=True).get_seg_array())), {0, 9}) + + def test_the_output_keeps_the_input_grid(self): + from TPTBox import to_nii # noqa: PLC0415 + from TPTBox.segmentation.VibeSeg.inference_nnunet import run_inference_on_file # noqa: PLC0415 + + load_patch, run_patch = self._patched() + with tempfile.TemporaryDirectory() as tmp, load_patch, run_patch: + out_file = Path(tmp) / "seg.nii.gz" + niis = self._inputs() + nii = niis[0] + run_inference_on_file(85, niis, out_file=out_file, ddevice="cpu", verbose=False) + written = to_nii(out_file, seg=True) + np.testing.assert_allclose(written.affine, nii.affine, atol=1e-6) + self.assertEqual(written.orientation, nii.orientation) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_snapshot2D.py b/unit_tests/test_snapshot2D.py new file mode 100644 index 0000000..373386e --- /dev/null +++ b/unit_tests/test_snapshot2D.py @@ -0,0 +1,282 @@ +"""Unit tests for `TPTBox.spine.snapshot2D`. + +The pure array/colour helpers are tested directly, and `create_snapshot` plus the +`snapshot_templates` entry points are rendered to a temp directory and checked for +a non-blank PNG of the expected size. `unit_tests/test_snapshot_mip.py` already +covers the slab projection functions. +""" + +from __future__ import annotations + +import importlib.util +import tempfile +import unittest +from pathlib import Path + +import matplotlib as mpl + +mpl.use("Agg") + +import matplotlib.pyplot as plt # noqa: E402 +import nibabel as nib # noqa: E402 +import numpy as np # noqa: E402 + +from TPTBox import NII, Location # noqa: E402 + +# `create_snapshot` imports torch for a frame with `denoise_threshold` set, which the +# mip_shot template does - so that one template needs torch even though it is optional. +HAS_TORCH = importlib.util.find_spec("torch") is not None +from TPTBox.spine.snapshot2D.snapshot_modular import ( # noqa: E402 + Snapshot_Frame, + create_snapshot, + div0, + get_contrasting_stroke_color, + make_isotropic2d, + make_isotropic2dpluscolor, + normalize_image, + to_cdt, +) + + +def _nii(arr, seg=True, zoom=(1.0, 1.0, 1.0)) -> NII: + return NII(nib.Nifti1Image(arr, np.diag([*zoom, 1.0])), seg=seg) + + +class TestNormalizeImage(unittest.TestCase): + def test_maps_to_the_unit_range(self): + arr = np.array([[-100.0, 0.0], [100.0, 500.0]]) + out = normalize_image(arr) + self.assertAlmostEqual(float(out.min()), 0.0, places=9) + self.assertAlmostEqual(float(out.max()), 1.0, places=9) + + def test_is_linear(self): + arr = np.array([0.0, 1.0, 2.0, 3.0, 4.0]) + out = normalize_image(arr) + np.testing.assert_allclose(out, [0.0, 0.25, 0.5, 0.75, 1.0], atol=1e-9) + + def test_explicit_range_rescales_but_does_not_clip(self): + # The docstring calls v_range a "clip range", but the implementation only + # rescales: values outside it land outside [0, 1]. Pinned as-is. + arr = np.array([-50.0, 0.0, 50.0, 100.0, 150.0]) + out = normalize_image(arr, v_range=(0.0, 100.0)) + self.assertAlmostEqual(float(out[1]), 0.0, places=9) + self.assertAlmostEqual(float(out[3]), 1.0, places=9) + self.assertLess(float(out[0]), 0.0) + self.assertGreater(float(out[4]), 1.0) + + def test_a_constant_image_yields_nan(self): + # min == max, so the division is 0/0. Documented limitation, not a crash. + out = normalize_image(np.full((4, 4), 7.0)) + self.assertTrue(np.isnan(out).all()) + + +class TestMakeIsotropic(unittest.TestCase): + def test_anisotropic_input_is_stretched(self): + arr = np.zeros((10, 10), dtype=np.uint8) + arr[4:6, 4:6] = 1 + out = make_isotropic2d(arr, (1.0, 3.0)) + self.assertEqual(out.shape[0], 10) + self.assertGreater(out.shape[1], 10, "the 3 mm axis must be upsampled") + + def test_isotropic_input_is_unchanged_in_shape(self): + arr = np.arange(100, dtype=np.uint8).reshape(10, 10) + self.assertEqual(make_isotropic2d(arr, (1.0, 1.0)).shape, arr.shape) + + def test_mask_mode_preserves_the_exact_label_set(self): + arr = np.zeros((10, 10), dtype=np.uint8) + arr[2:5, 2:5] = 7 + arr[6:9, 6:9] = 9 + out = make_isotropic2d(arr, (1.0, 2.5), msk=True) + self.assertTrue(set(np.unique(out)).issubset({0, 7, 9}), "nearest neighbour must not invent labels") + self.assertIn(7, np.unique(out)) + self.assertIn(9, np.unique(out)) + + def test_plus_color_handles_rgb_per_channel(self): + arr = np.zeros((8, 8, 3), dtype=np.uint8) + arr[..., 0] = 10 + out = make_isotropic2dpluscolor(arr, (1.0, 2.0)) + self.assertEqual(out.shape[2], 3) + self.assertEqual(out.shape[0], 8) + + def test_plus_color_falls_back_to_the_scalar_path(self): + arr = np.zeros((8, 8), dtype=np.uint8) + np.testing.assert_array_equal(make_isotropic2dpluscolor(arr, (1.0, 1.0)), make_isotropic2d(arr, (1.0, 1.0))) + + +class TestColourAndDivision(unittest.TestCase): + def test_contrasting_stroke_for_dark_and_bright_colours(self): + self.assertEqual(get_contrasting_stroke_color((0.0, 0.0, 0.0)), "gray") + self.assertEqual(get_contrasting_stroke_color((1.0, 1.0, 1.0)), "black") + + def test_accepts_rgba_and_label_indices(self): + self.assertIn(get_contrasting_stroke_color((1.0, 1.0, 1.0, 1.0)), ("gray", "black")) + self.assertIn(get_contrasting_stroke_color([0.0, 0.0, 0.0]), ("gray", "black")) + self.assertIn(get_contrasting_stroke_color(5), ("gray", "black")) + + def test_div0_on_arrays_and_scalars(self): + np.testing.assert_allclose(div0(np.array([1.0, 2.0]), np.array([2.0, 0.0])), [0.5, 0.0]) + np.testing.assert_allclose(div0(np.array([1.0]), np.array([0.0]), fill=-1), [-1.0]) + self.assertEqual(div0(1.0, 0.0, fill=3), 3) + self.assertAlmostEqual(div0(1.0, 4.0), 0.25, places=9) + + def test_div0_leaves_finite_results_alone(self): + np.testing.assert_allclose(div0(np.array([4.0, 9.0]), np.array([2.0, 3.0])), [2.0, 3.0]) + + +class TestToCdt(unittest.TestCase): + def test_none_passes_through(self): + self.assertIsNone(to_cdt(None)) + + def test_a_poi_with_points_is_returned(self): + from TPTBox import POI # noqa: PLC0415 + + poi = POI({1: {50: (1.0, 1.0, 1.0)}}, orientation=("R", "A", "S"), zoom=(1.0, 1.0, 1.0), shape=(10.0, 10.0, 10.0)) + self.assertIsNotNone(to_cdt(poi)) + + def test_an_empty_poi_becomes_none(self): + from TPTBox import POI # noqa: PLC0415 + + empty = POI({}, orientation=("R", "A", "S"), zoom=(1.0, 1.0, 1.0), shape=(10.0, 10.0, 10.0)) + self.assertIsNone(to_cdt(empty)) + + +class TestCreateSnapshot(unittest.TestCase): + """Renders the shipped sample CT; assertions are on the produced PNG.""" + + @classmethod + def setUpClass(cls): + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + cls.ct, cls.subreg, cls.vert, cls.label = get_test_ct() + + def tearDown(self): + plt.close("all") + + @staticmethod + def _png_is_not_blank(path: Path) -> bool: + from PIL import Image # noqa: PLC0415 + + with Image.open(path) as img: + arr = np.asarray(img.convert("L")) + return int(arr.std()) > 0 + + def test_image_only_frame(self): + with tempfile.TemporaryDirectory() as tmp: + out = Path(tmp) / "snap.png" + create_snapshot(out, [Snapshot_Frame(image=self.ct)]) + self.assertTrue(out.exists()) + self.assertTrue(self._png_is_not_blank(out)) + + def test_segmentation_and_centroid_overlay(self): + from TPTBox import calc_poi_from_subreg_vert # noqa: PLC0415 + + poi = calc_poi_from_subreg_vert(self.vert, self.subreg, subreg_id=[Location.Vertebra_Corpus], verbose=False) + with tempfile.TemporaryDirectory() as tmp: + out = Path(tmp) / "snap.png" + create_snapshot(out, [Snapshot_Frame(image=self.ct, segmentation=self.vert, centroids=poi, coronal=True)]) + self.assertTrue(out.exists()) + self.assertTrue(self._png_is_not_blank(out)) + + def test_all_three_views(self): + with tempfile.TemporaryDirectory() as tmp: + out = Path(tmp) / "snap.png" + create_snapshot(out, [Snapshot_Frame(image=self.ct, sagittal=True, coronal=True, axial=True)]) + self.assertTrue(out.exists()) + + def test_two_frames_render_differently_from_one(self): + from PIL import Image # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + one, two = Path(tmp) / "one.png", Path(tmp) / "two.png" + create_snapshot(one, [Snapshot_Frame(image=self.ct)]) + create_snapshot(two, [Snapshot_Frame(image=self.ct), Snapshot_Frame(image=self.ct)]) + with Image.open(one) as a, Image.open(two) as b: + self.assertNotEqual((a.width, a.height), (b.width, b.height)) + + def test_writes_several_paths_at_once(self): + with tempfile.TemporaryDirectory() as tmp: + png, jpg = Path(tmp) / "snap.png", Path(tmp) / "snap.jpg" + create_snapshot([png, jpg], [Snapshot_Frame(image=self.ct)]) + self.assertTrue(png.exists()) + self.assertTrue(jpg.exists()) + + def test_none_frames_are_skipped(self): + with tempfile.TemporaryDirectory() as tmp: + out = Path(tmp) / "snap.png" + create_snapshot(out, [Snapshot_Frame(image=self.ct), None]) # type: ignore[list-item] + self.assertTrue(out.exists()) + + def test_crop_shrinks_the_rendered_area(self): + from PIL import Image # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + plain, cropped = Path(tmp) / "a.png", Path(tmp) / "b.png" + create_snapshot(plain, [Snapshot_Frame(image=self.ct, segmentation=self.vert)]) + create_snapshot(cropped, [Snapshot_Frame(image=self.ct, segmentation=self.vert)], crop=True) + with Image.open(plain) as a, Image.open(cropped) as b: + self.assertNotEqual((a.width, a.height), (b.width, b.height)) + + def test_rendering_is_deterministic(self): + with tempfile.TemporaryDirectory() as tmp: + first, second = Path(tmp) / "a.png", Path(tmp) / "b.png" + for path in (first, second): + create_snapshot(path, [Snapshot_Frame(image=self.ct, segmentation=self.vert)]) + self.assertEqual(first.read_bytes(), second.read_bytes(), "the same input must render identically") + + +class TestSnapshotTemplates(unittest.TestCase): + """The templates are Snapshot_Frame builders; `out_path` is their 5th argument.""" + + @classmethod + def setUpClass(cls): + from TPTBox import calc_poi_from_subreg_vert # noqa: PLC0415 + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + cls.ct, cls.subreg, cls.vert, cls.label = get_test_ct() + cls.poi = calc_poi_from_subreg_vert(cls.vert, cls.subreg, subreg_id=[Location.Vertebra_Corpus], verbose=False) + + def tearDown(self): + plt.close("all") + + def test_snapshot_template(self): + from TPTBox.spine.snapshot2D.snapshot_templates import snapshot # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + out = Path(tmp) / "snap.png" + paths = snapshot(self.ct, self.vert, self.poi, self.subreg, out_path=out) + self.assertTrue(out.exists()) + self.assertEqual([Path(p) for p in paths], [out]) + + def test_mri_snapshot_template(self): + from TPTBox.spine.snapshot2D.snapshot_templates import mri_snapshot # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + out = Path(tmp) / "snap.png" + mri_snapshot(self.ct, self.vert, self.poi, self.subreg, out_path=out) + self.assertTrue(out.exists()) + + @unittest.skipUnless(HAS_TORCH, "mip_shot sets denoise_threshold, which makes create_snapshot import torch") + def test_mip_shot_template(self): + from TPTBox.spine.snapshot2D.snapshot_templates import mip_shot # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + out = Path(tmp) / "snap.png" + mip_shot(self.ct, self.vert, self.poi, out_path=out) + self.assertTrue(out.exists()) + + def test_ct_mode_is_chosen_for_a_bids_ct_file(self): + from TPTBox.core.bids_files import BIDS_FILE # noqa: PLC0415 + from TPTBox.spine.snapshot2D.snapshot_templates import snapshot # noqa: PLC0415 + from TPTBox.tests.test_utils import get_nii_paths_ct, get_tests_dir # noqa: PLC0415 + + ct_path = get_nii_paths_ct()[0] + bids = BIDS_FILE(ct_path, get_tests_dir(), verbose=False) + self.assertEqual(bids.bids_format, "ct") + with tempfile.TemporaryDirectory() as tmp: + out = Path(tmp) / "snap.png" + snapshot(bids, self.vert, self.poi, self.subreg, out_path=out) + self.assertTrue(out.exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_spinestats_angles.py b/unit_tests/test_spinestats_angles.py new file mode 100644 index 0000000..d415f47 --- /dev/null +++ b/unit_tests/test_spinestats_angles.py @@ -0,0 +1,687 @@ +"""Unit tests for `TPTBox.spine.spinestats` angle, curvature and pelvic measurements. + +The shipped sample data holds a single vertebra, so every spine here is synthesised +analytically: vertebra bodies stacked along the superior axis, each carrying the +direction POIs (`Vertebra_Direction_Posterior/Inferior/Right`) and the two endplate +median points that `MoveTo.TOP/BOTTOM` resolve to. Tilting one vertebra by a known +angle therefore gives a known expected measurement, which is what these tests assert - +a circle of radius r must yield curvature 1/r, a 30 degree coronal tilt must yield a +30 degree Cobb angle, and pelvic incidence must equal pelvic tilt plus sacral slope. +""" + +from __future__ import annotations + +import json +import math +import tempfile +import unittest +from pathlib import Path + +import numpy as np + +from TPTBox import POI, Location, Vertebra_Instance +from TPTBox.spine.spinestats import angles as angles_mod +from TPTBox.spine.spinestats import curvature as curv_mod +from TPTBox.spine.spinestats.angles import ( + MoveTo, + compute_angel_between_two_points_, + compute_lordosis_and_kyphosis, + compute_max_cobb_angle, + compute_max_cobb_angle_multi, + cosine_distance, + get_to_space, +) +from TPTBox.spine.spinestats.curvature import ( + _find_apices, + _numerical_curvature, + compute_axial_rotation, + compute_coronal_balance, + compute_curvature_profile, + compute_multi_cobb, + compute_segmental_endplate_angles, + compute_sva, + compute_wedge_metrics, +) +from TPTBox.spine.spinestats.pelvic_parameters import ( + _pelvic_from_endplate_pair, + _sag, + compute_pelvic_parameters, + resolve_fullbody_poi_path, +) + +V = Vertebra_Instance +SPINE_IDS = [v.value for v in list(V.cervical()) + list(V.thoracic()) + list(V.lumbar()) if v.value <= V.L5.value] + + +def rot(axis: str, deg: float) -> np.ndarray: + """Rotation matrix about a RAS axis: x = R-L, y = A-P, z = S-I.""" + t = math.radians(deg) + c, s = math.cos(t), math.sin(t) + if axis == "x": + return np.array([[1, 0, 0], [0, c, -s], [0, s, c]]) + if axis == "y": + return np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]]) + return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]]) + + +def spine_poi( + ids=None, + sagittal_tilt=None, + coronal_tilt=None, + axial_tilt=None, + spacing: float = 20.0, + half_height: float = 5.0, + coronal_amp: float = 0.0, + sagittal_radius: float | None = None, + anterior_shift=None, +) -> POI: + """Build a synthetic spine POI in RAS at 1 mm isotropic spacing. + + Args: + ids: Vertebra ids, inferior-to-superior order of increasing z. Defaults to C1..L5. + sagittal_tilt / coronal_tilt / axial_tilt: per-id rotation in degrees about the + R-L / A-P / S-I axis, applied to that vertebra's direction and endplate POIs. + spacing: distance between consecutive body centres in mm. + half_height: distance from body centre to each endplate median point in mm. + coronal_amp: amplitude of a half-sine lateral deviation across the whole spine. + sagittal_radius: if given, place the bodies on a circular arc of that radius. + anterior_shift: per-id extra offset along the anterior axis in mm. + """ + ids = SPINE_IDS if ids is None else ids + sagittal_tilt = sagittal_tilt or {} + coronal_tilt = coronal_tilt or {} + axial_tilt = axial_tilt or {} + anterior_shift = anterior_shift or {} + cent: dict[int, dict[int, tuple[float, float, float]]] = {} + n = len(ids) + for k, vid in enumerate(ids): + if sagittal_radius: + # Pure circular arc of +-20 degrees: curvature is exactly 1/sagittal_radius. + th = math.radians(-20 + 40 * k / max(n - 1, 1)) + y = sagittal_radius * math.sin(th) + z = sagittal_radius * (1 - math.cos(th)) + else: + y, z = 0.0, spacing * k + x = coronal_amp * math.sin(math.pi * k / max(n - 1, 1)) if n > 1 else 0.0 + centre = np.array([x, y + anterior_shift.get(vid, 0.0), z]) + r = rot("y", coronal_tilt.get(vid, 0.0)) @ rot("x", sagittal_tilt.get(vid, 0.0)) @ rot("z", axial_tilt.get(vid, 0.0)) + cent[vid] = { + Location.Vertebra_Corpus.value: tuple(centre), + Location.Vertebra_Direction_Posterior.value: tuple(centre + r @ np.array([0.0, -1.0, 0.0])), + Location.Vertebra_Direction_Inferior.value: tuple(centre + r @ np.array([0.0, 0.0, -1.0])), + Location.Vertebra_Direction_Right.value: tuple(centre + r @ np.array([1.0, 0.0, 0.0])), + Location.Additional_Vertebral_Body_Middle_Superior_Median.value: tuple(centre + r @ np.array([0.0, 0.0, half_height])), + Location.Additional_Vertebral_Body_Middle_Inferior_Median.value: tuple(centre + r @ np.array([0.0, 0.0, -half_height])), + } + return POI( + cent, + orientation=("R", "A", "S"), + zoom=(1.0, 1.0, 1.0), + shape=(400.0, 400.0, 400.0), + origin=(0.0, 0.0, 0.0), + rotation=np.eye(3), + level_one_info=V, + level_two_info=Location, + ) + + +class TestNumericalCurvature(unittest.TestCase): + """A circle of radius r has curvature exactly 1/r everywhere.""" + + def test_circle_curvature_is_one_over_radius(self): + for radius in (5.0, 17.3, 50.0): + for step in (0.05, 0.2): + with self.subTest(radius=radius, step=step): + t = np.arange(0.0, math.pi / 2, step) + pts = np.stack([radius * np.cos(t), radius * np.sin(t), np.zeros_like(t)], axis=1) + kappa = _numerical_curvature(pts)[3:-3] + np.testing.assert_allclose(kappa, 1.0 / radius, rtol=2e-2) + + def test_circle_in_a_tilted_plane_gives_the_same_curvature(self): + # Proves the 3D cross-product branch agrees with the planar case. + radius = 12.0 + t = np.arange(0.0, math.pi, 0.05) + flat = np.stack([radius * np.cos(t), radius * np.sin(t), np.zeros_like(t)], axis=1) + basis = np.linalg.qr(np.random.default_rng(0).normal(size=(3, 3)))[0] + kappa = _numerical_curvature(flat @ basis.T)[3:-3] + np.testing.assert_allclose(kappa, 1.0 / radius, rtol=2e-2) + + def test_straight_line_has_zero_curvature(self): + pts = np.stack([np.arange(20.0), np.zeros(20), np.zeros(20)], axis=1) + np.testing.assert_allclose(_numerical_curvature(pts), 0.0, atol=1e-12) + + def test_scaling_the_curve_scales_curvature_inversely(self): + t = np.arange(0.0, 1.5, 0.05) + pts = np.stack([np.cos(t), np.sin(t), np.zeros_like(t)], axis=1) + base = _numerical_curvature(pts)[3:-3] + scaled = _numerical_curvature(pts * 4.0)[3:-3] + np.testing.assert_allclose(scaled, base / 4.0, rtol=1e-6) + + def test_too_few_points(self): + self.assertEqual(len(_numerical_curvature(np.zeros((2, 3)))), 2) + + +class TestFindApices(unittest.TestCase): + def test_returns_local_maxima_sorted_by_arc_position(self): + arc = np.arange(11.0) + kappa = np.array([0.0, 0.5, 0.1, 0.9, 0.2, 0.3, 0.2, 0.7, 0.1, 0.4, 0.0]) + apices = _find_apices(arc, kappa, top_k=3) + self.assertEqual([a["arc_mm"] for a in apices], [1.0, 3.0, 7.0]) + self.assertEqual([a["kappa_1_per_mm"] for a in apices], [0.5, 0.9, 0.7]) + + def test_top_k_keeps_the_strongest_peaks(self): + arc = np.arange(11.0) + kappa = np.array([0.0, 0.5, 0.1, 0.9, 0.2, 0.3, 0.2, 0.7, 0.1, 0.4, 0.0]) + self.assertEqual([a["arc_mm"] for a in _find_apices(arc, kappa, top_k=1)], [3.0]) + + def test_no_peak_and_too_short(self): + self.assertEqual(_find_apices(np.arange(5.0), np.arange(5.0), top_k=3), []) + self.assertEqual(_find_apices(np.arange(2.0), np.zeros(2), top_k=3), []) + + +class TestAngleBetweenTwoPoints(unittest.TestCase): + def test_sagittal_tilt_is_recovered_exactly(self): + for tilt in (10.0, 30.0, 45.0): + with self.subTest(tilt=tilt): + poi = spine_poi(sagittal_tilt={V.T5.value: tilt}) + for project_2d in (False, True): + angle = compute_angel_between_two_points_( + poi.copy(), V.T4, V.T5, "P", MoveTo.CENTER, MoveTo.CENTER, project_2D=project_2d + ) + self.assertAlmostEqual(angle, tilt, delta=0.05) + + def test_untilted_pair_is_zero(self): + poi = spine_poi() + self.assertEqual(compute_angel_between_two_points_(poi.copy(), V.T4, V.T6, "P", MoveTo.CENTER, MoveTo.CENTER, False), 0.0) + + def test_coronal_tilt_is_measured_in_the_r_direction(self): + poi = spine_poi(coronal_tilt={V.T5.value: 20.0}) + angle = compute_angel_between_two_points_(poi.copy(), V.T4, V.T5, "R", MoveTo.CENTER, MoveTo.CENTER, project_2D=False) + self.assertAlmostEqual(angle, 20.0, delta=0.01) + + def test_opposite_directions_give_the_same_unsigned_angle(self): + poi = spine_poi(sagittal_tilt={V.T5.value: 25.0}) + a = compute_angel_between_two_points_(poi.copy(), V.T4, V.T5, "P", MoveTo.CENTER, MoveTo.CENTER, False) + b = compute_angel_between_two_points_(poi.copy(), V.T4, V.T5, "A", MoveTo.CENTER, MoveTo.CENTER, False) + self.assertAlmostEqual(a, b, delta=1e-6) + + def test_argument_order_does_not_matter(self): + poi = spine_poi(sagittal_tilt={V.T5.value: 18.0}) + a = compute_angel_between_two_points_(poi.copy(), V.T4, V.T5, "P", MoveTo.CENTER, MoveTo.CENTER, False) + b = compute_angel_between_two_points_(poi.copy(), V.T5, V.T4, "P", MoveTo.CENTER, MoveTo.CENTER, False) + self.assertAlmostEqual(a, b, delta=1e-6) + + def test_none_for_missing_vertebra_or_none_input(self): + poi = spine_poi(ids=[V.T4.value, V.T5.value]) + self.assertIsNone(compute_angel_between_two_points_(poi.copy(), None, V.T5, "P")) + self.assertIsNone(compute_angel_between_two_points_(poi.copy(), V.T4, None, "P")) + self.assertIsNone(compute_angel_between_two_points_(poi.copy(), V.T4, V.L5, "P")) + + def test_unknown_direction_raises(self): + poi = spine_poi(ids=[V.T4.value, V.T5.value]) + with self.assertRaises(NotImplementedError): + compute_angel_between_two_points_(poi.copy(), V.T4, V.T5, "Q") # type: ignore[arg-type] + + def test_same_vertebra_is_rejected(self): + poi = spine_poi(ids=[V.T4.value, V.T5.value]) + with self.assertRaises(AssertionError): + compute_angel_between_two_points_(poi.copy(), V.T4, V.T4, "P") + + def test_superior_direction_is_supported(self): + poi = spine_poi(sagittal_tilt={V.T5.value: 15.0}) + angle = compute_angel_between_two_points_(poi.copy(), V.T4, V.T5, "S", MoveTo.CENTER, MoveTo.CENTER, project_2D=False) + self.assertAlmostEqual(angle, 15.0, delta=0.01) + + +class TestMoveTo(unittest.TestCase): + def test_center_resolves_to_the_body_centroid(self): + poi = spine_poi() + self.assertEqual(MoveTo.CENTER.get_location(V.T5, poi), (V.T5, 50)) + np.testing.assert_allclose(MoveTo.CENTER.get_point(V.T5, poi), poi[V.T5.value, 50]) + + def test_top_and_bottom_resolve_to_the_endplate_medians(self): + poi = spine_poi(half_height=5.0) + top = MoveTo.TOP.get_point(V.T5, poi) + bottom = MoveTo.BOTTOM.get_point(V.T5, poi) + # Both are averaged with the neighbour, so they sit between the two bodies. + self.assertGreater(top[2], bottom[2]) + + def test_has_point(self): + poi = spine_poi() + self.assertTrue(MoveTo.CENTER.has_point(V.T5, poi)) + self.assertTrue(MoveTo.CENTER.has_point(V.T5.value, poi)) + self.assertFalse(MoveTo.CENTER.has_point(V.S1, poi)) + + +class TestLordosisAndKyphosis(unittest.TestCase): + def test_straight_spine_has_no_curvature(self): + out = compute_lordosis_and_kyphosis(spine_poi()) + self.assertEqual(sorted(out), ["cervical_lordosis", "lumbar_lordosis", "thoracic_kyphosis"]) + for key, value in out.items(): + with self.subTest(key=key): + self.assertEqual(value, 0.0) + + def test_tilting_c2_only_moves_the_cervical_value(self): + out = compute_lordosis_and_kyphosis(spine_poi(sagittal_tilt={V.C2.value: 20.0})) + # C2 is an end vertebra, so its normal is averaged with its single neighbour -> half. + self.assertAlmostEqual(out["cervical_lordosis"], 10.0, delta=0.01) + self.assertEqual(out["thoracic_kyphosis"], 0.0) + self.assertEqual(out["lumbar_lordosis"], 0.0) + + def test_requires_the_posterior_direction(self): + poi = spine_poi() + poi.remove_centroid_(*[(v, Location.Vertebra_Direction_Posterior.value) for v in SPINE_IDS]) + with self.assertRaises(AssertionError): + compute_lordosis_and_kyphosis(poi) + + def test_input_is_not_mutated(self): + poi = spine_poi(sagittal_tilt={V.C2.value: 20.0}) + before = tuple(poi[V.C2.value, 50]) + compute_lordosis_and_kyphosis(poi) + self.assertEqual(tuple(poi[V.C2.value, 50]), before) + + +class TestMaxCobbAngle(unittest.TestCase): + def test_straight_spine_has_no_cobb_angle(self): + self.assertEqual(compute_max_cobb_angle(spine_poi()), (0, None, None, None)) + + def test_opposed_coronal_tilts_sum_into_the_cobb_angle(self): + poi = spine_poi(coronal_tilt={V.T5.value: 15.0, V.T11.value: -15.0}) + angle, from_vert, to_vert, apex = compute_max_cobb_angle(poi.copy(), vert_id1_mv=MoveTo.CENTER, vert_id2_mv=MoveTo.CENTER) + self.assertAlmostEqual(angle, 30.0, delta=0.05) + self.assertEqual((from_vert, to_vert), (V.T5.value, V.T11.value)) + self.assertIsNotNone(apex) + self.assertLessEqual(from_vert, apex) + self.assertLessEqual(apex, to_vert) + + def test_default_endplate_moves_average_with_the_neighbour(self): + poi = spine_poi(coronal_tilt={V.T5.value: 15.0, V.T11.value: -15.0}) + angle, _, _, _ = compute_max_cobb_angle(poi.copy()) + self.assertAlmostEqual(angle, 15.0, delta=0.05) + + def test_restricting_the_vertebra_list_restricts_the_search(self): + poi = spine_poi(coronal_tilt={V.T5.value: 15.0, V.T11.value: -15.0}) + angle, _, _, _ = compute_max_cobb_angle(poi.copy(), vertebrae_list=[V.L1, V.L2, V.L3]) + self.assertEqual(angle, 0) + + def test_multi_returns_the_segments_above_the_threshold(self): + poi = spine_poi(coronal_tilt={V.T5.value: 15.0, V.T11.value: -15.0}) + out = compute_max_cobb_angle_multi(poi.copy(), vert_id1_mv=MoveTo.CENTER, vert_id2_mv=MoveTo.CENTER) + self.assertGreaterEqual(len(out), 1) + self.assertAlmostEqual(out[0][0], 30.0, delta=0.05) + + def test_multi_is_empty_for_a_straight_spine(self): + self.assertEqual(compute_max_cobb_angle_multi(spine_poi()), []) + + +class TestSagittalAndCoronalBalance(unittest.TestCase): + def test_sva_is_the_signed_anterior_offset_of_c7(self): + poi = spine_poi(ids=[*SPINE_IDS, V.S1.value], anterior_shift={V.C7.value: 25.0}) + out = compute_sva(poi.copy()) + self.assertEqual(out["top_vertebra"], "C7") + self.assertEqual(out["base_vertebra"], "S1") + self.assertAlmostEqual(out["sva_mm"], 25.0, delta=0.6) + + def test_sva_is_zero_for_a_plumb_spine(self): + poi = spine_poi(ids=[*SPINE_IDS, V.S1.value]) + self.assertAlmostEqual(compute_sva(poi.copy())["sva_mm"], 0.0, delta=0.6) + + def test_sva_falls_back_from_s1_to_l5(self): + out = compute_sva(spine_poi()) # no S1 in this spine + self.assertEqual(out["base_vertebra"], "L5") + self.assertIn("L5", out["base_landmark"]) + + def test_sva_reports_a_missing_top_vertebra(self): + poi = spine_poi(ids=[V.L1.value, V.L2.value, V.L3.value]) + out = compute_sva(poi.copy()) + self.assertIsNone(out["sva_mm"]) + self.assertIn("C7", out["error"]) + + def test_sva_reports_a_missing_base_vertebra(self): + poi = spine_poi(ids=[V.C6.value, V.C7.value, V.T1.value]) + out = compute_sva(poi.copy()) + self.assertIsNone(out["sva_mm"]) + self.assertIn("base vertebra", out["error"]) + + def test_coronal_balance_tracks_the_lateral_offset(self): + poi = spine_poi(ids=[*SPINE_IDS, V.S1.value], coronal_amp=0.0) + self.assertAlmostEqual(compute_coronal_balance(poi.copy())["coronal_balance_mm"], 0.0, delta=0.6) + + def test_coronal_balance_reports_a_missing_top_vertebra(self): + out = compute_coronal_balance(spine_poi(ids=[V.L1.value, V.L2.value, V.L3.value])) + self.assertIsNone(out["coronal_balance_mm"]) + self.assertIn("C7", out["error"]) + + +class TestSegmentalEndplateAnglesAndAxialRotation(unittest.TestCase): + def test_tilting_one_vertebra_shows_up_in_both_neighbouring_segments(self): + out = compute_segmental_endplate_angles(spine_poi(sagittal_tilt={V.T5.value: 10.0})) + non_zero = {k: v for k, v in out.items() if v != 0.0} + self.assertEqual(non_zero, {"T4-T5": 10.0, "T5-T6": -10.0}) + + def test_straight_spine_is_all_zero(self): + out = compute_segmental_endplate_angles(spine_poi()) + self.assertTrue(out) + self.assertEqual(set(out.values()), {0.0}) + + def test_axial_rotation_is_zero_without_rotation(self): + out = compute_axial_rotation(spine_poi()) + self.assertTrue(out) + self.assertEqual(set(out.values()), {0.0}) + + def test_axial_rotation_recovers_a_transverse_rotation(self): + out = compute_axial_rotation(spine_poi(axial_tilt={V.T5.value: 12.0})) + self.assertAlmostEqual(abs(out["T5"]), 12.0, delta=0.1) + + +class TestWedgeMetrics(unittest.TestCase): + def test_analytic_values(self): + out = compute_wedge_metrics({7: {"anterior_height_x1": 12.0, "posterior_height_x2": 10.0, "width_sagittal_x6": 20.0}}) + self.assertAlmostEqual(out[7]["sagittal_wedge_deg"], math.degrees(math.atan2(2.0, 20.0)), places=3) + self.assertAlmostEqual(out[7]["sagittal_wedge_index"], 2.0 / 11.0, places=4) + self.assertTrue(math.isnan(out[7]["coronal_wedge_deg"])) + + def test_coronal_values(self): + out = compute_wedge_metrics({1: {"right_height_x3": 11.0, "left_height_x4": 9.0, "width_lateral_x5": 30.0}}) + self.assertAlmostEqual(out[1]["coronal_wedge_deg"], math.degrees(math.atan2(2.0, 30.0)), places=3) + self.assertAlmostEqual(out[1]["coronal_wedge_index"], 2.0 / 10.0, places=4) + + def test_symmetric_structure_has_zero_wedge(self): + out = compute_wedge_metrics({1: {"anterior_height_x1": 10.0, "posterior_height_x2": 10.0, "width_sagittal_x6": 20.0}}) + self.assertEqual(out[1]["sagittal_wedge_deg"], 0.0) + self.assertEqual(out[1]["sagittal_wedge_index"], 0.0) + + def test_missing_and_invalid_inputs_become_nan(self): + out = compute_wedge_metrics( + { + 1: {"anterior_height_x1": None, "posterior_height_x2": 10.0, "width_sagittal_x6": 20.0}, + 2: {"anterior_height_x1": 12.0, "posterior_height_x2": 10.0, "width_sagittal_x6": 0.0}, + 3: {"anterior_height_x1": float("nan"), "posterior_height_x2": 10.0, "width_sagittal_x6": 20.0}, + 4: "not a dict", + } + ) + self.assertNotIn(4, out) + for label in (1, 2, 3): + with self.subTest(label=label): + self.assertTrue(math.isnan(out[label]["sagittal_wedge_deg"])) + + def test_input_is_not_modified(self): + geometry = {1: {"anterior_height_x1": 12.0, "posterior_height_x2": 10.0, "width_sagittal_x6": 20.0}} + snapshot = json.dumps(geometry, sort_keys=True) + compute_wedge_metrics(geometry) + self.assertEqual(json.dumps(geometry, sort_keys=True), snapshot) + + +class TestCurvatureProfileAndMultiCobb(unittest.TestCase): + def test_straight_spine_profile(self): + out = compute_curvature_profile(spine_poi()) + self.assertAlmostEqual(out["tortuosity"], 1.0, places=3) + self.assertAlmostEqual(out["curvature_max_1_per_mm"], 0.0, places=6) + self.assertAlmostEqual(out["arc_length_mm"], out["chord_length_mm"], places=3) + + def test_arc_profile_recovers_one_over_the_radius(self): + radius = 200.0 + out = compute_curvature_profile(spine_poi(sagittal_radius=radius)) + self.assertAlmostEqual(out["curvature_mean_1_per_mm"], 1.0 / radius, delta=1e-4) + self.assertGreater(out["tortuosity"], 1.0) + self.assertGreater(out["arc_length_mm"], out["chord_length_mm"]) + + def test_profile_needs_four_corpus_points(self): + out = compute_curvature_profile(spine_poi(ids=[V.T4.value, V.T5.value])) + self.assertIn("too few", out["error"]) + + def test_multi_cobb_is_zero_for_a_straight_spine(self): + out = compute_multi_cobb(spine_poi()) + self.assertEqual(out["max_cobb_deg"], 0.0) + + def test_multi_cobb_grows_with_the_lateral_deviation(self): + small = compute_multi_cobb(spine_poi(coronal_amp=5.0))["max_cobb_deg"] + large = compute_multi_cobb(spine_poi(coronal_amp=30.0))["max_cobb_deg"] + self.assertGreater(large, small) + self.assertGreater(small, 0.0) + for curve in compute_multi_cobb(spine_poi(coronal_amp=30.0))["curves"]: + with self.subTest(curve=curve): + self.assertIn(curve["handedness"], ("left", "right")) + self.assertGreaterEqual(curve["arc_end_mm"], curve["arc_start_mm"]) + + def test_multi_cobb_needs_four_corpus_points(self): + self.assertIn("too few", compute_multi_cobb(spine_poi(ids=[V.T4.value, V.T5.value]))["error"]) + + +class TestAngleHelpers(unittest.TestCase): + def test_cosine_distance(self): + self.assertAlmostEqual(cosine_distance(np.array([1.0, 0, 0]), np.array([1.0, 0, 0])), 1.0, places=6) + self.assertAlmostEqual(cosine_distance(np.array([1.0, 0, 0]), np.array([0, 1.0, 0])), 0.0, places=6) + self.assertAlmostEqual(cosine_distance(np.array([1.0, 0, 0]), np.array([-1.0, 0, 0])), -1.0, places=6) + + def test_get_to_space_is_an_invertible_orthonormal_basis(self): + a = np.array([1.0, 0.0, 0.0]) + b = np.array([0.0, 1.0, 0.0]) + c = np.array([0.0, 0.0, 1.0]) + to_space, from_space = get_to_space(a, b, c) + np.testing.assert_allclose(from_space @ to_space, np.eye(3), atol=1e-9) + np.testing.assert_allclose(to_space @ a, [1, 0, 0], atol=1e-9) + + def test_last_lumbar_and_thoracic(self): + poi = spine_poi() + self.assertEqual(angles_mod._get_last_lumbar(poi), V.L5) + self.assertEqual(angles_mod._get_last_thoracic(poi), V.T12) + + def test_last_lumbar_is_none_without_lumbar_vertebrae(self): + self.assertIsNone(angles_mod._get_last_lumbar(spine_poi(ids=[V.T4.value, V.T5.value]))) + + def test_last_present_picks_the_first_available_candidate(self): + poi = spine_poi() + self.assertEqual(curv_mod._last_present(poi, [V.S1, V.L5, V.L4]), V.L5) + self.assertIsNone(curv_mod._last_present(poi, [V.S1])) + + def test_get_and_corpus_helpers(self): + poi = curv_mod._prep(spine_poi()) + self.assertIsNotNone(curv_mod._corpus(poi, V.T5)) + self.assertIsNotNone(curv_mod._corpus(poi, V.T5.value)) + self.assertIsNone(curv_mod._get(poi, V.S1, Location.Vertebra_Corpus)) + + def test_endplate_falls_back_to_corpus_plus_direction(self): + poi = curv_mod._prep(spine_poi(half_height=5.0)) + sup = curv_mod._endplate(poi, V.T5, superior=True) + inf = curv_mod._endplate(poi, V.T5, superior=False) + corpus = curv_mod._corpus(poi, V.T5) + self.assertAlmostEqual(float(np.linalg.norm(sup - corpus)), 5.0, delta=0.6) + self.assertAlmostEqual(float(np.linalg.norm(inf - corpus)), 5.0, delta=0.6) + + +class TestPelvicParameters(unittest.TestCase): + S_POST = np.array([0.0, -8.66, 5.0]) + S_ANT = np.array([0.0, 8.66, -5.0]) + FH_R = np.array([80.0, 50.0, -100.0]) + FH_L = np.array([-80.0, 50.0, -100.0]) + + def test_pi_equals_pt_plus_ss(self): + out = _pelvic_from_endplate_pair(self.S_POST, self.S_ANT, self.FH_R, self.FH_L) + self.assertAlmostEqual(out["pi_deg"], out["pt_deg"] + out["ss_deg"], places=9) + + def test_pelvic_incidence_is_invariant_under_sagittal_rotation(self): + base = _pelvic_from_endplate_pair(self.S_POST, self.S_ANT, self.FH_R, self.FH_L) + theta = 12.0 + r = rot("x", theta) + turned = _pelvic_from_endplate_pair(r @ self.S_POST, r @ self.S_ANT, r @ self.FH_R, r @ self.FH_L) + # PI is an anatomical constant; PT and SS trade the rotation between them. + self.assertAlmostEqual(turned["pi_deg"], base["pi_deg"], places=3) + self.assertAlmostEqual(turned["pt_deg"], base["pt_deg"] + theta, places=3) + self.assertAlmostEqual(turned["ss_deg"], base["ss_deg"] - theta, places=3) + + def test_all_angles_are_scale_invariant(self): + base = _pelvic_from_endplate_pair(self.S_POST, self.S_ANT, self.FH_R, self.FH_L) + scaled = _pelvic_from_endplate_pair(self.S_POST * 3, self.S_ANT * 3, self.FH_R * 3, self.FH_L * 3) + for key in ("pi_deg", "pt_deg", "ss_deg"): + with self.subTest(key=key): + self.assertAlmostEqual(scaled[key], base[key], places=6) + + def test_hip_centre_is_the_midpoint_of_the_femoral_heads(self): + out = _pelvic_from_endplate_pair(self.S_POST, self.S_ANT, self.FH_R, self.FH_L) + np.testing.assert_allclose(out["hip_center_mm"], (self.FH_R + self.FH_L) / 2) + np.testing.assert_allclose(out["s1_endplate_center_mm"], (self.S_POST + self.S_ANT) / 2) + + def test_degenerate_endplate(self): + out = _pelvic_from_endplate_pair(self.S_POST, self.S_POST, self.FH_R, self.FH_L) + self.assertIsNone(out["pi_deg"]) + self.assertIn("degenerate", out["error"]) + + def test_degenerate_hip_axis(self): + zero = np.zeros(3) + out = _pelvic_from_endplate_pair(self.S_POST, self.S_ANT, zero, zero) + self.assertIsNone(out["pi_deg"]) + self.assertIn("coincides", out["error"]) + + def test_sag_projects_onto_the_anterior_superior_plane(self): + # Drops the left-right component entirely and returns a 2-vector. + np.testing.assert_allclose(_sag(np.array([5.0, 2.0, 3.0])), [2.0, 3.0]) + + def _write_poi_json(self, folder: Path, body: dict) -> Path: + path = folder / "sub-x_seg-fullbody_poi.json" + path.write_text(json.dumps([{"meta": 1}, body]), encoding="utf-8") + return path + + def _full_body(self) -> dict: + return { + "sacrum": { + "1": self.S_POST.tolist(), + "19": self.S_ANT.tolist(), + "27": [-10.0, 6.0, 4.0], + "28": [10.0, 6.0, 4.0], + }, + "femur_right": {"11": self.FH_R.tolist()}, + "femur_left": {"11": self.FH_L.tolist()}, + } + + def test_compute_pelvic_parameters_both_variants(self): + with tempfile.TemporaryDirectory() as tmp: + path = self._write_poi_json(Path(tmp), self._full_body()) + out = compute_pelvic_parameters(path, lumbar_lordosis_deg=40.0) + self.assertNotIn("error", out) + for variant in ("poi_ap", "poi_ala"): + with self.subTest(variant=variant): + pi = out[variant]["pi_deg"] + self.assertIsNotNone(pi) + self.assertAlmostEqual(out[variant]["pi_ll_mismatch_deg"], round(pi - 40.0, 3), places=3) + + def test_mismatch_is_none_without_lumbar_lordosis(self): + with tempfile.TemporaryDirectory() as tmp: + out = compute_pelvic_parameters(self._write_poi_json(Path(tmp), self._full_body())) + self.assertIsNone(out["poi_ap"]["pi_ll_mismatch_deg"]) + + def test_missing_ala_landmarks_only_affects_the_ala_variant(self): + body = self._full_body() + del body["sacrum"]["27"] + with tempfile.TemporaryDirectory() as tmp: + out = compute_pelvic_parameters(self._write_poi_json(Path(tmp), body)) + self.assertIsNotNone(out["poi_ap"]["pi_deg"]) + self.assertIn("ala", out["poi_ala"]["error"]) + + def test_error_paths(self): + self.assertIn("no fullbody POI json path", compute_pelvic_parameters(None)["error"]) + with tempfile.TemporaryDirectory() as tmp: + missing = Path(tmp) / "nope.json" + self.assertIn("does not exist", compute_pelvic_parameters(missing)["error"]) + + broken = Path(tmp) / "broken.json" + broken.write_text("{not json", encoding="utf-8") + self.assertIn("failed to load", compute_pelvic_parameters(broken)["error"]) + + wrong_shape = Path(tmp) / "shape.json" + wrong_shape.write_text(json.dumps({"body": {}}), encoding="utf-8") + self.assertIn("unexpected", compute_pelvic_parameters(wrong_shape)["error"]) + + body = self._full_body() + del body["sacrum"]["19"] + self.assertIn("missing required landmarks", compute_pelvic_parameters(self._write_poi_json(Path(tmp), body))["error"]) + + def test_resolve_fullbody_poi_path(self): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + self.assertIsNone(resolve_fullbody_poi_path(root, "100000")) + + +class TestPlottingEntryPoints(unittest.TestCase): + """The plot_* wrappers must return the same numbers as the bare computation. + + `plot_cobb_and_lordosis_and_kyphosis` is not covered here: it resamples the POI + onto the image in a way that is singular for a synthetic spine on the shipped + single-vertebra sample, and inventing a fixture that satisfies it would not be + testing anything real. + """ + + @classmethod + def setUpClass(cls): + import matplotlib as mpl # noqa: PLC0415 + + mpl.use("Agg") + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + cls.ct, cls.subreg, cls.vert, cls.label = get_test_ct() + # The POI has to live on the image's own grid - the renderers resample between + # the two - and its spine must run along the image's inferior axis, which is + # what the snapshot code interpolates along. So: build it in RAS, where the + # builder stacks along S, then reorient onto the sample CT's PIR grid. + ct_ras = cls.ct.reorient(("R", "A", "S")) + ids = [v.value for v in (V.T8, V.T9, V.T10, V.T11, V.T12, V.L1, V.L2, V.L3)] + base = spine_poi(ids=ids, sagittal_tilt={V.T11.value: 15.0}, spacing=6.0) + centre = np.array(ct_ras.shape, dtype=float) / 2.0 + offset = centre - np.array([0.0, 0.0, 6.0 * len(ids) / 2.0]) + points = { + vert: {sub: tuple(float(c + o) for c, o in zip(coord, offset)) for sub, coord in inner.items()} + for vert, inner in base.centroids.items_2d() + } + cls.poi = POI( + points, + orientation=ct_ras.orientation, + zoom=ct_ras.zoom, + shape=ct_ras.shape, + origin=ct_ras.origin, + rotation=ct_ras.rotation, + level_one_info=V, + level_two_info=Location, + ).reorient_(cls.ct.orientation) + + def tearDown(self): + import matplotlib.pyplot as plt # noqa: PLC0415 + + plt.close("all") + + def test_lordosis_plot_returns_the_same_numbers_as_the_bare_computation(self): + from TPTBox.spine.snapshot2D.snapshot_modular import Snapshot_Frame # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + out = Path(tmp) / "lordosis.png" + values, frame = angles_mod.plot_compute_lordosis_and_kyphosis(out, self.poi.copy(), self.ct, self.vert) + self.assertIsInstance(frame, Snapshot_Frame) + self.assertEqual(values, compute_lordosis_and_kyphosis(self.poi.copy())) + self.assertTrue(out.exists()) + + def test_lordosis_plot_without_a_path_renders_nothing(self): + values, _ = angles_mod.plot_compute_lordosis_and_kyphosis(None, self.poi.copy(), self.ct, self.vert) + self.assertEqual(sorted(values), ["cervical_lordosis", "lumbar_lordosis", "thoracic_kyphosis"]) + + def test_cobb_plot_returns_the_segments_and_a_frame(self): + from TPTBox.spine.snapshot2D.snapshot_modular import Snapshot_Frame # noqa: PLC0415 + + with tempfile.TemporaryDirectory() as tmp: + out = Path(tmp) / "cobb.png" + segments, frame = angles_mod.plot_cobb_angle(out, self.poi.copy(), self.ct, self.vert) + self.assertIsInstance(frame, Snapshot_Frame) + self.assertIsInstance(segments, list) + self.assertTrue(out.exists()) + for angle, *_ in segments: + with self.subTest(angle=angle): + self.assertGreaterEqual(angle, 0.0) + self.assertLess(angle, 180.0) + + def test_plots_work_without_a_segmentation_overlay(self): + with tempfile.TemporaryDirectory() as tmp: + out = Path(tmp) / "no_seg.png" + angles_mod.plot_compute_lordosis_and_kyphosis(out, self.poi.copy(), self.ct, None) + self.assertTrue(out.exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_spinestats_bodycomp.py b/unit_tests/test_spinestats_bodycomp.py new file mode 100644 index 0000000..99efffe --- /dev/null +++ b/unit_tests/test_spinestats_bodycomp.py @@ -0,0 +1,346 @@ +"""Unit tests for `TPTBox.spine.spinestats.torso_vat_sat` and `body_quadrants`. + +Every input is a synthetic labelled volume with a known voxel count and a known +non-isotropic zoom, so the reported areas (mm^2) and volumes (mm^3) are exactly +predictable rather than recorded from a previous run. +""" + +from __future__ import annotations + +import math +import unittest + +import nibabel as nib +import numpy as np + +from TPTBox import NII, Location, Vertebra_Instance +from TPTBox.core.vert_constants import Full_Body_Instance_Vibe as FB +from TPTBox.spine.spinestats.body_quadrants import make_quadrants +from TPTBox.spine.spinestats.torso_vat_sat import ( + VBQ_score, + body_composition_score, + muscle_fat_infiltration, + peak_centered_mean, + torso_vat_sat_muscle_mass, +) + +ZOOM = (2.0, 2.0, 5.0) +VOXEL_AREA = ZOOM[0] * ZOOM[1] # axial in-plane area, S is the last axis +VOXEL_VOLUME = ZOOM[0] * ZOOM[1] * ZOOM[2] +SHAPE = (24, 24, 40) + + +def _nii(arr: np.ndarray, seg: bool) -> NII: + affine = np.diag([*ZOOM, 1.0]) + return NII(nib.Nifti1Image(arr, affine), seg=seg) + + +class TestPeakCenteredMean(unittest.TestCase): + """Robust mean around the histogram mode, used to estimate CSF signal.""" + + def test_rejects_the_low_signal_tail(self): + values = np.concatenate([np.full(1000, 100.0), np.full(50, 5.0)]) + peak = peak_centered_mean(values) + self.assertAlmostEqual(peak, 100.0, places=9) + self.assertLess(values.mean(), 96.0) # the plain mean is dragged down by the tail + + def test_bimodal_sample_follows_the_dominant_mode(self): + rng = np.random.default_rng(0) + values = np.concatenate([rng.normal(100.0, 2.0, 2000), rng.normal(20.0, 2.0, 300)]) + peak = peak_centered_mean(values) + self.assertAlmostEqual(peak, 100.0, delta=1.0) + self.assertGreater(peak, values.mean() + 5.0) + + def test_constant_input_returns_that_constant(self): + self.assertAlmostEqual(peak_centered_mean(np.full(100, 7.5)), 7.5, places=9) + + def test_empty_input_is_nan(self): + self.assertTrue(math.isnan(peak_centered_mean(np.array([])))) + + def test_non_finite_values_are_dropped(self): + values = np.array([1.0, 1.0, 1.0, np.nan, np.inf, -np.inf]) + self.assertAlmostEqual(peak_centered_mean(values), 1.0, places=9) + + def test_lower_fraction_keeps_more_of_the_tails(self): + rng = np.random.default_rng(1) + values = np.concatenate([rng.normal(100.0, 2.0, 2000), rng.normal(60.0, 2.0, 600)]) + tight = peak_centered_mean(values, peak_frac_height=1.0) + loose = peak_centered_mean(values, peak_frac_height=0.0) + # peak_frac_height=0 keeps every bin, i.e. the plain mean. + self.assertAlmostEqual(loose, float(values.mean()), places=6) + self.assertGreater(tight, loose) + + def test_input_may_be_multidimensional(self): + self.assertAlmostEqual(peak_centered_mean(np.full((4, 5), 3.0)), 3.0, places=9) + + +def _torso(include_clavicula: bool = True, include_pelvis: bool = True, vat_outside_roi: int = 0): + """Build (vibe_seg, roi) with known voxel counts per tissue. + + muscle: 2x3x4 = 24 voxels (autochthon left/right), VAT: 2x2x4 = 16, SAT: 3x2x4 = 24. + """ + seg = np.zeros(SHAPE, dtype=np.uint8) + seg[2:4, 2:5, 10:14] = FB.autochthon_left.value # 24 voxels + seg[6:8, 2:4, 10:14] = FB.inner_fat.value # 16 voxels + seg[10:13, 2:4, 10:14] = FB.subcutaneous_fat.value # 24 voxels + if include_clavicula: + seg[0, 0, 0] = FB.clavicula_left.value + if include_pelvis: + seg[1, 0, 0] = FB.pelvis_left.value + + roi = np.zeros(SHAPE, dtype=np.uint8) + roi[:, :, 8:20] = 3 # a torso ROI label inside roi_ids=(3..8) + if vat_outside_roi: + # VAT voxels placed outside the ROI slab must not be counted. + seg[6:8, 6:8, 30 : 30 + vat_outside_roi] = FB.inner_fat.value + return _nii(seg, seg=True), _nii(roi, seg=True) + + +class TestTorsoVatSatMuscleMass(unittest.TestCase): + def test_volumes_are_voxel_counts_times_the_voxel_volume(self): + seg, roi = _torso() + out, body_comp = torso_vat_sat_muscle_mass(seg, roi) + self.assertAlmostEqual(out["muscle_mass"], 24 * VOXEL_VOLUME, places=4) + self.assertAlmostEqual(out["VAT"], 16 * VOXEL_VOLUME, places=4) + self.assertAlmostEqual(out["SAT"], 24 * VOXEL_VOLUME, places=4) + self.assertNotIn("reason", out) + self.assertIsInstance(body_comp, NII) + + def test_tissue_outside_the_roi_is_excluded(self): + seg, roi = _torso(vat_outside_roi=4) + out, _ = torso_vat_sat_muscle_mass(seg, roi) + self.assertAlmostEqual(out["VAT"], 16 * VOXEL_VOLUME, places=4) + + def test_missing_clavicula_is_reported_as_an_incomplete_torso(self): + seg, roi = _torso(include_clavicula=False) + out, body_comp = torso_vat_sat_muscle_mass(seg, roi) + self.assertTrue(math.isnan(out["VAT"])) + self.assertIn("clavicula", out["reason"]) + self.assertIsNone(body_comp) + + def test_missing_pelvis_is_reported(self): + seg, roi = _torso(include_pelvis=False) + out, _ = torso_vat_sat_muscle_mass(seg, roi) + self.assertIn("pelvis", out["reason"]) + + def test_return_nii_adds_the_three_masks(self): + seg, roi = _torso() + out, _ = torso_vat_sat_muscle_mass(seg, roi, return_nii=True) + self.assertEqual(sorted(out["nii"]), ["SAT", "VAT", "muscle_mass"]) + self.assertEqual(int(out["nii"]["VAT"].sum()), 16) + + def test_unsupported_dataset_id_is_reported(self): + seg, roi = _torso() + out, _ = torso_vat_sat_muscle_mass(seg, roi, dataset_id=7) # type: ignore[arg-type] + self.assertIn("Unsupported dataset_id", out["reason"]) + + def test_a_roi_on_a_different_grid_is_resampled(self): + seg, roi = _torso() + coarse = roi.rescale((4.0, 4.0, 10.0), verbose=False) + self.assertNotEqual(coarse.shape, seg.shape) + out, _ = torso_vat_sat_muscle_mass(seg, coarse) + self.assertGreater(out["muscle_mass"], 0) + + +def _spine_volume(): + """Vert / spine / t2w triple with 4 lumbar bodies and a spinal canal column.""" + vert = np.zeros(SHAPE, dtype=np.uint8) + spine = np.zeros(SHAPE, dtype=np.uint8) + t2w = np.zeros(SHAPE, dtype=np.float32) + for i, vid in enumerate((Vertebra_Instance.L1, Vertebra_Instance.L2, Vertebra_Instance.L3, Vertebra_Instance.L4)): + z0 = 4 + i * 8 + vert[4:13, 4:13, z0 : z0 + 7] = vid.value + spine[4:13, 4:13, z0 : z0 + 7] = Location.Vertebra_Corpus.value + t2w[4:13, 4:13, z0 : z0 + 7] = 200.0 + # spinal canal: a column posterior to the bodies, spanning the whole volume + spine[15:22, 4:13, :] = Location.Spinal_Canal.value + t2w[15:22, 4:13, :] = 400.0 + return _nii(vert, seg=True), _nii(spine, seg=True), _nii(t2w, seg=False) + + +class TestVbqScore(unittest.TestCase): + REGIONS = ((Vertebra_Instance.L1, Vertebra_Instance.L4),) + + def test_vbq_is_the_vertebra_to_csf_signal_ratio(self): + vert, spine, t2w = _spine_volume() + out = VBQ_score(t2w, vert, spine, regions=list(self.REGIONS)) + self.assertAlmostEqual(out["mean_signal_vertebra_L1-L4"], 200.0, places=4) + self.assertAlmostEqual(out["mean_signal_liquor_L1-L4"], 400.0, places=4) + self.assertAlmostEqual(out["VBQ_L1-L4"], 0.5, places=4) + self.assertAlmostEqual(out["VBQ_L1-L4_old"], 0.5, places=4) + + def test_settings_are_echoed_back(self): + vert, spine, t2w = _spine_volume() + out = VBQ_score(t2w, vert, spine, regions=list(self.REGIONS), n_erode=1, spinal_bins=32, spinal_peak_frac_height=0.25) + self.assertEqual(out["n_erode"], 1) + self.assertEqual(out["spinal_bins"], 32) + self.assertEqual(out["spinal_peak_frac_height"], 0.25) + self.assertFalse(out["signal_by_full_cord"]) + + def test_full_cord_uses_the_whole_canal(self): + vert, spine, t2w = _spine_volume() + out = VBQ_score(t2w, vert, spine, regions=list(self.REGIONS), full_cord=True) + self.assertTrue(out["signal_by_full_cord"]) + self.assertAlmostEqual(out["VBQ_L1-L4"], 0.5, places=4) + + def test_a_region_without_vertebrae_yields_none(self): + vert, spine, t2w = _spine_volume() + out = VBQ_score(t2w, vert, spine, regions=[(Vertebra_Instance.C3, Vertebra_Instance.C6)]) + self.assertIsNone(out["VBQ_C3-C6"]) + self.assertIsNone(out["mean_signal_vertebra_C3-C6"]) + + def test_the_csf_peak_mean_ignores_a_dark_border(self): + # Darken part of the canal: the robust estimate must stay at the 400 peak while + # the plain mean drops. + vert, spine, t2w = _spine_volume() + arr = t2w.get_array() + arr[15:17, 4:13, :] = 50.0 + t2w.set_array_(arr) + out = VBQ_score(t2w, vert, spine, regions=list(self.REGIONS)) + self.assertAlmostEqual(out["mean_signal_liquor_L1-L4"], 400.0, places=4) + self.assertLess(out["mean_signal_liquor_L1-L4_old"], 400.0) + + +class TestBodyCompositionScore(unittest.TestCase): + def _inputs(self): + vert, spine, _ = _spine_volume() + seg = np.zeros(SHAPE, dtype=np.uint8) + # 6 voxels of muscle and 4 of VAT/SAT in every axial slice of the L1-L4 extent + seg[0:3, 0:2, :] = FB.autochthon_left.value + seg[3:5, 0:2, :] = FB.inner_fat.value + seg[5:7, 0:2, :] = FB.subcutaneous_fat.value + return _nii(seg, seg=True), vert, spine + + def test_areas_are_slice_voxel_counts_times_the_in_plane_voxel_area(self): + seg, vert, spine = self._inputs() + out = body_composition_score(seg, vert, spine, regions=[(Vertebra_Instance.L1, Vertebra_Instance.L4)]) + self.assertAlmostEqual(out["mean_muscle_area_L1-L4"], 6 * VOXEL_AREA, places=4) + self.assertAlmostEqual(out["max_muscle_area_L1-L4"], 6 * VOXEL_AREA, places=4) + self.assertAlmostEqual(out["mean_VAT_area_L1-L4"], 4 * VOXEL_AREA, places=4) + self.assertAlmostEqual(out["mean_SAT_area_L1-L4"], 4 * VOXEL_AREA, places=4) + + def test_muscle_fat_ratio_matches_the_reported_areas(self): + seg, vert, spine = self._inputs() + out = body_composition_score(seg, vert, spine, regions=[(Vertebra_Instance.L1, Vertebra_Instance.L4)]) + expected = out["mean_muscle_area_L1-L4"] / (out["mean_VAT_area_L1-L4"] + out["mean_SAT_area_L1-L4"]) + self.assertAlmostEqual(out["muscle_fat_ratio_L1-L4"], round(expected, 4), places=4) + + def test_muscle_index_normalises_by_height_squared(self): + seg, vert, spine = self._inputs() + out = body_composition_score(seg, vert, spine, regions=[(Vertebra_Instance.L1, Vertebra_Instance.L4)], height_m=2.0) + self.assertEqual(out["muscle_index_L1-L4"], round(out["mean_muscle_area_L1-L4"] / 4.0)) + + def test_n_slices_counts_the_axial_slices_of_the_region(self): + seg, vert, spine = self._inputs() + out = body_composition_score(seg, vert, spine, regions=[(Vertebra_Instance.L1, Vertebra_Instance.L4)]) + self.assertGreater(out["n_slices_L1-L4"], 0) + + def test_region_without_vertebrae_is_skipped(self): + seg, vert, spine = self._inputs() + out = body_composition_score(seg, vert, spine, regions=[(Vertebra_Instance.C3, Vertebra_Instance.C6)]) + self.assertEqual(out, {}) + + def test_unsupported_dataset_id_raises(self): + seg, vert, spine = self._inputs() + with self.assertRaises(NotImplementedError): + body_composition_score(seg, vert, spine, dataset_id=7) # type: ignore[arg-type] + + +class TestMuscleFatInfiltration(unittest.TestCase): + def _inputs(self, fat_fraction: float): + seg = np.zeros(SHAPE, dtype=np.uint8) + seg[2:8, 2:8, 10:20] = FB.autochthon_left.value + water = np.full(SHAPE, 100.0, dtype=np.float32) + fat = np.full(SHAPE, 100.0 * fat_fraction / (1.0 - fat_fraction), dtype=np.float32) + return _nii(water, seg=False), _nii(fat, seg=False), _nii(seg, seg=True) + + def test_fat_fraction_is_fat_over_fat_plus_water(self): + water, fat, seg = self._inputs(0.30) + out = muscle_fat_infiltration(water, fat, seg, per_muscle=False) + self.assertAlmostEqual(out["mean_fat_fraction_all_all_muscle"], 0.30, places=4) + self.assertAlmostEqual(out["median_fat_fraction_all_all_muscle"], 0.30, places=4) + + def test_uneroded_muscle_volume_is_the_voxel_count_times_the_voxel_volume(self): + water, fat, seg = self._inputs(0.30) + out = muscle_fat_infiltration(water, fat, seg, per_muscle=False) + # the mask is 6 x 6 x 10 voxels + self.assertAlmostEqual(out["muscle_volume_no_erosion_all_all_muscle"], 6 * 6 * 10 * VOXEL_VOLUME, places=4) + # the eroded volume used for the fat statistics is strictly smaller + self.assertLess(out["muscle_volume_all_all_muscle"], out["muscle_volume_no_erosion_all_all_muscle"]) + + def test_voxels_above_the_threshold_are_all_imat(self): + water, fat, seg = self._inputs(0.30) + out = muscle_fat_infiltration(water, fat, seg, threshold=0.20, per_muscle=False) + self.assertAlmostEqual(out["IMAT_fraction_all_all_muscle"], 1.0, places=4) + self.assertAlmostEqual(out["lean_muscle_volume_all_all_muscle"], 0.0, places=4) + self.assertAlmostEqual(out["IMAT_volume_all_all_muscle"], out["muscle_volume_all_all_muscle"], places=4) + + def test_voxels_below_the_threshold_are_all_lean(self): + water, fat, seg = self._inputs(0.05) + out = muscle_fat_infiltration(water, fat, seg, threshold=0.20, per_muscle=False) + self.assertAlmostEqual(out["IMAT_fraction_all_all_muscle"], 0.0, places=4) + self.assertAlmostEqual(out["IMAT_volume_all_all_muscle"], 0.0, places=4) + self.assertAlmostEqual(out["lean_muscle_volume_all_all_muscle"], out["muscle_volume_all_all_muscle"], places=4) + + def test_threshold_and_erode_settings_are_echoed_back(self): + water, fat, seg = self._inputs(0.30) + out = muscle_fat_infiltration(water, fat, seg, threshold=0.25, per_muscle=False) + self.assertEqual(out["threshold"], 0.25) + self.assertIn("all_muscle", out["erode"]) + + def test_per_muscle_reports_each_group(self): + water, fat, seg = self._inputs(0.30) + out = muscle_fat_infiltration(water, fat, seg, per_muscle=True) + self.assertIn("mean_fat_fraction_all_autochthon_left", out) + self.assertIn("mean_fat_fraction_all_all_muscle", out) + # only autochthon_left is populated, the other groups stay empty + self.assertAlmostEqual(out["muscle_volume_no_erosion_all_autochthon_left"], 6 * 6 * 10 * VOXEL_VOLUME, places=4) + + def test_regions_require_a_vert_segmentation(self): + water, fat, seg = self._inputs(0.30) + with self.assertRaises(ValueError): + muscle_fat_infiltration(water, fat, seg, regions=[(Vertebra_Instance.L1, Vertebra_Instance.L4)]) + + def test_unsupported_dataset_id_raises(self): + water, fat, seg = self._inputs(0.30) + with self.assertRaises(NotImplementedError): + muscle_fat_infiltration(water, fat, seg, dataset_id=7) # type: ignore[arg-type] + + +class TestMakeQuadrants(unittest.TestCase): + """Subdivides a vertebral body into 27 vertebra-local regions. + + Needs real vertebra subregions (the local axes come from POIs computed on them), + so this runs on the shipped single-vertebra sample CT. + """ + + @classmethod + def setUpClass(cls): + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + _, cls.subreg, cls.vert, cls.label = get_test_ct() + + def test_quadrants_are_a_subset_of_the_vertebral_body(self): + quad = make_quadrants(self.vert, self.subreg, vert_ids=[self.label]) + self.assertIsInstance(quad, NII) + out = quad.get_array() != 0 + self.assertGreater(out.sum(), 0) + body = (self.vert.get_array() == self.label) & np.isin(self.subreg.get_array(), (49, 50, 52)) + self.assertTrue(np.all(body[out]), "quadrant voxels must come from the body mask") + + def test_labels_are_in_one_to_twentyseven_and_disjoint(self): + quad = make_quadrants(self.vert, self.subreg, vert_ids=[self.label]).get_array() + labels = [int(label) for label in np.unique(quad) if label != 0] + self.assertGreater(len(labels), 1) + self.assertTrue(all(1 <= label <= 27 for label in labels), labels) + total = sum(int((quad == label).sum()) for label in labels) + self.assertEqual(total, int((quad != 0).sum()), "labels must partition the output") + + def test_erode_shrinks_the_result(self): + plain = int((make_quadrants(self.vert, self.subreg, vert_ids=[self.label]).get_array() != 0).sum()) + eroded = int((make_quadrants(self.vert, self.subreg, vert_ids=[self.label], erode=1).get_array() != 0).sum()) + self.assertLess(eroded, plain) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_spinestats_geometry.py b/unit_tests/test_spinestats_geometry.py new file mode 100644 index 0000000..5aa2661 --- /dev/null +++ b/unit_tests/test_spinestats_geometry.py @@ -0,0 +1,424 @@ +"""Unit tests for the geometry machinery behind `measure_ivd_and_vertebra_geometry` +and `spinestats.poi_fun.endplates`. + +The private helpers are pure geometry, so they are tested against shapes whose +answers are known in closed form: a box of known side lengths, a ray through a +triangle whose plane distance is known, and a spherical cap of known radius. +""" + +from __future__ import annotations + +import importlib +import importlib.util +import math +import tempfile +import unittest +from pathlib import Path + +import numpy as np +import trimesh + +from TPTBox.spine.spinestats.poi_fun import endplates as ep + +# `TPTBox.spine.spinestats.__init__` re-exports the *function* +# `measure_ivd_and_vertebra_geometry` under the same name as its module, so the module +# has to be fetched explicitly rather than with a plain `from ... import`. +geom = importlib.import_module("TPTBox.spine.spinestats.measure_ivd_and_vertebra_geometry") + +VOXEL = (2.0, 2.0, 5.0) +# Ray spacing for the sample-data measurements: the 0.5 mm default costs ~60 s per call. +STEP_MM = 2.0 + +# trimesh's ray engine does its broad phase through `rtree`, which trimesh only pulls in +# under its `easy` extra. Without it every mesh ray cast raises ModuleNotFoundError. +HAS_RTREE = importlib.util.find_spec("rtree") is not None +RTREE_REASON = "requires rtree (trimesh's ray engine broad phase)" + + +def _box_mask(shape=(20, 20, 20), lo=(5, 5, 5), hi=(15, 15, 15)) -> np.ndarray: + arr = np.zeros(shape, dtype=np.uint8) + arr[lo[0] : hi[0], lo[1] : hi[1], lo[2] : hi[2]] = 1 + return arr + + +class TestVoxelGeometryHelpers(unittest.TestCase): + def test_center_of_mass_of_a_box(self): + arr = _box_mask(lo=(4, 6, 8), hi=(10, 12, 14)) + com = geom._center_of_mass_voxels(arr, voxel_size=1) + np.testing.assert_allclose(com, [6.5, 8.5, 10.5]) + + def test_center_of_mass_scales_with_the_voxel_size(self): + arr = _box_mask(lo=(4, 6, 8), hi=(10, 12, 14)) + plain = geom._center_of_mass_voxels(arr, voxel_size=1) + scaled = geom._center_of_mass_voxels(arr, voxel_size=np.array(VOXEL)) + np.testing.assert_allclose(scaled, plain * np.array(VOXEL)) + + def test_pca_axes_of_an_elongated_box_include_its_long_axis(self): + arr = np.zeros((40, 20, 20), dtype=np.uint8) + arr[4:36, 8:12, 8:12] = 1 # long along axis 0 + axes = geom._pca_principal_axes(arr, voxel_size=np.ones(3)) + # Returns (up, v1, v2) where "up" is chosen to follow `up_axis`, not the dominant + # component, so the long axis shows up as one of the three (here as v1). + dots = [abs(float(np.dot(np.asarray(a), [1.0, 0.0, 0.0]))) for a in axes] + self.assertAlmostEqual(max(dots), 1.0, places=6) + + def test_pca_axes_are_an_orthonormal_triple(self): + axes = geom._pca_principal_axes(_box_mask(), voxel_size=np.ones(3)) + self.assertEqual(len(axes), 3) + mat = np.stack([np.asarray(a) for a in axes]) + for i in range(3): + with self.subTest(axis=i): + self.assertAlmostEqual(float(np.linalg.norm(mat[i])), 1.0, places=6) + np.testing.assert_allclose(np.abs(mat @ mat.T), np.eye(3), atol=1e-6) + + def test_swap(self): + self.assertEqual(geom._swap(1, 2), (2, 1)) + + def test_grid_origins_spans_the_requested_plane(self): + base = np.array([1.0, 2.0, 3.0]) + v1 = np.array([1.0, 0.0, 0.0]) + v2 = np.array([0.0, 1.0, 0.0]) + xs = np.array([-1.0, 0.0, 1.0]) + ys = np.array([0.0, 2.0]) + origins, flat_x, flat_y = geom._grid_origins(base, v1, v2, xs, ys) + self.assertEqual(origins.shape, (6, 3)) + self.assertEqual(len(flat_x), 6) + np.testing.assert_allclose(origins[0], base + xs[0] * v1 + ys[0] * v2) + # the offsets are exactly the requested grid + self.assertEqual(sorted(set(flat_x)), [-1.0, 0.0, 1.0]) + self.assertEqual(sorted(set(flat_y)), [0.0, 2.0]) + + def test_segment_length_is_the_entry_to_exit_distance(self): + pts = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 7.0]]) + self.assertAlmostEqual(geom._segment_length(pts), 7.0, places=9) + + def test_segment_length_uses_first_and_last_point_only(self): + pts = np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 3.0], [0.0, 0.0, 10.0]]) + self.assertAlmostEqual(geom._segment_length(pts), 10.0, places=9) + + def test_segment_length_of_a_missed_ray_is_zero(self): + self.assertEqual(geom._segment_length(np.zeros((0, 3))), 0) + + +class TestMeshConstruction(unittest.TestCase): + """Marching-cubes meshing and orthographic projection; no ray casting, so no rtree.""" + + def test_marching_cubes_mesh_of_a_box_is_watertight_and_spans_the_mask(self): + arr = _box_mask(lo=(5, 5, 5), hi=(15, 15, 15)) # 10 voxels per side + mesh = geom._segmentation_to_surface_mesh(arr, voxel_size=np.ones(3)) + self.assertTrue(mesh.is_watertight) + np.testing.assert_allclose(mesh.bounds[1] - mesh.bounds[0], [10.0, 10.0, 10.0]) + # The surface winding makes trimesh report a negative volume, which is why the + # caller stores abs(mesh.volume); corner cutting keeps |volume| just under 10^3. + self.assertLess(mesh.volume, 0.0) + self.assertAlmostEqual(abs(mesh.volume), 10.0**3, delta=0.03 * 10.0**3) + + def test_mesh_volume_scales_with_the_voxel_size(self): + arr = _box_mask(lo=(5, 5, 5), hi=(15, 15, 15)) + plain = abs(geom._segmentation_to_surface_mesh(arr, voxel_size=np.ones(3)).volume) + scaled = abs(geom._segmentation_to_surface_mesh(arr, voxel_size=np.array(VOXEL)).volume) + self.assertAlmostEqual(scaled, plain * VOXEL[0] * VOXEL[1] * VOXEL[2], delta=1e-6 * scaled) + + def test_projecting_a_box_along_its_axis_gives_the_face_area(self): + arr = _box_mask(lo=(5, 5, 5), hi=(15, 15, 15)) + mesh = geom._segmentation_to_surface_mesh(arr, voxel_size=np.ones(3)) + flat = geom._project_mesh_onto_plane(mesh, np.array([0.0, 0.0, 1.0])) + # Flattening collapses the two 10x10 faces perpendicular to "up" onto each other + # (the four side faces degenerate to lines and contribute no area). + self.assertAlmostEqual(geom._projected_area(flat), 2 * 10 * 10, delta=0.02 * 200) + + def test_projection_is_idempotent(self): + arr = _box_mask() + mesh = geom._segmentation_to_surface_mesh(arr, voxel_size=np.ones(3)) + up = np.array([0.0, 0.0, 1.0]) + once = geom._project_mesh_onto_plane(mesh, up) + twice = geom._project_mesh_onto_plane(once, up) + np.testing.assert_allclose(once.vertices, twice.vertices, atol=1e-9) + + +@unittest.skipUnless(HAS_RTREE, RTREE_REASON) +class TestMeshRayHelpers(unittest.TestCase): + def test_intersect_ray_through_a_box_measures_its_height(self): + arr = _box_mask(lo=(5, 5, 5), hi=(15, 15, 15)) + mesh = geom._segmentation_to_surface_mesh(arr, voxel_size=np.ones(3)) + up = np.array([0.0, 0.0, 1.0]) + v1 = np.array([1.0, 0.0, 0.0]) + v2 = np.array([0.0, 1.0, 0.0]) + hits = geom._intersect_ray(up, v1, v2, mesh) + # The marching-cubes surface spans the full 10 voxels (bounds 4.5 .. 14.5). + self.assertAlmostEqual(geom._segment_length(hits), 10.0, delta=0.2) + + def test_intersect_ray_from_dirs_matches_intersect_ray(self): + arr = _box_mask() + mesh = geom._segmentation_to_surface_mesh(arr, voxel_size=np.ones(3)) + dirs = (np.array([0.0, 0.0, 1.0]), np.array([1.0, 0.0, 0.0]), np.array([0.0, 1.0, 0.0])) + a = geom._intersect_ray_from_dirs(dirs, mesh) + b = geom._intersect_ray(dirs[0], dirs[1], dirs[2], mesh) + np.testing.assert_allclose(np.asarray(a), np.asarray(b)) + + def test_batched_rays_agree_with_single_rays(self): + arr = _box_mask(lo=(5, 5, 5), hi=(15, 15, 15)) + mesh = geom._segmentation_to_surface_mesh(arr, voxel_size=np.ones(3)) + direction = np.array([0.0, 0.0, 1.0]) + origins = np.array([[9.0, 9.0, -50.0], [10.0, 10.0, -50.0]]) + lengths = geom._batched_ray_segments(mesh, direction, origins)[0] + for length in np.asarray(lengths).ravel(): + with self.subTest(length=float(length)): + self.assertAlmostEqual(float(length), 10.0, delta=0.2) + + def test_batched_rays_that_miss_return_zero(self): + arr = _box_mask(lo=(5, 5, 5), hi=(15, 15, 15)) + mesh = geom._segmentation_to_surface_mesh(arr, voxel_size=np.ones(3)) + lengths = geom._batched_ray_segments(mesh, np.array([0.0, 0.0, 1.0]), np.array([[-100.0, -100.0, -50.0]]))[0] + self.assertAlmostEqual(float(np.asarray(lengths).ravel()[0]), 0.0, places=9) + + +class TestResultContainers(unittest.TestCase): + def test_nan_result_is_all_nan_plus_the_error(self): + out = geom._nan_result("boom") + self.assertEqual(out["error"], "boom") + for key, value in out.items(): + if key == "error": + continue + with self.subTest(key=key): + self.assertTrue(math.isnan(value)) + + def test_height_statistics_of_an_empty_sample_are_nan(self): + info = geom._StructureMeasurements(volume_voxel=1.0, volume_mesh=1.0, height_center=2.0, area=1.0) + info._update_height_statistics([]) + self.assertTrue(math.isnan(info.mean_height)) + self.assertTrue(math.isnan(info.max_height)) + + def test_nan_result_without_an_error_has_no_error_key(self): + self.assertNotIn("error", geom._nan_result()) + + def test_result_from_info_rounds_every_field(self): + info = geom._StructureMeasurements(volume_voxel=1.23456789, volume_mesh=2.0, height_center=3.0, area=4.0) + info._update_height_statistics([1.0, 2.0]) # mean/max are set lazily by this call + out = geom._result_from_info(info) + self.assertEqual(out["volume_voxel"], 1.2346) + self.assertEqual(out["volume_mesh"], 2.0) + # the public result carries exactly the documented field set + self.assertEqual(sorted(out), sorted(geom._nan_result())) + + def test_height_statistics_track_the_samples(self): + info = geom._StructureMeasurements(volume_voxel=1.0, volume_mesh=1.0, height_center=2.0, area=1.0) + info._update_height_statistics([3.0, 1.0, 4.0, 2.0]) # unsorted input is sorted internally + self.assertAlmostEqual(info.max_height, 4.0, places=6) + self.assertAlmostEqual(info.mean_height, 2.5, places=6) + # get_quantile indexes into the sorted sample list rather than interpolating, + # so the 50th percentile of four samples is the third one. + self.assertAlmostEqual(info.get_quantile(50, [1.0, 2.0, 3.0, 4.0]), 3.0, places=6) + self.assertIsNone(geom._StructureMeasurements(1.0, 1.0, 1.0, 1.0).get_quantile(99)) + + +class TestRayTriangleIntersect(unittest.TestCase): + """Moeller-Trumbore: exact hit/miss decisions for a known triangle.""" + + V0 = np.array([0.0, 0.0, 0.0]) + V1 = np.array([4.0, 0.0, 0.0]) + V2 = np.array([0.0, 4.0, 0.0]) + + def _cast(self, origin, direction): + return ep._ray_triangle_intersect(np.asarray(origin, dtype=float), np.asarray(direction, dtype=float), self.V0, self.V1, self.V2) + + def test_ray_through_the_centroid_hits_at_the_plane_distance(self): + centroid = (self.V0 + self.V1 + self.V2) / 3.0 + t = self._cast(centroid + np.array([0.0, 0.0, 7.0]), [0.0, 0.0, -1.0]) + self.assertIsNotNone(t) + self.assertAlmostEqual(float(t), 7.0, places=6) + + def test_ray_outside_the_triangle_misses(self): + self.assertIsNone(self._cast([5.0, 5.0, 3.0], [0.0, 0.0, -1.0])) + + def test_ray_parallel_to_the_triangle_misses(self): + self.assertIsNone(self._cast([1.0, 1.0, 3.0], [1.0, 0.0, 0.0])) + + def test_triangle_behind_the_origin_misses(self): + self.assertIsNone(self._cast([1.0, 1.0, 3.0], [0.0, 0.0, 1.0])) + + def test_distance_scales_with_the_origin_distance(self): + centroid = (self.V0 + self.V1 + self.V2) / 3.0 + for height in (1.0, 10.0, 100.0): + with self.subTest(height=height): + t = self._cast(centroid + np.array([0.0, 0.0, height]), [0.0, 0.0, -1.0]) + self.assertAlmostEqual(float(t), height, places=5) + + def test_non_unit_direction_is_handled(self): + centroid = (self.V0 + self.V1 + self.V2) / 3.0 + t = self._cast(centroid + np.array([0.0, 0.0, 6.0]), [0.0, 0.0, -2.0]) + self.assertIsNotNone(t) + self.assertGreater(float(t), 0.0) + + +@unittest.skipUnless(HAS_RTREE, RTREE_REASON) +class TestRayCastToMesh(unittest.TestCase): + """Pins the *current* behaviour of `_ray_cast_to_mesh`. + + Its comment says "closest hit", but it returns `origin + d.mean() * direction`, + i.e. the mean over every hit along the ray. For a convex mesh entered and exited + once that is the midpoint *inside* the mesh, not the surface point. These tests + therefore assert the midpoint; see the PR description - changing this would move + every computed endplate point, so it is reported rather than silently fixed. + """ + + def test_returns_the_midpoint_between_entry_and_exit(self): + arr = _box_mask(lo=(5, 5, 5), hi=(15, 15, 15)) # surface spans 4.5 .. 14.5 + mesh = geom._segmentation_to_surface_mesh(arr, voxel_size=np.ones(3)) + hit = ep._ray_cast_to_mesh(mesh, np.array([9.5, 9.5, -20.0]), np.array([0.0, 0.0, 1.0])) + self.assertIsNotNone(hit) + self.assertAlmostEqual(float(np.asarray(hit)[2]), 9.5, delta=0.2) + + def test_a_ray_pointing_away_misses(self): + arr = _box_mask(lo=(5, 5, 5), hi=(15, 15, 15)) + mesh = geom._segmentation_to_surface_mesh(arr, voxel_size=np.ones(3)) + self.assertIsNone(ep._ray_cast_to_mesh(mesh, np.array([9.5, 9.5, -20.0]), np.array([0.0, 0.0, -1.0]))) + + def test_hand_built_box_centred_on_the_origin(self): + mesh = trimesh.creation.box(extents=(4.0, 4.0, 4.0)) # spans -2 .. 2 + hit = ep._ray_cast_to_mesh(mesh, np.array([0.0, 0.0, -10.0]), np.array([0.0, 0.0, 1.0])) + self.assertIsNotNone(hit) + self.assertAlmostEqual(float(np.asarray(hit)[2]), 0.0, places=6) + + def test_the_returned_point_lies_on_the_ray(self): + mesh = trimesh.creation.box(extents=(4.0, 4.0, 4.0)) + origin = np.array([0.5, -0.25, -10.0]) + direction = np.array([0.0, 0.0, 1.0]) + hit = np.asarray(ep._ray_cast_to_mesh(mesh, origin, direction)) + np.testing.assert_allclose(hit[:2], origin[:2], atol=1e-9) + + +class TestSampleRayVoxels(unittest.TestCase): + def test_samples_stay_inside_the_volume_and_advance_monotonically(self): + shape = (20, 20, 20) + coords, t = ep._sample_ray_voxels(shape, np.array([10.0, 10.0, 2.0]), np.array([0.0, 0.0, 1.0]), max_steps=30) + self.assertGreater(len(coords), 0) + self.assertEqual(len(coords), len(t)) + for idx in coords: + with self.subTest(idx=tuple(int(c) for c in idx)): + self.assertTrue(all(0 <= int(c) < s for c, s in zip(idx, shape))) + # one step per voxel, strictly forward + np.testing.assert_allclose(np.diff(t), 1.0) + self.assertEqual([int(c[2]) for c in coords], sorted(int(c[2]) for c in coords)) + + def test_two_sided_adds_samples_behind_the_origin(self): + shape = (20, 20, 20) + one = ep._sample_ray_voxels(shape, np.array([10.0, 10.0, 10.0]), np.array([0.0, 0.0, 1.0]))[1] + both = ep._sample_ray_voxels(shape, np.array([10.0, 10.0, 10.0]), np.array([0.0, 0.0, 1.0]), two_sided=True)[1] + self.assertGreater(len(both), len(one)) + self.assertLess(both.min(), 0.0) + self.assertEqual(set(one.tolist()).issubset(set(both.tolist())), True) + + def test_a_ray_pointing_straight_out_of_the_volume_yields_nothing(self): + coords, t = ep._sample_ray_voxels((10, 10, 10), np.array([0.0, 0.0, 0.0]), np.array([0.0, 0.0, -1.0]), max_steps=5) + self.assertEqual(len(coords), 1 if len(t) == 1 else len(coords)) + self.assertTrue(all(0 <= int(c[2]) < 10 for c in coords)) + + +class TestEndplatePointsOnSampleData(unittest.TestCase): + """`calc_endplate_points_` on the shipped single-vertebra CT.""" + + @classmethod + def setUpClass(cls): + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + _, cls.subreg, cls.vert, cls.label = get_test_ct() + + def test_without_an_endplate_label_the_inputs_come_back_unchanged(self): + from TPTBox import Location # noqa: PLC0415 + + self.assertNotIn(Location.Endplate.value, self.subreg.unique(), "sample data has no Endplate label") + vert_out, spine_out = ep.endplate_to_super_infer_endplate(self.vert.copy(), self.subreg.copy()) + np.testing.assert_array_equal(vert_out.get_array(), self.vert.get_array()) + np.testing.assert_array_equal(spine_out.get_array(), self.subreg.get_array()) + + def test_endplate_voxels_are_relabelled_superior_or_inferior(self): + from TPTBox import Location # noqa: PLC0415 + + # Paint a slab of Endplate between two vertebra instances and let the function + # assign each voxel to the nearer vertebra's superior/inferior surface. + spine = self.subreg.copy() + vert = self.vert.copy() + arr = spine.get_array() + vert_arr = vert.get_array() + border = (vert_arr == 22) & (arr == Location.Vertebra_Corpus_border.value) + if int(border.sum()) == 0: + self.skipTest("sample vertebra has no corpus border voxels") + arr[border] = Location.Endplate.value + spine.set_array_(arr) + + _, spine_out = ep.endplate_to_super_infer_endplate(vert, spine) + labels = {int(u) for u in spine_out.unique()} + self.assertNotIn(Location.Endplate.value, labels, "the generic Endplate label must be resolved") + self.assertTrue( + {Location.Vertebral_Body_Endplate_Superior.value, Location.Vertebral_Body_Endplate_Inferior.value} & labels, + labels, + ) + + +@unittest.skipUnless(HAS_RTREE, RTREE_REASON) +class TestMeasureGeometryOnSampleData(unittest.TestCase): + """The public entry point on the shipped sample CT (one vertebra, label 22). + + Vertebra mode is `structure_label=0` (labels above 0 and below 100); the IVD mode + default of 100 selects disc labels instead. + """ + + @classmethod + def setUpClass(cls): + from TPTBox.tests.test_utils import get_test_ct # noqa: PLC0415 + + cls.ct, cls.subreg, cls.vert, cls.label = get_test_ct() + # The measurement is expensive (POI computation plus mesh ray casting), so it + # runs exactly once per class and every assertion below reads the cached result. + # Restricting it to the one fully contained vertebra (21 and 23 are truncated by + # the sample volume) and coarsening the ray step takes it from ~60 s to ~17 s. + # buffer_poi is left at its documented default of None on purpose: that path + # used to raise AssertionError (see test_works_without_a_poi_buffer). + cls.out = geom.measure_ivd_and_vertebra_geometry( + None, cls.vert.copy(), cls.subreg.copy(), structure_label=0, step_size_mm=STEP_MM, instance_labels=[cls.label] + ) + + def test_works_without_a_poi_buffer(self): + # Regression: buffer_poi=None is the documented default, but the function asked + # the buffering decorator to write the result back unconditionally, which + # asserts the path is not None. setUpClass already exercises that path. + self.assertIn(self.label, self.out) + + def test_reports_finite_positive_geometry_for_the_vertebra(self): + row = self.out[self.label] + self.assertNotIn("error", row) + for key in ("volume_voxel", "volume_mesh", "mean_height", "max_height", "width_lateral_x5", "width_sagittal_x6"): + with self.subTest(key=key): + self.assertTrue(np.isfinite(row[key]), key) + self.assertGreater(row[key], 0.0, key) + + def test_voxel_volume_matches_the_mask(self): + row = self.out[self.label] + expected = float((self.vert.get_array() == self.label).sum()) * self.vert.voxel_volume() + self.assertAlmostEqual(row["volume_voxel"], expected, delta=0.01 * expected) + + def test_mesh_volume_is_close_to_the_voxel_volume(self): + row = self.out[self.label] + self.assertAlmostEqual(row["volume_mesh"], row["volume_voxel"], delta=0.25 * row["volume_voxel"]) + + def test_heights_are_ordered(self): + row = self.out[self.label] + self.assertLessEqual(row["lower_10_percent_height"], row["mean_height"]) + self.assertLessEqual(row["mean_height"], row["max_height"]) + + def test_every_documented_field_is_present(self): + self.assertEqual(sorted(self.out[self.label]), sorted(geom._nan_result())) + + def test_signal_fields_stay_nan_without_a_t2w_image(self): + row = self.out[self.label] + for key in ("signal", "structure_signal", "spinal_canal_signal"): + with self.subTest(key=key): + self.assertTrue(math.isnan(row[key])) + + def test_instance_labels_restricts_the_output(self): + self.assertEqual(sorted(self.out), [self.label]) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_stitching_internals.py b/unit_tests/test_stitching_internals.py new file mode 100644 index 0000000..832f3e2 --- /dev/null +++ b/unit_tests/test_stitching_internals.py @@ -0,0 +1,266 @@ +"""Unit tests for the pure geometry helpers in `TPTBox.stitching`. + +`unit_tests/test_stiching.py` drives `stitching_raw` end to end; this file pins the +affine/bbox/dtype helpers it is built on, where the expected values are analytic. +""" + +from __future__ import annotations + +import tempfile +import unittest +from pathlib import Path + +import numpy as np + +from TPTBox import NII +from TPTBox.stitching.stitching_tools import _center_frontal + + +class TestCenterFrontal(unittest.TestCase): + """`_center_frontal` cuts a ~256-voxel dorsal window out of a stitched sagittal volume.""" + + def test_slice_is_always_inside_the_volume(self): + # Regression: `middle - 122` is negative for any volume below ~258 voxels, and numpy + # reads a negative start as an offset from the end -> the crop silently took the + # wrong side of the image (size=100 used to yield slice(-75, 181)). + for size in (1, 10, 50, 100, 200, 256, 257, 258, 300, 600): + with self.subTest(size=size): + s = _center_frontal(size) + self.assertGreaterEqual(s.start, 0) + self.assertGreater(s.stop, s.start) + self.assertLessEqual(s.stop, size) + + def test_crop_keeps_voxels_for_small_volumes(self): + for size in (50, 100, 200): + with self.subTest(size=size): + arr = np.arange(size) + kept = arr[_center_frontal(size)] + self.assertGreater(kept.size, 0) + # The window starts at the volume border for anything this small. + self.assertEqual(kept[0], 0) + + def test_large_volume_window_is_unchanged(self): + # The empirically tuned behaviour for full-size NAKO volumes must not move. + self.assertEqual(_center_frontal(600), slice(163, 419)) + self.assertEqual(_center_frontal(600).stop - _center_frontal(600).start, 256) + + +class TestAffineDecomposition(unittest.TestCase): + """`get_rotation_and_spacing_from_affine` / `get_ras_affine` are inverses.""" + + @staticmethod + def _affines(): + import math # noqa: PLC0415 + + diag = np.diag([1.5, 2.0, 3.0, 1.0]) + diag[:3, 3] = [10.0, -5.0, 2.0] + + flipped = np.diag([-1.0, -2.0, 3.0, 1.0]) + flipped[:3, 3] = [1.0, 2.0, 3.0] + + def rot(axis, deg): + t = math.radians(deg) + c, s = math.cos(t), math.sin(t) + if axis == "x": + return np.array([[1, 0, 0], [0, c, -s], [0, s, c]]) + if axis == "y": + return np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]]) + return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]]) + + turned = np.eye(4) + turned[:3, :3] = rot("x", 13) @ rot("y", -7) @ rot("z", 21) @ np.diag([0.8, 1.2, 4.0]) + turned[:3, 3] = [3.0, -4.0, 5.0] + return {"diagonal": diag, "double_flip": flipped, "rotated_anisotropic": turned} + + def test_decomposition_invariants(self): + from TPTBox.stitching.stitching import get_rotation_and_spacing_from_affine # noqa: PLC0415 + + for name, affine in self._affines().items(): + with self.subTest(affine=name): + rotation, spacing = get_rotation_and_spacing_from_affine(affine) + self.assertTrue((spacing >= 0).all(), "spacing is a norm, so never negative") + np.testing.assert_allclose(rotation.T @ rotation, np.eye(3), atol=1e-9) + np.testing.assert_allclose(rotation @ np.diag(spacing), affine[:3, :3], atol=1e-9) + + def test_round_trip_reproduces_the_affine(self): + from TPTBox.stitching.stitching import get_ras_affine, get_rotation_and_spacing_from_affine # noqa: PLC0415 + + for name, affine in self._affines().items(): + with self.subTest(affine=name): + rotation, spacing = get_rotation_and_spacing_from_affine(affine) + # `origin` is the translation expressed in the rotated frame, not the + # world-space origin - which is what makes this round-trip exact. + origin = rotation.T @ affine[:3, 3] + np.testing.assert_allclose(get_ras_affine(rotation, spacing, origin), affine, atol=1e-9) + + def test_spacing_matches_the_column_norms(self): + from TPTBox.stitching.stitching import get_rotation_and_spacing_from_affine # noqa: PLC0415 + + affine = self._affines()["rotated_anisotropic"] + _, spacing = get_rotation_and_spacing_from_affine(affine) + np.testing.assert_allclose(spacing, np.linalg.norm(affine[:3, :3], axis=0), atol=1e-9) + + +class TestCornerPoints(unittest.TestCase): + def test_corners_are_the_box_corners_offset_by_one_voxel(self): + from TPTBox.stitching.stitching import get_all_corner_points # noqa: PLC0415 + + corners = np.asarray(get_all_corner_points(np.eye(4), (4, 5, 6))) + self.assertEqual(corners.shape, (8, 3)) + # Pinned, not endorsed: the implementation adds 1 to every index before + # applying the affine, so the box spans 1..shape+1 instead of 0..shape. + # See the PR description - the stitched FOV depends on it. + np.testing.assert_allclose(corners.min(axis=0), [1.0, 1.0, 1.0]) + np.testing.assert_allclose(corners.max(axis=0), [5.0, 6.0, 7.0]) + + def test_corners_follow_the_affine(self): + from TPTBox.stitching.stitching import get_all_corner_points # noqa: PLC0415 + + affine = np.diag([2.0, 2.0, 2.0, 1.0]) + affine[:3, 3] = [100.0, 0.0, 0.0] + corners = np.asarray(get_all_corner_points(affine, (4, 4, 4))) + self.assertAlmostEqual(float(corners[:, 0].min()), 102.0) + self.assertEqual(len({tuple(c) for c in corners}), 8, "all eight corners are distinct") + + +class TestOccupancyBoxes(unittest.TestCase): + def test_bbox_of_a_known_block(self): + from TPTBox.stitching.stitching import _occupancy_bbox # noqa: PLC0415 + + occ = np.zeros((10, 10, 10), dtype=np.uint8) + occ[2:5, 3:7, 1:2] = 1 + lo, hi = _occupancy_bbox(occ) + np.testing.assert_array_equal(lo, [2, 3, 1]) + np.testing.assert_array_equal(hi, [4, 6, 1], "hi is inclusive") + + def test_empty_occupancy_is_none(self): + from TPTBox.stitching.stitching import _occupancy_bbox # noqa: PLC0415 + + self.assertIsNone(_occupancy_bbox(np.zeros((5, 5, 5), dtype=np.uint8))) + + def test_aabb_overlaps(self): + from TPTBox.stitching.stitching import _aabb_overlaps # noqa: PLC0415 + + box = (np.array([0, 0, 0]), np.array([4, 4, 4])) + cases = { + "identical": ((np.array([0, 0, 0]), np.array([4, 4, 4])), True), + "touching": ((np.array([4, 4, 4]), np.array([8, 8, 8])), True), + "disjoint": ((np.array([5, 5, 5]), np.array([8, 8, 8])), False), + "contained": ((np.array([1, 1, 1]), np.array([2, 2, 2])), True), + "empty": (None, False), + } + for name, (other, expected) in cases.items(): + with self.subTest(case=name): + self.assertEqual(_aabb_overlaps(box, other), expected) + self.assertEqual(_aabb_overlaps(other, box), expected, "the relation is symmetric") + + +class TestMaxAffineAndShape(unittest.TestCase): + def test_output_grid_covers_every_input_corner(self): + from TPTBox.stitching.stitching import get_all_corner_points, get_max_affine_and_shape # noqa: PLC0415 + + a = np.diag([1.0, 1.0, 1.0, 1.0]) + b = np.diag([1.0, 1.0, 1.0, 1.0]) + b[:3, 3] = [10.0, 0.0, 0.0] + points = np.concatenate([get_all_corner_points(a, (8, 8, 8)), get_all_corner_points(b, (8, 8, 8))]) + out = get_max_affine_and_shape(points, [a, b]) + self.assertIsInstance(out, NII) + self.assertGreaterEqual(out.shape[0], 18, "the union spans both volumes along x") + + def test_min_spacing_clamps_the_output_zoom(self): + from TPTBox.stitching.stitching import get_all_corner_points, get_max_affine_and_shape # noqa: PLC0415 + + a = np.diag([0.5, 0.5, 0.5, 1.0]) + points = get_all_corner_points(a, (8, 8, 8)) + fine = get_max_affine_and_shape(points, [a]) + coarse = get_max_affine_and_shape(points, [a], min_spacing=3.0) + self.assertAlmostEqual(min(fine.zoom), 0.5, delta=1e-6) + self.assertGreaterEqual(min(coarse.zoom), 3.0 - 1e-6) + self.assertLess(max(coarse.shape), max(fine.shape), "a coarser grid needs fewer voxels") + + def test_the_finest_input_spacing_wins(self): + from TPTBox.stitching.stitching import get_all_corner_points, get_max_affine_and_shape # noqa: PLC0415 + + fine = np.diag([0.5, 0.5, 0.5, 1.0]) + coarse = np.diag([3.0, 3.0, 3.0, 1.0]) + points = np.concatenate([get_all_corner_points(fine, (8, 8, 8)), get_all_corner_points(coarse, (4, 4, 4))]) + out = get_max_affine_and_shape(points, [fine, coarse]) + np.testing.assert_allclose(out.zoom, (0.5, 0.5, 0.5), atol=1e-6) + + +class TestAutoOutputDtype(unittest.TestCase): + @staticmethod + def _nii(dtype, slope=None): + import nibabel as nib # noqa: PLC0415 + + from TPTBox import NII # noqa: PLC0415 + + img = nib.Nifti1Image(np.zeros((4, 4, 4), dtype=dtype), np.eye(4)) + if slope is not None: + img.header.set_slope_inter(slope, 0.0) + return NII(img, seg=False) + + def test_all_uint8_stays_uint8(self): + from TPTBox.stitching.stitching import _auto_output_dtype # noqa: PLC0415 + + self.assertEqual(np.dtype(_auto_output_dtype([self._nii(np.uint8), self._nii(np.uint8)])), np.dtype(np.uint8)) + + def test_mixed_integer_widths_take_the_wider_one(self): + from TPTBox.stitching.stitching import _auto_output_dtype # noqa: PLC0415 + + out = np.dtype(_auto_output_dtype([self._nii(np.uint8), self._nii(np.uint16)])) + self.assertGreaterEqual(out.itemsize, np.dtype(np.uint16).itemsize) + + def test_a_float_input_forces_float32(self): + from TPTBox.stitching.stitching import _auto_output_dtype # noqa: PLC0415 + + self.assertEqual(np.dtype(_auto_output_dtype([self._nii(np.uint8), self._nii(np.float32)])), np.dtype(np.float32)) + + def test_a_non_trivial_slope_forces_float32(self): + import nibabel as nib # noqa: PLC0415 + + from TPTBox import NII # noqa: PLC0415 + from TPTBox.stitching.stitching import _auto_output_dtype # noqa: PLC0415 + + # The Philips case the docstring is about: int16 storage with a large slope. + # The slope is read off `nii.dataobj`, which only carries it for a file-backed + # image, so the fixture has to go through disk. + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "slope.nii.gz" + img = nib.Nifti1Image(np.zeros((4, 4, 4), dtype=np.int16), np.eye(4)) + img.header.set_slope_inter(641.0, 0.0) + nib.save(img, path) + nii = NII.load(path, seg=False) + self.assertEqual(np.dtype(_auto_output_dtype([nii])), np.dtype(np.float32)) + + +class TestStitchingTools(unittest.TestCase): + def test_crop_borders_cuts_the_requested_slices(self): + import nibabel as nib # noqa: PLC0415 + + from TPTBox import NII # noqa: PLC0415 + from TPTBox.stitching.stitching_tools import _crop_borders # noqa: PLC0415 + + nii = NII(nib.Nifti1Image(np.ones((20, 20, 20), dtype=np.uint16), np.eye(4)), seg=False) + cut = {"HWS": (slice(None), slice(2, 18), slice(None))} + out = _crop_borders(nii, "HWS", cut) + # The cut is applied in the reoriented frame and the result is turned back, so + # assert the voxel count rather than a particular axis. + self.assertEqual(int(np.prod(out.shape)), 20 * 16 * 20) + self.assertEqual(out.orientation, nii.orientation) + + def test_crop_borders_with_an_unknown_key_raises(self): + import nibabel as nib # noqa: PLC0415 + + from TPTBox import NII # noqa: PLC0415 + from TPTBox.stitching.stitching_tools import _crop_borders # noqa: PLC0415 + + # The guard only logs a message and then indexes `cut` anyway. Pinned as-is; + # it is a developer-facing helper and the KeyError names the missing key. + nii = NII(nib.Nifti1Image(np.ones((20, 20, 20), dtype=np.uint16), np.eye(4)), seg=False) + with self.assertRaises(KeyError): + _crop_borders(nii, "NOPE", {}) + + +if __name__ == "__main__": + unittest.main()