diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 91343a6..9ba09fd 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -14,7 +14,11 @@ jobs: fail-fast: false matrix: os: [ubuntu-latest, windows-latest] - python-version: ["3.9", "3.10", "3.11", "3.12"] + python-version: ["3.10", "3.11", "3.12"] + include: + # Python 3.9 is not available on ubuntu-latest (26.04); pin that cell to 22.04. + - os: ubuntu-22.04 + python-version: "3.9" steps: - uses: actions/checkout@v4 diff --git a/.github/workflows/tests_mr.yml b/.github/workflows/tests_mr.yml index 4945b6a..619dd3f 100644 --- a/.github/workflows/tests_mr.yml +++ b/.github/workflows/tests_mr.yml @@ -14,7 +14,11 @@ jobs: fail-fast: false matrix: os: [ubuntu-latest] - python-version: ["3.9","3.10","3.12"] + python-version: ["3.10", "3.12"] + include: + # Python 3.9 is not available on ubuntu-latest (26.04); pin that cell to 22.04. + - os: ubuntu-22.04 + python-version: "3.9" steps: - uses: actions/checkout@v4 diff --git a/TPTBox/core/dicom/dicom_extract.py b/TPTBox/core/dicom/dicom_extract.py index 9290b61..66350a6 100644 --- a/TPTBox/core/dicom/dicom_extract.py +++ b/TPTBox/core/dicom/dicom_extract.py @@ -158,6 +158,89 @@ def _make_fname(k: dict): return fname.file["json"], fname +# DICOM anatomical direction letter → unit vector in LPS (DICOM) space. +# L = +X (patient left), P = +Y (posterior), H = +Z (superior/head). +_DCM_LETTER_TO_LPS_VEC = { + "L": np.array([1.0, 0.0, 0.0]), + "R": np.array([-1.0, 0.0, 0.0]), + "P": np.array([0.0, 1.0, 0.0]), + "A": np.array([0.0, -1.0, 0.0]), + "H": np.array([0.0, 0.0, 1.0]), + "F": np.array([0.0, 0.0, -1.0]), +} +# Negate X and Y to turn DICOM LPS cosines into NIfTI RAS cosines. +_LPS_TO_RAS = np.array([-1.0, -1.0, 1.0]) + + +def _cosines_from_patient_orientation(patient_orientation) -> tuple[np.ndarray, np.ndarray] | tuple[None, None]: + """DICOM PatientOrientation (0020,0020) → RAS row/column cosines. + + The tag's two values are single-letter anatomical codes (L/R, A/P, H/F) + describing the direction the row and column axes point in patient space. + Returns (None, None) if the tag is absent, has blanks, or uses an + unknown letter — the caller then falls back to the next heuristic. + """ + if not patient_orientation or len(patient_orientation) < 2: + return None, None + row_letter = str(patient_orientation[0]).strip().upper()[:1] + col_letter = str(patient_orientation[1]).strip().upper()[:1] + row = _DCM_LETTER_TO_LPS_VEC.get(row_letter) + col = _DCM_LETTER_TO_LPS_VEC.get(col_letter) + if row is None or col is None: + return None, None + return row * _LPS_TO_RAS, col * _LPS_TO_RAS + + +_VIEW_TOKEN_TO_LETTERS = { + "ap": ("L", "F"), + "pa": ("L", "F"), + "ll": ("P", "F"), + "llat": ("P", "F"), + "lateral": ("P", "F"), # unspecified side defaults to left-lateral (clinical norm). + "lat": ("P", "F"), + "rl": ("A", "F"), + "rlat": ("A", "F"), +} + + +def _cosines_from_view_position(view_position) -> tuple[np.ndarray, np.ndarray] | tuple[None, None]: + """Fallback for projection X-rays whose PatientOrientation is blank. + + Uses the canonical radiography convention for the given ViewPosition + (0018,5101): rows run horizontally in the image, columns run top-to- + bottom (head→feet). Returns (None, None) for unknown or ambiguous + codes (e.g. a bare ``LAT`` that does not say which side). + """ + if not view_position: + return None, None + view = str(view_position).strip().lower() + letters = _VIEW_TOKEN_TO_LETTERS.get(view) + if letters is None: + return None, None + return _cosines_from_patient_orientation(list(letters)) + + +def _cosines_from_text_tokens(*texts) -> tuple[np.ndarray, np.ndarray] | tuple[None, None]: + """Last-ditch fallback: scan SeriesDescription / StudyDescription for AP/PA/LL/LAT tokens. + + Some scanners drop both PatientOrientation and ViewPosition for projection + X-rays, leaving only free-text descriptions. Scan them as whitespace- / + underscore- / dash-separated tokens so we only match standalone words, + not substrings like 'ap' inside 'scaphoid'. Returns (None, None) if no + recognisable projection token is found. + """ + import re + + for text in texts: + if not text: + continue + for token in re.split(r"[\s_\-]+", str(text).lower()): + letters = _VIEW_TOKEN_TO_LETTERS.get(token) + if letters is not None: + return _cosines_from_patient_orientation(list(letters)) + return None, None + + def dicom_to_nifti_multiframe_2d(ds, nii_path, pixel_array) -> None: """Convert a 2D multiframe DICOM dataset to a NIfTI file using its affine metadata. @@ -170,18 +253,48 @@ def dicom_to_nifti_multiframe_2d(ds, nii_path, pixel_array) -> None: dy, dx = map(float, ds.PixelSpacing) affine = np.eye(4) - if hasattr(ds, "ImageOrientationPatient"): - orientation = list(map(float, ds.ImageOrientationPatient)) - row_cosines = np.array(orientation[:3]) - col_cosines = np.array(orientation[3:]) - affine[:3, 0] = row_cosines * dx - affine[:3, 1] = col_cosines * dy + row_cos = col_cos = None + iop = getattr(ds, "ImageOrientationPatient", None) + if iop is not None: + orientation = list(map(float, iop)) + # ImageOrientationPatient is in DICOM LPS; flip to NIfTI RAS. + row_cos = np.array(orientation[:3]) * _LPS_TO_RAS + col_cos = np.array(orientation[3:]) * _LPS_TO_RAS + + # Projection X-rays (CR/DX) typically ship without ImageOrientationPatient + # but set PatientOrientation (0020,0020) or ViewPosition (0018,5101) + # instead. Without this fallback the affine stayed identity and viewers + # showed each series rotated or mirrored depending on how the modality + # happened to lay out the pixel buffer. + if row_cos is None: + row_cos, col_cos = _cosines_from_patient_orientation(getattr(ds, "PatientOrientation", None)) + if row_cos is None: + row_cos, col_cos = _cosines_from_view_position(getattr(ds, "ViewPosition", None)) + if row_cos is None: + row_cos, col_cos = _cosines_from_text_tokens( + getattr(ds, "SeriesDescription", None), + getattr(ds, "StudyDescription", None), + getattr(ds, "ProtocolName", None), + ) + + if row_cos is not None: + affine[:3, 0] = row_cos * dx + affine[:3, 1] = col_cos * dy + # The 2D image lies in the row x col plane; the "slice" axis must be + # the normal to that plane so the affine stays non-singular. + # Leaving np.eye's default (0,0,1) produces a degenerate affine + # whenever row/col already span X-Z or Y-Z (e.g. AP X-rays with + # row=L, col=F) — ITK refuses to open such a NIfTI. + slice_cos = np.cross(row_cos, col_cos) + dz = float(getattr(ds, "SliceThickness", 1.0) or 1.0) + affine[:3, 2] = slice_cos * dz else: affine[0, 0] = dx affine[1, 1] = dy - if hasattr(ds, "ImagePositionPatient"): - affine[:3, 3] = np.array(list(map(float, ds.ImagePositionPatient))) + ipp = getattr(ds, "ImagePositionPatient", None) + if ipp is not None: + affine[:3, 3] = np.array(list(map(float, ipp))) * _LPS_TO_RAS elif hasattr(ds, "ImagerPixelSpacing"): dy, dx = map(float, ds.ImagerPixelSpacing) @@ -413,6 +526,16 @@ def _extract_nii_from_dicom(dicom_out_path, nii_path): return False except Exception as e: logger.on_debug("Multi-Frame DICOM did not work:", e) + # Multi-echo Philips DIXON (e.g. mDIX quant): N slices x M echoes at + # the same ImagePositionPatient. dicom2nifti's own stack detector + # collapses to N/2 fake groups of 2 slices, then fails inside + # `four_d_to_nifti → create_affine` with `NOT_A_VOLUME`. Convert + # per-echo to 3-D and stack into a 4-D NIfTI ourselves; the + # downstream `_split_multi_echo_dixon` will then split it into + # per-echo `-eco` files. + if _dicoms_are_multi_echo(dicom_out_path) and _convert_multi_echo_dicoms_to_4d(dicom_out_path, nii_path): + logger.print("Save ", nii_path, Log_Type.SAVE) + return True ## The PDF dicom lands here convert_dicom.dicom_array_to_nifti(dicom_out_path, nii_path, True) else: @@ -646,6 +769,84 @@ def _from_dicom_to_nii( return nii_path if add_grid else None +def _dicoms_are_multi_echo(dcm_list) -> bool: + """True when ``dcm_list`` carries slices from more than one ``EchoNumbers`` value.""" + if not isinstance(dcm_list, list) or len(dcm_list) < 2: + return False + seen: set[int] = set() + for d in dcm_list: + try: + en = int(getattr(d, "EchoNumbers", 0) or 0) + except (TypeError, ValueError): + continue + if en > 0: + seen.add(en) + if len(seen) > 1: + return True + return False + + +def _convert_multi_echo_dicoms_to_4d(dcm_list: list, nii_path: str | Path) -> bool: + """Convert a multi-echo DICOM list to a 4-D NIfTI (echoes on axis 3). + + Per-echo buckets are converted separately with ``dicom2nifti`` (each is a + normal 3-D stack with no repeated ImagePositionPatient), then stacked along + a new fourth axis in ascending ``EchoNumbers``. Written to *nii_path*. The + caller's post-processing (``_split_multi_echo_dixon``) then splits the 4-D + file into per-echo ``-eco`` files. + + Returns True on success, False if any per-echo conversion fails or the + per-echo grids don't match. + """ + per_echo: dict[int, list] = {} + for d in dcm_list: + try: + en = int(getattr(d, "EchoNumbers", 0) or 0) + except (TypeError, ValueError): + continue + if en > 0: + per_echo.setdefault(en, []).append(d) + if len(per_echo) < 2: + return False + + nii_path = Path(nii_path) + tmp_dir = Path(tempfile.mkdtemp(prefix=f"{nii_path.stem}.echoes.", dir=str(nii_path.parent))) + try: + arrays: list[np.ndarray] = [] + ref_affine: np.ndarray | None = None + ref_header = None + ref_shape: tuple | None = None + for en in sorted(per_echo): + tmp_out = tmp_dir / f"echo{en}.nii.gz" + try: + convert_dicom.dicom_array_to_nifti(per_echo[en], str(tmp_out), True) + except Exception as e: # noqa: BLE001 + logger.on_warning(f"multi-echo split: per-echo conversion failed for EchoNumbers={en}: {type(e).__name__}: {e}") + return False + if not tmp_out.exists(): + logger.on_warning(f"multi-echo split: per-echo NIfTI missing for EchoNumbers={en}") + return False + n = nib.load(str(tmp_out)) + arr = np.asanyarray(n.dataobj, dtype=n.get_data_dtype()) + if ref_affine is None: + ref_affine = n.affine + ref_header = n.header + ref_shape = arr.shape + elif arr.shape != ref_shape or not np.allclose(n.affine, ref_affine, atol=1e-3): + logger.on_warning( + f"multi-echo split: grid mismatch between echoes (echo {en}: shape={arr.shape}); aborting per-echo stack." + ) + return False + arrays.append(arr) + stacked = np.stack(arrays, axis=-1) + out_img = nib.Nifti1Image(stacked, ref_affine, ref_header) + out_img.set_data_dtype(arrays[0].dtype) + out_img.to_filename(str(nii_path)) + return True + finally: + shutil.rmtree(tmp_dir, ignore_errors=True) + + def _split_multi_echo_dixon(nii_path: Path, json_path: Path, dcm_data_l) -> list[Path] | None: """Split a 4-D multi-echo DIXON (magnitude/phase) NIfTI into per-echo 3-D files. diff --git a/TPTBox/core/internal/nii_help.py b/TPTBox/core/internal/nii_help.py index ba79b0d..e814884 100644 --- a/TPTBox/core/internal/nii_help.py +++ b/TPTBox/core/internal/nii_help.py @@ -133,7 +133,6 @@ def _convert(obj): def _save_json(data, filepath: str | Path | bids_files.BIDS_FILE, indent=4, convert=_convert): - if isinstance(filepath, bids_files.BIDS_FILE): if "json" in filepath.file: filepath = filepath.file["json"] diff --git a/TPTBox/core/internal/train_nnUnet/prepere_dataset.py b/TPTBox/core/internal/train_nnUnet/prepere_dataset.py index bf9a4da..ced21f4 100644 --- a/TPTBox/core/internal/train_nnUnet/prepere_dataset.py +++ b/TPTBox/core/internal/train_nnUnet/prepere_dataset.py @@ -199,6 +199,13 @@ def _validate_config(cfg: DatasetConfig) -> None: "Use 'nnUNetTrainerNoMirroring', a SmaugLab DAExt trainer (mirror is stripped from its " "params JSON), or drop the mirror pairs." ) + # SmaugLab DAExt trainers assume the SmaugLab (non-CT) preprocessing path; combining them with + # is_ct=True mixes CT normalization with augmentations tuned for the SmaugLab regime. + if cfg.is_ct and "DAExt" in cfg.nn_trainer: + logger.on_warning( + f"is_ct=True combined with nn_trainer='{cfg.nn_trainer}': SmaugLab DAExt trainers are " + "not intended to be used with CT normalization. Set is_ct=False or pick a non-DAExt trainer." + ) if errors: logger.on_fail("Config validation failed:") for e in errors: diff --git a/TPTBox/core/internal/train_nnUnet/train.py b/TPTBox/core/internal/train_nnUnet/train.py index a9677dc..126ede4 100644 --- a/TPTBox/core/internal/train_nnUnet/train.py +++ b/TPTBox/core/internal/train_nnUnet/train.py @@ -248,7 +248,6 @@ def _run_training( enable_deep_supervision=True, save_every=1, # 50 ): - try: from nnunetv2.run.run_training import get_trainer_from_args, join, maybe_load_checkpoint except ImportError as e: @@ -409,7 +408,6 @@ def _patch_plans(self): # ---------------------------------------------------------- def _train_fold(self, fold: int | str): - # from nnunetv2.run.run_training import run_training print(f"Training fold {fold}") @@ -440,7 +438,6 @@ def _train_fold(self, fold: int | str): # ---------------------------------------------------------- def _train(self): - if self.cfg.single_gpu: os.environ["CUDA_VISIBLE_DEVICES"] = self.cfg.gpus[0] diff --git a/TPTBox/core/nii_wrapper.py b/TPTBox/core/nii_wrapper.py index 625f141..00c4ba0 100755 --- a/TPTBox/core/nii_wrapper.py +++ b/TPTBox/core/nii_wrapper.py @@ -2618,7 +2618,8 @@ def to_stl( to_world: bool = True, include_normals: bool = False, number_path=False, - _crop = True + _crop = True, + _raise_error=False ) -> Mesh: """Convert a binary segmentation label into an STL surface mesh using marching cubes. @@ -2672,7 +2673,13 @@ def to_stl( seg = self.extract_label(label) # Prepare binary mask if _crop and to_world: # this speed up the stl generation - crop = seg.compute_crop(0,1) + try: + crop = seg.compute_crop(0,1) + except ValueError: + log.on_fail(f"to_stl: empty segmentation. {label=}; Nothing will be saved") + if _raise_error: + raise + return None # type: ignore seg = seg.apply_crop(crop) seg_arr = np.pad(seg.clamp(0, 1).get_array(), 1) # Marching cubes (voxel coordinates) @@ -2687,7 +2694,7 @@ def to_stl( verts += np.array([b.start for b in bb]) # Convert to world coordinates using affine if to_world: - affine = self.affine # (4, 4) + affine = seg.affine # (4, 4) verts_h = np.c_[verts, np.ones(len(verts))] # homogeneous coords verts = (affine @ verts_h.T).T[:, :3] diff --git a/TPTBox/core/poi.py b/TPTBox/core/poi.py index 26013fb..470586a 100755 --- a/TPTBox/core/poi.py +++ b/TPTBox/core/poi.py @@ -45,6 +45,16 @@ ### CURRENT TYPE DEFINITIONS C = TypeVar("C", bound="POI") + +# Direction-vector auto-transform helpers live in poi_fun.vector_fields. +# Re-exported here for backwards compatibility with existing callers. +from TPTBox.core.poi_fun.vector_fields import ( # noqa: E402 + POI_INFO_VECTOR_FIELDS_KEY, + _remap_vector_field_keys_inplace, + _rotate_direction_vectors_inplace, + _transform_direction_vectors_inplace, +) + POI_Reference = Union[ bids_files.BIDS_FILE, Path, @@ -527,8 +537,11 @@ def reorient( self.shape = shape self.origin = origin self.rotation = rotation + _transform_direction_vectors_inplace(self.info, trans) return self - return self.copy(orientation=axcodes_to, centroids=points, zoom=zoom, shape=shape, origin=origin, rotation=rotation) + new_poi = self.copy(orientation=axcodes_to, centroids=points, zoom=zoom, shape=shape, origin=origin, rotation=rotation) + _transform_direction_vectors_inplace(new_poi.info, trans) + return new_poi def reorient_(self, axcodes_to: AX_CODES = ("P", "I", "R"), decimals=3, verbose: logging = False, _shape=None) -> Self: """In-place variant of :meth:`reorient`.""" @@ -537,6 +550,11 @@ def reorient_(self, axcodes_to: AX_CODES = ("P", "I", "R"), decimals=3, verbose: def rescale(self, voxel_spacing: ZOOMS = (1, 1, 1), decimals=ROUNDING_LVL, verbose: logging = True, inplace=False) -> Self: """Rescale the POI coordinates to a new voxel spacing in the current x-y-z-orientation. + Direction-vector fields registered under ``info[POI_INFO_VECTOR_FIELDS_KEY]`` + are left untouched: they are assumed to live in mm-space aligned with the + current voxel axes (see :func:`_transform_direction_vectors_inplace`), and + a spacing change relabels axes without rotating them. + Args: voxel_spacing (tuple[float, float, float], optional): New voxel spacing in millimeters. Defaults to (1, 1, 1). decimals (int, optional): Number of decimal places to round the rescaled coordinates to. Defaults to ROUNDING_LVL. @@ -617,7 +635,11 @@ def to_global(self, itk_coords=False) -> POI_Global: ) def resample_from_to(self, ref: Has_Grid) -> POI: - """Resample this POI to the grid of another image by converting to global and back. + """Resample this POI to the voxel grid of ``ref``. + + Registered direction-vector fields (see :data:`POI_INFO_VECTOR_FIELDS_KEY`) + are rotated by ``R_ref^T @ R_self`` so their components stay aligned with + the target grid's axes. Args: ref (Has_Grid): Target image grid (any object providing affine/orientation info). @@ -625,7 +647,9 @@ def resample_from_to(self, ref: Has_Grid) -> POI: Returns: POI: A new POI in the voxel space of ``ref``. """ - return self.to_global().to_other(ref) + out = self.to_global().to_other(ref) + _rotate_direction_vectors_inplace(out.info, self.rotation, getattr(ref, "rotation", None)) + return out def resample_from_to_(self, ref: Has_Grid) -> Self: """In-place variant of :meth:`resample_from_to`.""" @@ -1377,6 +1401,9 @@ def calc_centroids( ctd_list[first_stage, int(i)] = out_coord else: ctd_list[int(i), second_stage] = out_coord + if extend_to is not None: + args.setdefault("level_one_info", extend_to.level_one_info) + args.setdefault("level_two_info", extend_to.level_two_info) return POI(ctd_list, **msk_nii._extract_affine(), **args) diff --git a/TPTBox/core/poi_fun/poi_abstract.py b/TPTBox/core/poi_fun/poi_abstract.py index afcc90c..ebece0f 100755 --- a/TPTBox/core/poi_fun/poi_abstract.py +++ b/TPTBox/core/poi_fun/poi_abstract.py @@ -154,28 +154,30 @@ def normalize_label_name(d: dict | None) -> dict[int, dict]: Target: ``{region:int -> {subregion:int -> name:str, "name": group_name:str}}``. Accepts and migrates: the old flat ``{"(1, 2)": "C2"}`` form, a JSON-loaded nested - form with string keys ``{"1": {"2": "C2", "name": "Spine"}}``, and the already-nested - form (idempotent). + form with string keys ``{"1": {"2": "C2", "name": "Spine"}}``, the already-nested + form (idempotent), and any mix of the three in the same dict (produced e.g. when + ``join_left_`` merges a nested POI into a dict that still receives flat + ``"(region, subregion)"`` string entries). """ import ast if not d: return {} out: dict[int, dict] = {} - # old flat format: every key is a "(region, subregion)" tuple string - if all(isinstance(k, str) and k.strip().startswith("(") for k in d): - for k, name in d.items(): + for k, v in d.items(): + # legacy flat entry: "(region, subregion)" -> "name" + if isinstance(k, str) and k.strip().startswith("("): region, subregion = ast.literal_eval(k) - out.setdefault(int(region), {})[int(subregion)] = name - return out - # nested format (region keys may be JSON strings) - for region, sub in d.items(): - target = out.setdefault(int(region), {}) - if isinstance(sub, dict): - for s, name in sub.items(): + out.setdefault(int(region), {})[int(subregion)] = v + continue + # nested entry: region -> {subregion: name, "name": group_name} (or degenerate region -> name) + region_i = int(k) + target = out.setdefault(region_i, {}) + if isinstance(v, dict): + for s, name in v.items(): target[_GROUP_NAME_KEY if s == _GROUP_NAME_KEY else int(s)] = name - else: # degenerate {region: name} -> treat as the region group name - target[_GROUP_NAME_KEY] = sub + else: + target[_GROUP_NAME_KEY] = v return out @@ -637,11 +639,16 @@ def map_labels( poi_new[region:subreg] = value new_values = poi_new if new_values is None: + # No mapping happened (all input maps were empty) -> info stays as is. return self if inplace else self.copy() + + from TPTBox.core.poi_fun.vector_fields import _remap_vector_field_keys_inplace + + target = self if inplace else self.copy(centroids=new_values) if inplace: self.centroids = new_values - return self - return self.copy(centroids=new_values) + _remap_vector_field_keys_inplace(target.info, label_map_region_, label_map_subregion_) + return target def map_labels_( self, @@ -865,6 +872,19 @@ def set_label_name(self, region, subregion, name: str) -> None: d.setdefault(_id_of(region), {})[_id_of(subregion)] = name self.info[LABEL_NAME] = d + def iter_label_names(self) -> Iterator[tuple[int, int, str]]: + """Yield ``(region, subregion, name)`` triples for every custom per-point name stored in ``info["label_name"]``. + + The per-region group name (stored under the reserved ``"name"`` key) is + skipped; use :meth:`level_one_name` / :meth:`iter_level_one_names` for + that. Normalizes any legacy flat ``{"(r, s)": "name"}`` form on the fly. + """ + for region, subs in label_name_dict(self.info).items(): + for sub, name in subs.items(): + if sub == _GROUP_NAME_KEY: + continue + yield region, sub, name + def level_one_name(self, region) -> str | None: """Return the group (level-one) name of ``region``. diff --git a/TPTBox/core/poi_fun/poi_global.py b/TPTBox/core/poi_fun/poi_global.py index 97aa98a..935919f 100755 --- a/TPTBox/core/poi_fun/poi_global.py +++ b/TPTBox/core/poi_fun/poi_global.py @@ -158,7 +158,8 @@ def to_cord_system(self, itk_coords: bool, inplace: bool = False) -> Self: """Convert between ITK (LPS) and NIfTI (RAS) coordinate systems. Flips the first two coordinate axes when switching between the two - systems (LPS ↔ RAS only differs in the sign of x and y). + systems (LPS ↔ RAS only differs in the sign of x and y). Registered + direction-vector fields in ``info`` are flipped along the same axes. Args: itk_coords: ``True`` for ITK/LPS output, ``False`` for NIfTI/RAS. @@ -167,12 +168,19 @@ def to_cord_system(self, itk_coords: bool, inplace: bool = False) -> Self: Returns: ``POI_Global`` in the requested coordinate system. """ + import numpy as np + + from TPTBox.core.poi_fun.vector_fields import _transform_direction_vectors_inplace + out = self if inplace else self.copy() if self.itk_coords == itk_coords: return out out.itk_coords = itk_coords for k1, k2, v in self.items(): out[k1, k2] = (-v[0], -v[1], v[2]) + # x and y are negated, z is kept -- express as a signed-permutation trans. + trans = np.array([[0, -1], [1, -1], [2, 1]], dtype=int) + _transform_direction_vectors_inplace(out.info, trans) return out def to_other(self, msk: Has_Grid, verbose=False) -> poi.POI: diff --git a/TPTBox/core/poi_fun/vector_fields.py b/TPTBox/core/poi_fun/vector_fields.py new file mode 100644 index 0000000..e3271d4 --- /dev/null +++ b/TPTBox/core/poi_fun/vector_fields.py @@ -0,0 +1,182 @@ +"""Helpers for auto-transforming direction-vector fields stored in ``POI.info``. + +Producers write direction vectors into ``poi.info[""]`` as +``{key: (x, y, z)}`` dicts and register ```` in +``poi.info[POI_INFO_VECTOR_FIELDS_KEY]``. The helpers here are then called by +``POI.reorient`` / ``POI.resample_from_to`` / ``POI_Global.to_cord_system`` / +``POI.map_labels`` to keep the vectors and their keys aligned with the POI +points as the POI is transformed. + +Vectors are assumed to live in mm-space aligned with the POI's current voxel +axes; ``rescale`` is therefore a no-op for them. +""" + +from __future__ import annotations + +import numpy as np + +from TPTBox.core.vert_constants import Vertebra_Instance + +# poi.info key naming the info dicts whose values are direction vectors +# in the POI's current axis frame (i.e. mm-space aligned with poi.orientation). +POI_INFO_VECTOR_FIELDS_KEY = "_vector_fields" + +# poi.info key naming *additional* info dicts (typically scalar-valued, e.g. +# per-vertebra angles or curvature) whose top-level keys are region labels or +# ``Vertebra_Instance`` names. They participate in :func:`map_labels` key +# remapping but not in reorient / resample / to_cord_system vector transforms. +POI_INFO_LABEL_KEYED_FIELDS_KEY = "_label_keyed_fields" + + +def _transform_direction_vectors_inplace(info: dict, trans: np.ndarray) -> None: + """Reorient every registered direction-vector field in ``info`` according to ``trans``. + + ``trans`` is the output of ``nibabel.orientations.ornt_transform`` applied to + the POI's source and target axcodes. Each value stored under a registered + field must be a 3-tuple / 3-list (or ``None`` / non-3-tuple, which is skipped) + interpreted as a unit direction in the POI's current mm-space voxel-axis + frame. Values are updated in place. + """ + field_names = info.get(POI_INFO_VECTOR_FIELDS_KEY) + if not field_names: + return + perm = np.asarray(trans[:, 0], dtype=int) + flip = np.asarray(trans[:, 1], dtype=int) + for name in field_names: + vectors = info.get(name) + if not isinstance(vectors, dict): + continue + for k, v in list(vectors.items()): + if v is None: + continue + try: + v_arr = np.asarray(v, dtype=float) + except (TypeError, ValueError): + continue + if v_arr.shape != (3,): + continue + new_v = np.zeros(3, dtype=float) + new_v[perm] = v_arr * flip + vectors[k] = tuple(float(x) for x in new_v) + + +def _map_one_key(k, key_map: dict | None): + """Map a single dict key via ``key_map``. + + ``int`` keys are looked up directly. ``str`` keys resolve through + ``Vertebra_Instance`` (name -> value) and, on hit, the mapped value is + turned back into the corresponding ``Vertebra_Instance`` name (or kept + as ``str(id)`` when the target isn't a known instance). + """ + if not key_map: + return k + if isinstance(k, int) and k in key_map: + return key_map[k] + if isinstance(k, str): + try: + label = Vertebra_Instance[k].value + except KeyError: + return k + if label in key_map: + try: + return Vertebra_Instance(key_map[label]).name + except ValueError: + return str(key_map[label]) + return k + + +def _remap_vector_field_keys_inplace(info: dict, region_map: dict | None, subregion_map: dict | None = None) -> None: + """Remap keys of every registered label-keyed field in ``info``. + + Considers fields registered under both :data:`POI_INFO_VECTOR_FIELDS_KEY` + (direction vectors) and :data:`POI_INFO_LABEL_KEYED_FIELDS_KEY` (scalar + per-label fields), plus ``label_name`` (always handled; migrated to the + nested form on the fly). + + Behaviour is dispatched by *value type*: + + - **Flat fields** (values are tuples / scalars): only outer keys are + remapped via ``region_map``. + - **Nested fields** (values are dicts, e.g. ``label_name`` / + ``{region: {subregion: name, "name": group}}``): outer keys are + remapped via ``region_map``, inner keys via ``subregion_map``, and the + special ``"name"`` group entry is preserved. When two source regions + collide onto one target, their inner dicts merge (last-write-wins on + overlapping keys). + + Keys may be integer labels or ``Vertebra_Instance``-name strings; both + are matched against the int-keyed maps. No-op if there is nothing to do. + """ + from TPTBox.core.poi_fun.poi_abstract import _GROUP_NAME_KEY, LABEL_NAME, label_name_dict + + if not region_map and not subregion_map: + return + field_names: list[str] = [] + for key in (POI_INFO_VECTOR_FIELDS_KEY, POI_INFO_LABEL_KEYED_FIELDS_KEY): + names = info.get(key) + if names: + field_names.extend(names) + # label_name is always handled -- ensure the nested-form migration runs, then include it. + if LABEL_NAME in info and LABEL_NAME not in field_names: + label_name_dict(info) + field_names.append(LABEL_NAME) + if not field_names: + return + for name in field_names: + vectors = info.get(name) + if not isinstance(vectors, dict): + continue + remapped: dict = {} + for k, v in vectors.items(): + new_k = _map_one_key(k, region_map) + if new_k is None: + continue # drop entries whose region is mapped to None + new_v = v + # Nested field: recurse into inner dict. + if isinstance(v, dict): + new_inner: dict = {} + for ik, iv in v.items(): + if ik == _GROUP_NAME_KEY: + new_inner[_GROUP_NAME_KEY] = iv + continue + new_ik = _map_one_key(ik, subregion_map) + if new_ik is None: + continue # drop entries whose subregion is mapped to None + new_inner[new_ik] = iv + new_v = new_inner + # Merge on outer-key collision when both values are dicts (label_name-style). + if new_k in remapped and isinstance(remapped[new_k], dict) and isinstance(new_v, dict): + remapped[new_k].update(new_v) + else: + remapped[new_k] = new_v + vectors.clear() + vectors.update(remapped) + + +def _rotate_direction_vectors_inplace(info: dict, src_rot, tgt_rot) -> None: + """Rotate registered direction-vector fields from ``src_rot`` to ``tgt_rot``. + + Composes ``tgt_rot.T @ src_rot`` and applies it in place. No-op when either + rotation is None or no vector fields are registered. + """ + if src_rot is None or tgt_rot is None: + return + field_names = info.get(POI_INFO_VECTOR_FIELDS_KEY) + if not field_names: + return + R = np.asarray(tgt_rot, dtype=float).T @ np.asarray(src_rot, dtype=float) + for name in field_names: + vectors = info.get(name) + if not isinstance(vectors, dict): + continue + for k, v in list(vectors.items()): + if v is None: + continue + try: + v_arr = np.asarray(v, dtype=float) + except (TypeError, ValueError): + continue + if v_arr.shape != (3,): + continue + new_v = R @ v_arr + vectors[k] = tuple(float(x) for x in new_v) diff --git a/TPTBox/core/poi_fun/vertebra_direction.py b/TPTBox/core/poi_fun/vertebra_direction.py index 7a51081..b307ed5 100644 --- a/TPTBox/core/poi_fun/vertebra_direction.py +++ b/TPTBox/core/poi_fun/vertebra_direction.py @@ -147,15 +147,26 @@ def calc_orientation_of_vertebra_PIR( if last_vert == 20: last_vert = None break - max_vert_key = max(vert_keys) + # Only pre-sacral vertebrae influence the "next free label" search — S1 (and higher) + # may be present in poi_iso via `calc_endplate_points_` (Sacrum_Endplate landmark), + # but their label must NOT be counted here, otherwise spline anchors would land at + # labels 27/28 and get promoted to phantom vertebrae by the downstream + # `poi_iso.extract_subregion(source_subreg_point_id)` loop (which would then write + # spurious direction POIs at those labels and fool `_get_last_thoracic` into + # returning T13 / COCC). + presacral_keys = [k for k in vert_keys if k < Vertebra_Instance.S1.value] + max_vert_key = max(presacral_keys) if presacral_keys else max(vert_keys) + anchor_labels: set[int] = set() if (Vertebra_Instance.S1.value, Location.Vertebral_Body_Endplate_Superior.value) in poi_iso: poi_iso[max_vert_key + 1, spline_subreg_point_id] = poi_iso[ (Vertebra_Instance.S1.value, Location.Vertebral_Body_Endplate_Superior.value) ] max_vert_key += 1 + anchor_labels.add(max_vert_key) if last_vert is not None and (last_vert, Location.Vertebral_Body_Endplate_Inferior.value) in poi_iso: poi_iso[max_vert_key + 1, spline_subreg_point_id] = poi_iso[(last_vert, Location.Vertebral_Body_Endplate_Inferior.value)] max_vert_key += 1 + anchor_labels.add(max_vert_key) ##### # spline: body_spline, body_spline_der = poi_iso.fit_spline(location=spline_subreg_point_id, vertebra=True) @@ -184,6 +195,10 @@ def calc_orientation_of_vertebra_PIR( fill_back = out.copy() if do_fill_back else None # Draw a plain with the up_vector an cut it with intersection_target for reg_label, _, cords in poi_iso.extract_subregion(source_subreg_point_id).items(): + # Spline-anchor labels (added above to influence the spline fit) must not be + # processed here — they'd otherwise get written back to `ret` as pseudo-vertebrae. + if reg_label in anchor_labels: + continue # calculate_normal_vector if reg_label in down_vector: normal_vector_down = down_vector[reg_label] @@ -222,13 +237,25 @@ def calc_orientation_of_vertebra_PIR( arr = subreg_sar.set_array(fill_back).reorient(poi.orientation).rescale_(poi.zoom).get_array() fill_back_nii.set_array_(arr) - ret = calc_centroids(subreg_iso.set_array(out), second_stage=subreg_id, extend_to=poi_iso.copy(), inplace=True) + # Strip spline anchors from the working POI before extending: they were only + # needed to influence `fit_spline` above, and must not leak into the returned POI + # as pseudo-vertebrae (they'd otherwise show up as phantom regions labelled 25/26 + # / 27/28, fooling `_get_last_thoracic` etc. into treating them as real). + poi_iso_clean = poi_iso.copy() + for _a in anchor_labels: + if (_a, spline_subreg_point_id.value) in poi_iso_clean: + del poi_iso_clean.centroids[_a, spline_subreg_point_id.value] + ret = calc_centroids(subreg_iso.set_array(out), second_stage=subreg_id, extend_to=poi_iso_clean, inplace=True) poi._vert_orientation_pir = {} if save_normals_in_info: poi.info["vert_orientation_PIR"] = poi._vert_orientation_pir # calc posterior vector and the crossproduct for vert_id, normal_down in down_vector.items(): + # Missing plane/source POIs at this level (e.g. sacrum without a spinal-canal + # intersection) are legitimate — skip silently instead of logging a failure. + if (vert_id, subreg_id.value) not in ret or (vert_id, source_subreg_point_id.value) not in ret: + continue try: # get two points and compute the direction: a = np.array(ret[vert_id : subreg_id.value]) - 1 diff --git a/TPTBox/core/poi_fun/vertebra_pois_non_centroids.py b/TPTBox/core/poi_fun/vertebra_pois_non_centroids.py index 8f0583a..beb4d04 100755 --- a/TPTBox/core/poi_fun/vertebra_pois_non_centroids.py +++ b/TPTBox/core/poi_fun/vertebra_pois_non_centroids.py @@ -365,7 +365,12 @@ def compute_non_centroid_pois( # noqa: C901 log.on_text("Compute Vertebra Endplate DIRECTIONS", verbose=verbose) sub_regions = poi.keys_subregion() - if any(a.value not in sub_regions for a in endplate[:2]): # skip if all exists + # Also (re)run when S1 is present in the segmentation but its Sacrum_Endplate-derived + # superior-endplate landmark is not in the POI yet -- otherwise the sacrum block inside + # calc_endplate_points_ never gets a chance to populate poi.info["angle_superior_endplate"]["S1"]. + s1 = Vertebra_Instance.S1.value + sacrum_endplate_missing = s1 in _vert_ids and (s1, Location.Vertebral_Body_Endplate_Superior.value) not in poi + if any(a.value not in sub_regions for a in endplate[:2]) or sacrum_endplate_missing: # skip if all exists poi, *_ = calc_endplate_points_(poi, vert, subreg, _vert_ids=_vert_ids, log=log) ### STEP 1 Vert Direction### if Location.Vertebra_Direction_Inferior in locations: diff --git a/TPTBox/registration/_deformable/multilabel_segmentation.py b/TPTBox/registration/_deformable/multilabel_segmentation.py index d4b7949..2468d48 100644 --- a/TPTBox/registration/_deformable/multilabel_segmentation.py +++ b/TPTBox/registration/_deformable/multilabel_segmentation.py @@ -4,7 +4,9 @@ from pathlib import Path from typing import TypeVar -from TPTBox import NII, POI +import numpy as np + +from TPTBox import NII, POI, Print_Logger from TPTBox.core.internal.deep_learning_utils import DEVICES from TPTBox.core.poi import calc_centroids from TPTBox.core.poi_fun.poi_global import POI_Global @@ -14,6 +16,7 @@ from TPTBox.registration._ridged_points.point_registration import Point_Registration _NIIOrPOI = TypeVar("_NIIOrPOI", NII, POI) +logger = Print_Logger() def _r_axis_slicer(axis: int) -> tuple[slice, ...]: @@ -435,6 +438,40 @@ def transform_poi_inverse(self, poi_target: POI_Global | POI): return poi +def _filter_pois(data_target: NII, nii_atlas_fov: NII, poi_target_cms, poi_atlas_cms): + u = data_target.unique() + cms_ids = [int(l) for l in u if not data_target.extract_label(l).is_segmentation_in_border()] + dropped = [int(l) for l in u if int(l) not in cms_ids] + if dropped: + logger.on_warning(f"labels {dropped} touch target FOV border - excluded from rigid CMS") + poi_target_cms_fit: POI | None = None + if not cms_ids: + # All target labels truncated. Fall back to the innermost point of each + # mask (argmax of the EDT) so the rigid stage at least matches the + # robust interior of the visible region, not the CMS of a truncated blob. + from scipy.ndimage import distance_transform_edt + + cms_ids = [int(l) for l in u] + logger.on_warning(" all labels truncated - using innermost-point fallback for rigid anchor") + poi_target_cms_fit = data_target.make_empty_POI() + poi_atlas_cms_fit = nii_atlas_fov.make_empty_POI() + for lbl in cms_ids: + tmask = data_target.extract_label(lbl).get_seg_array().astype(bool) + amask = nii_atlas_fov.extract_label(lbl).get_seg_array().astype(bool) + if not tmask.any() or not amask.any(): + continue + ti = np.unravel_index(int(np.argmax(distance_transform_edt(tmask))), tmask.shape) + ai = np.unravel_index(int(np.argmax(distance_transform_edt(amask))), amask.shape) + poi_target_cms_fit[lbl, 40] = tuple(float(x) for x in ti) + poi_atlas_cms_fit[lbl, 40] = tuple(float(x) for x in ai) + return (poi_target_cms_fit, poi_atlas_cms_fit) + else: + return ( + poi_target_cms.extract_region(cms_ids), + poi_atlas_cms.extract_region(cms_ids), + ) + + class Template_Registration2: """Atlas-to-target multi-stage registration with a pluggable pre-registration. @@ -542,16 +579,17 @@ def __init__( # noqa: C901 if pre_registration is not None: self.reg_point = pre_registration else: - if poi_target_cms is None: - x_seg = target_seg.extract_label(cms_ids, keep_label=True) if cms_ids else target_seg - poi_target = calc_centroids(x_seg, second_stage=40, bar=True) - else: - poi_target = poi_target_cms.resample_from_to(target_seg) if poi_cms is None: x_seg = atlas_seg.extract_label(cms_ids, keep_label=True) if cms_ids else atlas_seg poi_cms_local = calc_centroids(x_seg, second_stage=40, bar=True) else: poi_cms_local = poi_cms + if poi_target_cms is None: + x_seg = target_seg.extract_label(cms_ids, keep_label=True) if cms_ids else target_seg + poi_target = calc_centroids(x_seg, second_stage=40, bar=True) + poi_target, poi_cms_local = _filter_pois(target_seg, atlas_seg, poi_target, poi_cms_local) + else: + poi_target = poi_target_cms.resample_from_to(target_seg) if not poi_cms_local.assert_affine(atlas_seg, raise_error=False): poi_cms_local = poi_cms_local.resample_from_to(atlas_seg) self.reg_point = Deepali_Point_Registration(poi_target, poi_cms_local, verbose=False, ddevice=ddevice, gpu=gpu) diff --git a/TPTBox/segmentation/rib/_rib_assign.py b/TPTBox/segmentation/rib/_rib_assign.py index c88406a..d667fec 100644 --- a/TPTBox/segmentation/rib/_rib_assign.py +++ b/TPTBox/segmentation/rib/_rib_assign.py @@ -13,12 +13,9 @@ from __future__ import annotations -import os from collections.abc import Iterable from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass -from enum import Enum -from functools import partial import numpy as np from tqdm import tqdm @@ -59,7 +56,12 @@ def _split_ccs_by_side(rib_cc: NII, cms_cc: POI, cms_cc2: POI, cms_vert: POI) -> continue center = cms_cc[cc, 50] center2 = cms_cc2[cc, 50] if (cc, 50) in cms_cc2 else center # noqa: SIM401 - cand = RibCandidate(cc_label=cc, volume=vols.get(cc, 0), z=center[inf_axis], x=center2[right_axis]) + cand = RibCandidate( + cc_label=cc, + volume=vols.get(cc, 0), + z=center[inf_axis], + x=center2[right_axis], + ) distances = cms_vert.calculate_distances_cord(center2) min_key = min(distances, key=distances.get) # type: ignore # x < vertebra center => patient right in RAS-like orientation @@ -81,7 +83,11 @@ def _touching_surface(mask_a: np.ndarray, mask_b: np.ndarray, nii: NII, axis_wei All other axes have weight 1.0. """ if axis_weights is None: - axis_weights = {nii.get_axis("S"): 0.01, nii.get_axis("A"): 0.1, nii.get_axis("R"): 1} + axis_weights = { + nii.get_axis("S"): 0.01, + nii.get_axis("A"): 0.1, + nii.get_axis("R"): 1, + } surface = 0.0 for axis in range(mask_a.ndim): @@ -102,7 +108,12 @@ def _touching_surface(mask_a: np.ndarray, mask_b: np.ndarray, nii: NII, axis_wei def _try_erosion_split( - cc_label: int, binary_cc: NII, erosion_pixels: int, min_volume: int, _pass=0, verbose=True + cc_label: int, + binary_cc: NII, + erosion_pixels: int, + min_volume: int, + _pass=0, + verbose=True, ) -> tuple[int, list[np.ndarray]] | None: """Try to split one CC by erosion. @@ -157,7 +168,10 @@ def _try_erosion_split( target = remaining[0] # print("merge", target) else: - target = max(remaining, key=lambda candidate: _touching_surface(sub_mask, infected_arr == candidate, infected)) + target = max( + remaining, + key=lambda candidate: _touching_surface(sub_mask, infected_arr == candidate, infected), + ) # print("merge of many", target) infected_arr[infected_arr == sub] = target sub_labels.remove(sub) @@ -240,7 +254,14 @@ def split_touching_rib_ccs( with ThreadPoolExecutor(max_workers=num_workers) as executor: futures = [ - executor.submit(_try_erosion_split, cc_label, binary_cc, erosion_pixels, min_volume, _pass) + executor.submit( + _try_erosion_split, + cc_label, + binary_cc, + erosion_pixels, + min_volume, + _pass, + ) for cc_label, binary_cc in binary_ccs ] @@ -347,7 +368,12 @@ def assign_ribs_to_vert_segmentation( cms_vert = calc_centroids(vert_seg) if split_touching: rib_cc = split_touching_rib_ccs( - rib_cc, erosion_pixels=erosion_pixels, min_volume=min_volume, vert_ids=vert_ids, short_cut=short_cut, verbose=verbose + rib_cc, + erosion_pixels=erosion_pixels, + min_volume=min_volume, + vert_ids=vert_ids, + short_cut=short_cut, + verbose=verbose, ) cms_cc = calc_centroids(rib_cc * vert_pred.calc_convex_hull(None).dilate_msk_euclid(5)) cms_cc2 = calc_centroids(rib_cc) # .dilate_msk_euclid(5) @@ -406,8 +432,7 @@ def assign_ribs_to_vert_segmentation( # Merge rib assignments back into the original vert / sem segmentations # without disturbing existing (non-rib) labels. Skip unmatched CCs (sentinel error_value=255). # 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. + # NII.__and__ rejects non-integer arrays, so the NII-level `&` raised TypeError. rib_inst_arr = rib_inst.get_seg_array() matched = (rib_inst_arr != 0) & (rib_inst_arr != error_value) vert_seg[matched] = rib_inst[matched] diff --git a/TPTBox/segmentation/spineps.py b/TPTBox/segmentation/spineps.py index dca938f..0df6fb7 100644 --- a/TPTBox/segmentation/spineps.py +++ b/TPTBox/segmentation/spineps.py @@ -74,7 +74,7 @@ def run_spineps( save_debug_data: bool = False, verbose: bool = False, save_raw: bool = False, - ignore_compatibility_issues: bool = False, + ignore_compatibility_issues: bool = True, use_cpu: bool = False, **args, ) -> dict[ diff --git a/TPTBox/spine/spinestats/_load_nako.py b/TPTBox/spine/spinestats/_load_nako.py index b0462ee..e2a98c6 100644 --- a/TPTBox/spine/spinestats/_load_nako.py +++ b/TPTBox/spine/spinestats/_load_nako.py @@ -7,6 +7,7 @@ from TPTBox import Print_Logger from TPTBox.core.bids_files import BIDS_FILE, BIDS_Family, Buffered_BIDS_Global_info from TPTBox.core.nii_wrapper import to_nii +from TPTBox.spine.spinestats._load_nako_wh import get_current_best_VERIDAH as _get_current_best_VERIDAH log = Print_Logger() @@ -320,6 +321,8 @@ def loop_over_repaired_nako( subj_dict["vert"] = vert subj_dict["spine"] = spine subj_dict["poi"] = poi + veridah = _get_current_best_VERIDAH(sub) + subj_dict["veridah"] = str(veridah) if veridah is not None else None yield subj_dict diff --git a/TPTBox/spine/spinestats/_load_nako_wh.py b/TPTBox/spine/spinestats/_load_nako_wh.py index 3925965..939a610 100644 --- a/TPTBox/spine/spinestats/_load_nako_wh.py +++ b/TPTBox/spine/spinestats/_load_nako_wh.py @@ -154,6 +154,14 @@ def resolve_pick( # (``main-conflict:pd``, ``mevibe-conflict:pd``, ``vibe-conflict:pd``, …). idx = max(range(len(labels)), key=labels.__getitem__) return candidates[idx] # transient: do not save + if key == "main:T2haste" or key.endswith(":T2haste"): + # auto-accept: when duplicates differ only in the presence of a ``sequ`` entity + # (one raw file without sequ, one or more with an explicit sequ number for the same + # acquisition), prefer the sequ-numbered candidate. Skips the prompt only when this + # split is unambiguous (exactly one sequ'd candidate); otherwise falls through. + with_sequ = [c for c in candidates if getattr(c, "get", lambda *_: None)("sequ", None) is not None] + if len(with_sequ) == 1 and len(with_sequ) < len(candidates): + return with_sequ[0] # transient: do not save choice, reason = _prompt_choice(sub, key, question, labels, allow_discard=allow_discard) if choice == "__skip__": return candidates[0] if candidates else None # transient: do not save @@ -455,44 +463,69 @@ def get_corrected_mevibe(fam: BIDS_Family, compute_PDFF=True): # TODO return di # PDFF is recomputed out = {key: _check(fam[f"mevibe_part-{key}"]) for key in ["eco0-opp1", "eco1-pip1", "eco2-opp2", "eco3-in1", "eco4-pop1", "eco5-arb1"]} - pdff = _check(fam["mevibe_part-fat-fraction"]) if "mevibe_part-water_desc-reconstructed" in fam: - # if "mevibe_part-fat-fraction_desc-reconstructed" not in fam: fat = _check(fam["mevibe_part-fat_desc-reconstructed"]) water = _check(fam["mevibe_part-water_desc-reconstructed"]) - - else: + elif "mevibe_part-water" in fam: fat = _check(fam["mevibe_part-fat"]) water = _check(fam["mevibe_part-water"]) - out["mevibe_part-fat"] = fat - out["mevibe_part-fat"] = water - pdff = water.get_changed_bids( - "nii.gz", bids_format=water.bids_format, parent=water.parent, info={"part": "fat-fraction", "desc": "reconstructed"} + else: + fat = None + water = None + + # Anchor for BIDS derivations: prefer a raw/reconstructed water image, fall back to + # whatever fraction file the family exposes. + if water is not None: + anchor = water + elif "mevibe_part-water-fraction" in fam: + anchor = _check(fam["mevibe_part-water-fraction"]) + elif "mevibe_part-fat-fraction" in fam: + anchor = _check(fam["mevibe_part-fat-fraction"]) + else: + return out + + pdff = anchor.get_changed_bids( + "nii.gz", bids_format=anchor.bids_format, parent=anchor.parent, info={"part": "fat-fraction", "desc": "reconstructed"} ) - pdwf = water.get_changed_bids( - "nii.gz", bids_format=water.bids_format, parent=water.parent, info={"part": "water-fraction", "desc": "reconstructed"} + pdwf = anchor.get_changed_bids( + "nii.gz", bids_format=anchor.bids_format, parent=anchor.parent, info={"part": "water-fraction", "desc": "reconstructed"} ) if compute_PDFF and (not pdff.exists() or not pdwf.exists()): - water_nii = to_nii(water) - fat_nii = to_nii(fat) - water_nii.set_dtype_() - fat_nii.set_dtype_() - if not pdff.exists(): - nii = fat_nii / (water_nii + fat_nii) - nii[water_nii + fat_nii == 0] = 0 - nii *= 1000 - nii.set_dtype_("smallest_int") - nii.save(pdff) - if not pdwf.exists(): - nii = water_nii / (water_nii + fat_nii) - nii[water_nii + fat_nii == 0] = 0 - nii *= 1000 - nii.set_dtype_("smallest_int") - nii.save(pdwf) + if fat is not None and water is not None: + water_nii = to_nii(water) + fat_nii = to_nii(fat) + water_nii.set_dtype_() + fat_nii.set_dtype_() + if not pdff.exists(): + nii = fat_nii / (water_nii + fat_nii) + nii[water_nii + fat_nii == 0] = 0 + nii *= 1000 + nii.set_dtype_("smallest_int") + nii.save(pdff) + if not pdwf.exists(): + nii = water_nii / (water_nii + fat_nii) + nii[water_nii + fat_nii == 0] = 0 + nii *= 1000 + nii.set_dtype_("smallest_int") + nii.save(pdwf) + else: + # No raw fat/water — derive the missing fraction from the one the scanner shipped. + if not pdff.exists() and "mevibe_part-water-fraction" in fam: + wf_nii = to_nii(_check(fam["mevibe_part-water-fraction"])) + nii = 1000 - wf_nii + nii.set_dtype_("smallest_int") + nii.save(pdff) + if not pdwf.exists() and "mevibe_part-fat-fraction" in fam: + ff_nii = to_nii(_check(fam["mevibe_part-fat-fraction"])) + nii = 1000 - ff_nii + nii.set_dtype_("smallest_int") + nii.save(pdwf) + + # Downstream (see _apply_corrections_to_subj_dict) expects "mevibe_part-fat" to hold PDFF. if pdff.exists(): out["mevibe_part-fat"] = pdff - if pdff.exists(): + if pdwf.exists(): out["mevibe_part-fat"] = pdwf # else: # pdff = _check(fam["mevibe_part-fat-fraction_desc-reconstructed"]) @@ -500,6 +533,25 @@ def get_corrected_mevibe(fam: BIDS_Family, compute_PDFF=True): # TODO return di return out +def get_current_best_VERIDAH(sub) -> Path | None: + """Return the newest VERIDAH-label JSON for ``sub`` (V2 preferred), or ``None`` if missing. + + The file lives alongside the T2w segmentation under + ``derivatives_spine_inference_162_sacrumfix///T2w/`` and holds the + ``orig_label -> fpath`` remapping used by :func:`compute_veridah_variants`. + """ + sub = str(sub).split("_")[0].replace("sub-", "") + for folder in ("derivatives_spine_inference_162_sacrumfix",): + for suffix in ("VERIDAH-label-V2", "VERIDAH-label"): + p = Path( + f"/DATA/NAS/datasets_processed/NAKO/dataset-nako/{folder}/{sub[:3]}/{sub}/T2w/" + f"sub-{sub}_sequ-stitched_acq-sag_mod-T2w_seg-vert_desc-{suffix}_stat.json" + ) + if p.exists(): + return p + return None + + def get_current_best_T2w_seg(sub, black_list_t2w=None): if black_list_t2w is None: black_list_t2w = [ @@ -569,10 +621,11 @@ def loop_over_repaired_nako( test=False, verbose=False, sort=True, - test_key="/100/10", # path matching. if you want on specific us a 6 digits + test_key="/102/", # path matching. if you want on specific us a 6 digits decision_cache: DecisionCache | Path | str | None = None, corrected_index: dict | Path | str | None = None, skip_subject=None, + vibe_mismatch_snap_dir: Path | str | None = None, ): """Iterate over the repaired NAKO dataset yielding per-subject file dicts. @@ -595,6 +648,12 @@ def loop_over_repaired_nako( test_key: Path substring passed to the BIDS scanner's ``filter_file`` when ``test=True``; only paths containing this substring are indexed. Defaults to a hard-coded example subject. baseline_metadata: Path to the NAKO baseline CSV used to look up height metadata. + vibe_mismatch_snap_dir: When set, subjects whose VIBE parts don't share a shape are skipped + (not yielded); a review snapshot is written to this directory as + ``sub-_vibe-shape-mismatch.jpg`` when that file doesn't already exist. + Two sub-folders ``accept/`` and ``reject/`` are also created on demand: if the reviewer + moves the jpg into ``accept/`` the next run auto-answers the grid prompt with "y" + (resample); if moved into ``reject/`` it auto-answers "n" (drop mismatched keys). Yields: Dict mapping short keys to ``BIDS_FILE`` entries for one subject. @@ -703,6 +762,18 @@ def loop_over_repaired_nako( q.filter_format("mevibe") # q.filter("sequ", "me1") mevibe_fams = list(q.loop_dict(key_addendum=["mod", "part", "desc"])) + # Drop derivative-only families and incomplete acquisitions: get_corrected_mevibe + # unconditionally indexes all six echoes, so a family missing any of them would crash. + _echo_keys = [f"mevibe_part-{k}" for k in ("eco0-opp1", "eco1-pip1", "eco2-opp2", "eco3-in1", "eco4-pop1", "eco5-arb1")] + _kept = [] + for f in mevibe_fams: + missing = [k for k in _echo_keys if k not in f] + if missing: + if "mevibe_part-eco0-opp1" in f: + log.on_warning(f"sub-{sub}: incomplete mevibe family {f.family_id!r} (missing {missing}); skipping") + continue + _kept.append(f) + mevibe_fams = _kept if len(mevibe_fams) > 1: labels = [str(f.get("mevibe_part-eco0-opp1", f)) for f in mevibe_fams] cached_pick = _cached_pick(cache.get(sub, "mevibe_fam")) @@ -792,6 +863,7 @@ def loop_over_repaired_nako( else: cache.set(sub, "vibe_fam", {"pick": labels[choice], "reason": reason}) vibe_fams = [vibe_fams[choice]] + _vibe_skip_subject = False for fam in vibe_fams: vibe_by_key: dict = {} for _k in ( @@ -804,6 +876,34 @@ def loop_over_repaired_nako( ): if fam.get(_k): vibe_by_key[_k] = fam[_k][0] + if vibe_mismatch_snap_dir is not None: + sigs = _vibe_grid_sigs(vibe_by_key) + if len(set(sigs.values())) > 1: + snap_dir = Path(vibe_mismatch_snap_dir) + accept_dir = snap_dir / "accept" + reject_dir = snap_dir / "reject" + for d in (snap_dir, accept_dir, reject_dir): + d.mkdir(parents=True, exist_ok=True) + snap_name = f"sub-{sub}_vibe-shape-mismatch.jpg" + snap_path = snap_dir / snap_name + if (accept_dir / snap_name).exists(): + # User verified this mismatch is fine — auto-answer "y" (resample). + cache.set(sub, "grid_mismatch:vibe", {"decision": "resample", "reason": "accepted-via-snap"}) + elif (reject_dir / snap_name).exists(): + # User rejected this subject's VIBE — auto-answer "n" (drop mismatched keys). + cache.set(sub, "grid_mismatch:vibe", {"decision": "remove", "reason": "rejected-via-snap"}) + else: + if not snap_path.exists(): + try: + _save_vibe_shape_mismatch_snapshot(vibe_by_key, snap_path) + except Exception as e: # noqa: BLE001 + log.on_warning(f"sub-{sub}: failed to save vibe mismatch snapshot: {e}") + log.on_warning( + f"sub-{sub}: VIBE grid mismatch {sigs} — awaiting review " + f"(move {snap_name} into accept/ or reject/); skipping subject" + ) + _vibe_skip_subject = True + break check_same_grid(cache, sub, "vibe", vibe_by_key, inphase_key="vibe_part-inphase") # Propagate the check's outcome back to ``fam`` so the downstream unpack # loop below picks up resampled files (or skips removed keys). @@ -835,10 +935,14 @@ def loop_over_repaired_nako( for k, k2 in mapp.items(): if k in fam: subj_dict[k2] = fam[k][0] + if _vibe_skip_subject: + continue vert, spine, poi = get_current_best_T2w_seg(sub) subj_dict["vert"] = vert subj_dict["spine"] = spine subj_dict["poi"] = poi + veridah = get_current_best_VERIDAH(sub) + subj_dict["veridah"] = str(veridah) if veridah is not None else None if corrected_index: _apply_corrections_to_subj_dict(str(sub), subj_dict, corrected_index) verify_missing_images(cache, sub, subj_dict) @@ -865,6 +969,69 @@ def _is_grid_only_json(path: Path | str) -> bool: _CANONICAL_DONE_ROOT = Path("/DATA/NAS/datasets_processed/NAKO/dataset-nako-canonical/.hardlink_done") +_VIBE_MISMATCH_SNAP_DIR = Path("/DATA/NAS/datasets_processed/NAKO/dataset-nako-canonical/snaps/vibe-missmatch") + + +def _vibe_grid_sigs(vibe_by_key: dict) -> dict[str, str]: + """Return ``{key: grid-signature-string}`` for each present VIBE part. + + Mirrors ``check_same_grid``'s comparison: reads ``bf.get_grid_info()`` and + stringifies it, skipping msk entries and files that have no NIfTI. Two + signatures being unequal is exactly what would cause ``check_same_grid`` + to prompt the user. + """ + sigs: dict[str, str] = {} + for k, bf in vibe_by_key.items(): + if bf is None: + continue + if getattr(bf, "format", None) == "msk" or k.startswith("msk"): + continue + get_nii_file = getattr(bf, "get_nii_file", None) + if callable(get_nii_file) and get_nii_file() is None: + continue + try: + g = bf.get_grid_info() + except Exception: # noqa: BLE001 + continue + if g is None: + continue + sigs[k] = str(g) + return sigs + + +def _save_vibe_shape_mismatch_snapshot(vibe_by_key: dict, out_path: Path) -> None: + """Save a review snapshot of VIBE parts whose grids don't agree. + + Renders one sagittal+coronal frame per present VIBE part (in/out/water/fat), + titled with the part's shape, so a reviewer can eyeball what's off. + """ + from TPTBox.spine.snapshot2D import Snapshot_Frame, create_snapshot + + frames = [] + for k in ("vibe_part-inphase", "vibe_part-outphase", "vibe_part-water", "vibe_part-fat"): + bf = vibe_by_key.get(k) + if bf is None: + continue + try: + shape = tuple(bf.get_grid_info().shape) # type: ignore[union-attr] + except Exception: # noqa: BLE001 + shape = None + frames.append( + Snapshot_Frame( + image=bf, + mode="MRI", + sagittal=True, + coronal=True, + axial=False, + crop_msk=False, + title=f"{k} shape={shape}", + ) + ) + if not frames: + return + out_path.parent.mkdir(parents=True, exist_ok=True) + create_snapshot(snp_path=[out_path], frames=frames) + def _hard_link_done_marker(sub: str) -> Path: return _CANONICAL_DONE_ROOT / f"{sub}.done" @@ -875,10 +1042,42 @@ def is_hard_linked(sub: str) -> bool: return _hard_link_done_marker(str(sub)).exists() +def _stale_hardlink_targets(srcs: dict[str, Path], new_path: Path) -> list[Path]: + """Return targets that exist on disk but no longer share an inode with their source. + + A hardlink pair shares an inode (``os.path.samefile``); a stale target either points at + a previous inode (source was rewritten, not overwritten in place) or at a different + source entirely (e.g. flipped between ``dataset-nako`` and ``rawdata-corrected``). + Callers unlink the returned paths before re-running ``bf.symlink_files``. + """ + base = str(new_path)[: -len(".nii.gz")] + stale: list[Path] = [] + for ext, src in srcs.items(): + t = Path(base + "." + ext) + if not t.exists(): + continue + try: + if not t.samefile(Path(src)): + stale.append(t) + except OSError: + stale.append(t) + return stale + + def hard_link( d: dict, dataset="/DATA/NAS/datasets_processed/NAKO/dataset-nako/", + refresh_stale: bool = False, ): + """Materialize a subject's file dict into ``dataset-nako-canonical`` via hardlinks. + + When ``refresh_stale=True`` every already-linked target is checked against its + current source with :func:`os.path.samefile`. Mismatches (source was rewritten, + or the ``subj_dict`` slot flipped between ``dataset-nako`` and + ``rawdata-corrected``) are unlinked and re-linked so ``-canonical`` matches + the current state. Default is ``False`` — cheaper, but stale references + survive between runs. + """ subj = d.pop("id") d.pop("dataset") log.on_log(subj) @@ -895,7 +1094,14 @@ def hard_link( info={"ses": "baseline"}, dataset_path="/DATA/NAS/datasets_processed/NAKO/dataset-nako-canonical", ) - if not new_path.exists(): + stale: list[Path] = [] + if refresh_stale and new_path.exists(): + srcs_t2w = {ext: Path(src) for ext, src in bf.file.items() if Path(src).exists()} + stale = _stale_hardlink_targets(srcs_t2w, new_path) + for t in stale: + t.unlink() + print("Refreshing stale hardlink:", t) + if stale or not new_path.exists(): new_path.parent.mkdir(parents=True, exist_ok=True) bf.symlink_files(new_path, hard_link=True) # exist_ok=True, print("Hard linked:", new_path) @@ -923,6 +1129,7 @@ def hard_link( "vert", # t2w (stiched) "spine", # t2w (stiched) "poi", # t2w (stiched) + "veridah", # VERIDAH enumeration-anomaly relabeling (V2 preferred) ] imgs = [ "pd", @@ -948,6 +1155,7 @@ def hard_link( info = {"run": None} if isinstance(bf, str): bf = BIDS_FILE(bf, dataset) + bf.info.pop("run", None) assert len([k for k, v in bf.loop_keys() if k not in allowed_keys]) == 0, ( [k for k, v in bf.loop_keys() if k not in allowed_keys], bf, @@ -969,7 +1177,11 @@ def hard_link( # Compute the real target paths (per extension) and skip if all already exist. base = str(new_path)[: -len(".nii.gz")] targets = {ext: Path(base + "." + ext) for ext in srcs} - if all(t.exists() for t in targets.values()): + stale = _stale_hardlink_targets(srcs, new_path) if refresh_stale else [] + for t in stale: + t.unlink() + print("Refreshing stale hardlink:", t) + if not stale and all(t.exists() for t in targets.values()): continue new_path.parent.mkdir(parents=True, exist_ok=True) # Temporarily restrict bf.file to just the sources we want linked, then restore. @@ -1119,9 +1331,6 @@ def _restitch_vibe_water_fat( result: dict[str, Path] = {} for part in ("water", "fat"): out_path = out_dir / f"sub-{sub}_sequ-stitched_acq-ax_part-{part}_desc-{tag}_vibe.nii.gz" - if out_path.exists(): - result[part] = out_path - continue images: list[Path] = [] for c in chunks_sorted: if c in corrected_set: @@ -1132,23 +1341,59 @@ def _restitch_vibe_water_fat( p = raw_dir / f"sub-{sub}_acq-ax_chunk-{c}_part-{part}_vibe.nii.gz" if p.exists(): images.append(p) - if len(images) < 2: - log.on_warning(f"sub-{sub} part-{part}: only {len(images)} chunk(s) available; cannot stitch") + # Re-stitch when a chunk (corrected or raw) is newer than the cached output; + # covers the case where nako_export regenerated a corrected chunk in place. + needs_stitch = True + if out_path.exists() and images: + out_mtime = out_path.stat().st_mtime + if all(p.stat().st_mtime <= out_mtime for p in images): + needs_stitch = False + _delete_stale_corrected_stitches(out_dir, sub, part, keep=out_path) + if needs_stitch: + if len(images) < 2: + log.on_warning(f"sub-{sub} part-{part}: only {len(images)} chunk(s) available; cannot stitch") + continue + try: + _stitching_fn( + [str(p) for p in images], + str(out_path), + is_seg=False, + bias_field=False, + verbose=False, + verbose_stitching=False, + ) + except Exception as e: # noqa: BLE001 + log.on_warning(f"sub-{sub} part-{part}: stitching failed: {type(e).__name__}: {e}") + continue + result[part] = out_path + return result + + +def _delete_stale_corrected_stitches(out_dir: Path, sub: str, part: str, keep: Path) -> None: + """Remove older ``desc-corrected*`` re-stitches whose tag no longer matches ``keep``. + + Called after resolving the current tag: whenever a new set of corrected chunks + yields a different ``desc-{tag}`` filename, the file(s) written for the previous + set are no longer reachable through the index and must not shadow later runs. + """ + keep_resolved = keep.resolve() if keep.exists() else keep + pattern = f"sub-{sub}_sequ-stitched_acq-ax_part-{part}_desc-corrected*_vibe.nii.gz" + for stale in out_dir.glob(pattern): + if stale.exists() and stale.resolve() == keep_resolved: continue try: - _stitching_fn( - [str(p) for p in images], - str(out_path), - is_seg=False, - bias_field=False, - verbose=False, - verbose_stitching=False, - ) - except Exception as e: # noqa: BLE001 - log.on_warning(f"sub-{sub} part-{part}: stitching failed: {type(e).__name__}: {e}") + stale.unlink() + except OSError as e: + log.on_warning(f"sub-{sub} part-{part}: could not delete stale {stale.name}: {e}") continue - result[part] = out_path - return result + base = stale.name[: -len(".nii.gz")] + for sidecar in stale.parent.glob(base + ".*"): + if sidecar.name == stale.name: + continue + try: + sidecar.unlink() + except OSError: + pass def _corrected_bids_file(path: Path, dataset_root: Path) -> BIDS_FILE: @@ -1407,8 +1652,28 @@ def _drain(fs): help="Trace hard_link()'s planned links for one subject; check src exists + same fs as target.", ) parser.add_argument("--workers", type=int, default=None, help="Number of worker processes (default: cpu_count-1).") + parser.add_argument( + "--vibe-mismatch-snaps", + nargs="?", + const=str(_VIBE_MISMATCH_SNAP_DIR), + # default=None, + default=str(_VIBE_MISMATCH_SNAP_DIR), + metavar="DIR", + help="Skip subjects whose VIBE parts have grid mismatches; write a review .jpg to DIR " + f"(defaults to {_VIBE_MISMATCH_SNAP_DIR}) unless one already exists there. " + "Pass '' to disable.", + ) + parser.add_argument( + "--no-refresh-stale", + dest="refresh_stale", + action="store_false", + help="Disable per-target samefile check. When set, subjects already flagged .hardlink_done are " + "skipped and existing targets in -canonical are never re-linked, even if the source flipped " + "between dataset-nako and rawdata-corrected or was rewritten with a new inode.", + ) + parser.set_defaults(refresh_stale=True) args = parser.parse_args() - test = True + test = False if args.build_corrected_index: build_corrected_index() @@ -1418,5 +1683,11 @@ def _drain(fs): precompute_grid_info_parallel(num_workers=args.workers, test=test) else: corrected = load_corrected_index() - for d in loop_over_repaired_nako(test=test, corrected_index=corrected, skip_subject=is_hard_linked): - hard_link(d) + skip = None if args.refresh_stale else is_hard_linked + for d in loop_over_repaired_nako( + test=test, + corrected_index=corrected, + skip_subject=skip, + vibe_mismatch_snap_dir=args.vibe_mismatch_snaps or None, + ): + hard_link(d, refresh_stale=args.refresh_stale) diff --git a/TPTBox/spine/spinestats/_run_all.py b/TPTBox/spine/spinestats/_run_all.py index 236d189..4098e6f 100644 --- a/TPTBox/spine/spinestats/_run_all.py +++ b/TPTBox/spine/spinestats/_run_all.py @@ -23,8 +23,7 @@ from tqdm import tqdm -from TPTBox import Print_Logger -from TPTBox.core.bids_files import BIDS_FILE +from TPTBox import BIDS_FILE, POI, Print_Logger from TPTBox.core.dicom.dicom2nii_utils import load_json from TPTBox.core.internal.nii_help import save_json from TPTBox.core.nii_wrapper import to_nii @@ -32,6 +31,27 @@ DATASET_ROOT = Path("/DATA/NAS/datasets_processed/NAKO/dataset-nako") logger = Print_Logger() +# Version stamp written into every aggregate `_stat.json` at +# ``_provenance.version`` (built by :func:`_build_provenance`). Bump this whenever +# the produced numbers change in a way that requires already-cached files to be +# recomputed. Current bumps: +# 1 -- initial release (WK-direction based cobb/lordosis; no S1 sacrum endplate). +# 2 -- endplate-plane based lordosis / kyphosis (average of the two flanking +# endplates at each disc), S1 sacrum-endplate landmarks, and correct +# cranio-caudal ordering for the T13 annotation label. +# 3 -- endplate-plane based cobb (chain-closed, matches lordosis convention), +# C2 excluded from Cobb search (starts at C3), non-overlapping multi-cobb +# curves, anatomic (not numeric) vertebra ordering in +# ``compute_angel_between_two_points_`` so T13 pairs measure the right +# endplates, and per-curve apex vertebra keys in ``curv`` +# (``{cervical_lordosis,thoracic_kyphosis,lumbar_lordosis}_apex``). +CURRENT_VERSION = 3 +# Highest ``CURRENT_VERSION`` bump that required *new POI landmarks*. Stat files +# older than this were written before those landmarks existed, so the sibling +# POI buffer must be rebuilt (not just the angle keys). Bumps that only change +# how numbers are derived from an unchanged POI set do not update this. +POI_INVALIDATING_VERSION = 2 + # Top-level keys we require inside a finished json before we consider a # subject "done" and skip recomputation. cobb/curv are optional and only # added when run_all is called with cobb=True. @@ -96,6 +116,15 @@ def get_nako_paths(nako_id: str) -> dict[str, Path | None]: fullbody_poi = ( DATASET_ROOT / f"derivatives-fullbody-poi/{pfx}/{sub}/vibe/sub-{sub}_sequ-stitched_acq-ax_part-water_seg-fullbody_poi.json" ) + veridah = None + for suffix in ("VERIDAH-label-V2", "VERIDAH-label"): + p = ( + DATASET_ROOT / f"derivatives_spine_inference_162_sacrumfix/{pfx}/{sub}/T2w/" + f"sub-{sub}_sequ-stitched_acq-sag_mod-T2w_seg-vert_desc-{suffix}_stat.json" + ) + if p.exists(): + veridah = p + break roi = ( DATASET_ROOT / f"derivatives_Abdominal-Segmentation/{pfx}/{sub}/vibe/sub-{nako_id}_sequ-stitched_acq-ax_mod-vibe_seg-ROI_msk.nii.gz" ) @@ -112,6 +141,7 @@ def get_nako_paths(nako_id: str) -> dict[str, Path | None]: "roi": roi, "vibeseg100": vibeseg100 if vibeseg100.exists() else None, "fullbody_poi": fullbody_poi if fullbody_poi.exists() else None, + "veridah": veridah, "dataset": DATASET_ROOT, } @@ -126,6 +156,33 @@ def _segmentation_inputs(file_dict: dict) -> list[Path]: ] +def _stat_version(loaded_stat: dict | None) -> int: + """Return the ``_provenance.version`` of a loaded stat dict (defaults to 1).""" + if not isinstance(loaded_stat, dict): + return 1 + prov = loaded_stat.get("_provenance") + if isinstance(prov, dict): + try: + return int(prov.get("version", 1)) + except (TypeError, ValueError): + return 1 + return 1 + + +def _poi_is_stale_wrt_stat(stat_path: Path, loaded_stat: dict) -> bool: + """Return True when the POI buffer sitting next to ``stat_path`` should be rebuilt. + + Staleness is derived from the stat json's ``_provenance.version``: whenever + that reads < :data:`POI_INVALIDATING_VERSION`, the sibling POI buffer is + treated as outdated (older stat + older POI travelled together). + ``loaded_stat`` is the already-loaded stat dict — pass ``{}`` if none exists + (then nothing is stale since there's no version marker to invalidate against). + """ + if not stat_path.exists() or not loaded_stat: + return False + return _stat_version(loaded_stat) < POI_INVALIDATING_VERSION + + def _is_cache_valid(json_path: Path, seg_files: list[Path], required_keys: tuple[str, ...]) -> tuple[bool, dict | None]: """Return (valid, loaded_dict). @@ -135,6 +192,9 @@ def _is_cache_valid(json_path: Path, seg_files: list[Path], required_keys: tuple - json fails to parse, - any required main key is missing from the loaded dict. """ + # TODO(hash-invalidation): also compare each entry in data["_provenance"]["inputs"] + # against the current file's sha1 (see _build_provenance) and force recompute on + # mismatch. Deferred while we assume inputs are unchanged. if not json_path.exists(): return False, None json_mtime = json_path.stat().st_mtime @@ -150,9 +210,99 @@ def _is_cache_valid(json_path: Path, seg_files: list[Path], required_keys: tuple for k in required_keys: if k not in data: return False, None + if _stat_version(data) < CURRENT_VERSION: + return False, None return True, data +_PROVENANCE_INPUT_KEYS: tuple[str, ...] = ( + "t2w", + "vibe_part-water", + "vibe_part-fat", + "vibe_part-inphase", + "vibe_part-outphase", + "vibe-water", + "vibe-fat", + "vibe-inphase", + "vibe-outphase", + "vert", + "spine", + "vibeseg100", + "roi", + "fullbody_poi", + "veridah", +) + + +def _resolve_prov_path(v) -> Path | None: + """Best-effort ``file_dict`` value → filesystem Path for provenance recording.""" + if v is None: + return None + if isinstance(v, BIDS_FILE): + nii = v.get_nii_file() + if nii is not None: + return Path(nii) + j = v.file.get("json") if hasattr(v, "file") else None + return Path(j) if j is not None else None + if isinstance(v, (str, Path)): + s = str(v) + return Path(s) if s else None + return None + + +def _file_provenance(path: Path, prior: dict | None = None) -> dict: + """Return provenance dict for ``path``. + + Shape is either ``{"path", "mtime_ns", "sha1"}`` or, when the file is + missing, ``{"path", "missing": True}``. When ``prior`` has the same + ``mtime_ns`` as the current file, its ``sha1`` is reused to avoid + re-hashing — this is what keeps reruns cheap while the "assume inputs + unchanged" mode is in effect. + """ + import hashlib + + p = Path(path) + if not p.exists(): + return {"path": str(p), "missing": True} + st = p.stat() + if isinstance(prior, dict) and prior.get("mtime_ns") == st.st_mtime_ns and isinstance(prior.get("sha1"), str): + return {"path": str(p), "mtime_ns": st.st_mtime_ns, "sha1": prior["sha1"]} + h = hashlib.sha1() + with p.open("rb") as f: + for chunk in iter(lambda: f.read(1 << 20), b""): + h.update(chunk) + return {"path": str(p), "mtime_ns": st.st_mtime_ns, "sha1": h.hexdigest()} + + +def _build_provenance(file_dict: dict, poi_out: Path | str | None, prior: dict | None) -> dict: + """Build the ``_provenance`` block for an aggregated per-subject JSON. + + Records ``path`` / ``mtime_ns`` / ``sha1`` for every known input in + ``file_dict`` plus the POI json at ``poi_out``. Reuses + ``prior["inputs"][k]`` to skip re-hashing files whose mtime is unchanged. + """ + from datetime import datetime, timezone + + prior_inputs = (prior or {}).get("inputs", {}) if isinstance(prior, dict) else {} + inputs: dict[str, dict] = {} + for key in _PROVENANCE_INPUT_KEYS: + if key not in file_dict: + continue + p = _resolve_prov_path(file_dict[key]) + if p is None: + continue + inputs[key] = _file_provenance(p, prior_inputs.get(key)) + if poi_out is not None: + p = Path(poi_out) + if p.exists(): + inputs["poi"] = _file_provenance(p, prior_inputs.get("poi")) + return { + "version": CURRENT_VERSION, + "written_at": datetime.now(timezone.utc).isoformat(timespec="seconds"), + "inputs": inputs, + } + + def run_all( file_dict, override: bool = False, @@ -247,6 +397,26 @@ def run_all( out = loaded except Exception: out = {} + # Version stale -> drop the ANGLE-related keys so the corresponding _need() + # checks trigger a recompute of just the affected metrics + JPGs, without + # invalidating the expensive body-composition / VBQ / muscle_fat / torso + # blocks that don't depend on the endplate-based lordosis fix. + _cur_ver = _stat_version(out) + if _cur_ver < CURRENT_VERSION: + logger.on_warning("version bump", _cur_ver, "->", CURRENT_VERSION, ": recomputing angle keys") + for _k in ( + "cobb", + "curv", + "curv_veridah", + "sva", + "coronal_balance", + "axial_rotation", + "segmental_endplate_angles", + "curvature_profile", + "multi_cobb", + "endplate_internal_angle", + ): + out.pop(_k, None) def _need(*keys: str, compute: bool) -> bool: return override or (any(k not in out for k in keys) and compute) @@ -308,7 +478,8 @@ def _need(*keys: str, compute: bool) -> bool: _merge_per_vertebra_metrics(out) save = True #### - need_poi = need_cobb or need_ivd or need_vert or need_curvature + need_veridah = override or "curv_veridah" not in out + need_poi = need_cobb or need_ivd or need_vert or need_curvature or need_veridah need_t2w = need_ivd or need_vert or need_vbq need_vert_nii = need_poi or need_vbq or need_bcs or need_mfi need_spine_nii = need_vert_nii or need_vbq @@ -316,8 +487,20 @@ def _need(*keys: str, compute: bool) -> bool: need_roi = need_mfi or need_torso need_vibe_wf = need_mfi - if not (need_cobb or need_ivd or need_vert or need_vbq or need_bcs or need_mfi or need_torso or need_curvature or need_pelvic): + if not ( + need_cobb + or need_ivd + or need_vert + or need_vbq + or need_bcs + or need_mfi + or need_torso + or need_curvature + or need_pelvic + or need_veridah + ): if _merge_endplate_angles(out, Path(poi_out)) or save: + out["_provenance"] = _build_provenance(file_dict, poi_out, out.get("_provenance")) save_json(final_out, out) return out @@ -334,13 +517,18 @@ def _need(*keys: str, compute: bool) -> bool: poi = None if need_poi: logger.on_debug("calc_poi_from_subreg_vert") - poi = calc_poi_from_subreg_vert( - vert, - spine, - subreg_id=[Location.Vertebra_Corpus, Location.Vertebra_Direction_Posterior, Location.Endplate, Location.Vertebra_Disc], - buffer_file=poi_out, - save_buffer_file=True, - ) + if _poi_is_stale_wrt_stat(final_out, out): + Path(poi_out).unlink(missing_ok=True) + if poi_out.exists(): + poi = POI.load(poi_out) + else: + poi = calc_poi_from_subreg_vert( + vert, + spine, + subreg_id=[Location.Vertebra_Corpus, Location.Vertebra_Direction_Posterior, Location.Endplate, Location.Vertebra_Disc], + buffer_file=poi_out, + save_buffer_file=True, + ) if need_cobb: try: project_2D = False @@ -404,6 +592,27 @@ def _need(*keys: str, compute: bool) -> bool: logger.on_fail("curvature error caught") logger.print_error() + if need_veridah and poi is not None: + try: + from TPTBox.spine.spinestats.veridah_angles import compute_veridah_variants, plot_veridah_variants + + logger.on_debug("veridah variants") + out["curv_veridah"] = compute_veridah_variants(poi, file_dict.get("veridah")) + if file_dict.get("veridah") is not None and file_dict.get("vert") is not None: + veridah_jpg_out = t2w_bf.get_changed_path( + "jpg", "snp", "derivatives_spine_inference_162_sacrumfix_subregionmeasures-v2", info={"seg": "cobb-veridah"} + ) + plot_veridah_variants( + veridah_jpg_out, + poi, + file_dict["t2w"], + file_dict["vert"], + file_dict.get("veridah"), + ) + except Exception: + logger.on_fail("veridah variants error caught") + logger.print_error() + # Merge wedge metrics directly into the per-label vert_geometry / ivd_geometry # entries so they land in per_vertebra.xlsx / per_ivd.xlsx automatically. for geom_key in ("vert_geometry", "ivd_geometry"): @@ -451,6 +660,8 @@ def _need(*keys: str, compute: bool) -> bool: logger.print_error() _merge_endplate_angles(out, Path(poi_out)) + out["_provenance"] = _build_provenance(file_dict, poi_out, out.get("_provenance")) + out.pop("_version", None) # TODO can beremoved logger.on_save("save", final_out.name) save_json(final_out, out) return out @@ -604,6 +815,7 @@ def _rows_from_json(subject_id: str, data: dict) -> tuple[dict[str, Any], list[d "axial_rotation", "endplate_internal_angle", "segmental_endplate_angles", + "_provenance", ) subject_view = {k: v for k, v in data.items() if k not in _PER_SUBJECT_EXCLUDE} _flatten("", subject_view, per_subject) @@ -844,7 +1056,7 @@ def _run_one(args: tuple[dict, bool, bool]) -> tuple[str, str | None, dict]: os.nice(20) OUT_FOLDER = Path("/DATA/NAS/ongoing_projects/robert/test/NAKO-stats") OUT_FOLDER.mkdir(parents=True, exist_ok=True) - N_CPUS = 1 # set >1 to parallelize + N_CPUS = 20 # set >1 to parallelize OVERRIDE = False aggregate = True do_not_update = False @@ -858,14 +1070,14 @@ def _run_one(args: tuple[dict, bool, bool]) -> tuple[str, str | None, dict]: try: if test: subjects = loop_over_repaired_nako(test=True) - total = 15 + total = 10 aggregate = False elif aggregate: subjects = loop_over_repaired_nako(test=False, sort=aggregate) else: # subjects = tqdm(loop_over_repaired_nako(test=False, sort=aggregate), total=30645) - l = loop_over_repaired_nako(test=False, sort=aggregate) - # total = 1000 + l = loop_over_repaired_nako(test=False, sort=True) # aggregate + # total = 15 # subjects = iter([next(l) for _ in range(total)]) subjects = l @@ -883,7 +1095,7 @@ def _run_one(args: tuple[dict, bool, bool]) -> tuple[str, str | None, dict]: from itertools import islice with ProcessPoolExecutor(max_workers=N_CPUS, max_tasks_per_child=100) as ex: - batch_size = 100 + batch_size = 500 l = tqdm(total=total) while True: gc.collect() diff --git a/TPTBox/spine/spinestats/angles.py b/TPTBox/spine/spinestats/angles.py index 606f56a..2ac3553 100644 --- a/TPTBox/spine/spinestats/angles.py +++ b/TPTBox/spine/spinestats/angles.py @@ -15,7 +15,7 @@ from TPTBox.spine.snapshot2D.snapshot_modular import Snapshot_Frame, create_snapshot IVD_MORE_ACCURATE = 15 -VERT_START_COBB = Vertebra_Instance.C2 +VERT_START_COBB = Vertebra_Instance.C3 class MoveTo(Enum): @@ -66,10 +66,10 @@ def get_location(self, v: Vertebra_Instance | int, poi: POI) -> tuple: subreg = Location.Additional_Vertebral_Body_Middle_Inferior_Median if (v, subreg) in poi: return (v, subreg) - # Test if it has next POINT + # Fall back to averaging v's centroid with the next vertebra's centroid. next_vert = v.get_next_poi(poi) - if next_vert is not None and (next_vert, 50) in poi: - return (v, subreg, next_vert, subreg) + if next_vert is not None and (v, 50) in poi and (next_vert, 50) in poi: + return (v, 50, next_vert, 50) elif self == self.TOP: prev_vert = v.get_previous_poi(poi) # Test IVD @@ -84,9 +84,9 @@ def get_location(self, v: Vertebra_Instance | int, poi: POI) -> tuple: subreg = Location.Additional_Vertebral_Body_Middle_Superior_Median if (v, subreg) in poi: return (v, subreg) - # Test if it has next POINT - if prev_vert is not None and (prev_vert, 50) in poi: - return (v, subreg, prev_vert, subreg) + # Fall back to averaging v's centroid with the previous vertebra's centroid. + if prev_vert is not None and (v, 50) in poi and (prev_vert, 50) in poi: + return (v, 50, prev_vert, 50) return (v, 50) def get_point(self, v: Vertebra_Instance | int, poi: POI) -> np.ndarray: @@ -254,11 +254,15 @@ def compute_angel_between_two_points_( return None assert id1 != id2, id1 - # Ensure id1 is less than id2 - if id1 > id2: + # Ensure id1 is anatomically above id2 (not just numerically smaller). + # T13 has label value 28 but sits between T12 (19) and L1 (20) in anatomical + # order; using the label value directly would misplace it and swap the wrong + # MoveTo semantics onto TOP/BOTTOM. + _order = Vertebra_Instance.order_dict() + if _order.get(id1, id1) > _order.get(id2, id2): id1, id2 = id2, id1 # Reorient and rescale the POI data - poi.reorient_().rescale_() + poi.reorient_().rescale_(verbose=False) recompute_use_ivd_direction = False # Determine direction-specific settings location2 = None @@ -388,17 +392,244 @@ def compute_lordosis_and_kyphosis(poi: POI, project_2D=True) -> dict[str, float poi = poi.copy() for k, i in curvature_definition.items(): - angle = compute_angel_between_two_points_( - poi, i.get_start_vert(poi), i.get_stop_vert(poi), "P", i.start_move, i.stop_move, project_2D - ) + start = i.get_start_vert(poi) + stop = i.get_stop_vert(poi) + angle = compute_angel_between_two_points_(poi, start, stop, "P", i.start_move, i.stop_move, project_2D) out[k] = round(angle, 4) if angle is not None else None + out[f"{k}_apex"] = _find_curve_apex(poi, start, stop, i.start_move, i.stop_move, Location.Vertebra_Direction_Posterior) return out +def _find_curve_apex( + poi: POI, + from_vert: Vertebra_Instance | int | None, + to_vert: Vertebra_Instance | int | None, + from_mv: MoveTo, + to_mv: MoveTo, + location: Location, +) -> int | None: + """Return the vertebra between ``from_vert`` and ``to_vert`` whose direction is closest to the endpoint bisector. + + Same apex heuristic as :func:`compute_max_cobb_angle`, just parameterised by + the direction location so it also works for lordosis / kyphosis + (``Vertebra_Direction_Posterior``). Returns ``None`` when either endpoint + direction is unavailable or no intermediate vertebra is present. + """ + if from_vert is None or to_vert is None: + return None + from_v = from_vert if isinstance(from_vert, Vertebra_Instance) else Vertebra_Instance(from_vert) + to_v = to_vert if isinstance(to_vert, Vertebra_Instance) else Vertebra_Instance(to_vert) + a = _get_norm(poi, from_v, from_mv, location, 1) + b = _get_norm(poi, to_v, to_mv, location, 1) + if a is None or b is None: + return None + apex_v = (a + b) / 2 + order = Vertebra_Instance.order() + try: + i_from = order.index(from_v) + i_to = order.index(to_v) + except ValueError: + return None + if i_from > i_to: + i_from, i_to = i_to, i_from + apex: int | None = None + cos_dis = -np.inf + for v in order[i_from : i_to + 1]: + n = _get_norm(poi, v, to_mv, location, 1) + if n is None: + continue + cos_new = cosine_distance(n, apex_v) + if cos_new > cos_dis: + cos_dis = cos_new + apex = v.value + return apex + + +def _single_endplate_ap_direction(poi: POI, vert: Vertebra_Instance, side: MoveTo) -> np.ndarray | None: + """A/P direction of *one* endplate (superior for TOP, inferior for BOTTOM). + + Uses ``Vertebra_Corpus``, ``Vertebral_Body_Endplate_Superior/Inferior`` and + ``Vertebra_Direction_Right`` to build a unit vector that lies in the endplate + plane and points *anteriorly* (matches :func:`_get_norm`'s ``inv=1`` sign + convention). Returns ``None`` if any required POI landmark is missing. + """ + if side == MoveTo.TOP: + endplate_loc = Location.Vertebral_Body_Endplate_Superior + elif side == MoveTo.BOTTOM: + endplate_loc = Location.Vertebral_Body_Endplate_Inferior + else: + return None + if (vert, 50) not in poi or (vert, endplate_loc) not in poi or (vert, Location.Vertebra_Direction_Right) not in poi: + return None + corpus = np.array(poi[vert, 50], dtype=float) + ep = np.array(poi[vert, endplate_loc], dtype=float) + r_pt = np.array(poi[vert, Location.Vertebra_Direction_Right], dtype=float) + n = ep - corpus + if side == MoveTo.BOTTOM: + n = -n # flip inferior endplate so both cases point superior + n_norm = np.linalg.norm(n) + r_vec = r_pt - corpus + r_norm = np.linalg.norm(r_vec) + if n_norm < 1e-8 or r_norm < 1e-8: + return None + n /= n_norm + r_vec /= r_norm + # cross(right, superior-pointing) lies in the endplate plane and points posterior; + # negate to match _get_norm's default sign (anterior for inv=1). + p = -np.cross(r_vec, n) + p_norm = np.linalg.norm(p) + if p_norm < 1e-8: + return None + return p / p_norm + + +def _endplate_ap_direction(poi: POI, vert: Vertebra_Instance, mv: MoveTo) -> np.ndarray | None: + """Return the endplate-plane A/P direction at a vertebra-disc *boundary*. + + For a lordosis / kyphosis chain to close at every transition (e.g. so + ``thoracic_kyphosis + lumbar_lordosis`` measures the same T4-top→L5-inferior + end-to-end angle as the total), the "bottom of upper" and "top of lower" at + each disc must use the *same* reference direction. This function averages the + two flanking endplate P/A directions at the disc: + + - :attr:`MoveTo.BOTTOM` at vertebra ``V`` averages ``V``'s inferior endplate + with the superior endplate of the next vertebra in the POI (``V.get_next_poi``). + - :attr:`MoveTo.TOP` at vertebra ``V`` averages ``V``'s superior endplate + with the inferior endplate of the previous vertebra (``V.get_previous_poi``). + + Falls back to the single-endplate direction when the neighbour is missing. + Returns ``None`` when even the local endplate is unavailable. + """ + if isinstance(vert, int): + vert = Vertebra_Instance(vert) + if mv == MoveTo.BOTTOM: + neighbour = vert.get_next_poi(poi) + neighbour_side = MoveTo.TOP + elif mv == MoveTo.TOP: + neighbour = vert.get_previous_poi(poi) + neighbour_side = MoveTo.BOTTOM + else: + return _single_endplate_ap_direction(poi, vert, mv) + own = _single_endplate_ap_direction(poi, vert, mv) + other = _single_endplate_ap_direction(poi, neighbour, neighbour_side) if neighbour is not None else None + # If the local endplate landmark is missing (calc_endplate_points_ ray-cast + # sometimes fails to hit the mask), mirror across the disc: use the neighbour's + # endplate as the direction proxy so both sides of the boundary agree. + if own is None: + return other + if other is None: + return own + avg = own + other + n = np.linalg.norm(avg) + if n < 1e-8: + return own + return avg / n + + +def _single_endplate_r_direction(poi: POI, vert: Vertebra_Instance, side: MoveTo) -> np.ndarray | None: + """R/L direction of *one* endplate (superior for TOP, inferior for BOTTOM). + + Builds a unit vector that lies in the endplate plane along the vertebra's + right/left axis: takes ``corpus - Vertebra_Direction_Right`` (matching the + WK path's ``inv=1`` sign convention, which points *left*) and projects it + onto the plane perpendicular to the endplate normal (``endplate_point - + corpus``). Returns ``None`` if any required POI landmark is missing. + """ + if side == MoveTo.TOP: + endplate_loc = Location.Vertebral_Body_Endplate_Superior + elif side == MoveTo.BOTTOM: + endplate_loc = Location.Vertebral_Body_Endplate_Inferior + else: + return None + if (vert, 50) not in poi or (vert, endplate_loc) not in poi or (vert, Location.Vertebra_Direction_Right) not in poi: + return None + corpus = np.array(poi[vert, 50], dtype=float) + ep = np.array(poi[vert, endplate_loc], dtype=float) + r_pt = np.array(poi[vert, Location.Vertebra_Direction_Right], dtype=float) + n = ep - corpus + if side == MoveTo.BOTTOM: + n = -n # flip inferior endplate so both cases point superior + n_norm = np.linalg.norm(n) + # WK path uses ``corpus - right_pt`` (points anatomical LEFT); mirror that + # convention so both paths agree when they meet in + # ``compute_angel_between_two_points_``. + r_vec = corpus - r_pt + r_norm = np.linalg.norm(r_vec) + if n_norm < 1e-8 or r_norm < 1e-8: + return None + n /= n_norm + r_vec /= r_norm + # Project the vertebra right/left vector onto the endplate plane + r_ep = r_vec - np.dot(r_vec, n) * n + r_ep_norm = np.linalg.norm(r_ep) + if r_ep_norm < 1e-8: + return None + return r_ep / r_ep_norm + + +def _endplate_r_direction(poi: POI, vert: Vertebra_Instance, mv: MoveTo) -> np.ndarray | None: + """R/L direction in the endplate plane at a vertebra-disc *boundary*. + + Analogue of :func:`_endplate_ap_direction` for the coronal (right) direction, + used to obtain a classical endplate-line orientation for Cobb angles. Averages + the two flanking endplate R directions at a disc: + + - :attr:`MoveTo.BOTTOM` at vertebra ``V`` averages ``V``'s inferior endplate + with the superior endplate of the next vertebra. + - :attr:`MoveTo.TOP` at vertebra ``V`` averages ``V``'s superior endplate + with the inferior endplate of the previous vertebra. + + Falls back to the single-endplate direction when the neighbour is missing. + Returns ``None`` when even the local endplate is unavailable. + """ + if isinstance(vert, int): + vert = Vertebra_Instance(vert) + if mv == MoveTo.BOTTOM: + neighbour = vert.get_next_poi(poi) + neighbour_side = MoveTo.TOP + elif mv == MoveTo.TOP: + neighbour = vert.get_previous_poi(poi) + neighbour_side = MoveTo.BOTTOM + else: + return _single_endplate_r_direction(poi, vert, mv) + own = _single_endplate_r_direction(poi, vert, mv) + other = _single_endplate_r_direction(poi, neighbour, neighbour_side) if neighbour is not None else None + if own is None: + return other + if other is None: + return own + avg = own + other + n = np.linalg.norm(avg) + if n < 1e-8: + return own + return avg / n + + def _get_norm(poi: POI, id1: int | Vertebra_Instance, mv: MoveTo, location: Location, inv: int = 1) -> np.ndarray | None: # noqa: ARG001 - """Return the normalised direction vector from a location POI to the vertebra centroid.""" + """Return the normalised direction vector from a location POI to the vertebra centroid. + + When ``location`` is :attr:`Location.Vertebra_Direction_Posterior` and ``mv`` targets + an endplate (:attr:`MoveTo.TOP` / :attr:`MoveTo.BOTTOM`), the buffered per-endplate + landmark (``Vertebral_Body_Endplate_Superior`` / ``_Inferior``) is preferred over + the averaged vertebral-body posterior direction — this yields the classical + endplate-line orientation used in Cobb-style lordosis/kyphosis measurements. + The same automatic switch applies to ``Location.Vertebra_Direction_Right``: it + is projected into the endplate plane so Cobb (scoliosis) angles are measured + between endplate lines, analogous to the sagittal case. The endplate direction + is used only when both the relevant endplate point and + ``Vertebra_Direction_Right`` are present in ``poi`` for that vertebra; otherwise + the code falls back to the WK-based averaged direction below. + """ if isinstance(id1, int): id1 = Vertebra_Instance(id1) + if location == Location.Vertebra_Direction_Posterior and mv in (MoveTo.TOP, MoveTo.BOTTOM): + ep_norm = _endplate_ap_direction(poi, id1, mv) + if ep_norm is not None: + return ep_norm * inv + if location == Location.Vertebra_Direction_Right and mv in (MoveTo.TOP, MoveTo.BOTTOM): + ep_norm = _endplate_r_direction(poi, id1, mv) + if ep_norm is not None: + return ep_norm * inv subreg = 50 if location in [Location.Vertebra_Disc_Inferior, Location.Vertebra_Disc_Superior]: subreg = 100 @@ -520,6 +751,8 @@ def compute_max_cobb_angle( assert b is not None apex_v = (a + b) / 2 for i in vertebrae_list[vertebrae_list.index(Vertebra_Instance(from_vert)) : vertebrae_list.index(Vertebra_Instance(to_vert)) + 1]: + if i.value not in poi.keys_region(): + continue try: a = _get_norm(poi, i, vert_id2_mv, Location.Vertebra_Direction_Right, 1) if a is None: @@ -615,8 +848,13 @@ def compute_max_cobb_angle_multi( assert from_vert is not None assert to_vert is not None out_list.append((max_angle, from_vert, to_vert, apex)) + # Exclusive split: neither endpoint of the just-found curve may participate + # in a subsequent curve. Prevents overlaps like (T1-T5) + (T5-T7); a + # sibling curve below the current one starts strictly caudal to to_vert, + # a sibling above ends strictly cranial to from_vert. Drop the ``+ 1`` + # on the ``below`` slice to restore the textbook (endpoint-shared) split. above = vertebrae_list[: vertebrae_list.index(Vertebra_Instance(from_vert))] - below = vertebrae_list[vertebrae_list.index(Vertebra_Instance(to_vert)) :] + below = vertebrae_list[vertebrae_list.index(Vertebra_Instance(to_vert)) + 1 :] compute_max_cobb_angle_multi( poi, above, @@ -686,6 +924,7 @@ def plot_compute_lordosis_and_kyphosis( seg: Image_Reference | None = None, line_len=100, project_2D=True, + curvature_definition=curvature_definition, ) -> tuple[dict[str, float | None], Snapshot_Frame]: """Plots and computes the angles of lordosis and kyphosis on a spinal image. @@ -701,6 +940,12 @@ def plot_compute_lordosis_and_kyphosis( seg (Image_Reference | None): The segmentation image reference. Optional, can be None. line_len (int): The length of the lines representing the vertebrae directions (default is 100). project_2D (bool, optional): If True, the angles are computed in the 2D sagittal projection; otherwise in 3D. Defaults to True. + curvature_definition (dict[str, Def_Curvature], optional): Mapping of output-key name + → :class:`Def_Curvature` describing which vertebra pair defines each angle. + Defaults to the module-level ``curvature_definition`` (cervical_lordosis, + thoracic_kyphosis, lumbar_lordosis). Pass a custom dict to compute a different + set of segmental angles or to override the ``last_thoracic`` / ``last_lumbar`` + resolution — the output dict's keys mirror this mapping's keys. Returns: tuple: A tuple containing: @@ -724,7 +969,7 @@ def plot_compute_lordosis_and_kyphosis( >>> print(angles) {'cervical_lordosis': 34.5, 'thoracic_kyphosis': 42.7, 'lumbar_lordosis': 50.3} """ - poi = poi.reorient().rescale_() + poi = poi.reorient().rescale_(verbose=False) poi = _add_artificial_ivd(poi) out = [] text_out = [] @@ -744,12 +989,28 @@ def plot_compute_lordosis_and_kyphosis( out.append((id1.value, s, (-a[0] * line_len * 3, -a[1] * line_len * 3))) out2 = compute_lordosis_and_kyphosis(poi, project_2D=project_2D) for name, v in out2.items(): - if v is None: + if v is None or name not in curvature_definition: + # Skip auxiliary keys like ``*_apex`` that live alongside the angles + # in the same dict but have no curve definition of their own. continue + # Apex annotation: mark the apex vertebra body with a star + label so + # the reader can see which vertebra the ``*_apex`` json key refers to. + apex_v = out2.get(f"{name}_apex") + if apex_v is not None and (apex_v, 50) in poi: + text_out.append((apex_v, ("*apex", -60))) id1 = curvature_definition[name].get_start_vert(poi) id2 = curvature_definition[name].get_stop_vert(poi) - vert = round((id1.value + id2.value) / 2) + # Cranio-caudal midpoint via the anatomical order — arithmetic mean of + # `.value` breaks for T13 (value 28, ordered after T12 but numbered after + # S1/COCC), landing the annotation on L1 instead of somewhere thoracic. + order = Vertebra_Instance.order() + try: + i1, i2 = order.index(id1), order.index(id2) + mid_inst = order[(min(i1, i2) + max(i1, i2)) // 2] + vert = mid_inst.value + except ValueError: + vert = round((id1.value + id2.value) / 2) while (vert, 50) not in poi and vert != 0: vert -= 1 text_out.append((vert, (f"{v:.1f}° - {str(name).split('_')[-1]}", 25))) @@ -806,7 +1067,7 @@ def plot_cobb_angle( >>> plot_cobb_angle("output.png", poi, img, seg, line_len=100, threshold_deg=10) """ - poi = poi.reorient().rescale_() + poi = poi.reorient().rescale_(verbose=False) poi = _add_artificial_ivd(poi) out = [] @@ -824,34 +1085,36 @@ def plot_cobb_angle( for id1, mv in zip([from_vert, to_vert], [vert_id1_mv, vert_id2_mv]): c = mv.get_location(id1, poi) - if use_ivd_direction and id1 > IVD_MORE_ACCURATE: - norm1_post = _get_norm(poi, id1, mv, Location.Vertebra_Direction_Posterior) - a = _get_norm(poi, id1, mv, Location.Vertebra_Disc_Inferior) - a = np.cross(a, norm1_post) - else: - a = _get_norm(poi, id1, mv, Location.Vertebra_Direction_Right) - - # print(a, id1, mv, c) + # Always go through Vertebra_Direction_Right so _get_norm routes + # to the endplate-plane right direction (chain-closed across the + # shared disc: `_endplate_r_direction` averages both flanking + # endplates). The old ``use_ivd_direction`` branch pulled + # ``Vertebra_Disc_Inferior`` from each vertebra separately — + # T9-BOTTOM used the T9/T10 disc but T10-TOP used the T10/T11 + # disc, so the same anatomic boundary got two different lines. + a = _get_norm(poi, id1, mv, Location.Vertebra_Direction_Right) assert a is not None out.append((apex, c, (-a[2] * line_len, a[1] * line_len))) out.append((apex, c, (a[2] * line_len, -a[1] * line_len))) - # a = _get_norm(poi, id1, mv, Location.Vertebra_Disc_Inferior) - # out.append((apex, c, (a[2] * line_len, -a[1] * line_len))) if apex is not None: - cord = poi[apex, 50] - s = f"copp angle: {max_angle:.1f}° {Vertebra_Instance(from_vert)} - {Vertebra_Instance(to_vert)}" - text_out.append((apex, (s, 25, cord[1]))) + # Align the label with the disc below the apex vertebra (IVD height) + # rather than the vertebra body centre, so the text sits at the same + # cranio-caudal level as the drawn Cobb line at the apex. + cord = poi[apex, Location.Vertebra_Disc.value] if (apex, Location.Vertebra_Disc.value) in poi else poi[apex, 50] + s = f"copp angle\n{max_angle:.1f}° {Vertebra_Instance(from_vert)} - {Vertebra_Instance(to_vert)}" + text_out.append((apex, (s, 35, cord[1]))) poi.info["line_segments_cor"] = out + poi.info.get("line_segments_cor", []) poi.info["text_cor"] = text_out + poi.info.get("text_cor", []) axis = poi.get_axis("R") width = poi.shape[axis] / poi.zoom[axis] / 2 - if width < 50: + min_half_width_mm = 80 + if width < min_half_width_mm: padd = [(0, 0) for _ in range(3)] - padd[axis] = (int(50 - width), int(50 - width)) - img = to_nii(img).apply_pad(padd) - seg = to_nii(seg, True).apply_pad(padd) + padd[axis] = (int(min_half_width_mm - width), int(min_half_width_mm - width)) + img = to_nii(img).apply_pad(padd, verbose=False) + seg = to_nii(seg, True).apply_pad(padd, verbose=False) poi = poi.resample_from_to(seg) frame = Snapshot_Frame( img, diff --git a/TPTBox/spine/spinestats/curvature.py b/TPTBox/spine/spinestats/curvature.py index 6b40bd6..de577a1 100644 --- a/TPTBox/spine/spinestats/curvature.py +++ b/TPTBox/spine/spinestats/curvature.py @@ -70,7 +70,7 @@ def _prep(poi: POI) -> POI: """Return the POI in the internal ``(P, I, R)`` orientation at 1 mm scale.""" - return poi.reorient().rescale() + return poi.reorient().rescale(verbose=False) def _get(poi: POI, vert: int | Vertebra_Instance, loc: Location) -> np.ndarray | None: diff --git a/TPTBox/spine/spinestats/poi_fun/endplates.py b/TPTBox/spine/spinestats/poi_fun/endplates.py index 48c8730..17a9e02 100644 --- a/TPTBox/spine/spinestats/poi_fun/endplates.py +++ b/TPTBox/spine/spinestats/poi_fun/endplates.py @@ -399,10 +399,26 @@ def calc_endplate_points_( poi.info.setdefault("angle_superior_endplate", {}) poi.info.setdefault("angle_inferior_endplate", {}) + # Register direction-vector fields (auto-transformed by reorient / resample / + # to_cord_system) and additional per-label scalar fields (whose keys are + # remapped by map_labels). + from TPTBox.core.poi_fun.vector_fields import POI_INFO_LABEL_KEYED_FIELDS_KEY, POI_INFO_VECTOR_FIELDS_KEY + + _vec_fields = poi.info.setdefault(POI_INFO_VECTOR_FIELDS_KEY, []) + for _f in ("angle_superior_endplate", "angle_inferior_endplate"): + if _f not in _vec_fields: + _vec_fields.append(_f) + _lbl_fields = poi.info.setdefault(POI_INFO_LABEL_KEYED_FIELDS_KEY, []) + for _f in ("endplate_internal_angle",): + if _f not in _lbl_fields: + _lbl_fields.append(_f) if compute_curvature: poi.info.setdefault("curvature_superior_endplate", {}) poi.info.setdefault("curvature_inferior_endplate", {}) + for _f in ("curvature_superior_endplate", "curvature_inferior_endplate"): + if _f not in _lbl_fields: + _lbl_fields.append(_f) # Collect normals per vertebra so we can compute the inter-endplate # angle once both superior and inferior have been processed. @@ -429,7 +445,7 @@ def calc_endplate_points_( # Location.Sacrum_Endplate, endplate_nii = spine.extract_label(Location.Sacrum_Endplate) # vert * - endplate_nii = endplate_nii.apply_crop(endplate_nii.compute_crop(0, 2, raise_error=False)) + endplate_nii = endplate_nii.apply_crop(endplate_nii.compute_crop(0, 3, raise_error=False)) if endplate_nii.max() > 0: c = endplate_nii.dilate_msk(2).get_connected_components() if c.max() != 1: @@ -438,9 +454,16 @@ def calc_endplate_points_( endplate_nii = c.extract_label(superior_label) cms_local_override = None - last_vert = max(vert_ids) + # A vertebra can appear in the segmentation (vert_ids) but have no POI + # entry (e.g. too few voxels for a centroid). Pick the highest vertebra + # that actually has a centroid in poi so the fallback to Vertebra_Corpus + # below cannot KeyError. + vert_ids_in_poi = [v for v in vert_ids if (v, Location.Vertebra_Corpus.value) in poi] + last_vert = max(vert_ids_in_poi) if vert_ids_in_poi else None - if (last_vert, Location.Vertebral_Body_Endplate_Inferior.value) in poi: + if last_vert is None: + pass # no anchor available; skip the sacrum endplate override + elif (last_vert, Location.Vertebral_Body_Endplate_Inferior.value) in poi: cms_local_override = poi[last_vert, Location.Vertebral_Body_Endplate_Inferior] elif (last_vert, Location.Vertebra_Disc.value) in poi: cms_local_override = poi[last_vert, Location.Vertebra_Disc.value] @@ -448,17 +471,19 @@ def calc_endplate_points_( cms_local_override = vert.extract_label(100 + last_vert).center_of_masses()[1] else: cms_local_override = poi[last_vert, Location.Vertebra_Corpus] - _endplate( - poi, - endplate_nii, - Location.Vertebral_Body_Endplate_Superior, - Vertebra_Instance.S1.value, - log, - normals_by_vert, - cms_local_override=cms_local_override, - flip_direction=True, - compute_curvature=compute_curvature, - ) + + if last_vert is not None: + _endplate( + poi, + endplate_nii, + Location.Vertebral_Body_Endplate_Superior, + Vertebra_Instance.S1.value, + log, + normals_by_vert, + cms_local_override=cms_local_override, + flip_direction=True, + compute_curvature=compute_curvature, + ) # Angle between superior and inferior endplate normals, per vertebra. for vert_id, normals in normals_by_vert.items(): n_sup = normals.get(Location.Vertebral_Body_Endplate_Superior) @@ -508,7 +533,7 @@ def endplate_to_super_infer_endplate(vert: NII, spine: NII) -> tuple[NII, NII]: spine = spine.copy() vert_org = vert.copy() vert[vert >= 40] = 0 - vert[spine.extract_label([Location.Vertebra_Corpus, Location.Vertebra_Corpus_border]) != 1] = 0 + vert[spine.extract_label([Location.Vertebra_Corpus, Location.Vertebra_Corpus_border, Vertebra_Instance.S1]) != 1] = 0 vert %= 100 v = vert.infect( spine.extract_label( @@ -522,8 +547,10 @@ def endplate_to_super_infer_endplate(vert: NII, spine: NII) -> tuple[NII, NII]: verbose=False, ) endplate_nii = v * endplate_nii + spine[endplate_nii == Vertebra_Instance.S1.value] = Location.Sacrum_Endplate.value spine[np.logical_and(endplate_nii == vert_org % 100, endplate_nii != 0)] = Location.Vertebral_Body_Endplate_Inferior.value spine[spine == Location.Endplate.value] = Location.Vertebral_Body_Endplate_Superior.value + vert_org[endplate_nii != 0] = v[endplate_nii != 0] + 200 return vert_org, spine diff --git a/TPTBox/spine/spinestats/veridah_angles.py b/TPTBox/spine/spinestats/veridah_angles.py new file mode 100644 index 0000000..2fa366d --- /dev/null +++ b/TPTBox/spine/spinestats/veridah_angles.py @@ -0,0 +1,194 @@ +"""VERIDAH-aware lordosis / kyphosis variants. + +Two extra angle computations on top of the standard ``curv`` block in +:mod:`TPTBox.spine.spinestats.angles`: + +- ``anomaly``: relabel the POI using the VERIDAH ``orig_label -> fpath`` + map (so a supernumerary T13 lands on label 28 instead of shifting the + entire lumbar enumeration down by one) and recompute all three + regional angles (cervical / thoracic / lumbar). +- ``k4 / k5 / k6 / k7``: force the last ``k`` present vertebrae above + the sacrum to be interpreted as L1..Lk, drop everything cranial, and + compute **only** ``lumbar_lordosis``. Independent of VERIDAH. + +Both variants live under the top-level JSON key ``curv_veridah`` (see +``_run_all.py``); the standard ``curv`` output is untouched. + +:func:`plot_veridah_variants` writes a matching snapshot JPG. It re-uses +the already-computed POI points (no re-segmentation, no re-POI-derivation) +by driving the label re-key through :meth:`POI.map_labels` — which also +carries the registered ``label_name`` and vector-field metadata along — +and applies the same integer mapping to the vertebra NIfTI via +:meth:`NII.map_labels` so the overlay matches the relabeled points. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +from TPTBox.core.nii_wrapper import NII, Image_Reference, to_nii +from TPTBox.core.poi import POI, POI_Descriptor +from TPTBox.core.vert_constants import Vertebra_Instance +from TPTBox.spine.snapshot2D.snapshot_modular import create_snapshot +from TPTBox.spine.spinestats.angles import compute_lordosis_and_kyphosis, plot_compute_lordosis_and_kyphosis + + +def _load_veridah(path: Path) -> dict | None: + """Return the single dict inside a VERIDAH ``_stat.json``, or ``None`` on error.""" + try: + data = json.loads(Path(path).read_text()) + except (OSError, json.JSONDecodeError): + return None + if isinstance(data, list) and data and isinstance(data[0], dict): + return data[0] + if isinstance(data, dict): + return data + return None + + +def _relabel_poi(poi: POI, mapping: dict[int, int]) -> POI: + return poi.map_labels(label_map_region=mapping) + + +def _last_k_relabel(poi: POI, k: int) -> POI: + """Keep the ``k`` most caudal non-sacral vertebrae; relabel them L1..Lk. + + Uses :meth:`Vertebra_Instance.order` to walk cranio-caudal, ignores + sacrum members, and keeps the last ``k`` present labels. The bottom + one becomes ``Lk`` (label ``20 + k - 1``), stepping up by one to + ``L1`` (label 20). + """ + order = Vertebra_Instance.order() + sacrum_vals = {v.value for v in Vertebra_Instance.sacrum()} + present = {r for r, _s, _c in poi.centroids.items()} + non_sacral_in_order = [v.value for v in order if v.value in present and v.value not in sacrum_vals] + if len(non_sacral_in_order) < k: + return poi.copy(centroids=POI_Descriptor()) + last_k = non_sacral_in_order[-k:] # cranio → caudal + # last_k[0] should be L1 (20), last_k[-1] should be Lk (20+k-1). + mapping = {orig: 20 + i for i, orig in enumerate(last_k)} + return _relabel_poi(poi, mapping) + + +def _compute_anomaly_variant(poi: POI, veridah_json_path: Path | None, project_2D: bool = False) -> dict[str, float | None] | None: + """Recompute the three regional angles after applying the VERIDAH label correction. + + ``project_2D`` must match the setting used for the standard ``curv`` block + in :mod:`_run_all` (default ``False``, i.e. 3D angles). Returns ``None`` + when the VERIDAH file is missing or unusable. + """ + if veridah_json_path is None: + return None + veridah = _load_veridah(Path(veridah_json_path)) + if veridah is None: + return None + orig = veridah.get("orig_label") + fpath = veridah.get("fpath") + if not isinstance(orig, list) or not isinstance(fpath, list) or len(orig) != len(fpath): + return None + mapping = {int(o): int(f) for o, f in zip(orig, fpath) if int(o) != int(f)} + return compute_lordosis_and_kyphosis(_relabel_poi(poi, mapping), project_2D=project_2D) + + +def _compute_k_variant(poi: POI, k: int, project_2D: bool = False) -> dict[str, float | None]: + """Compute lumbar-lordosis only, with the last ``k`` vertebrae treated as L1..Lk.""" + relabeled = _last_k_relabel(poi, k) + if not relabeled.centroids: + return {"lumbar_lordosis": None} + full = compute_lordosis_and_kyphosis(relabeled, project_2D=project_2D) + return {"lumbar_lordosis": full.get("lumbar_lordosis"), "lumbar_lordosis_apex": full.get("lumbar_lordosis_apex")} + + +def _veridah_region_map(veridah_json_path: Path) -> dict[int, int] | None: + """Return the ``orig_label -> fpath`` region mapping from a VERIDAH stat json.""" + veridah = _load_veridah(Path(veridah_json_path)) + if veridah is None: + return None + orig = veridah.get("orig_label") + fpath = veridah.get("fpath") + if not isinstance(orig, list) or not isinstance(fpath, list) or len(orig) != len(fpath): + return None + return {int(o): int(f) for o, f in zip(orig, fpath) if int(o) != int(f)} + + +def plot_veridah_variants( + jpg_path: str | Path | None, + poi: POI, + img: Image_Reference, + seg_vert: Image_Reference, + veridah_json_path: Path | str | None, + line_len: int = 100, + project_2D: bool = False, +) -> tuple[dict[str, float | None] | None, str | None]: + """Render the VERIDAH-corrected lordosis / kyphosis snapshot next to the standard one. + + Re-uses the *already computed* POI points -- the angle numbers are cheap + trig on those points, and no re-segmentation / re-``calc_poi_from_subreg_vert`` + happens. The re-key is driven through :meth:`POI.map_labels`, so any + registered ``label_name`` / vector-field metadata is carried along + automatically. The vertebra segmentation is remapped with the same + integer table via :meth:`NII.map_labels` so the overlay matches. + + Args: + jpg_path: Output path for the snapshot JPG. If ``None``, the image is not saved. + poi: The (already computed) POI object. + img: The reference image on which to plot. + seg_vert: Vertebra segmentation NIfTI, remapped alongside the POI. + veridah_json_path: Path to the VERIDAH ``_stat.json`` giving + ``{orig_label, fpath}`` lists. + line_len: Length of the direction lines drawn on the snapshot. + project_2D: Compute the angles as 2D sagittal projections. + + Returns: + ``(anomaly_angles, saved_path)`` -- ``anomaly_angles`` is the same + dict :func:`compute_lordosis_and_kyphosis` returns (or ``None`` + when the VERIDAH file is missing/malformed), ``saved_path`` is + the resolved output path (or ``None`` when nothing was written). + """ + if veridah_json_path is None: + return None, None + mapping = _veridah_region_map(Path(veridah_json_path)) + if mapping is None: + return None, None + + relabeled_poi = _relabel_poi(poi, mapping) + if not relabeled_poi.centroids: + return None, None + # Remap the vertebra segmentation with the same integer table so the overlay + # numbering matches the (already-remapped) POI. NII.map_labels leaves labels + # that are not listed in ``mapping`` unchanged. + seg_nii = seg_vert if isinstance(seg_vert, NII) else to_nii(seg_vert, seg=True) + seg_relabeled = seg_nii.map_labels(mapping, verbose=False) # type: ignore[arg-type] + + angles, _frame = plot_compute_lordosis_and_kyphosis( + None, + relabeled_poi, + img, + seg_relabeled, + line_len=line_len, + project_2D=project_2D, + ) + saved: str | None = None + if jpg_path is not None: + create_snapshot(jpg_path, [_frame]) + saved = str(jpg_path) + return angles, saved + + +def compute_veridah_variants(poi: POI, veridah_json_path: Path | str | None, project_2D: bool = False) -> dict: + """Return the ``curv_veridah`` block for one subject. + + Shape:: + + { + "anomaly": {"cervical_lordosis", "thoracic_kyphosis", "lumbar_lordosis"} | None, + "k4": {"lumbar_lordosis": …}, + "k5": {...}, "k6": {...}, "k7": {...}, + } + """ + veridah_path = Path(veridah_json_path) if veridah_json_path is not None else None + out: dict = {"anomaly": _compute_anomaly_variant(poi, veridah_path, project_2D=project_2D)} + for k in (4, 5, 6, 7): + out[f"k{k}"] = _compute_k_variant(poi, k, project_2D=project_2D) + return out diff --git a/docs/api/poi_fun.md b/docs/api/poi_fun.md index d1050ca..675c4fb 100644 --- a/docs/api/poi_fun.md +++ b/docs/api/poi_fun.md @@ -44,3 +44,123 @@ from segmentation volumes. options: show_source: true filters: ["!^_"] + +## Direction-Vector Fields in `POI.info` + +`POI.info` can hold *auxiliary* per-vertebra data alongside the main POI points — +direction vectors (e.g. endplate PCA normals) and scalar metadata (e.g. per-vertebra +wedge angles or curvature). Producers register field names under two well-known +keys in `poi.info`, and `POI.reorient` / `POI.resample_from_to` / +`POI_Global.to_cord_system` / `POI.map_labels` then keep those fields aligned +with the POI points automatically. + +**Two registries:** + +- `poi.info["_vector_fields"]` (`POI_INFO_VECTOR_FIELDS_KEY`) — a list of field + names whose values are `{key: (x, y, z)}` dicts holding **unit direction + vectors in mm-space aligned with `poi.orientation`**. These fields are + auto-transformed by `reorient`, `resample_from_to`, and + `POI_Global.to_cord_system`, and their keys are also remapped by `map_labels`. + `rescale` is a no-op for them. +- `poi.info["_label_keyed_fields"]` (`POI_INFO_LABEL_KEYED_FIELDS_KEY`) — a + list of field names whose values are `{key: }` dicts (typically + scalars). Only their **keys** are remapped by `map_labels`; no + orientation-based transform is applied. + +Keys in either type of field may be integer region labels (`20`) or +`Vertebra_Instance`-name strings (`"L1"`); both forms are supported. + +**Producer pattern:** + +```python +from TPTBox.core.poi_fun.vector_fields import ( + POI_INFO_VECTOR_FIELDS_KEY, + POI_INFO_LABEL_KEYED_FIELDS_KEY, +) + +# Direction vectors (auto-rotated on reorient / resample): +vec_fields = poi.info.setdefault(POI_INFO_VECTOR_FIELDS_KEY, []) +if "my_vector_field" not in vec_fields: + vec_fields.append("my_vector_field") +poi.info["my_vector_field"] = {"L1": (0.09, -0.99, -0.02)} # ... + +# Scalar per-vertebra metadata (key-remapped by map_labels only): +lbl_fields = poi.info.setdefault(POI_INFO_LABEL_KEYED_FIELDS_KEY, []) +if "my_scalar_field" not in lbl_fields: + lbl_fields.append("my_scalar_field") +poi.info["my_scalar_field"] = {"L1": 3.14} # ... +``` + +**Assigning names (`info["label_name"]`):** + +Human-readable per-point and per-region names live in +`poi.info["label_name"]` as `{region: {subregion: name, "name": group_name}}`. +Use the accessors on `Abstract_POI` (available on both `POI` and `POI_Global`) +instead of writing the dict directly: + +```python +poi.set_label_name(region=2, subregion=10, name="FLCPC") # per-point label +poi.set_level_one_name(region=2, name="Femur") # region group name + +poi.label_name(2, 10) # -> "FLCPC" +poi.level_one_name(2) # -> "Femur" +``` + +`region` / `subregion` accept `int`, numeric string, or `Enum` members. A +custom name in `label_name` always takes priority over the auto-derived name +from `level_one_info` / `level_two_info`; if none is set, `.label_name(...)` +falls back to the enum name, and finally to the raw id as a string. A warning +is emitted when a custom name conflicts with the `level_two_info` enum name +for the same id. + +`map_labels` remaps this field automatically: the same +`_remap_vector_field_keys_inplace` helper handles both flat fields and the +nested `label_name` structure by dispatching on value type. For `label_name`, +`label_map_region` remaps the top-level region keys, `label_map_subregion` +remaps the inner subregion keys, and the `"name"` group entry is preserved. +Two source regions colliding onto one target merge inner dicts with +*last-write-wins* on overlapping keys. No explicit registration required — +`label_name` is always handled. + +**Names in 3D Slicer:** + +`POI_Global.save_mrk(...)` writes a `.mrk.json` markup file whose control +points and groups inherit these names directly: + +```python +poi_global = poi.to_global() +poi_global.save_mrk("points.mrk.json", pointLabelsVisibility=True) +``` + +Per control point, `save_mkr.get_desc(poi, region, subregion)` looks up: + +- `label` — from `poi.info["label_name"][region][subregion]`; falls back to + the `level_two_info` enum name (or the raw subregion id). +- `name2` (group label shown in the markup tree) — from + `poi.info["label_name"][region]["name"]`; falls back to + `poi.info["label_group_name"][region]` and finally to the + `level_one_info` enum name. + +So `poi.set_label_name(...)` and `poi.set_level_one_name(...)` are all you +need: Slicer displays those strings on hover, in the markup tree, and (when +`pointLabelsVisibility=True`) as 3D annotations. Enable +`split_by_region=True` on `save_mrk` to get one Slicer group per region +(named via `level_one_name`). + +**Caveats:** + +- `resample_from_to` uses the real `R_ref.T @ R_self` rotation of the two + affines (safe for skewed grids); axcode-only heuristics are deliberately + avoided. +- `map_labels` handles duplicates with a *last-write-wins* policy after remap. +- `map_labels`' `label_map_full` (mapping `(region, subreg)` tuples) does NOT + trigger the key-remap of these fields — only `label_map_region` / + `label_map_subregion` do. + +::: TPTBox.core.poi_fun.vector_fields + options: + show_source: true + filters: ["!^_"] + members: + - POI_INFO_VECTOR_FIELDS_KEY + - POI_INFO_LABEL_KEYED_FIELDS_KEY diff --git a/unit_tests/test_poi_vector_fields.py b/unit_tests/test_poi_vector_fields.py new file mode 100644 index 0000000..9892251 --- /dev/null +++ b/unit_tests/test_poi_vector_fields.py @@ -0,0 +1,471 @@ +"""Unit tests for auto-transform of direction-vector and label-keyed fields in ``POI.info``. + +Covers: +- ``POI.reorient`` — signed permutation of direction vectors. +- ``POI.resample_from_to`` — full rotation (``R_ref.T @ R_self``). +- ``POI.rescale`` — no-op for mm-space vectors. +- ``POI_Global.to_cord_system`` — RAS <-> LPS axis flips. +- ``POI.map_labels`` — key remap for both vector and label-keyed fields. +""" + +from __future__ import annotations + +import sys +import unittest +from pathlib import Path + +file = Path(__file__).resolve() +sys.path.append(str(file.parents[2])) + +from typing import ClassVar # noqa: E402 + +import numpy as np # noqa: E402 + +from TPTBox.core.poi import POI # noqa: E402 +from TPTBox.core.poi_fun.vector_fields import ( # noqa: E402 + POI_INFO_LABEL_KEYED_FIELDS_KEY, + POI_INFO_VECTOR_FIELDS_KEY, +) + + +def _make_poi(orientation=("P", "I", "R"), zoom=(1.0, 1.0, 1.0)) -> POI: + """Build a small POI with identity rotation and one L1 corpus point.""" + centroids: dict[int, dict[int, tuple[float, float, float]]] = {20: {50: (1.0, 2.0, 3.0)}} + return POI( + centroids, + orientation=orientation, + zoom=zoom, + shape=(100, 100, 100), + origin=(0.0, 0.0, 0.0), + rotation=np.eye(3), + ) + + +def _register(poi: POI, vec_fields=(), lbl_fields=()) -> None: + if vec_fields: + poi.info.setdefault(POI_INFO_VECTOR_FIELDS_KEY, []).extend(vec_fields) + if lbl_fields: + poi.info.setdefault(POI_INFO_LABEL_KEYED_FIELDS_KEY, []).extend(lbl_fields) + + +class Test_Vector_Field_Reorient(unittest.TestCase): + def test_pir_to_las_signed_permutation(self): + # In PIR, +x=P, +y=I, +z=R. In LAS, +x=L=-R, +y=A=-P, +z=S=-I. + # v_PIR = (P=0.020, I=-0.092, R=0.996). Same physical direction in LAS is: + # L=-R=-0.996, A=-P=-0.020, S=-I=0.092 -> (-0.996, -0.020, 0.092). + poi = _make_poi(orientation=("P", "I", "R")) + _register(poi, vec_fields=["v"]) + v_in = np.array([0.02, -0.09, 0.996]) + v_in = v_in / np.linalg.norm(v_in) + poi.info["v"] = {"L1": tuple(v_in)} + out = poi.reorient(("L", "A", "S")) + got = np.asarray(out.info["v"]["L1"]) + expected = np.array([-v_in[2], -v_in[0], -v_in[1]]) + np.testing.assert_allclose(got, expected, atol=1e-9) + # Norm is preserved by a signed permutation. + self.assertAlmostEqual(float(np.linalg.norm(got)), 1.0, places=6) + # Full round-trip: PIR -> LAS -> PIR equals original. + back = out.reorient(("P", "I", "R")) + np.testing.assert_allclose(back.info["v"]["L1"], v_in, atol=1e-9) + + def test_unregistered_field_untouched(self): + poi = _make_poi() + # not registered + poi.info["orphan_vec"] = {"L1": (1.0, 0.0, 0.0)} + out = poi.reorient(("L", "A", "S")) + self.assertEqual(tuple(out.info["orphan_vec"]["L1"]), (1.0, 0.0, 0.0)) + + def test_none_or_malformed_values_skipped(self): + poi = _make_poi() + _register(poi, vec_fields=["v"]) + poi.info["v"] = {"L1": (0.0, 1.0, 0.0), "L2": None, "L3": (1.0, 2.0)} + out = poi.reorient(("L", "A", "S")) + self.assertIsNone(out.info["v"]["L2"]) + # 2-tuple stays unchanged, since shape mismatch is silently skipped + self.assertEqual(tuple(out.info["v"]["L3"]), (1.0, 2.0)) + # 3-tuple got transformed + self.assertEqual(len(out.info["v"]["L1"]), 3) + + +class Test_Vector_Field_Rescale(unittest.TestCase): + def test_rescale_leaves_vectors_untouched(self): + poi = _make_poi(zoom=(0.5, 0.5, 3.0)) + _register(poi, vec_fields=["v"]) + v = (0.1, -0.9, 0.05) + poi.info["v"] = {"L1": v} + out = poi.rescale((1.0, 1.0, 1.0)) + np.testing.assert_allclose(out.info["v"]["L1"], v, atol=1e-12) + + +class Test_Vector_Field_ResampleFromTo(unittest.TestCase): + def test_identity_grid_is_noop(self): + poi = _make_poi() + _register(poi, vec_fields=["v"]) + v = (0.09, -0.99, -0.02) + poi.info["v"] = {"L1": v} + # ref = a copy of the same POI -> same rotation -> R_ref.T @ R_self = I + ref = poi.copy() + out = poi.resample_from_to(ref) + np.testing.assert_allclose(out.info["v"]["L1"], v, atol=1e-6) + + def test_rotated_ref_applies_rotation(self): + poi = _make_poi() + _register(poi, vec_fields=["v"]) + v = (1.0, 0.0, 0.0) + poi.info["v"] = {"L1": v} + # Target POI has a 90-degree z-rotation + R_z90 = np.array([[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]) + ref = poi.copy() + ref.rotation = R_z90 + out = poi.resample_from_to(ref) + # v_target = R_ref.T @ R_self @ v = R_z90.T @ I @ [1,0,0] = R_z90.T @ [1,0,0] = [0, -1, 0] + np.testing.assert_allclose(out.info["v"]["L1"], (0.0, -1.0, 0.0), atol=1e-9) + + +class Test_ToCordSystem(unittest.TestCase): + def test_ras_to_lps_flips_xy_only(self): + poi = _make_poi() + _register(poi, vec_fields=["v"]) + poi.info["v"] = {"L1": (0.1, 0.2, 0.3)} + g = poi.to_global() # RAS by default + self.assertFalse(g.itk_coords) + self.assertEqual(tuple(g.info["v"]["L1"]), (0.1, 0.2, 0.3)) + g_itk = g.to_cord_system(itk_coords=True) + np.testing.assert_allclose(g_itk.info["v"]["L1"], (-0.1, -0.2, 0.3), atol=1e-12) + # roundtrip + g_back = g_itk.to_cord_system(itk_coords=False) + np.testing.assert_allclose(g_back.info["v"]["L1"], (0.1, 0.2, 0.3), atol=1e-12) + + +class Test_MapLabels(unittest.TestCase): + def test_vector_field_keys_remapped(self): + poi = _make_poi() + _register(poi, vec_fields=["v"]) + v = (0.1, -0.9, 0.4) + poi.info["v"] = {"L1": v} + out = poi.map_labels(label_map_region={20: 2}) # L1 -> C2 + self.assertNotIn("L1", out.info["v"]) + self.assertIn("C2", out.info["v"]) + np.testing.assert_allclose(out.info["v"]["C2"], v, atol=1e-12) + + def test_label_keyed_scalar_field_remapped(self): + poi = _make_poi() + _register(poi, lbl_fields=["endplate_internal_angle"]) + poi.info["endplate_internal_angle"] = {"L1": 3.14} + out = poi.map_labels(label_map_region={20: 2}) + self.assertNotIn("L1", out.info["endplate_internal_angle"]) + self.assertAlmostEqual(out.info["endplate_internal_angle"]["C2"], 3.14) + + def test_int_keys_also_remapped(self): + poi = _make_poi() + _register(poi, vec_fields=["v"]) + poi.info["v"] = {20: (0.0, 1.0, 0.0), 21: (1.0, 0.0, 0.0)} + out = poi.map_labels(label_map_region={20: 2}) + self.assertNotIn(20, out.info["v"]) + self.assertIn(2, out.info["v"]) + # unaffected key + self.assertIn(21, out.info["v"]) + np.testing.assert_allclose(out.info["v"][2], (0.0, 1.0, 0.0), atol=1e-12) + + def test_unknown_string_key_kept(self): + poi = _make_poi() + _register(poi, lbl_fields=["s"]) + poi.info["s"] = {"NOT_A_VERT": 42.0, "L1": 1.0} + out = poi.map_labels(label_map_region={20: 2}) + self.assertIn("NOT_A_VERT", out.info["s"]) + self.assertNotIn("L1", out.info["s"]) + self.assertIn("C2", out.info["s"]) + + def test_empty_map_is_noop(self): + poi = _make_poi() + _register(poi, vec_fields=["v"]) + v = (0.1, -0.9, 0.4) + poi.info["v"] = {"L1": v} + out = poi.map_labels(label_map_region={}) + self.assertEqual(tuple(out.info["v"]["L1"]), v) + + +class Test_LabelName_MapLabels(unittest.TestCase): + """label_name (nested {region: {subregion: name, "name": group}}) rides map_labels.""" + + def _poi_with_label_name(self) -> POI: + poi = _make_poi() + # nested format directly + poi.info["label_name"] = { + 20: {50: "L1_corpus", 100: "L1_disc", "name": "Spine"}, + 21: {50: "L2_corpus"}, + } + return poi + + def test_region_remap(self): + poi = self._poi_with_label_name() + out = poi.map_labels(label_map_region={20: 2}) + ln = out.info["label_name"] + self.assertNotIn(20, ln) + self.assertIn(2, ln) + # inner subregion keys unchanged, group name preserved + self.assertEqual(ln[2][50], "L1_corpus") + self.assertEqual(ln[2][100], "L1_disc") + self.assertEqual(ln[2]["name"], "Spine") + # unaffected region still there + self.assertEqual(ln[21][50], "L2_corpus") + + def test_subregion_remap(self): + poi = self._poi_with_label_name() + out = poi.map_labels(label_map_subregion={50: 51}) + ln = out.info["label_name"] + # region keys unchanged + self.assertIn(20, ln) + self.assertNotIn(50, ln[20]) + self.assertIn(51, ln[20]) + self.assertEqual(ln[20][51], "L1_corpus") + # group name key survives + self.assertEqual(ln[20]["name"], "Spine") + # unrelated subregion 100 preserved + self.assertEqual(ln[20][100], "L1_disc") + + def test_region_and_subregion_remap(self): + poi = self._poi_with_label_name() + out = poi.map_labels(label_map_region={20: 2}, label_map_subregion={50: 51}) + ln = out.info["label_name"] + self.assertEqual(ln[2][51], "L1_corpus") + self.assertEqual(ln[2]["name"], "Spine") + + def test_region_collision_last_write_wins_on_inner(self): + # Both 20 and 21 map onto 2: their inner dicts should merge; overlapping + # inner keys let the later region's value win. + poi = _make_poi() + poi.info["label_name"] = { + 20: {50: "A", "name": "grp20"}, + 21: {50: "B", 60: "unique"}, + } + out = poi.map_labels(label_map_region={20: 2, 21: 2}) + ln = out.info["label_name"] + self.assertIn(2, ln) + # value at key 50 should come from the second insert (21 -> "B") + self.assertEqual(ln[2][50], "B") + # inner keys from both are preserved + self.assertEqual(ln[2][60], "unique") + # group name from region 20 comes along + self.assertEqual(ln[2].get("name"), "grp20") + + def test_flat_legacy_format_migrated_and_remapped(self): + # Ensure the migration path in label_name_dict / normalize_label_name is + # triggered by the remap, so old flat "(region, subreg)" strings work too. + poi = _make_poi() + poi.info["label_name"] = {"(20, 50)": "L1_corpus"} + out = poi.map_labels(label_map_region={20: 2}) + ln = out.info["label_name"] + self.assertIn(2, ln) + self.assertEqual(ln[2][50], "L1_corpus") + + +class Test_LabelName_LegacyMigration(unittest.TestCase): + """The old flat ``{"(region, subreg)": name}`` format must migrate cleanly. + + Real-world fixture: a leg-atlas POI with 38 flat entries across 5 regions, + including multi-digit subregion ids like ``(2, 10)`` — those must survive + ``ast.literal_eval`` parsing. + """ + + _ATLAS_FLAT: ClassVar[dict[str, str]] = { + "(1, 1)": "TGT", + "(1, 2)": "FHC", + "(1, 3)": "FNC", + "(1, 4)": "FAAP", + "(2, 1)": "FLCD", + "(2, 2)": "FMCD", + "(2, 3)": "FLCP", + "(2, 4)": "FMCP", + "(2, 5)": "FNP", + "(2, 6)": "FADP", + "(2, 7)": "TGPP", + "(2, 8)": "TGCP", + "(2, 9)": "FMCPC", + "(2, 10)": "FLCPC", + "(2, 11)": "TRMP", + "(2, 12)": "TRLP", + "(3, 1)": "TLCL", + "(3, 2)": "TMCM", + "(3, 3)": "TKC", + "(3, 4)": "TLCA", + "(3, 5)": "TLCP", + "(3, 6)": "TMCA", + "(3, 7)": "TMCP", + "(3, 8)": "TTP", + "(3, 9)": "TAAP", + "(3, 10)": "TMIT", + "(3, 11)": "TLIT", + "(4, 1)": "FLM", + "(4, 2)": "TMM", + "(4, 3)": "TAC", + "(4, 4)": "TADP", + "(5, 1)": "PPP", + "(5, 2)": "PDP", + "(5, 3)": "PMP", + "(5, 4)": "PLP", + "(5, 5)": "PRPP", + "(5, 6)": "PRDP", + "(5, 7)": "PRHP", + } + + def test_normalize_label_name_migrates_flat_atlas(self): + from TPTBox.core.poi_fun.poi_abstract import normalize_label_name + + nested = normalize_label_name(dict(self._ATLAS_FLAT)) + # region keys are ints + self.assertEqual(set(nested.keys()), {1, 2, 3, 4, 5}) + # multi-digit inner keys survive parsing + self.assertEqual(nested[2][10], "FLCPC") + self.assertEqual(nested[2][12], "TRLP") + self.assertEqual(nested[3][11], "TLIT") + # inner keys are ints too, no leftover string keys + for region, inner in nested.items(): + for k in inner: + self.assertIsInstance(k, int, f"inner key {k!r} in region {region} is not int") + # count is preserved + total = sum(len(v) for v in nested.values()) + self.assertEqual(total, len(self._ATLAS_FLAT)) + + def test_normalize_is_idempotent(self): + from TPTBox.core.poi_fun.poi_abstract import normalize_label_name + + once = normalize_label_name(dict(self._ATLAS_FLAT)) + twice = normalize_label_name({**once}) + self.assertEqual(once, twice) + + def test_normalize_empty_and_none(self): + from TPTBox.core.poi_fun.poi_abstract import normalize_label_name + + self.assertEqual(normalize_label_name(None), {}) + self.assertEqual(normalize_label_name({}), {}) + + def test_label_name_dict_caches_migration_in_info(self): + from TPTBox.core.poi_fun.poi_abstract import label_name_dict + + info = {"label_name": dict(self._ATLAS_FLAT)} + out = label_name_dict(info) + # returned dict is the normalized form + self.assertEqual(out[1][1], "TGT") + # and the migrated form is cached back into info + self.assertIs(info["label_name"], out) + # a second call is a no-op (still nested) + out2 = label_name_dict(info) + self.assertIs(out2, info["label_name"]) + + def test_load_poi_migrates_atlas_from_disk(self): + """End-to-end: write the legacy JSON, load it, expect the nested form.""" + import json + import tempfile + + from TPTBox.core.poi import POI + + payload = [ + { + "direction": ["R", "A", "S"], + "zoom": [1.0, 1.0, 1.0], + "origin": [0.0, 0.0, 0.0], + "shape": [10, 10, 10], + "rotation": [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], + "format": "POI", + "label_name": dict(self._ATLAS_FLAT), + }, + # one dummy point so the file loads as a POI: {region: {subregion: (x,y,z)}} + {"1": {"1": [1.0, 2.0, 3.0]}}, + ] + with tempfile.NamedTemporaryFile(mode="w", suffix="_poi.json", delete=False) as f: + json.dump(payload, f) + path = f.name + try: + poi = POI.load(path) + ln = poi.info["label_name"] + self.assertEqual(ln[1][1], "TGT") + self.assertEqual(ln[2][10], "FLCPC") + # no more flat "(...)"-style keys + self.assertFalse(any(isinstance(k, str) and k.startswith("(") for k in ln)) + finally: + Path(path).unlink(missing_ok=True) + + def test_migration_then_map_labels(self): + """A flat-format POI still remaps correctly through map_labels.""" + poi = _make_poi() + poi.info["label_name"] = dict(self._ATLAS_FLAT) # legacy flat + # rename region 2 -> 20 and region 5 -> 50 in one shot + out = poi.map_labels(label_map_region={2: 20, 5: 50}) + ln = out.info["label_name"] + self.assertIn(20, ln) + self.assertIn(50, ln) + self.assertNotIn(2, ln) + self.assertNotIn(5, ln) + self.assertEqual(ln[20][10], "FLCPC") + self.assertEqual(ln[50][7], "PRHP") + + +class Test_LabelName_Accessors(unittest.TestCase): + """`set_label_name` / `set_level_one_name` write into ``info['label_name']`` + and get read back by ``label_name`` / ``level_one_name`` and by the + Slicer/mkr exporter. + """ + + def _poi_with_enums(self) -> POI: + from TPTBox.core.vert_constants import Location, Vertebra_Instance + + poi = _make_poi() + poi.level_one_info = Vertebra_Instance + poi.level_two_info = Location + return poi + + def test_set_and_get_label_name(self): + poi = self._poi_with_enums() + poi.set_label_name(region=20, subregion=50, name="L1_corpus") + self.assertEqual(poi.label_name(20, 50), "L1_corpus") + # unset points fall back to the level_two_info enum name (or the raw id when + # no enum entry matches). + fallback = poi.label_name(21, 50) + self.assertIsInstance(fallback, str) + self.assertNotEqual(fallback, "L1_corpus") + + def test_set_and_get_level_one_name(self): + poi = self._poi_with_enums() + poi.set_level_one_name(region=20, name="Vertebra L1 custom") + self.assertEqual(poi.level_one_name(20), "Vertebra L1 custom") + # unset region falls back to the level_one_info enum name (L1 -> "L1") + self.assertEqual(poi.level_one_name(21), "L2") + + def test_set_label_name_persists_in_info(self): + poi = _make_poi() + poi.set_label_name(20, 50, "L1_corpus") + poi.set_level_one_name(20, "Femur") + ln = poi.info["label_name"] + self.assertEqual(ln[20][50], "L1_corpus") + self.assertEqual(ln[20]["name"], "Femur") + + def test_names_flow_into_slicer_export(self): + """`get_desc` (used by save_mrk) reads label/group name from label_name.""" + from TPTBox.core.poi_fun.save_mkr import get_desc + + poi = _make_poi() + poi.set_label_name(20, 50, "L1_corpus") + poi.set_level_one_name(20, "Spine") + g = poi.to_global() + name, name2, label = get_desc(g, region=20, subregion=50) + # `label` is the per-point custom name; `name2` is the region group name. + self.assertEqual(label, "L1_corpus") + self.assertEqual(name2, "Spine") + + +class Test_Composition(unittest.TestCase): + def test_reorient_and_map_labels_commute(self): + # For a vector-transform + key-remap, order shouldn't matter. + poi_a = _make_poi() + poi_b = _make_poi() + for poi in (poi_a, poi_b): + _register(poi, vec_fields=["v"]) + poi.info["v"] = {"L1": (0.02, -0.09, 0.996)} + out_a = poi_a.reorient(("L", "A", "S")).map_labels(label_map_region={20: 2}) + out_b = poi_b.map_labels(label_map_region={20: 2}).reorient(("L", "A", "S")) + np.testing.assert_allclose(out_a.info["v"]["C2"], out_b.info["v"]["C2"], atol=1e-9) + + +if __name__ == "__main__": + unittest.main() diff --git a/unit_tests/test_spinestats_angles.py b/unit_tests/test_spinestats_angles.py index d415f47..1b66df9 100644 --- a/unit_tests/test_spinestats_angles.py +++ b/unit_tests/test_spinestats_angles.py @@ -256,10 +256,11 @@ def test_has_point(self): 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(): + angle_keys = sorted(k for k in out if not k.endswith("_apex")) + self.assertEqual(angle_keys, ["cervical_lordosis", "lumbar_lordosis", "thoracic_kyphosis"]) + for key in angle_keys: with self.subTest(key=key): - self.assertEqual(value, 0.0) + self.assertEqual(out[key], 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})) @@ -660,7 +661,8 @@ def test_lordosis_plot_returns_the_same_numbers_as_the_bare_computation(self): 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"]) + angle_keys = sorted(k for k in values if not k.endswith("_apex")) + self.assertEqual(angle_keys, ["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