diff --git a/TPTBox/core/bids_files.py b/TPTBox/core/bids_files.py index 801748e0..9071a0ec 100755 --- a/TPTBox/core/bids_files.py +++ b/TPTBox/core/bids_files.py @@ -142,7 +142,7 @@ def get_values_from_name(path: Path | str, verbose: bool) -> tuple[str, dict[str (e.g. ``"sub-001_ses-01_T1w"``), * ``file_type`` is the extension (e.g. ``"nii.gz"``). """ - name = Path(path).name + name = path.rpartition("/")[2] if isinstance(path, str) else path.name bids_key, file_type = name.split(".", maxsplit=1) @@ -179,7 +179,6 @@ def Buffered_BIDS_Global_info( parents: Sequence[str] | str = ["rawdata", "derivatives"], additional_key: Sequence[str] = ["sequ", "seg", "ovl"], verbose: bool = True, - file_name_manipulation: typing.Callable[[str], str] | None = None, sequence_splitting_keys: list[str] | None = None, filter_file: typing.Callable[[Path], bool] | None = None, max_age_days: int = 30, @@ -295,7 +294,6 @@ def save_buffer(f: Path, buffer_name: str) -> list[Path]: parents, additional_key, verbose=verbose, - file_name_manipulation=file_name_manipulation, sequence_splitting_keys=sequence_splitting_keys, filter_folder=lambda _x, _y: False, additional_file_list=files, @@ -305,7 +303,7 @@ def save_buffer(f: Path, buffer_name: str) -> list[Path]: _cont = 0 -def _scan_tree(path, lvl=1, filter_folder=lambda _x, _y: True, verbose=False): +def _scan_tree(path, lvl=1, filter_folder=None, verbose=False): """Recursively yield DirEntry objects for given directory.""" global _cont # noqa: PLW0603 for entry in os.scandir(path): @@ -330,7 +328,6 @@ def __init__( parents: Sequence[str] | str = ["rawdata", "derivatives"], additional_key: Sequence[str] = ["sequ", "seg", "ovl"], verbose: bool = True, - file_name_manipulation: typing.Callable[[str], str] | None = None, sequence_splitting_keys: list[str] | None = None, filter_folder: typing.Callable[[Path, int], bool] | None = None, additional_file_list: dict[str | Path, list[Path]] | None = None, @@ -346,6 +343,7 @@ def __init__( filter_folder = lambda p, lvl: True if (lvl != 2 or p.name in ["sub-123","sub-456"]) else False """ self.count_file = 0 + self._p_counter = 0 if sequence_splitting_keys is None: from TPTBox.core.bids_constants import sequence_splitting_keys @@ -360,7 +358,6 @@ def __init__( assert isinstance(parents, Sequence), "parents is not a list" self.__bids_list: dict = {} - self.file_name_manipulation = file_name_manipulation # Validate for ds in datasets: ds_path = Path(ds) if isinstance(ds, str) else ds @@ -431,39 +428,28 @@ def add_file_2_subject(self, bids: BIDS_FILE | Path, ds: Path | str | None = Non ds = bids.dataset else: raise AssertionError("Dataset-path required") - if isinstance(bids, (Path, str)): - try: - bids_key, file_type = str(bids).rsplit("/", maxsplit=1)[-1].split(".", maxsplit=1) - # print(bids_key) - except Exception: + name = bids.rpartition("/")[2] if isinstance(bids, str) else bids.name + bids_key, sep, file_type = name.partition(".") + if not sep: print("[!] skip file with out a type declaration:", bids.name) - # raise e return if bids_key in self._global_bids_list: self._global_bids_list[bids_key].add_file(bids) return - bids = BIDS_FILE( - bids, - ds, - verbose=self.verbose, - file_name_manipulation=self.file_name_manipulation, - ) - + bids = BIDS_FILE(bids, ds, verbose=self.verbose) subject = bids.info.get("sub", "unsorted") if subject not in self.subjects: self.subjects[subject] = Subject_Container(subject, self.sequence_splitting_keys) self.count_file += 1 - ( - print( - f"Found: {subject}, total file keys {(self.count_file)}, total subjects = {len(self.subjects)} ", - end="\r", - ) - if self.verbose - else None - ) + self._p_counter -= 1 + if self.verbose and self._p_counter < 0: + print(f"Found: {subject}, total file keys {(self.count_file)}, total subjects = {len(self.subjects)}", end="\r") + self._p_counter += random.randint(10, 250) + self.subjects[subject].add(bids) + self._global_bids_list[bids.BIDS_key] = bids def enumerate_subjects(self, sort: bool = False, shuffle: bool = False) -> list[tuple[str, Subject_Container]]: """Return all subject identifiers together with their :class:`Subject_Container`. @@ -641,14 +627,7 @@ def get_sequence_files( class BIDS_FILE: """Representation of a single BIDS-compliant file with parsed entities and dataset context.""" - def __init__( - self, - file: Path | str, - dataset: Path | str, - verbose=True, - bids_ds: BIDS_Global_info | None = None, - file_name_manipulation: typing.Callable[[str], str] | None = None, - ): + def __init__(self, file: Path | str, dataset: Path | str, verbose=True, bids_ds: BIDS_Global_info | None = None): """Multi-file BIDS record sharing the same identifier (all extensions of one file stem). Holds references to `.nii.gz`, `.json`, etc. simultaneously. @@ -671,24 +650,33 @@ def __init__( file = Path(file) if not isinstance(file, Path) else file self.dataset = Path(dataset) if not isinstance(dataset, Path) else dataset self.verbose = verbose - if file_name_manipulation is not None: - if "WS_" in str(file): - file.rename(file.parent / Path(file_name_manipulation(file.name))) - name = file_name_manipulation(file.name) - else: - name = file.name + name = file.name self.format, self.info, self.BIDS_key, file_type = get_values_from_name(name, verbose) if bids_ds is not None: bids_ds.add_file_2_subject(bids=self, ds=self.dataset) - self.file = {file_type: file} - bids_key, _ = file.name.split(".", maxsplit=1) - for file_type in ["nii.gz", "json", "png"]: - if file_type in self.file: - continue - if os.path.exists(os.path.join(file.parent, bids_key + "." + file_type)): - self.file[file_type] = Path(file.parent, bids_key + "." + file_type) - self.file = dict(sorted(self.file.items())) + self._file = {file_type: file} + self._checked = False + + @property + def file(self) -> dict[str, Path]: + """Returns a dict mapping file types to paths. ["nii.gz", "json", "png"] are automatic searched for. + + Returns: + dict[str, Path]: _description_ + """ + if not self._checked: + files = {p.parent for p in self._file.values()} + for f in files: + bids_key = self.BIDS_key + for file_type in ["nii.gz", "json", "png"]: + if file_type in self._file: + continue + if os.path.exists(os.path.join(f, bids_key + "." + file_type)): + self._file[file_type] = Path(f, bids_key + "." + file_type) + self._file = dict(sorted(self._file.items())) + self._checked = True + return self._file def get_file(self, ending: str = "json", default: Path | None = None) -> Path | None: """Return the path for a given file extension, or *default* if absent. @@ -811,11 +799,7 @@ def remove(self, key: str) -> str: assert key != "sub", "not allowed to remove subject name" return self.info.pop(key) - def add_file( - self, - path: Path, - bids_ds: BIDS_Global_info | None = None, - ) -> None: + def add_file(self, path: Path, bids_ds: BIDS_Global_info | None = None) -> None: """Associate an additional file extension with this BIDS entry. Used to register companion files (e.g. a ``.json`` sidecar alongside @@ -834,12 +818,16 @@ def add_file( bids_key, file_type = Path(path).name.split(".", maxsplit=1) assert bids_key == self.BIDS_key, f"only aligned data aka same name different file type: {bids_key} != {self.BIDS_key}" - bids_dic_file = self.file - if file_type not in self.file: + bids_dic_file = self._file + if file_type not in bids_dic_file: bids_dic_file[file_type] = path if bids_ds is not None: - bids_ds._global_bids_list[bids_key].file = dict(sorted(bids_dic_file.items())) - self.file = dict(sorted(bids_dic_file.items())) + bids_ds._global_bids_list[bids_key]._file = dict(sorted(bids_dic_file.items())) + self._file = bids_ds._global_bids_list[bids_key]._file + else: + self._file = dict(sorted(bids_dic_file.items())) + elif bids_dic_file[file_type] != path: + print("BIDS_Key conflict!", path, "<-->", bids_dic_file) def rename_files(self, path: Path | str, ending: str = ".nii.gz") -> None: """Rename all associated files on disk to a new base path. @@ -863,36 +851,50 @@ def rename_files(self, path: Path | str, ending: str = ".nii.gz") -> None: p = Path(path + "." + key) value.rename(p) - def symlink_files(self, path: Path | str, ending: str = ".nii.gz", exist_ok: bool = False) -> None: - """Create symbolic links for all associated files at a new base path. + def symlink_files(self, path: Path | str, ending: str = ".nii.gz", exist_ok: bool = False, hard_link: bool = False) -> None: + """Create symbolic or hard links for all associated files at a new base path. - Equivalent to :meth:`rename_files` but creates symlinks rather than - moving files. Existing correct symlinks are silently skipped. + Equivalent to :meth:`rename_files` but creates links rather than moving + files. Existing correct symlinks/hard links are silently skipped. Args: path: Target path including the primary extension (e.g. ``/out/sub-001_T1w.nii.gz``). ending: Extension used to compute the base stem; a leading dot is added automatically if absent. + exist_ok: If ``True``, skip existing files. + hard_link: If ``True``, create hard links using :func:`os.link` + instead of symbolic links. Raises: AssertionError: If *path* does not end with *ending*, or if an - existing symlink at the target points elsewhere. + existing link at the target points elsewhere. """ - ending = ending if ending[0] == "." else "." + ending + ending = ending if ending.startswith(".") else "." + ending path = str(path) - assert path.endswith(ending), f"set 'ending' to the part after the '.'\n {path} does not end with {ending}" + assert path.endswith(ending), f"set 'ending' to the part after the '.'\n{path} does not end with {ending}" path = path.replace(ending, "") + for key, value in self.file.items(): p = Path(path + "." + key) - if os.path.islink(p): - assert Path(os.readlink(p)) == value, f"{p} exists" - continue - if exist_ok and p.exists(): - continue - - os.symlink(value, p) + if hard_link: + if p.exists(): + same = p.stat().st_ino == value.stat().st_ino and p.stat().st_dev == value.stat().st_dev + if exist_ok: + p.unlink(missing_ok=True) + else: + assert same, f"{p} exists" + continue + os.link(value, p) + else: + if os.path.islink(p): + if exist_ok and p.exists(): + p.unlink(missing_ok=True) + else: + assert Path(os.readlink(p)) == value, f"{p} exists" + continue + os.symlink(value, p) def get_path_decomposed(self, file_type: str | None = None) -> tuple[Path, str, str, str]: """Decompose the file path relative to the dataset root. @@ -917,7 +919,6 @@ def get_path_decomposed(self, file_type: str | None = None) -> tuple[Path, str, parent = folder_list[0] subpath = folder_list[1:-1] filename = folder_list[-1] - # print(parent, subpath, filename) return self.dataset, parent, str.join("/", subpath), filename @property @@ -1989,7 +1990,7 @@ def loop_dict( for sequ, values in self.candidates.items() ) if sort: - l = sorted(l) # type: ignore + l = sorted(l, key=lambda x: x.family_id) # type: ignore return l diff --git a/TPTBox/core/dicom/dicom_extract.py b/TPTBox/core/dicom/dicom_extract.py index 459c3969..42bb5116 100644 --- a/TPTBox/core/dicom/dicom_extract.py +++ b/TPTBox/core/dicom/dicom_extract.py @@ -7,6 +7,7 @@ import zipfile from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor, as_completed +from datetime import datetime from pathlib import Path import dicom2nifti @@ -52,27 +53,39 @@ def _next_letter_suffix(s: str, inc: int = 1) -> str: return "".join(reversed(result)) -def _inc_key(keys: dict, inc: int = 1, k="sequ") -> None: - """Increment the sequence key inside *keys* by appending letter suffixes.""" - if k not in keys: - keys[k] = "0" - value = str(keys[k]) - try: - # Pure number: 100 -> 100-a - int(value) - keys[k] = f"{value}-a" - return # noqa: TRY300 - except ValueError: - pass +def _inc_key(keys: dict, inc: int = 1, k: str = "sequ", path_exists: Callable[[dict], bool] | None = None) -> None: + """Increment the sequence key inside *keys* by appending letter suffixes. - try: - base, suffix = value.rsplit("-", maxsplit=1) - if suffix.isalpha(): - keys[k] = f"{base}-{_next_letter_suffix(suffix, inc)}" - else: - keys[k] = f"{base}-a" - except ValueError: - keys[k] = f"{value}-a" + When ``path_exists`` is given, keep incrementing until it returns ``False`` — i.e. + until the filename generated from *keys* no longer collides with an existing file + on disk. This guarantees the caller never receives keys that would produce a + duplicate filename. + """ + + def _step() -> None: + if k not in keys: + keys[k] = "0" + value = str(keys[k]) + try: + # Pure number: 100 -> 100-a + int(value) + keys[k] = f"{value}-a" + return # noqa: TRY300 + except ValueError: + pass + + try: + base, suffix = value.rsplit("-", maxsplit=1) + if suffix.isalpha(): + keys[k] = f"{base}-{_next_letter_suffix(suffix, inc)}" + else: + keys[k] = f"{base}-a" + except ValueError: + keys[k] = f"{value}-a" + + _step() + while path_exists is not None and path_exists(keys): + _step() def _generate_bids_path( @@ -105,10 +118,17 @@ def _generate_bids_path( ses, # Session, if exist ) args = {"file_type": "json", "parent": parent, "make_parent": True, "additional_folder": mri_format, "bids_format": mri_format} - fname = BIDS_FILE(Path(p, "sub-000_ct.nii.gz"), dataset_nifti_dir).get_changed_bids(**args, info=keys, non_strict_mode=True) - while test_name_conflict(simp_json, fname.file["json"]): - _inc_key(keys) - fname = BIDS_FILE(Path(p, "sub-000_ct.nii.gz"), dataset_nifti_dir).get_changed_bids(**args, info=keys, non_strict_mode=True) + + def _make_fname(k: dict): + return BIDS_FILE(Path(p, "sub-000_ct.nii.gz"), dataset_nifti_dir).get_changed_bids(**args, info=k, non_strict_mode=True) + + fname = _make_fname(keys) + # If a file already sits at this path, check whether its content matches ours + # (ignoring the "grid" key). Same content → reuse the existing filename. + # Different content → let _inc_key find a fresh, non-colliding filename. + if test_name_conflict(simp_json, fname.file["json"]): + _inc_key(keys, path_exists=lambda k: Path(_make_fname(k).file["json"]).exists()) + fname = _make_fname(keys) return fname.file["json"], fname @@ -547,10 +567,21 @@ def _add_grid_info_to_json(nii_path: Path | str, simp_json: Path | str, force_up Returns: The updated JSON dictionary including the ``"grid"`` key. """ - json_dict = load_json(simp_json) if Path(simp_json).exists() else {} - if "grid" in json_dict and not force_update: + nii_path = Path(nii_path) + simp_json = Path(simp_json) + + # Always preserve the existing JSON contents (DICOM metadata written by save_json). + # The mtime comparison is only used to short-circuit re-computing the grid when the + # sidecar is already up to date; it must NOT decide whether to keep the DICOM keys. + json_dict = load_json(simp_json) if simp_json.exists() else {} + json_up_to_date = ( + simp_json.exists() + and nii_path.exists() + and datetime.fromtimestamp(simp_json.stat().st_mtime) > datetime.fromtimestamp(nii_path.stat().st_mtime) + ) + if "grid" in json_dict and not force_update and json_up_to_date: return json_dict - print("Read Grid info", Path(simp_json).exists(), "grid" in json_dict) + print("Read Grid info") nii = NII.load(nii_path, False) gird = { "shape": nii.shape, diff --git a/TPTBox/core/internal/nii_help.py b/TPTBox/core/internal/nii_help.py index d51e636a..88d58d36 100644 --- a/TPTBox/core/internal/nii_help.py +++ b/TPTBox/core/internal/nii_help.py @@ -54,6 +54,9 @@ def wrapper(self, file: str | Path | bids_files.BIDS_FILE, *args, **kwargs): if file_type in file.file: file = file.file[file_type] break + else: + raise ValueError(f"No supported file type found in BIDS_FILE. Expected one of: {file_types}") + file = Path(file) if isinstance(file, str) else file # Ensure the file is a Path object backup_file = file.with_suffix(file.suffix + ".backup") file_existed = file.exists() diff --git a/TPTBox/core/internal/train_nnUnet/prepere_dataset.py b/TPTBox/core/internal/train_nnUnet/prepere_dataset.py index 8b50a9f9..c39dbb05 100644 --- a/TPTBox/core/internal/train_nnUnet/prepere_dataset.py +++ b/TPTBox/core/internal/train_nnUnet/prepere_dataset.py @@ -34,6 +34,7 @@ class DatasetConfig: # ── Preprocessing / spacing ─────────────────────────────────────────────── spacing: tuple[float, float, float] = (1, 1, 1) + orientation: tuple[str, str, str] = ("R", "A", "S") is_ct: bool = True num_input: int = 1 axis: str = "S" @@ -130,13 +131,11 @@ def _build_label_mapping( mapping_forward: dict[int, int] = {} labels_mapping_return: dict[str, str | int] = {} - for new_idx, (orig_idx, name) in enumerate( - sorted(dataset_mapping.items()), - start=1, - ): + for new_idx, (orig_idx, name) in enumerate(sorted(dataset_mapping.items()), start=1): labels_mapping[name] = new_idx - mapping_forward[orig_idx] = new_idx - labels_mapping_return[str(new_idx)] = enums.get(name, orig_idx) + if orig_idx != new_idx: + mapping_forward[orig_idx] = new_idx + labels_mapping_return[str(new_idx)] = enums.get(name, orig_idx) # ---------------------------------------------------------- # remap mirror pairs @@ -179,7 +178,7 @@ def build_dataset(cfg: DatasetConfig) -> None: from _prep_ds import add_file, finalize_ds, run, set_up_dataset labels_mapping, mapping_forward, mapping_back, mirror = _build_label_mapping(cfg) - logger.on_text(f"Label count : {len(mapping_forward)} classes") + logger.on_text(f"Label count : {len(labels_mapping) - 1} classes") logger.on_text(f"Mirror pairs : {len(mirror) if mirror else 0}") logger.on_text(f"Trainer : {cfg.nn_trainer}") logger.on_text(f"Spacing : {cfg.spacing}") @@ -199,8 +198,8 @@ def build_dataset(cfg: DatasetConfig) -> None: labels_found = set(seg_nii.unique()) labels_found.discard(0) # ignore background - expected_labels = set(mapping_forward.keys()) - + expected_labels = set(labels_mapping.values()) + expected_labels.remove(0) missing_mapping = labels_found - expected_labels unused_mapping = expected_labels - labels_found @@ -211,7 +210,7 @@ def build_dataset(cfg: DatasetConfig) -> None: logger.on_fail(f"Labels present in segmentation but missing in mapping: {sorted(missing_mapping)}") if unused_mapping: - logger.on_warning(f"Labels defined in mapping but not found in sample: {sorted(unused_mapping)}") + logger.on_ok(f"Unmapped labels {sorted(unused_mapping)}") # Test remapping out = seg_nii.map_labels(mapping_forward) @@ -219,8 +218,7 @@ def build_dataset(cfg: DatasetConfig) -> None: logger.on_text(f"Remapped labels : {remapped_labels}") - expected_remapped = set(mapping_forward.values()) - unexpected = set(remapped_labels) - expected_remapped - {0} + unexpected = set(remapped_labels) - expected_labels - {0} if unexpected: logger.on_fail(f"Unexpected labels after remapping: {sorted(unexpected)}") @@ -238,6 +236,7 @@ def build_dataset(cfg: DatasetConfig) -> None: num_input=cfg.num_input, is_ct=cfg.is_ct, base=cfg.nnunet_base, + orientation=cfg.orientation, ) dataset_settings["labels_mapping"] = mapping_back diff --git a/TPTBox/core/internal/train_nnUnet/train.py b/TPTBox/core/internal/train_nnUnet/train.py index ba4a847a..6a643131 100644 --- a/TPTBox/core/internal/train_nnUnet/train.py +++ b/TPTBox/core/internal/train_nnUnet/train.py @@ -104,11 +104,9 @@ def _run_training_highjack(self: nnUNetTrainer) -> None: self.on_validation_epoch_end(val_outputs) self.on_epoch_end() - l = list(self.dataset_json["labels"].keys()) - self.print_to_log_file( - "Dice", - ", ".join([f"{l[e]}:{i:.3f}" for e, i in enumerate(self.logger.my_fantastic_logging["dice_per_class_or_region"][-1], 1)]), - ) + id_to_name = {v: k for k, v in self.dataset_json["labels"].items()} + dice = self.logger.get_value("dice_per_class_or_region", step=-1) + self.print_to_log_file("Dice", ", ".join([f"{id_to_name[e]}:{i:.3f}" for e, i in enumerate(dice, 0)])) self.on_train_end() diff --git a/TPTBox/core/nii_wrapper.py b/TPTBox/core/nii_wrapper.py index d656c990..91fee579 100755 --- a/TPTBox/core/nii_wrapper.py +++ b/TPTBox/core/nii_wrapper.py @@ -2290,6 +2290,9 @@ def map_labels(self, label_map:LABEL_MAP , verbose:logging=True, inplace=False) If inplace is True, returns the current NIfTI image object with mapped labels. Otherwise, returns a new NIfTI image object with mapped labels. """ data_orig = self.get_seg_array() + if len(label_map) == 0: + log.print("Skip map_labels; map is empty", verbose=verbose) + return self if inplace else self.copy() # the before/after np_unique scans are only used for the verbose log line; skip them otherwise labels_before = [v for v in np_unique(data_orig) if v > 0] if verbose else None # enforce keys to be str to support both str and int diff --git a/TPTBox/core/poi.py b/TPTBox/core/poi.py index 057e541b..6f3d6c30 100755 --- a/TPTBox/core/poi.py +++ b/TPTBox/core/poi.py @@ -1089,7 +1089,7 @@ def calc_poi_from_subreg_vert( level_two_info=Location, ) if extend_to is None - else extend_to.apply_crop(crop, inplace=True) + else extend_to.resample_from_to_(vert_msk) ) if _vert_ids is None: diff --git a/TPTBox/core/poi_fun/poi_abstract.py b/TPTBox/core/poi_fun/poi_abstract.py index 7b1bd418..8b4bb86d 100755 --- a/TPTBox/core/poi_fun/poi_abstract.py +++ b/TPTBox/core/poi_fun/poi_abstract.py @@ -612,7 +612,7 @@ def fit_spline( self, smoothness: int = 10, samples_per_poi=20, - location: int | Abstract_lvl = Location.Vertebra_Corpus, + location: int | Enum | list[int] | list[Enum] | None = Location.Vertebra_Corpus, vertebra=False, ) -> tuple[np.ndarray, np.ndarray]: """Fits a spline interpolation through a set of centroids and calculates the first derivative of the spline curve. @@ -630,12 +630,11 @@ def fit_spline( - spline_1st_derivative: A 2D NumPy array representing the first derivative of the spline curve. shape: first dimension to select a cord, second dimension to select all X/Y/Z """ - if isinstance(location, Abstract_lvl): - location = location.value - if location not in self.keys_subregion() and not isinstance(location, Sequence): - raise ValueError(f"The location {location} is not computed in this POI class") # Extract subregion based on the provided location - poi = self.extract_subregion(*location) if isinstance(location, Sequence) else self.extract_subregion(location) + poi = self.copy() if location is None else self.extract_subregion(location) + if len(poi) == 0: + raise ValueError(f"The location {location} is not computed in this POI class") + # If vertebra sorting is requested, perform it poi = poi.sort(inplace=False, order_dict=Vertebra_Instance.order_dict() if vertebra else None) # Convert centroids to NumPy array for processing @@ -803,7 +802,7 @@ def remove(self, *label: tuple[int, int], inplace: bool = False) -> Self: obj.centroids.pop(loc, None) return obj - def extract_subregion(self, *location: int | list[int] | Enum, inplace: bool = False) -> Self: + def extract_subregion(self, *location: int | list[int] | list[Enum] | Enum, inplace: bool = False) -> Self: """Return a POI containing only the specified subregion(s). Args: @@ -813,9 +812,9 @@ def extract_subregion(self, *location: int | list[int] | Enum, inplace: bool = F Returns: Filtered POI. """ - location = _flatten(location) + location_ = _flatten(location) - location_values = tuple(l if isinstance(l, int) else l.value for l in location) + location_values = tuple(l if isinstance(l, int) else l.value for l in location_) extracted_centroids = POI_Descriptor() for x1, x2, y in self.centroids.items(): if x2 in location_values: diff --git a/TPTBox/core/poi_fun/poi_global.py b/TPTBox/core/poi_fun/poi_global.py index 0f798795..07b4c9ef 100755 --- a/TPTBox/core/poi_fun/poi_global.py +++ b/TPTBox/core/poi_fun/poi_global.py @@ -273,6 +273,9 @@ def save( grid before saving. verbose: Emit a save log message. Defaults to ``True``. """ + if Path(out_path).name.endswith("mrk.json"): + logging.on_warning("use save_mrk to save .mrk.json files") + return self.save_mrk(out_path) return save_poi( self, out_path, make_parents, additional_info, save_hint=save_hint, resample_reference=resample_reference, verbose=verbose ) diff --git a/TPTBox/core/poi_fun/vertebra_direction.py b/TPTBox/core/poi_fun/vertebra_direction.py index 94ac553a..b1214800 100644 --- a/TPTBox/core/poi_fun/vertebra_direction.py +++ b/TPTBox/core/poi_fun/vertebra_direction.py @@ -1,6 +1,7 @@ from __future__ import annotations from collections.abc import Sequence +from typing import Literal from warnings import warn import numpy as np @@ -8,7 +9,7 @@ from TPTBox import NII, POI, Print_Logger, calc_poi_from_subreg_vert from TPTBox.core.poi_fun._help import make_spine_plot, sacrum_w_o_direction -from TPTBox.core.vert_constants import DIRECTIONS, Location, _plane_dict, never_called +from TPTBox.core.vert_constants import DIRECTIONS, Location, Vertebra_Instance, _plane_dict, never_called Vertebra_Orientation = tuple[np.ndarray, np.ndarray, np.ndarray] _log = Print_Logger() @@ -63,8 +64,21 @@ def calc_orientation_of_vertebra_PIR( spine_plot_path: str | None = None, save_normals_in_info=False, _orientation_version=0, + method: Literal["spline", "endplate"] = "endplate", ) -> tuple[POI, NII | None]: - """Calculate the orientation of vertebrae using PIR (Posterior, Inferior, Right) DIRECTIONS. + """Compute the PIR (Posterior–Inferior–Right) orientation for each vertebra. + + The algorithm proceeds in four stages: + + 1. Estimate the inferior direction from vertebral endplates when available, + otherwise fall back to the local spline tangent. + 2. Construct a plane through the vertebral corpus orthogonal to the + inferior direction and intersect it with the posterior vertebral + structures. + 3. Compute the posterior direction from the corpus to the intersection + centroid. + 4. Construct the PIR coordinate frame and store the corresponding + direction landmarks. Args: poi (POI | None): Point of interest. If None, computed from `vert` and `subreg`. @@ -85,12 +99,61 @@ def calc_orientation_of_vertebra_PIR( if _orientation_version != 0: warn("out dated _orientation_version; Set is to 0", stacklevel=1) - # Step 1 compute the up direction + # --------------------------------------------------------------------- + # 1. Compute inferior directions + # --------------------------------------------------------------------- # check if label 50 is already computed in POI if poi is None or spline_subreg_point_id.value not in poi.keys_subregion(): poi = calc_poi_from_subreg_vert(vert, subreg, extend_to=poi, subreg_id=spline_subreg_point_id) # compute Spline in ISO space poi_iso = poi.rescale().reorient() + vert_keys = list(poi.keys_region()) + down_vector: dict[int, np.ndarray] = {} + + if method == "endplate": + for key in poi_iso.keys_region(): + if (key, Location.Vertebra_Corpus.value) not in poi_iso: + continue + + corpus = np.asarray(poi_iso[key, Location.Vertebra_Corpus], float) + directions = [] + + if (key, Location.Vertebral_Body_Endplate_Inferior.value) in poi_iso: + inf = np.asarray(poi_iso[key, Location.Vertebral_Body_Endplate_Inferior], float) + v = inf - corpus + if norm(v) > 1e-6: + directions.append(v / norm(v)) + + if (key, Location.Vertebral_Body_Endplate_Superior.value) in poi_iso: + sup = np.asarray(poi_iso[key, Location.Vertebral_Body_Endplate_Superior], float) + v = corpus - sup + if norm(v) > 1e-6: + directions.append(v / norm(v)) + + if len(directions) == 1: + down_vector[key] = directions[0] + elif len(directions) == 2: + down = directions[0] + directions[1] + if norm(down) > 1e-6: + down_vector[key] = down / norm(down) + ### Add last endplate to spline + last_vert = 25 + while last_vert not in vert_keys: + last_vert -= 1 + if last_vert == 20: + last_vert = None + break + max_vert_key = max(vert_keys) + 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 + 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 + ##### + # spline: body_spline, body_spline_der = poi_iso.fit_spline(location=spline_subreg_point_id, vertebra=True) # Step 2 compute the back direction by Spinosus_Process or arcus intersection_target = [Location.Spinosus_Process, Location.Arcus_Vertebrae] @@ -115,14 +178,17 @@ def calc_orientation_of_vertebra_PIR( out = target_labels * 0 fill_back_nii = subreg_iso.copy() if do_fill_back else None fill_back = out.copy() if do_fill_back else None - down_vector: dict[int, np.ndarray] = {} # 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(): # calculate_normal_vector - distances = np.sqrt(np.sum((body_spline - np.array(cords)) ** 2, -1)) - normal_vector_post = body_spline_der[np.argmin(distances)] - normal_vector_post /= np.linalg.norm(normal_vector_post) - down_vector[reg_label] = normal_vector_post.copy() + if reg_label in down_vector: + normal_vector_down = down_vector[reg_label] + else: + # spline fallback + distances = np.sqrt(np.sum((body_spline - np.array(cords)) ** 2, -1)) + normal_vector_down = body_spline_der[np.argmin(distances)] + normal_vector_down /= np.linalg.norm(normal_vector_down) + down_vector[reg_label] = normal_vector_down.copy() # create_plane_coords # The main axis will be treated differently idx = [_plane_dict[i] for i in subreg_iso.orientation] @@ -133,7 +199,7 @@ def calc_orientation_of_vertebra_PIR( dim1, dim2 = dims # Make a plane through start_point with the norm of "normal_vector", which is shifted by "shift" along the norm # create_subregion: 1 where the selected subreg is, else 0 - select = _create_plane_mask(subreg_iso.shape, np.array(cords), normal_vector_post, axis, dim1, dim2) # type: ignore + select = _create_plane_mask(subreg_iso.shape, np.array(cords), normal_vector_down, axis, dim1, dim2) # type: ignore out[out == 0] += (target_labels * select * reg_label)[out == 0] if fill_back is not None: diff --git a/TPTBox/core/poi_fun/vertebra_pois_non_centroids.py b/TPTBox/core/poi_fun/vertebra_pois_non_centroids.py index ef1e0250..8f0583a9 100755 --- a/TPTBox/core/poi_fun/vertebra_pois_non_centroids.py +++ b/TPTBox/core/poi_fun/vertebra_pois_non_centroids.py @@ -15,6 +15,7 @@ from TPTBox.core.poi_fun.vertebra_direction import calc_center_spinal_cord, calc_orientation_of_vertebra_PIR from TPTBox.core.vert_constants import Location, vert_directions from TPTBox.spine.spinestats import calculate_IVD_POI +from TPTBox.spine.spinestats.poi_fun.endplates import calc_endplate_points_ _log = Print_Logger() all_poi_functions: dict[int, Strategy_Pattern] = {} @@ -92,21 +93,30 @@ class Strategy_Pattern: Args: target (Location): The target location for which this strategy is defined. strategy (Callable): The strategy function that implements the desired behavior. - prerequisite (set[Location] | None, optional): A set of prerequisite locations that must be satisfied before applying this strategy. Defaults to None. - **args: Additional keyword arguments to be passed to the strategy function. + prerequisite (set[Location] | None, optional): A set of prerequisite locations that must be + satisfied before applying this strategy. Defaults to None. + prio (int, optional): Scheduling offset added to ``target.value`` in :meth:`prority`. + Strategies with a lower resulting priority run first. Defaults to 0. + sakrum (bool, optional): When ``True``, this strategy is also applied to sacrum vertebra IDs; + otherwise sacrum vertebrae are skipped in :func:`compute_non_centroid_pois`. Defaults to False. + **args: Additional keyword arguments to be passed to the strategy function. Any ``Location`` + values found in ``args`` (including inside sequences) are auto-added to ``prerequisite``, + and passing ``direction=...`` additionally adds ``Location.Vertebra_Direction_Inferior``. Attributes: target (Location): The target location for which this strategy is defined. args (dict): Additional keyword arguments to be passed to the strategy function. prerequisite (set[Location]): A set of prerequisite locations that must be satisfied before applying this strategy. strategy (Callable): The strategy function that implements the desired behavior. + sacrum (bool): Whether this strategy also applies to sacrum vertebrae. Note: The strategy function should accept the following arguments: - - poi (POI): The point of interest. - - current_subreg (NII): The current subregion. - - vert_id (int): The vertex ID. - - bb: The bounding box. + - poi (POI): The point of interest container being populated. + - current_subreg (NII): The current (cropped) subregion segmentation for this vertebra. + - location (Location): The target ``Location`` this strategy is computing. + - vert_id (int): The vertebra ID currently being processed. + - bb: The bounding box used to crop the vertebra. - log (Logger_Interface, optional): The logger interface. Defaults to _log, which should be defined globally. Example: @@ -297,7 +307,9 @@ def __init__(self, target: Location, *prerequisite: Location, **args) -> None: Strategy_Computed_Before(L.Spinal_Cord,L.Vertebra_Disc,L.Vertebra_Corpus,L.Dens_axis) Strategy_Computed_Before(L.Spinal_Canal,L.Vertebra_Corpus) Strategy_Computed_Before(L.Vertebra_Disc_Inferior,L.Vertebra_Disc_Inferior) - +Strategy_Computed_Before(L.Vertebral_Body_Endplate_Superior,L.Vertebra_Corpus) +Strategy_Computed_Before(L.Vertebral_Body_Endplate_Inferior,L.Vertebra_Corpus) +Strategy_Computed_Before(L.Endplate,L.Vertebra_Corpus) # fmt: on def compute_non_centroid_pois( # noqa: C901 @@ -314,10 +326,14 @@ def compute_non_centroid_pois( # noqa: C901 Runs the full non-centroid POI pipeline: + 0. Vertebral body endplates — computed via + :func:`~TPTBox.spine.spinestats.poi_fun.endplates.calc_endplate_points_` + when any of ``Vertebral_Body_Endplate_Inferior``, + ``Vertebral_Body_Endplate_Superior`` or ``Endplate`` is requested. 1. Vertebra orientation (PIR direction vectors) — always computed first if ``Location.Vertebra_Direction_Inferior`` is requested. - 2. Global landmarks: spinal canal / cord centres, intervertebral disc - POIs, dense-axis tip. + 2. Global landmarks: spinal canal / cord centres, spinal canal at IVD + level, dens-axis tip, and articular process midpoints (left/right). 3. Per-vertebra landmarks via the registered :class:`Strategy_Pattern` functions (extreme points, ray casts, corner finders, etc.). 4. Intervertebral disc (IVD) POIs. @@ -342,14 +358,21 @@ def compute_non_centroid_pois( # noqa: C901 _vert_ids = vert.unique() locations = list(locations) if isinstance(locations, Sequence) else [locations] - ### STEP 1 Vert Direction### - assert 52 not in poi.keys_region() + ### Step 0 Endplates ### + endplate = [Location.Vertebral_Body_Endplate_Inferior, Location.Vertebral_Body_Endplate_Superior, Location.Endplate] + if any(i in locations for i in endplate): + [locations.remove(i) for i in endplate if i in locations] + 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 + poi, *_ = calc_endplate_points_(poi, vert, subreg, _vert_ids=_vert_ids, log=log) + ### STEP 1 Vert Direction### if Location.Vertebra_Direction_Inferior in locations: log.on_text("Compute Vertebra DIRECTIONS", verbose=verbose) ### Calc vertebra direction; We always need them, so we just compute them. ### sub_regions = poi.keys_subregion() - if any(a.value not in sub_regions for a in vert_directions): + if any(a.value not in sub_regions for a in vert_directions): # skip if all exists poi, _ = calc_orientation_of_vertebra_PIR( poi, vert, subreg, do_fill_back=False, save_normals_in_info=False, _orientation_version=_orientation_version ) @@ -399,6 +422,11 @@ def compute_non_centroid_pois( # noqa: C901 poi = calc_center_spinal_cord( poi, subreg, source_subreg_point_id=Location.Vertebra_Disc, subreg_id=Location.Spinal_Canal_ivd_lvl, add_dense=True ) + if any(i in locations for i in [Location.Articular_Process_Midpoint_Left, Location.Articular_Process_Midpoint_Right]): + from TPTBox.spine.spinestats.poi_fun.articularis_midpoint import calc_all_facet_joint_pois + + p = calc_all_facet_joint_pois(vert, subreg) + poi.join_left_(p) # Step 3 Compute on individual Vertebras ivd_location = set() @@ -406,7 +434,7 @@ def compute_non_centroid_pois( # noqa: C901 if vert_id >= 39: continue current_vert = vert.extract_label(vert_id) - bb = current_vert.compute_crop() + bb = current_vert.compute_crop(raise_error=False) current_vert.apply_crop_(bb) current_subreg = subreg.apply_crop(bb) * current_vert for location in locations: diff --git a/TPTBox/core/vert_constants.py b/TPTBox/core/vert_constants.py index 4ee2afae..86cf5c30 100755 --- a/TPTBox/core/vert_constants.py +++ b/TPTBox/core/vert_constants.py @@ -264,7 +264,7 @@ class Full_Body_Instance_Vibe(Abstract_lvl): sternum = 63 costal_cartilages = 64 subcutaneous_fat = 65 - muscle = 66 + muscle_other = 66 inner_fat = 67 IVD = 68 vertebra_body = 69 @@ -272,6 +272,22 @@ class Full_Body_Instance_Vibe(Abstract_lvl): spinal_channel = 71 bone_other = 72 + @classmethod + def muscle_(cls) -> list[Full_Body_Instance_Vibe]: + """Return individually segmented muscle group instance labels.""" + return [ + Full_Body_Instance_Vibe.gluteus_maximus_right, + Full_Body_Instance_Vibe.gluteus_maximus_left, + Full_Body_Instance_Vibe.gluteus_medius_right, + Full_Body_Instance_Vibe.gluteus_medius_left, + Full_Body_Instance_Vibe.gluteus_minimus_right, + Full_Body_Instance_Vibe.gluteus_minimus_left, + Full_Body_Instance_Vibe.autochthon_right, + Full_Body_Instance_Vibe.autochthon_left, + Full_Body_Instance_Vibe.iliopsoas_right, + Full_Body_Instance_Vibe.iliopsoas_left, + ] + @classmethod def get_Full_Body_Instance_mapping(cls) -> dict[int, int]: """Return a mapping from ``Full_Body_Instance`` label values to ``Full_Body_Instance_Vibe`` label values.""" @@ -326,10 +342,9 @@ def get_Full_Body_Instance_mapping(cls) -> dict[int, int]: Full_Body_Instance.clavicula_right.value: cls.clavicula_right.value, # clavicula_right Full_Body_Instance.femur_left.value: cls.femur_left.value, # femur_left Full_Body_Instance.femur_right.value: cls.femur_right.value, # femur_right - # `hip_left`/`hip_right` do not exist on this enum; the labels are called - # pelvis_left/pelvis_right (50/51), matching `50: pelvis_left` in the inverse map. Full_Body_Instance.pelvis_left.value: cls.pelvis_left.value, # hip_left Full_Body_Instance.pelvis_right.value: cls.pelvis_right.value, # hip_right + Full_Body_Instance.channel.value: cls.spinal_cord.value, # spinal_cord Full_Body_Instance.gluteus_maximus_left.value: cls.gluteus_maximus_left.value, # gluteus_maximus_left Full_Body_Instance.gluteus_maximus_right.value: cls.gluteus_maximus_right.value, # gluteus_maximus_right Full_Body_Instance.gluteus_medius_left.value: cls.gluteus_medius_left.value, # gluteus_medius_left @@ -343,7 +358,7 @@ def get_Full_Body_Instance_mapping(cls) -> dict[int, int]: Full_Body_Instance.sternum.value: cls.sternum.value, # sternum Full_Body_Instance.costal_cartilage.value: cls.costal_cartilages.value, # costal_cartilages Full_Body_Instance.subcutaneous_fat.value: cls.subcutaneous_fat.value, # subcutaneous_fat - Full_Body_Instance.muscle_other.value: cls.muscle.value, # muscle + Full_Body_Instance.muscle_other.value: cls.muscle_other.value, # muscle Full_Body_Instance.inner_fat.value: cls.inner_fat.value, # inner_fat Full_Body_Instance.ivd.value: cls.IVD.value, # IVD Full_Body_Instance.vert_body.value: cls.vertebra_body.value, # vertebra_body @@ -1252,6 +1267,9 @@ def save_as_name(cls) -> bool: Implant_Target_Left = 92 Implant_Target_Right = 93 + Articular_Process_Midpoint_Left = 94 + Articular_Process_Midpoint_Right = 95 + # Muscle_Inserts_Rib_left = 90 # Muscle_Inserts_Rib_right = 91 # Ligament attachment points diff --git a/TPTBox/segmentation/VibeSeg/vibeseg.py b/TPTBox/segmentation/VibeSeg/vibeseg.py index 04d46156..bbe49aa4 100644 --- a/TPTBox/segmentation/VibeSeg/vibeseg.py +++ b/TPTBox/segmentation/VibeSeg/vibeseg.py @@ -98,6 +98,8 @@ def run_vibeseg( dataset_id: int = 100, padd: int = 5, keep_size: bool = False, + memory_max=9900000, # in MB + model_path=None, **args, ) -> NII: """Run the VibeSeg whole-body segmentation model on a single image. @@ -112,6 +114,7 @@ def run_vibeseg( padd: Number of voxels to pad the image before inference. keep_size: If True, keep the model's native output resolution instead of resampling back to the input image space. + memory_max: MAX GPU memory in MB. Changes the super-batches are used. Might speed up inference. At least 8000 **args: Additional keyword arguments forwarded to ``run_inference_on_file``. Returns: @@ -123,15 +126,17 @@ def run_vibeseg( args[k] = v return run_inference_on_file( dataset_id, - [to_nii(i)] if not isinstance(i, (list, tuple)) else [to_nii(j) for j in i], + [to_nii(i)] if not isinstance(i, list | tuple) else [to_nii(j) for j in i], out_file=out_seg, override=override, gpu=gpu, ddevice=ddevice, padd=padd, keep_size=keep_size, + memory_max=memory_max, + model_path=model_path, **args, - )[0] + )[0] # type: ignore def run_nnunet( diff --git a/TPTBox/segmentation/nnUnet_utils/predictor.py b/TPTBox/segmentation/nnUnet_utils/predictor.py index f61a317a..c760b084 100755 --- a/TPTBox/segmentation/nnUnet_utils/predictor.py +++ b/TPTBox/segmentation/nnUnet_utils/predictor.py @@ -3,12 +3,15 @@ # method for deep learning-based biomedical image segmentation. Nature methods, 18(2), 203-211. from __future__ import annotations +import itertools import os import time import traceback from collections.abc import Generator from dataclasses import dataclass, field from math import ceil, floor +from queue import Queue +from threading import Thread import numpy as np import torch @@ -550,24 +553,14 @@ def _internal_maybe_mirror_and_predict(self, x: torch.Tensor, network) -> torch. if mirror_axes is not None: # check for invalid numbers in mirror_axes # x should be 5d for 3d images and 4d for 2d. so the max value of mirror_axes cannot exceed len(x.shape) - 3 - assert max(mirror_axes) <= len(x.shape) - 3, "mirror_axes does not match the dimension of the input!" - - num_predictons = 2 ** len(mirror_axes) - if 0 in mirror_axes: - prediction += torch.flip(network(torch.flip(x, (2,))), (2,)) - if 1 in mirror_axes: - prediction += torch.flip(network(torch.flip(x, (3,))), (3,)) - if 2 in mirror_axes: - prediction += torch.flip(network(torch.flip(x, (4,))), (4,)) - if 0 in mirror_axes and 1 in mirror_axes: - prediction += torch.flip(network(torch.flip(x, (2, 3))), (2, 3)) - if 0 in mirror_axes and 2 in mirror_axes: - prediction += torch.flip(network(torch.flip(x, (2, 4))), (2, 4)) - if 1 in mirror_axes and 2 in mirror_axes: - prediction += torch.flip(network(torch.flip(x, (3, 4))), (3, 4)) - if 0 in mirror_axes and 1 in mirror_axes and 2 in mirror_axes: - prediction += torch.flip(network(torch.flip(x, (2, 3, 4))), (2, 3, 4)) - prediction /= num_predictons + assert max(mirror_axes) <= x.ndim - 3, "mirror_axes does not match the dimension of the input!" + + mirror_axes = [m + 2 for m in mirror_axes] + axes_combinations = [c for i in range(len(mirror_axes)) for c in itertools.combinations(mirror_axes, i + 1)] + for axes in axes_combinations: + prediction += torch.flip(self.network(torch.flip(x, axes)), axes) + prediction /= len(axes_combinations) + 1 + return prediction def predict_sliding_window_return_logits( @@ -787,30 +780,56 @@ def _allocate(self, data: torch.Tensor, results_device, pbar: tqdm, gauss: bool return predicted_logits, n_predictions, gaussian, results_device def _run_sub(self, data: torch.Tensor, network, results_device, slicers, pbar: tqdm, addendum: str = "", logger=logger): - """Iterate over slicers, run inference per tile (optionally batched), and accumulate results.""" + """Iterate over slicers, run inference in batches while asynchronously preparing the next batch.""" slicers = list(slicers) + + def producer(d, slicers, batch_size, q): + for batch_start in range(0, len(slicers), batch_size): + batch_slicers = slicers[batch_start : batch_start + batch_size] + + if batch_size == 1: + work_on = torch.clone(d[batch_slicers[0]][None], memory_format=torch.contiguous_format) + else: + work_on = torch.stack([torch.clone(d[sl], memory_format=torch.contiguous_format) for sl in batch_slicers], dim=0) + q.put((work_on.to(self.device, non_blocking=False), batch_slicers)) + q.put("end") + try: - data = data.to(self.device) # type: ignore + batch_size = max(1, self.tile_batch_size) + data = data.to(results_device) predicted_logits, n_predictions, gaussian, results_device = self._allocate(data, results_device, pbar, logger=logger) + pbar.desc = f"running prediction {addendum}" + queue = Queue(maxsize=2) + t = Thread(target=producer, args=(data, slicers, batch_size, queue), daemon=True) + t.start() prediction = None work_on = None - batch_size = max(1, self.tile_batch_size) - for batch_start in range(0, len(slicers), batch_size): - batch_slicers = slicers[batch_start : batch_start + batch_size] - # batch_size == 1 keeps the original view (no copy); larger batches stack tiles into a - # dense (B, C, *patch) tensor (valid because all tiles share the same patch_size). - work_on = data[batch_slicers[0]][None] if batch_size == 1 else torch.stack([data[sl] for sl in batch_slicers], dim=0) - work_on = work_on.to(self.device, non_blocking=False) + + while True: + item = queue.get() + if item == "end": + queue.task_done() + break + work_on, batch_slicers = item prediction = self._internal_maybe_mirror_and_predict(work_on, network=network).to(results_device) + for b, sl in enumerate(batch_slicers): - pbar.update(1) pred = prediction[b] if pred.shape[0] != predicted_logits.shape[0]: pred = pred.squeeze(0) - predicted_logits[sl] += pred * gaussian if self.use_gaussian else pred - n_predictions[sl[1:]] += gaussian if self.use_gaussian else 1 + if self.use_gaussian: + predicted_logits[sl] += pred * gaussian + n_predictions[sl[1:]] += gaussian + else: + predicted_logits[sl] += pred + n_predictions[sl[1:]] += 1 + pbar.update(1) + queue.task_done() + queue.join() + return predicted_logits, n_predictions # noqa: TRY300 + except RuntimeError: try: del predicted_logits @@ -820,6 +839,7 @@ def _run_sub(self, data: torch.Tensor, network, results_device, slicers, pbar: t del prediction except UnboundLocalError: pass + empty_cache(self.device) empty_cache(results_device) self.memory_base += 1000 diff --git a/TPTBox/segmentation/spineps.py b/TPTBox/segmentation/spineps.py index b7dd246a..dca938fd 100644 --- a/TPTBox/segmentation/spineps.py +++ b/TPTBox/segmentation/spineps.py @@ -64,10 +64,10 @@ def get_outpaths_spineps( def run_spineps( file_path: str | Path | BIDS_FILE, dataset: str | Path | None = None, - model_semantic: str | Path = "t2w", - model_instance: str | Path = "instance", - model_labeling: str | None = "t2w_labeling", - derivative_name: str = "derivative", + model_semantic: str | Path = "t2w", # t2w, vibe, ct + model_instance: str | Path = "instance", # instance, ct_instance + model_labeling: str | None = "t2w_labeling", # t2w_labeling, ct_labeling + derivative_name: str = "derivatives", override_semantic: bool = False, override_instance: bool = False, lambda_semantic=None, diff --git a/TPTBox/spine/snapshot2D/snapshot_modular.py b/TPTBox/spine/snapshot2D/snapshot_modular.py index c5618e50..dae553b1 100755 --- a/TPTBox/spine/snapshot2D/snapshot_modular.py +++ b/TPTBox/spine/snapshot2D/snapshot_modular.py @@ -717,6 +717,8 @@ def plot_sag_centroids( elif len(x) == 2: (text, a) = x b = zms[0] * ctd[color, curve_location][0] + if isinstance(color, int): + color = get_color_by_label(color).rgb / 255 # noqa: PLW2901 axs.text( a, b, diff --git a/TPTBox/spine/spinestats/README.md b/TPTBox/spine/spinestats/README.md index fce970e1..de4b078e 100644 --- a/TPTBox/spine/spinestats/README.md +++ b/TPTBox/spine/spinestats/README.md @@ -1,27 +1,379 @@ # Spine Statistics (`spine/spinestats`) -Clinical spine measurements computed from `POI` objects and `NII` segmentations. +Clinical spine and body-composition measurements computed from `POI` +objects and `NII` segmentations. ## Modules | Module | Description | |---|---| -| `distances.py` | Distances between anatomical landmarks (IVD height, canal diameter, …) | -| `angles.py` | Cobb angle and other spine curvature measurements | -| `ivd_pois.py` | Intervertebral disc (IVD) Point of Interest computation | +| `angles.py` | Cobb angle, cervical lordosis, thoracic kyphosis, lumbar lordosis | +| `measure_ivd_and_vertebra_geometry.py` | Per-structure geometry (heights, widths, x1–x6) and T2 signal ratio for vertebrae and IVDs | +| `torso_vat_sat.py` | VBQ score, body composition CSA, muscle fat infiltration, torso VAT/SAT/muscle volumes; also `peak_centered_mean` | +| `vertebra_anatomical_widths.py` | Anatomical distances per vertebra (IVD height, body height, LR/AP widths) stored on `POI.info` | | `body_quadrants.py` | Subdivides vertebra bodies into anatomical quadrants | -| `make_endplate.py` | Generates superior/inferior endplate surfaces from segmentations | +| `_run_all.py` | End-to-end NAKO pipeline: resolves paths, runs all analyses, writes one json per subject, streams Excel summaries | +| `poi_fun/` | Points-of-interest sub-package used by the geometry code | ## Key functions | Function | Module | Description | |---|---|---| -| `compute_cobb_angle(poi)` | `angles.py` | Cobb angle between vertebra pairs from POI coordinates | -| `compute_ivd_height(poi)` | `distances.py` | Inter-vertebral disc height at a given level | -| `calc_ivd_pois(vert_nii, subreg_nii)` | `ivd_pois.py` | Compute IVD POIs from vertebra + subregion segmentations | -| `make_endplate(nii, poi, ...)` | `make_endplate.py` | Fit an endplate surface to a vertebra body | +| `compute_max_cobb_angle` / `compute_max_cobb_angle_multi` | `angles.py` | Maximum Cobb angle (single value or list of scoliotic segments) | +| `compute_lordosis_and_kyphosis` | `angles.py` | Cervical / thoracic / lumbar curvature angles | +| `plot_cobb_and_lordosis_and_kyphosis` | `angles.py` | Combined computation + snapshot | +| `measure_ivd_and_vertebra_geometry` | `measure_ivd_and_vertebra_geometry.py` | Geometry + signal per label | +| `VBQ_score` | `torso_vat_sat.py` | Vertebral Bone Quality score (T2 vertebra / T2 CSF) | +| `body_composition_score` | `torso_vat_sat.py` | Per-level axial CSA of muscle / VAT / SAT / psoas / autochthon | +| `muscle_fat_infiltration` | `torso_vat_sat.py` | Dixon fat-fraction based muscle-quality metrics | +| `torso_vat_sat_muscle_mass` | `torso_vat_sat.py` | Whole-torso VAT / SAT / muscle volumes inside an ROI | +| `peak_centered_mean` | `torso_vat_sat.py` | Robust mean around the histogram peak (used to suppress non-CSF voxels) | +| `compute_all_distances` | `vertebra_anatomical_widths.py` | Compute IVD/vertebra distances and store them on `POI.info` | +| `run_all` | `_run_all.py` | Full pipeline: geometry + signal + composition + torso volumes, writes one json | +| `ExcelCollector` | `_run_all.py` | Background process that turns per-subject jsons into rolling Excel summaries | ## Coordinate convention -All measurement functions consume `POI` objects (voxel or world space) produced by -`calc_centroids` or `calc_poi_from_subreg_vert` from the `core` module. +All measurement functions consume `POI` objects (voxel or world space) +produced by `calc_centroids` or `calc_poi_from_subreg_vert` from the +`core` module. + +--- + +# Pipeline output reference (`run_all`) + +This section documents every key produced by `run_all(file_dict)` in +`_run_all.py`, together with its unit and important implementation +details. It is aimed at radiologists reviewing the numbers, so it +focuses on "what does this mean clinically" and "how was it computed", +not on the Python API. A standalone copy of this reference lives at +`all_output_reference.md` in the same folder. + +## How the pipeline is organised + +`run_all` writes a single json per subject with these top-level keys: + +| Key | Source function | What it covers | +|---|---|---| +| `ivd_geometry` | `measure_ivd_and_vertebra_geometry(..., structure_label=100)` | intervertebral discs | +| `vert_geometry` | `measure_ivd_and_vertebra_geometry(..., structure_label=50)` | vertebral bodies | +| `VBQ_score` | `VBQ_score` | vertebral bone quality (T2 signal ratio) | +| `body_composition_score` | `body_composition_score` | axial CSA per tissue at chosen vertebral levels | +| `muscle_fat_infiltration` | `muscle_fat_infiltration` | Dixon fat-fraction based muscle-quality metrics | +| `torso_vat_sat_muscle_mass` | `torso_vat_sat_muscle_mass` | whole-torso VAT / SAT / muscle volume | +| `cobb`, `curv` | `plot_cobb_and_lordosis_and_kyphosis` | only when called with `cobb=True` | + +Distance metrics from `vertebra_anatomical_widths.compute_all_distances` +are not currently written into the json by `run_all`; they live on the +returned `POI.info` dict (see the `vertebra_anatomical_widths.py` +section below). + +Angles are in **degrees**, lengths in **millimetres**, areas in **mm²**, +volumes in **mm³**, fat fractions are **unitless** in `[0, 1]`, MR +signal values are in **arbitrary units (a.u.)** and only meaningful as +ratios. + +Caching: `run_all(..., override=False)` (the default) reuses the json +when it exists, is newer than every input segmentation file, and +contains all of the required top-level keys. When some (but not all) +required keys are missing, the existing json is loaded and only the +missing top-level keys are recomputed; NII inputs that are not needed +for any missing key are skipped so partial reruns are cheap. Pass +`override=True` to force recomputation of every key. + +### Input requirements + +`run_all` assumes whole-body-style acquisitions: + +- **T2w image** — must cover the **full spine** (cervical through + sacrum). Curvature angles and per-vertebra geometry silently return + `None`/`NaN` for any level that is cropped away, and the VBQ ranges + (`C3-C6`, `T5-T8`, `L1-L4`) need every vertebra in the range to be + visible. +- **VIBE water/fat images** — must cover the **full torso**. Fat + fraction and muscle CSA are computed on whatever axial slices are + present, so a cropped VIBE silently biases per-region CSA and IMAT + volumes. +- **Segmentations** (`vert`, `spine`, `vibeseg100`, `roi`) — must + match the extent of their underlying image. In particular, + `torso_vat_sat_muscle_mass` explicitly verifies that both the + clavicula and the pelvis are present in the VIBESeg mask; if either + is missing it aborts, returns `NaN` volumes and stores the reason in + the `reason` key. Partially-covered scans should either be skipped + or handled outside this pipeline. + +### How to produce the segmentations + +All required segmentations can be produced from +`TPTBox.segmentation`: + +- **`vert` / `spine`** — run **SPINEPS** on the T2w image. Import from + `TPTBox.segmentation` (`run_spineps`, `get_outpaths_spineps`, + `_run_spineps_all`). SPINEPS returns both the per-vertebra instance + segmentation and the spine subregion segmentation used by every + spine-side function in this package. +- **`vibeseg100`** — run **VIBESegmentator** with + `run_vibeseg(..., dataset_id=100)` on the VIBE stack. Dataset **100** + is the general MR/CT body-composition model that `run_all` targets. + Dataset **12** is the 0.8 mm iso CT model and is also supported by + `body_composition_score`, `muscle_fat_infiltration` and + `torso_vat_sat_muscle_mass` (pass `dataset_id=12`). +- **`roi`** — run VIBESegmentator with dataset **278** on the VIBE + stack. Note: the raw dataset-278 ROI is **not perfect** and needs + postprocessing before it is fed into `run_all`; without cleanup the + torso extent used to gate VAT/SAT/muscle volumes and the per-region + muscle statistics will be off. + +## Signal-based conventions used everywhere + +Two things are worth understanding before reading the T2 signal keys: + +1. **Peak-centered mean.** Ordinary mean signal inside a mask is + sensitive to non-CSF voxels that leak into the spinal canal + segmentation (nerve roots, vessel walls). The pipeline instead + averages only voxels whose intensity falls in a window around the + histogram peak. When both a peak-centered and a plain-mean version + are stored, the plain-mean version is suffixed with `_old` for + comparison. See `peak_centered_mean` in `torso_vat_sat.py`. +2. **Erosion.** Muscle and vertebral-body masks are eroded by one or + two voxels before signal extraction to reduce partial-volume mixing + at the boundary. Volume metrics are reported both after erosion + (`*_volume_*`) and before erosion (`*_volume_no_erosion_*`) so the + effect of the erosion is auditable. + +--- + +## `ivd_geometry` and `vert_geometry` + +Both keys hold `dict[label_id, dict[metric_name, value]]`, where each +`label_id` is one intervertebral disc (`ivd_geometry`) or vertebra +(`vert_geometry`). If a label fails to evaluate, its entry contains +`error` (message string) plus all metrics set to `NaN`. + +| Key | Unit | Meaning | +|---|---|---| +| `volume_voxel` | mm³ | volume counted from the raw voxel mask | +| `volume_mesh` | mm³ | volume of the reconstructed surface mesh (same structure) | +| `height_center` | mm | height sampled through the structure's centroid | +| `mean_height` | mm | mean of the sampled heights over the structure surface | +| `max_height` | mm | maximum of the sampled heights | +| `lower_10_percent_height` | mm | 10th percentile of the sampled heights | +| `mean_diameter` | mm | diameter of the circle whose area equals the projected area | +| `anterior_height_x1` | mm | anterior height at the anterior point (x1) | +| `posterior_height_x2` | mm | posterior height at the posterior point (x2) | +| `right_height_x3` | mm | right-lateral height (x3) | +| `left_height_x4` | mm | left-lateral height (x4) | +| `width_lateral_x5` | mm | lateral width (x5) | +| `width_sagittal_x6` | mm | sagittal width (x6) | +| `signal` | unitless | peak-centered structure T2 signal / peak-centered spinal-canal T2 signal | +| `structure_signal` | a.u. | peak-centered T2 signal inside the eroded structure mask | +| `spinal_canal_signal` | a.u. | peak-centered T2 signal inside the eroded spinal canal reference | +| `signal_old` | unitless | same ratio computed with plain per-voxel means | +| `structure_signal_old` | a.u. | plain mean T2 signal inside the eroded structure mask | +| `spinal_canal_signal_old` | a.u. | plain mean T2 signal inside the eroded spinal canal | + +Implementation notes: +- Structure orientation for IVDs is estimated from the disc's own voxel + mask via PCA; vertebral orientation is read from precomputed POIs + (except C2/dens, which falls back to PCA). +- x1–x6 are the six clinically standard directional heights/widths (see + the geometry module docstring for the figure). +- Only labels present in the segmentation appear as keys. + +## `VBQ_score` + +`dict[str, float]`; one triple of entries per configured spinal range. +Default ranges are `C3-C6`, `T5-T8`, `L1-L4`. + +| Key template | Unit | Meaning | +|---|---|---| +| `mean_signal_vertebra_-` | a.u. | mean T2 signal inside the eroded vertebral body mask over the range | +| `mean_signal_liquor_-` | a.u. | **peak-centered** mean T2 signal inside the spinal canal over the same S/I extent | +| `mean_signal_liquor_-_old` | a.u. | plain-mean version, kept for backward comparison | +| `VBQ_-` | unitless | vertebral signal divided by the peak-centered CSF signal | +| `VBQ_-_old` | unitless | same ratio using the plain-mean CSF signal | + +Implementation notes: +- Vertebral body mask is eroded (default `n_erode=2`) to avoid the + cortical rim. +- The spinal canal is cropped to the same superior–inferior slab as the + vertebral bodies so the CSF reference matches the region of interest. +- Higher VBQ = darker vertebral bodies relative to CSF, associated in + the literature with lower bone quality. + +## `body_composition_score` + +`dict[str, float]`; per-region axial cross-sectional-area statistics of +five tissue classes. Default regions are `T12-L1` and `L3-L3`. + +Region tag: `{start.name}-{goal.name}` (e.g. `T12-L1`, `L3-L3`). + +| Key template | Unit | Meaning | +|---|---|---| +| `mean_muscle_area_{region}` | mm² | mean skeletal muscle CSA across the region | +| `max_muscle_area_{region}` | mm² | maximum skeletal muscle CSA in the region | +| `mean_VAT_area_{region}` | mm² | mean visceral adipose tissue CSA | +| `max_VAT_area_{region}` | mm² | maximum visceral adipose tissue CSA | +| `mean_SAT_area_{region}` | mm² | mean subcutaneous adipose tissue CSA | +| `max_SAT_area_{region}` | mm² | maximum subcutaneous adipose tissue CSA | +| `mean_psoas_area_{region}` | mm² | mean psoas CSA (left + right) | +| `max_psoas_area_{region}` | mm² | maximum psoas CSA | +| `mean_autochthon_area_{region}` | mm² | mean autochthonous back-muscle CSA | +| `max_autochthon_area_{region}` | mm² | maximum autochthonous back-muscle CSA | +| `n_slices_{region}` | count | number of axial slices contributing to the muscle statistic | +| `muscle_index_{region}` | mm²/m² | `mean_muscle_area / height_m²`; only present when `height_m` is supplied | +| `muscle_fat_ratio_{region}` | unitless | `mean_muscle_area / (mean_VAT_area + mean_SAT_area)`; `NaN` if the denominator is zero | + +Implementation notes: +- The superior–inferior extent of the vertebral bodies inside the region + defines the slice range. +- Axial voxel area is derived from the VIBE geometry; slices with zero + tissue are excluded before mean/max. +- If no vertebral body voxels are present for a region, that region is + silently skipped (no keys emitted). + +## `muscle_fat_infiltration` + +`dict[str, float]`; per (region, muscle group) Dixon-based fat +infiltration metrics. Muscle groups (dataset_id=100) are: +`all_muscle`, `iliopsoas_left`, `iliopsoas_right`, `autochthon_left`, +`autochthon_right`, `muscle_other`. Suffix is `{region}_{muscle}`; when +no region is given, the region tag is `all`. + +Fat fraction (FF) is computed voxel-wise as `FF = fat / (fat + water)`; +voxels with `FF >= threshold` (default 0.20) are IMAT, otherwise lean. + +| Key template | Unit | Meaning | +|---|---|---| +| `mean_fat_fraction_{suffix}` | [0, 1] | mean FF over the eroded muscle mask | +| `median_fat_fraction_{suffix}` | [0, 1] | median FF over the eroded muscle mask | +| `mean_lean_fat_fraction_{suffix}` | [0, 1] | mean FF of lean voxels (FF < threshold) | +| `mean_IMAT_fat_fraction_{suffix}` | [0, 1] | mean FF of IMAT voxels (FF ≥ threshold) | +| `muscle_volume_{suffix}` | mm³ | muscle volume after erosion | +| `muscle_volume_no_erosion_{suffix}` | mm³ | muscle volume **before** erosion (raw segmentation volume) | +| `lean_muscle_volume_{suffix}` | mm³ | lean-muscle volume within the eroded mask | +| `lean_muscle_volume_no_erosion_{suffix}` | mm³ | lean-muscle volume within the un-eroded mask | +| `IMAT_volume_{suffix}` | mm³ | IMAT volume within the eroded mask | +| `IMAT_volume_no_erosion_{suffix}` | mm³ | IMAT volume within the un-eroded mask | +| `IMAT_fraction_{suffix}` | [0, 1] | IMAT voxel fraction within the eroded mask | + +Implementation notes: +- Erosion iterations per muscle are configurable via the `erode` dict. + Defaults: `all_muscle=1`, `iliopsoas_*=1`, `autochthon_*=2`, + `muscle_other=1`. +- Fat-fraction statistics use the eroded mask; volumes are also + reported for the un-eroded mask so the caller can inspect the effect + of erosion. +- When `regions` are supplied, the analysis is restricted to the + superior–inferior extent of the vertebral bodies in each range and + the region tag becomes `{start.name}-{goal.name}`. + +## `torso_vat_sat_muscle_mass` + +`dict[str, float]`; whole-torso volumes restricted to the supplied ROI. +Only the results dict is stored in the json (the optional NII output is +dropped by `run_all` because it is not JSON-serializable). + +| Key | Unit | Meaning | +|---|---|---| +| `VAT` | mm³ | visceral adipose tissue volume inside the ROI | +| `SAT` | mm³ | subcutaneous adipose tissue volume inside the ROI | +| `muscle_mass` | mm³ | skeletal muscle volume inside the ROI | +| `reason` | string | present only if the computation failed; explains why | + +Implementation notes: +- The function checks that both the clavicula and the pelvis are + present in the segmentation to make sure the full torso is covered. + If either check fails, all three volumes are set to `NaN` and + `reason` is populated. +- ROI labels (default 3–8) select which sub-regions of the torso count. + +## `cobb` and `curv` (only when `cobb=True`) + +- `cobb`: `list[tuple[float, int, int, int | None]]` from + `compute_max_cobb_angle_multi` — one entry per detected scoliotic + segment: `(max_angle_deg, from_vertebra_id, to_vertebra_id, apex_id_or_none)`. + Angles are in **degrees**. +- `curv`: dict from `compute_lordosis_and_kyphosis`: + - `cervical_lordosis` (deg) — computed between C2 and C7 + - `thoracic_kyphosis` (deg) — computed between T4 and the last thoracic vertebra + - `lumbar_lordosis` (deg) — computed between L1 and the last lumbar vertebra + + Values can be `None` if the required vertebrae are missing from the + POI. + +--- + +## `vertebra_anatomical_widths.py` + +Not written into the json by `run_all`, but part of this package. +`compute_all_distances(poi, vert=..., subreg=...)` fills +``poi.info[key]`` for each of the four registered distances. Each entry +is a dict `{vertebra_region_id: distance_mm}`. + +| `poi.info` key | Unit | Endpoints (`Location`) | Meaning | +|---|---|---|---| +| `ivd_heights_center_mm` | mm | `Vertebra_Disc_Inferior` → `Vertebra_Disc_Superior` | IVD height at the disc centre | +| `vertebra_heights_center_mm` | mm | `Additional_Vertebral_Body_Middle_Superior_Median` → `Additional_Vertebral_Body_Middle_Inferior_Median` | Vertebral body height through the mid-body | +| `vertebra_width_LR_center_mm` | mm | `Muscle_Inserts_Vertebral_Body_Right` → `Muscle_Inserts_Vertebral_Body_Left` | Left–right (lateral) vertebral body width | +| `vertebra_width_AP_center_mm` | mm | `Additional_Vertebral_Body_Posterior_Central_Median` → `Additional_Vertebral_Body_Anterior_Central_Median` | Anterior–posterior (sagittal) vertebral body width | + +Implementation notes: +- `_compute_distance` short-circuits when the key already exists in + ``poi.info`` unless ``recompute=True`` is passed. +- POIs for the required endpoints are computed on demand via + ``calc_poi_from_subreg_vert(vert, subreg, ...)`` when + ``all_pois_computed=False`` and the endpoint locations are not yet in + the POI object. +- Distances are Euclidean in mm regardless of the input POI zoom + (`keep_zoom=False`). + +--- + +## Excel collector + +`ExcelCollector` in `_run_all.py` runs a background process that turns +each finished json into two rolling Excel files in a configurable +folder: + +- `per_subject.xlsx` — one row per subject with every scalar top-level + metric flattened to dotted keys + (e.g. `VBQ_score.VBQ_L1-L4`, `torso_vat_sat_muscle_mass.VAT`). +- `per_vertebra.xlsx` — one row per (subject, label), populated from + `vert_geometry` and `ivd_geometry`. The `source` column indicates + which of the two sections the row came from. + +Usage: + +```python +collector = ExcelCollector(out_folder="/tmp/nako_summary") +collector.start() +for nako_id in ids: + f = get_nako_paths(nako_id) + run_all(f) # writes the per-subject json + collector.submit(nako_id, _final_json_path(f)) +collector.close() # flushes and joins +``` + +The collector re-writes the Excel files every `flush_every` submissions +(default 25) and once more at shutdown, so partial runs still produce +usable summaries. + +## Batch entry point (`python -m ..._run_all`) + +Running the module directly loops over the NAKO cohort via +`loop_over_repaired_nako`, calls `run_all` per subject and streams the +results through `ExcelCollector`. Two knobs are exposed at the top of +the `__main__` block: + +- `N_CPUS` — set `>1` to run subjects in parallel through a + `ProcessPoolExecutor`; `1` keeps the sequential path. +- `OVERRIDE` — forwarded to `run_all` (see the caching note above). + +Before running, each subject is checked against `REQUIRED_INPUT_KEYS` +(ordered: `t2w`, `vert`, `spine`, `vibeseg100`, `roi`, +`vibe_part-water`, `vibe_part-fat`). Subjects with at least one missing +input are skipped and recorded in `missing_inputs.xlsx` under the +output folder, attributed to the **first** missing key in that order — +so a subject with several gaps still counts once. Subjects whose +`run_all` raises are also logged there with `error:`. diff --git a/TPTBox/spine/spinestats/__init__.py b/TPTBox/spine/spinestats/__init__.py index 785e1b5a..553b7a38 100644 --- a/TPTBox/spine/spinestats/__init__.py +++ b/TPTBox/spine/spinestats/__init__.py @@ -9,5 +9,6 @@ plot_compute_lordosis_and_kyphosis, ) from .body_quadrants import make_quadrants -from .ivd_pois import calculate_IVD_POI, calculate_pca_normal_np, compute_fake_ivd -from .make_endplate import endplate_extraction +from .measure_ivd_and_vertebra_geometry import measure_ivd_and_vertebra_geometry +from .poi_fun.ivd_pois import calculate_IVD_POI, calculate_pca_normal_np, compute_fake_ivd +from .torso_vat_sat import VBQ_score, torso_vat_sat_muscle_mass diff --git a/TPTBox/spine/spinestats/_load_nako.py b/TPTBox/spine/spinestats/_load_nako.py new file mode 100644 index 00000000..81b83645 --- /dev/null +++ b/TPTBox/spine/spinestats/_load_nako.py @@ -0,0 +1,324 @@ +import json +import os +from pathlib import Path + +import pandas as pd + +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 + +log = Print_Logger() + + +def _check(l: list[BIDS_FILE]): + """Pick the preferred BIDS file from a list of candidates. + + Prefers files with a ``rec`` entity (reconstruction variant, defaulting to ``"Hamilton"``). + If no such file is found, asserts that there is exactly one candidate and returns it. + + Args: + l: Candidate BIDS files sharing the same BIDS query key. + + Returns: + The chosen ``BIDS_FILE``. + """ + for i in l: + if i.get("rec", "Hamilton"): + return i + assert len(l) == 1, l + return l[0] + + +def get_corrected_mevibe(fam: BIDS_Family, compute_PDFF=True): # TODO return dict with literal + """Collect the six mevibe echo images plus fat/water/PDFF/PDWF for one subject family. + + The ``BIDS_Global_info`` used to build ``fam`` must include ``derivatives_mevibe`` as a + parent root, and its query key addendum must contain ``part`` and ``desc`` so the echo + images are addressable. If reconstructed fat/water images are present they are preferred + over the raw ones. When ``compute_PDFF`` is set and the reconstructed fat-fraction (PDFF) + or water-fraction (PDWF) maps are missing on disk, they are computed as + ``fat / (fat + water) * 1000`` (and the water equivalent), cast to the smallest int dtype, + and saved next to the reconstructed water image. + + Args: + fam: BIDS family for a single mevibe acquisition. + compute_PDFF: If True, generate and persist missing PDFF/PDWF maps. + + Returns: + Dict mapping mevibe part keys (``"eco0-opp1"`` … ``"eco5-arb1"``, ``"mevibe_part-fat"``) + to the chosen ``BIDS_FILE`` entries. + """ + # TODO figure out what to do when multiple present + # BIDS_GLOBAL_INFO needs to have "derivatives_mevibe" as an additional root + # additional keys must be part and desc + # 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: + 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"} + ) + pdwf = water.get_changed_bids( + "nii.gz", bids_format=water.bids_format, parent=water.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 pdff.exists(): + out["mevibe_part-fat"] = pdff + if pdff.exists(): + out["mevibe_part-fat"] = pdwf + # else: + # pdff = _check(fam["mevibe_part-fat-fraction_desc-reconstructed"]) + + return out + + +def get_current_best_T2w_seg(sub, black_list_t2w=None): + if black_list_t2w is None: + black_list_t2w = [ + # Head missing T2w + "106910", + "100470", + "105805", + "119399", # "Scoliosis, no head" + "125130", + ] + search_folders = [ + "derivatives_spine_vert_fixed", + "derivatives_spine_inference_combination162_148", + # "derivatives_spine_inference_combination", + # "derivatives_spine_inference_159_sacrumfix", + # "derivatives_spine_inference_148_preliminary", + # "derivatives_spine_inference_146_preliminary", # sub-128135_sequ-stitched_acq-sag_mod-T2w_seg-vert_msk.nii.gz + ] + sub = str(sub).split("_")[0].replace("sub-", "") + if sub in black_list_t2w: + return None, "", None + if sub in [ + # "100303", + # "109091", + "113612", + # "106991", + "102179", + # "102263", + "103730", + "103704", + "110618", + "123393", + "123222", + "124365", + "104249", + "104000", + ]: + search_folders = ["archive/derivatives_spine_inference_148_preliminary"] + for s in search_folders: + vert_T2w = f"/DATA/NAS/datasets_processed/NAKO/dataset-nako/{s}/{sub[:3]}/{sub}/T2w/sub-{sub}_sequ-stitched_acq-sag_mod-T2w_seg-vert_msk.nii.gz" + spine_T2w = f"/DATA/NAS/datasets_processed/NAKO/dataset-nako/{s}/{sub[:3]}/{sub}/T2w/sub-{sub}_sequ-stitched_acq-sag_mod-T2w_seg-spine_msk.nii.gz" + poi = f"/DATA/NAS/datasets_processed/NAKO/dataset-nako/{s}/{sub[:3]}/{sub}/T2w/sub-{sub}_sequ-stitched_acq-sag_mod-T2w_seg-spine_ctd.json" + + if Path(vert_T2w).exists(): + return vert_T2w, spine_T2w, poi + if not Path(vert_T2w).exists(): + T2w = Path( + f"/DATA/NAS/datasets_processed/NAKO/dataset-nako/rawdata_stitched/{sub[:3]}/{sub}/T2w/sub-{sub}_sequ-stitched_acq-sag_T2w.nii.gz" + ) + if T2w.exists(): + log.on_fail(f"Segmentation missing; {T2w.exists()=}", vert_T2w) + else: + T2w_org = list(Path(f"/DATA/NAS/datasets_processed/NAKO/dataset-nako/rawdata/{sub[:3]}/{sub}/T2w/").glob("*_T2w.nii.gz")) + if len(T2w_org) <= 2: + log.on_warning(f"Segmentation missing; {len(T2w_org)=}", Path(vert_T2w).name) + else: + log.on_debug(f"Segmentation missing; {(T2w_org)=}", Path(vert_T2w).name) + return None, "", None + return vert_T2w, spine_T2w, poi + + +def loop_over_repaired_nako( + add_mevibe=False, + add_vibe=True, + compute_PDFF=False, + raise_on_duplicate=True, + dataset="/DATA/NAS/datasets_processed/NAKO/dataset-nako/", + test=True, + verbose=False, + sort=True, + test_key="/110/11089", # path matching. if you want on specific us a 6 digits + baseline_metadata="/DATA/NAS/datasets_processed/NAKO/NAKO-732_Begleitdaten/NAKO-732_export_baseline.csv", +): + """Iterate over the repaired NAKO dataset yielding per-subject file dicts. + + Scans the NAKO BIDS dataset (including derivative roots for MEVIBE, inversion, and + abdominal segmentation), and for each subject collects a curated set of image and mask + files keyed by short names (e.g. ``"t2w"``, ``"MRSegmentator"``, ``"vibeseg100"``, + ``"roi"``). Optionally augments each subject with corrected MEVIBE outputs (see + :func:`get_corrected_mevibe`) and/or the four vibe part images (in-/out-phase, fat, + water), preferring reconstructed vibe fat/water when available. + + Args: + add_mevibe: Include corrected MEVIBE files (and optionally recompute PDFF/PDWF). + add_vibe: Include vibe part images. + compute_PDFF: Passed through to :func:`get_corrected_mevibe`. + raise_on_duplicate: Assert that each key resolves to exactly one file per subject. + dataset: Root path of the NAKO BIDS dataset. + test: If True, restrict scanning to a single hard-coded subject subtree for quick runs. + verbose: Log each subject id as it is processed. + + Yields: + Dict mapping short keys to ``BIDS_FILE`` entries for one subject. + """ + baseline = pd.read_csv(baseline_metadata, sep=";", decimal=",") + # ID is unique and matches `sub` + height_from_csv = baseline.set_index("ID")["a_anthro_groe_q"].replace(7777, pd.NA) + + gbi = Buffered_BIDS_Global_info( + datasets=dataset, + parents=[ + "rawdata", + "rawdata_stitched", + "derivatives_Abdominal-Segmentation", + # "derivatives_mevibe", #coopied into "derivatives_Abdominal-Segmentation" + "derivatives_inversion", + ], + # sequence_splitting_keys=["sub", "ses"], + # "/107/107472" + filter_file=(lambda x: test_key in str(x)) if test else None, # Figures sent to Paul "/117/117001"; "/113/113508" + ) + + for sub, subj in gbi.enumerate_subjects(sort=sort, shuffle=not sort): + subj_dict = {"id": sub, "dataset": dataset} + # Primary source: baseline CSV, height is in cm. + height_cm = height_from_csv.get(sub, pd.NA) + + if pd.notna(height_cm): + subj_dict["height_m"] = float(height_cm) / 100.0 # type: ignore + else: + q = subj.new_query(flatten=True) + q.filter_filetype("json") + for f in q.loop_list(): + try: + if not f.file["json"].exists(): + continue + js = f.open_json() + if "PatientSize" in js: + subj_dict["height_m"] = js["PatientSize"] + break + except json.decoder.JSONDecodeError: + log.on_fail(f, "json.decoder.JSONDecodeError") + if verbose: + log.on_log(sub) + mapping = {"T2w": "t2w"} + q = subj.new_query() + q.filter("chunk", lambda _: False, required=False) + # q.filter("mod", lambda x: str(x) not in "mevibe", required=False) + # q.filter_format(lambda x: str(x) not in ["mevibe"]) + keys = ["pd", "T2haste", *mapping.keys()] + mult_ok = ["pd", "T2haste"] + for fam in q.loop_dict(key_addendum=["mod", "part", "desc"]): + for k, v in fam.items(): + if k in keys: + k = mapping.get(k, k) # noqa: PLW2901 + + if raise_on_duplicate: + if k in mult_ok and k in subj_dict: + if int(v[0].get("sequ", 0)) < int(subj_dict[k].get("sequ", 0)): + continue + else: + assert len(v) == 1, v + assert k not in subj_dict, (k, subj_dict, v, subj_dict[k]) + subj_dict[k] = v[0] + + keys = ["msk_seg-body-composition_mod-mevibe"] + if add_mevibe: + q = subj.new_query() + q.filter_format("mevibe") + q.filter("sequ", "me1") + for fam in q.loop_dict(key_addendum=["mod", "part", "desc"]): + subj_dict = {**get_corrected_mevibe(fam, compute_PDFF=compute_PDFF), **subj_dict} + for k, v in fam.items(): + if k in keys: + k = mapping.get(k, k) # noqa: PLW2901 + if raise_on_duplicate: + assert len(v) == 1, v + assert k not in subj_dict, (k, subj_dict, v, subj_dict[k]) + subj_dict[k] = v[0] + + if add_vibe: + mapping = { + "msk_seg-MRSegmentator_part-inphase": "MRSegmentator", + "msk_seg-VibeSeg-100_mod-vibe_part-inphase": "vibeseg100", + "msk_seg-ROI_mod-vibe": "roi", + } + q = subj.new_query() + q.filter_format("vibe") + q.filter("chunk", lambda _: False, required=False) + # q.filter("run", lambda x: x != "2", required=False) + keys = [ + "vibe_part-inphase", + "vibe_part-outphase", + "vibe_part-fat", + "vibe_part-water", + "msk_seg-body-composition_mod-vibe", + *mapping.keys(), + ] + for fam in q.loop_dict(key_addendum=["mod", "part", "desc"]): + # print(fam) + for k, v in fam.items(): + # print("-", k) + if k in keys: + k = mapping.get(k, k) # noqa: PLW2901 + # print(k) + # print("*") + if raise_on_duplicate: + assert len(v) == 1, v + # assert k not in subj_dict, (k, subj_dict, v, subj_dict[k]) + if k in subj_dict and int(v[0].get("run", 0)) > int(subj_dict[k].get("run", 0)): + continue + subj_dict[k] = v[0] + + mapp = {"vibe_part-water_desc-reconstructed": "vibe_part-water", "vibe_part-fat_desc-reconstructed": "vibe_part-fat"} + for k, k2 in mapp.items(): + if k in fam: + subj_dict[k2] = fam[k][0] + vert, spine, poi = get_current_best_T2w_seg(sub) + subj_dict["vert"] = vert + subj_dict["spine"] = spine + subj_dict["poi"] = poi + yield subj_dict + + +if __name__ == "__main__": + from TPTBox import Print_Logger + + log = Print_Logger() + for d in loop_over_repaired_nako(): + print(d.keys()) + print(d["t2w"]) + break diff --git a/TPTBox/spine/spinestats/_run_all.py b/TPTBox/spine/spinestats/_run_all.py new file mode 100644 index 00000000..4ba2934c --- /dev/null +++ b/TPTBox/spine/spinestats/_run_all.py @@ -0,0 +1,643 @@ +"""End-to-end NAKO subject processing. + +For a given NAKO subject id, this module resolves the required inputs +(T2w + VIBE + spine/vertebra segmentation + VIBESeg-100 + ROI), runs the +spine and body-composition analyses, writes a single json with all +results, and (optionally) streams those results into Excel summaries in +parallel. + +The full documentation of the produced json keys, their units, and the +per-function conventions used by the pipeline lives in the folder +README (``TPTBox/spine/spinestats/README.md``) and its standalone copy +``all_output_reference.md`` — both are meant to be read by clinicians +reviewing the numbers. +""" + +from __future__ import annotations + +import gc +import multiprocessing as mp +import queue as _queue +from pathlib import Path +from typing import Any + +from tqdm import tqdm + +from TPTBox import Print_Logger +from TPTBox.core.bids_files import BIDS_FILE +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 +from TPTBox.spine.spinestats._load_nako import loop_over_repaired_nako + +DATASET_ROOT = Path("/DATA/NAS/datasets_processed/NAKO/dataset-nako") +logger = Print_Logger() +# 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. +REQUIRED_MAIN_KEYS: tuple[str, ...] = ( + "ivd_geometry", + "vert_geometry", + "VBQ_score", + "body_composition_score", + "muscle_fat_infiltration", + "torso_vat_sat_muscle_mass", +) + + +def get_nako_paths(nako_id: str) -> dict[str, Path | None]: + """Return a dict with all relevant paths for a given NAKO id. + + Keys: + t2w, vibe-water, vibe-fat, vibe-inphase, vibe-outphase, + vert, spine, roi, vibeseg100, dataset. + + Replaces the raw vibe stitched with the corrected version if both + corrected image + json exist. + """ + sub = str(nako_id).split("_")[0].replace("sub-", "") + pfx = sub[:3] + + t2w_stitched = DATASET_ROOT / f"rawdata_stitched/{pfx}/{sub}/T2w/sub-{sub}_sequ-stitched_acq-sag_T2w.nii.gz" + out = { + f"vibe-{a}": DATASET_ROOT / f"rawdata_stitched/{pfx}/{sub}/vibe/sub-{sub}_sequ-stitched_acq-ax_part-{a}_vibe.nii.gz" + for a in ["water", "fat", "inphase", "outphase"] + } + vibe_corr = ( + DATASET_ROOT + / f"derivatives_Abdominal-Segmentation/{pfx}/{sub}/vibe/sub-{sub}_sequ-stitched_acq-ax_part-water_desc-corrected_vibe.nii.gz" + ) + vibe_corr_json = vibe_corr.with_suffix("").with_suffix(".json") + if vibe_corr.exists() and vibe_corr_json.exists(): + out["vibe-water"] = vibe_corr + + # current best T2w spine seg (mirrors qa_spine_shift.get_current_best_T2w_seg) + search_folders = [ + "derivatives_spine_vert_fixed", + "derivatives_spine_inference_combination162_148", + ] + vert = spine = None + for s in search_folders: + base = DATASET_ROOT / f"{s}/{pfx}/{sub}/T2w" + v = base / f"sub-{sub}_sequ-stitched_acq-sag_mod-T2w_seg-vert_msk.nii.gz" + sp = base / f"sub-{sub}_sequ-stitched_acq-sag_mod-T2w_seg-spine_msk.nii.gz" + if v.exists(): + vert, spine = v, sp + 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" + ) + vibeseg100 = ( + DATASET_ROOT + / f"derivatives_Abdominal-Segmentation/{pfx}/{sub}/vibe/sub-{nako_id}_sequ-stitched_acq-ax_mod-vibe_part-inphase_seg-VibeSeg-100_msk.nii.gz" + ) + + return { + "t2w": t2w_stitched if t2w_stitched.exists() else None, + **out, + "vert": vert, + "spine": spine, + "roi": roi, + "vibeseg100": vibeseg100 if vibeseg100.exists() else None, + "dataset": DATASET_ROOT, + } + + +def _segmentation_inputs(file_dict: dict) -> list[Path]: + """Segmentation files whose mtime should invalidate a cached json.""" + keys = ("vert", "spine", "vibeseg100", "roi") + return [ + file_dict[k].file["nii.gz"] if isinstance(file_dict[k], BIDS_FILE) else Path(file_dict[k]) + for k in keys + if file_dict.get(k) is not None + ] + + +def _is_cache_valid(json_path: Path, seg_files: list[Path], required_keys: tuple[str, ...]) -> tuple[bool, dict | None]: + """Return (valid, loaded_dict). + + Cache is invalid (and needs recompute) if: + - json does not exist, + - json is older than any segmentation file, + - json fails to parse, + - any required main key is missing from the loaded dict. + """ + if not json_path.exists(): + return False, None + json_mtime = json_path.stat().st_mtime + for seg in seg_files: + if seg.exists() and seg.stat().st_mtime > json_mtime: + return False, None + try: + data = load_json(json_path) + except Exception: + return False, None + if not isinstance(data, dict): + return False, None + for k in required_keys: + if k not in data: + return False, None + return True, data + + +def run_all( + file_dict, + override: bool = False, + do_not_update=False, + need_cobb=False, + need_ivd=False, + need_vert=False, + need_vbq=True, + need_bcs=True, + need_mfi=True, + need_torso=True, +) -> dict[str, Any] | None: + """Run the full pipeline for one subject and return the results dict. + + Parameters + ---------- + file_dict : dict + Output of :func:`get_nako_paths`. + cobb : bool, default=False + If True, additionally compute Cobb / lordosis / kyphosis angles + (adds keys ``cobb`` and ``curv``). + override : bool, default=False + If False (default), skip recomputation and return the existing + json when it is still valid. The cache is considered valid when: + + - the target json exists, + - it is newer than every segmentation file listed in + ``_segmentation_inputs(file_dict)``, + - and it contains every key in :data:`REQUIRED_MAIN_KEYS` + (plus ``cobb``/``curv`` when ``cobb=True``). + + If True, all algorithms run even when a valid json already + exists and the json is overwritten. + + Returns: + ------- + dict + The full results dictionary (either freshly computed or loaded + from the cached json). See ``doc/all_output_reference.md`` for + the meaning of each key. + """ + from TPTBox import Location, calc_poi_from_subreg_vert + from TPTBox.spine.spinestats.angles import plot_cobb_and_lordosis_and_kyphosis + from TPTBox.spine.spinestats.measure_ivd_and_vertebra_geometry import ( + measure_ivd_and_vertebra_geometry, # structure_label: int = 100 and structure_label: int = 49 + ) + from TPTBox.spine.spinestats.torso_vat_sat import VBQ_score, body_composition_score, muscle_fat_infiltration, torso_vat_sat_muscle_mass + + if "t2w" not in file_dict: + return None + + t2w_bf = file_dict["t2w"] if isinstance(file_dict["t2w"], BIDS_FILE) else BIDS_FILE(file_dict["t2w"], file_dict["dataset"]) + poi_out = t2w_bf.get_changed_path( + "json", + "poi", + "derivatives_spine_inference_162_sacrumfix_subregionmeasures-v2", + info={"seg": "vert", "mod": "T2w", "desc": "vert-rotation"}, + ) + cobb_jpg_out = t2w_bf.get_changed_path( + "jpg", "snp", "derivatives_spine_inference_162_sacrumfix_subregionmeasures-v2", info={"seg": "cobb"} + ) + final_out = t2w_bf.get_changed_path( + "json", "stat", "derivatives_spine_inference_162_sacrumfix_subregionmeasures-v2", info={"seg": "all"} + ) + final_out = Path(final_out) + + required = (*REQUIRED_MAIN_KEYS, "cobb", "curv") + seg_files = _segmentation_inputs(file_dict) + + out: dict[str, Any] = {} + if not override: + valid, cached = _is_cache_valid(final_out, seg_files, required) + if valid and cached is not None and do_not_update: + return cached + # Reload existing json (if any) and only recompute the missing top-level keys. + if final_out.exists(): + try: + loaded = load_json(final_out) + if isinstance(loaded, dict): + out = loaded + except Exception: + out = {} + + def _need(*keys: str, compute: bool) -> bool: + return override or (any(k not in out for k in keys) and compute) + + need_cobb = _need("cobb", "curv", compute=need_cobb) + need_ivd = _need("ivd_geometry", compute=need_ivd) + need_vert = _need("vert_geometry", compute=need_vert) + need_vbq = _need("VBQ_score", compute=need_vbq) + need_bcs = _need("body_composition_score", compute=need_bcs) + need_mfi = _need("muscle_fat_infiltration", compute=need_mfi) + need_torso = _need("torso_vat_sat_muscle_mass", compute=need_torso) + # Recompute area + save = False + if "VBQ_score" in out and "VBQ_L1-L1_old" in out["VBQ_score"]: + logger.on_warning("redo vbq", t2w_bf.get("sub")) + need_vbq = True + del out["VBQ_score"] + save = True + if "torso_vat_sat_muscle_mass" in out and "Not a VIBESeg-100" in str(out.get("torso_vat_sat_muscle_mass", {}).get("reason", "")): + logger.on_warning("redo torso_vat_sat_muscle_mass", t2w_bf.get("sub")) + need_torso = True + del out["torso_vat_sat_muscle_mass"] + save = True + #### + need_poi = need_cobb or need_ivd or need_vert + 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 + need_vibe_seg = need_bcs or need_mfi or need_torso + 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): + if _merge_endplate_angles(out, Path(poi_out)) or save: + save_json(final_out, out) + return out + + logger.on_debug("load nii") + t2w = to_nii(file_dict["t2w"]) if need_t2w or need_cobb else None + vibe_water = to_nii(file_dict["vibe_part-water"], False) if need_vibe_wf else None + vibe_fat = to_nii(file_dict["vibe_part-fat"], False) if need_vibe_wf else None + vert = to_nii(file_dict["vert"], True) if need_vert_nii else None + spine = to_nii(file_dict["spine"], True) if need_spine_nii else None + vibe_seg = to_nii(file_dict["vibeseg100"], True) if need_vibe_seg else None + roi = to_nii(file_dict["roi"], True) if need_roi else None + height_m = file_dict.get("height_m") + + 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 need_cobb: + project_2D = False + threshold_deg = 10 + logger.on_debug("cobb") + cobb_val, curv, _ = plot_cobb_and_lordosis_and_kyphosis( + cobb_jpg_out, poi, file_dict["t2w"], file_dict["vert"], project_2D=project_2D, threshold_deg=threshold_deg + ) + out["cobb"] = cobb_val + out["curv"] = curv + out["project_2D"] = project_2D + out["min_coop_angle"] = threshold_deg + + if need_ivd: + logger.on_debug("measure_ivd_and_vertebra_geometry (ivd)") + out["ivd_geometry"] = measure_ivd_and_vertebra_geometry(t2w, vert, spine, buffer_poi=poi_out, structure_label=100) + if need_vert: + logger.on_debug("measure_ivd_and_vertebra_geometry (vert)") + out["vert_geometry"] = measure_ivd_and_vertebra_geometry(t2w, vert, spine, buffer_poi=poi_out, structure_label=0) + + if need_vbq: + logger.on_debug("VBQ_score") + out["VBQ_score"] = VBQ_score(t2w, vert, spine, full_cord=True) + + if need_bcs: + logger.on_debug("body_composition_score") + out["body_composition_score"] = body_composition_score(vibe_seg, vert, spine, dataset_id=100, height_m=height_m) + assert len(out["body_composition_score"]) != 0 + if need_mfi: + logger.on_debug("muscle_fat_infiltration") + out["muscle_fat_infiltration"] = muscle_fat_infiltration(vibe_water, vibe_fat, vibe_seg, vert, spine, roi=roi, dataset_id=100) + out["muscle_fat_infiltration"]["physics_model"] = "2-Point-Dixon" + if need_torso: + # torso_vat_sat_muscle_mass returns (results_dict, body_comp_nii). Keep + # only the serializable results dict so the whole json stays writable. + logger.on_debug("torso_vat_sat_muscle_mass") + torso_results, _body_comp = torso_vat_sat_muscle_mass(vibe_seg, roi, dataset_id=100) + out["torso_vat_sat_muscle_mass"] = torso_results + _merge_endplate_angles(out, Path(poi_out)) + logger.on_save("save", final_out.name) + save_json(final_out, out) + return out + + +def _read_endplate_internal_angles(poi_json_path: Path) -> dict[str, Any]: + """Read the ``endplate_internal_angle`` dict from a POI json (if present). + + The POI json is a list; the first element is the metadata dict where the + endplate-angle map (vertebra name -> angle in degrees) lives. + """ + if not poi_json_path.exists(): + return {} + try: + data = load_json(poi_json_path) + except Exception: + return {} + if isinstance(data, list): + for entry in data: + if isinstance(entry, dict) and isinstance(entry.get("endplate_internal_angle"), dict): + return entry["endplate_internal_angle"] + return {} + if isinstance(data, dict) and isinstance(data.get("endplate_internal_angle"), dict): + return data["endplate_internal_angle"] + return {} + + +def _merge_endplate_angles(out: dict[str, Any], poi_json_path: Path) -> bool: + """Attach the POI's per-vertebra endplate_internal_angle to ``out``. + + Adds a top-level ``endplate_internal_angle`` (vertebra-name -> angle) and, + for every entry in ``vert_geometry``, injects the matching angle as + ``endplate_internal_angle`` so it shows up in per-vertebra Excel rows. + Returns True iff ``out`` was modified. + """ + if out.get("endplate_internal_angle") is not None: + return False + angles = _read_endplate_internal_angles(poi_json_path) + if not angles: + return False + from TPTBox.core.vert_constants import Vertebra_Instance + + changed = False + if out.get("endplate_internal_angle") != angles: + out["endplate_internal_angle"] = angles + changed = True + vg = out.get("vert_geometry") + if isinstance(vg, dict): + for label, metrics in vg.items(): + if not isinstance(metrics, dict): + continue + try: + vname = Vertebra_Instance(int(label)).name + except Exception: + continue + angle = angles.get(vname) + if angle is None: + continue + if metrics.get("endplate_internal_angle") != angle: + metrics["endplate_internal_angle"] = angle + changed = True + return changed + + +# --------------------------------------------------------------------------- +# Excel collector (parallel, producer/consumer) +# --------------------------------------------------------------------------- + + +def _flatten(prefix: str, obj: Any, out: dict[str, Any]) -> None: + """Flatten nested dicts into dotted keys (leaves = scalars/None).""" + if isinstance(obj, dict): + for k, v in obj.items(): + new_key = f"{prefix}.{k}" if prefix else str(k) + _flatten(new_key, v, out) + else: + # Lists / tuples / NII placeholders end up here as-is; the writer + # will drop non-scalar values so the sheet stays clean. + out[prefix] = obj + + +def _rows_from_json(subject_id: str, data: dict) -> tuple[dict[str, Any], list[dict[str, Any]]]: + """Split one subject's json into (per-subject row, per-vertebra rows). + + Per-subject row: everything except the per-label geometry dicts, + flattened to dotted keys. + Per-vertebra rows: one row per label in ``ivd_geometry`` and + ``vert_geometry`` (source column indicates which). + """ + per_subject: dict[str, Any] = {"subject": subject_id} + subject_view = {k: v for k, v in data.items() if k not in ("ivd_geometry", "vert_geometry")} + _flatten("", subject_view, per_subject) + + per_vert: list[dict[str, Any]] = [] + for source_key in ("vert_geometry", "ivd_geometry"): + section = data.get(source_key) or {} + if not isinstance(section, dict): + continue + for label, metrics in section.items(): + if not isinstance(metrics, dict): + continue + row: dict[str, Any] = {"subject": subject_id, "source": source_key, "label": label} + row.update(metrics) + per_vert.append(row) + return per_subject, per_vert + + +def _collector_worker( + task_q: mp.Queue, + out_folder: Path, + per_subject_name: str, + per_vertebra_name: str, + flush_every: int, +) -> None: + import pandas as pd # local import so the main process starts fast + + out_folder = Path(out_folder) + out_folder.mkdir(parents=True, exist_ok=True) + subject_rows: list[dict[str, Any]] = [] + vertebra_rows: list[dict[str, Any]] = [] + seen: set[str] = set() + + def _flush() -> None: + if subject_rows: + pd.DataFrame(subject_rows).to_excel(out_folder / per_subject_name, index=False) + if vertebra_rows: + pd.DataFrame(vertebra_rows).to_excel(out_folder / per_vertebra_name, index=False) + + while True: + try: + item = task_q.get(timeout=1.0) + except _queue.Empty: + continue + if item is None: + _flush() + return + subject_id, json_path = item + if subject_id in seen: + continue + try: + data = load_json(Path(json_path)) + except Exception: + continue + per_subj, per_vert = _rows_from_json(str(subject_id), data) + subject_rows.append(per_subj) + vertebra_rows.extend(per_vert) + seen.add(subject_id) + if flush_every and len(seen) % flush_every == 0: + _flush() + + +class ExcelCollector: + """Background process that turns subject jsons into Excel summaries. + + Usage:: + + collector = ExcelCollector(out_folder="/tmp/nako_summary") + collector.start() + for nako_id in ids: + f = get_nako_paths(nako_id) + data = run_all(f) + collector.submit(nako_id, final_json_path_for(f)) + collector.close() # flushes and joins + """ + + def __init__( + self, + out_folder: str | Path, + per_subject_name: str = "per_subject.xlsx", + per_vertebra_name: str = "per_vertebra.xlsx", + flush_every: int = 200, + ) -> None: + self.out_folder = Path(out_folder) + self.per_subject_name = per_subject_name + self.per_vertebra_name = per_vertebra_name + self.flush_every = flush_every + self._queue: mp.Queue = mp.Queue() + self._proc: mp.Process | None = None + + def start(self) -> None: + if self._proc is not None: + return + self._proc = mp.Process( + target=_collector_worker, + args=(self._queue, self.out_folder, self.per_subject_name, self.per_vertebra_name, self.flush_every), + daemon=True, + ) + self._proc.start() + + def submit(self, subject_id: str, json_path: str | Path) -> None: + if self._proc is None: + raise RuntimeError("ExcelCollector not started") + self._queue.put((str(subject_id), str(json_path))) + + def close(self, join_timeout: float = 60.0) -> None: + if self._proc is None: + return + self._queue.put(None) + self._proc.join(timeout=join_timeout) + self._proc = None + + +def _final_json_path(file_dict: dict) -> Path: + """Recreate the json path run_all writes to, without re-running it.""" + t2w_bf = BIDS_FILE(file_dict["t2w"], file_dict["dataset"]) if not isinstance(file_dict["t2w"], BIDS_FILE) else file_dict["t2w"] + return Path( + t2w_bf.get_changed_path("json", "stat", "derivatives_spine_inference_162_sacrumfix_subregionmeasures-v2", info={"seg": "all"}) + ) + + +# Ordered list of required inputs for run_all. Order matters: for the +# missing-file report each subject is attributed to the FIRST missing +# key in this list, so a subject with several gaps is still counted once. +REQUIRED_INPUT_KEYS: tuple[str, ...] = ( + "t2w", + "vibe_part-water", + "vibe_part-fat", + "vert", + "spine", + "vibeseg100", + "roi", +) + + +def _first_missing_input(file_dict: dict) -> str | None: + """Return the first REQUIRED_INPUT_KEYS entry not present/on disk, else None.""" + for k in REQUIRED_INPUT_KEYS: + v = file_dict.get(k) + if v is None: + return k + p = v.file["nii.gz"] if isinstance(v, BIDS_FILE) else Path(v) + if not Path(p).exists(): + return k + return None + + +def _run_one(args: tuple[dict, bool, bool]) -> tuple[str, str | None, dict]: + """Worker: run_all for one subject; returns (subject_id, missing_key_or_None).""" + f, override, do_not_update = args + sub_id = str(f.get("id")) + missing = _first_missing_input(f) + if missing is not None: + return sub_id, missing, f + + try: + run_all(f, override=override, do_not_update=do_not_update) + except Exception as e: + logger.on_fail(f"run_all failed for {sub_id}: {e}") + logger.print_error() + return sub_id, f"error:{type(e).__name__}, {str(e)!s}", f + return sub_id, None, f + + +if __name__ == "__main__": + from concurrent.futures import ProcessPoolExecutor, as_completed + + import pandas as pd + + from TPTBox import No_Logger + + log = No_Logger() + + OUT_FOLDER = Path("/DATA/NAS/ongoing_projects/robert/test/NAKO-stats") + OUT_FOLDER.mkdir(parents=True, exist_ok=True) + N_CPUS = 40 # set >1 to parallelize + OVERRIDE = False + aggregate = True + do_not_update = False + test = False + if aggregate: + collector = ExcelCollector(out_folder=OUT_FOLDER) + collector.start() + missing_rows: list[dict[str, str]] = [] + total = 30645 + try: + if test: + subjects = loop_over_repaired_nako(test=True) + 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 + subjects = iter([next(l) for _ in range(total)]) + + if N_CPUS <= 1: + for f in subjects: + sub_id, missing, _ = _run_one((f, OVERRIDE, do_not_update)) + if missing is not None: + logger.on_fail("missing", list(f.keys()), missing) + missing_rows.append({"subject": sub_id, "missing": missing}) + continue + if aggregate: + collector.submit(sub_id, _final_json_path(f)) + else: + from itertools import islice + + with ProcessPoolExecutor(max_workers=N_CPUS) as ex: + batch_size = 1000 + l = tqdm(total=total) + while True: + gc.collect() + futs = [ex.submit(_run_one, (f, OVERRIDE, do_not_update)) for f in list(islice(subjects, batch_size))] + + if not futs: + break + for fut in as_completed(futs): + l.update(1) + sub_id, missing, f = fut.result() + if missing is not None: + logger.on_fail("missing", (sub_id), missing) + missing_rows.append({"subject": sub_id, "missing": missing}) + continue + if aggregate: + collector.submit(sub_id, _final_json_path(f)) + + finally: + if aggregate: + collector.close() + if missing_rows: + pd.DataFrame(missing_rows).to_excel(OUT_FOLDER / "missing_inputs.xlsx", index=False) diff --git a/TPTBox/spine/spinestats/all_output_reference.md b/TPTBox/spine/spinestats/all_output_reference.md new file mode 100644 index 00000000..acff5291 --- /dev/null +++ b/TPTBox/spine/spinestats/all_output_reference.md @@ -0,0 +1,275 @@ +# `_run_all.py` — Radiologist Reference + +This document lists every key produced by `run_all(file_dict)` in +`_run_all.py`, together with its unit and important implementation details. +It is aimed at radiologists reviewing the numbers, so it focuses on +"what does this mean clinically" and "how was it computed", not on the +Python API. + +## How the pipeline is organised + +`run_all` writes a single json per subject with these top-level keys: + +| Key | Source function | What it covers | +|---|---|---| +| `ivd_geometry` | `measure_ivd_and_vertebra_geometry(..., structure_label=100)` | intervertebral discs | +| `vert_geometry` | `measure_ivd_and_vertebra_geometry(..., structure_label=50)` | vertebral bodies | +| `VBQ_score` | `VBQ_score` | vertebral bone quality (T2 signal ratio) | +| `body_composition_score` | `body_composition_score` | axial CSA per tissue at chosen vertebral levels | +| `muscle_fat_infiltration` | `muscle_fat_infiltration` | Dixon fat-fraction based muscle-quality metrics | +| `torso_vat_sat_muscle_mass` | `torso_vat_sat_muscle_mass` | whole-torso VAT / SAT / muscle volume | +| `cobb`, `curv` | `plot_cobb_and_lordosis_and_kyphosis` | only when called with `cobb=True` | + +Angles are in **degrees**, lengths in **millimetres**, areas in **mm²**, +volumes in **mm³**, fat fractions are **unitless** in `[0, 1]`, MR signal +values are in **arbitrary units (a.u.)** and only meaningful as ratios. + +Caching: `run_all(..., override=False)` (the default) reuses the json +when it exists, is newer than every input segmentation file, and +contains all of the required top-level keys. Pass `override=True` to +force recomputation. + +### Input requirements + +`run_all` assumes whole-body-style acquisitions: + +- **T2w image** — must cover the **full spine** (cervical through + sacrum). Curvature angles and per-vertebra geometry silently return + `None`/`NaN` for any level that is cropped away, and the VBQ ranges + (`C3-C6`, `T5-T8`, `L1-L4`) need every vertebra in the range to be + visible. +- **VIBE water/fat images** — must cover the **full torso**. Fat + fraction and muscle CSA are computed on whatever axial slices are + present, so a cropped VIBE silently biases per-region CSA and IMAT + volumes. +- **Segmentations** (`vert`, `spine`, `vibeseg100`, `roi`) — must + match the extent of their underlying image. In particular, + `torso_vat_sat_muscle_mass` explicitly verifies that both the + clavicula and the pelvis are present in the VIBESeg mask; if either + is missing it aborts, returns `NaN` volumes and stores the reason in + the `reason` key. + +### How to produce the segmentations + +All required segmentations can be produced from +`TPTBox.segmentation`: + +- **`vert` / `spine`** — run **SPINEPS** on the T2w image + (`run_spineps`, `get_outpaths_spineps`, `_run_spineps_all`). +- **`vibeseg100`** — run **VIBESegmentator** with + `run_vibeseg(..., dataset_id=100)` on the VIBE stack. Dataset + **100** (MR and CT) is what `run_all` targets; dataset **12** + (0.8 mm iso CT) is also supported by the composition/infiltration + functions via `dataset_id=12`. +- **`roi`** — run VIBESegmentator with dataset **278** on the VIBE + stack. The raw dataset-278 ROI is **not perfect** and needs + postprocessing before it is fed into `run_all`. + +## Signal-based conventions used everywhere + +Two things are worth understanding before reading the T2 signal keys: + +1. **Peak-centered mean.** Ordinary mean signal inside a mask is + sensitive to non-CSF voxels that leak into the spinal canal + segmentation (nerve roots, vessel walls). The pipeline instead + averages only voxels whose intensity falls in a window around the + histogram peak. When both a peak-centered and a plain-mean version + are stored, the plain-mean version is suffixed with `_old` for + comparison. See `peak_centered_mean` in + `TPTBox/spine/spinestats/torso_vat_sat.py`. +2. **Erosion.** Muscle and vertebral-body masks are eroded by one or + two voxels before signal extraction to reduce partial-volume mixing + at the boundary. Volume metrics are reported both after erosion + (`*_volume_*`) and before erosion (`*_volume_no_erosion_*`) so the + effect of the erosion is auditable. + +--- + +## `ivd_geometry` and `vert_geometry` + +Both keys hold `dict[label_id, dict[metric_name, value]]`, where each +`label_id` is one intervertebral disc (`ivd_geometry`) or vertebra +(`vert_geometry`). If a label fails to evaluate, its entry contains +`error` (message string) plus all metrics set to `NaN`. + +| Key | Unit | Meaning | +|---|---|---| +| `volume_voxel` | mm³ | volume counted from the raw voxel mask | +| `volume_mesh` | mm³ | volume of the reconstructed surface mesh (same structure) | +| `height_center` | mm | height sampled through the structure's centroid | +| `mean_height` | mm | mean of the sampled heights over the structure surface | +| `max_height` | mm | maximum of the sampled heights | +| `lower_10_percent_height` | mm | 10th percentile of the sampled heights | +| `mean_diameter` | mm | diameter of the circle whose area equals the projected area | +| `anterior_height_x1` | mm | anterior height at the anterior point (x1) | +| `posterior_height_x2` | mm | posterior height at the posterior point (x2) | +| `right_height_x3` | mm | right-lateral height (x3) | +| `left_height_x4` | mm | left-lateral height (x4) | +| `width_lateral_x5` | mm | lateral width (x5) | +| `width_sagittal_x6` | mm | sagittal width (x6) | +| `signal` | unitless | peak-centered structure T2 signal / peak-centered spinal-canal T2 signal | +| `structure_signal` | a.u. | peak-centered T2 signal inside the eroded structure mask | +| `spinal_canal_signal` | a.u. | peak-centered T2 signal inside the eroded spinal canal reference | +| `signal_old` | unitless | same ratio computed with plain per-voxel means | +| `structure_signal_old` | a.u. | plain mean T2 signal inside the eroded structure mask | +| `spinal_canal_signal_old` | a.u. | plain mean T2 signal inside the eroded spinal canal | + +Implementation notes: +- Structure orientation for IVDs is estimated from the disc's own voxel + mask via PCA; vertebral orientation is read from precomputed POIs + (except C2/dens, which falls back to PCA). +- x1–x6 are the six clinically standard directional heights/widths (see + the geometry module docstring for the figure). +- Only labels present in the segmentation appear as keys. + +## `VBQ_score` + +`dict[str, float]`; one triple of entries per configured spinal range. +Default ranges are `C3-C6`, `T5-T8`, `L1-L4`. + +| Key template | Unit | Meaning | +|---|---|---| +| `mean_signal_vertebra_-` | a.u. | mean T2 signal inside the eroded vertebral body mask over the range | +| `mean_signal_liquor_-` | a.u. | **peak-centered** mean T2 signal inside the spinal canal over the same S/I extent | +| `mean_signal_liquor_-_old` | a.u. | plain-mean version, kept for backward comparison | +| `VBQ_-` | unitless | vertebral signal divided by the peak-centered CSF signal | +| `VBQ_-_old` | unitless | same ratio using the plain-mean CSF signal | + +Implementation notes: +- Vertebral body mask is eroded (default `n_erode=2`) to avoid the + cortical rim. +- The spinal canal is cropped to the same superior–inferior slab as the + vertebral bodies so the CSF reference matches the region of interest. +- Higher VBQ = darker vertebral bodies relative to CSF, associated in + the literature with lower bone quality. + +## `body_composition_score` + +`dict[str, float]`; per-region axial cross-sectional-area statistics of +five tissue classes. Default regions are `T12-L1` and `L3-L3`. + +Region tag: `{start.name}-{goal.name}` (e.g. `T12-L1`, `L3-L3`). + +| Key template | Unit | Meaning | +|---|---|---| +| `mean_muscle_area_{region}` | mm² | mean skeletal muscle CSA across the region | +| `max_muscle_area_{region}` | mm² | maximum skeletal muscle CSA in the region | +| `mean_VAT_area_{region}` | mm² | mean visceral adipose tissue CSA | +| `max_VAT_area_{region}` | mm² | maximum visceral adipose tissue CSA | +| `mean_SAT_area_{region}` | mm² | mean subcutaneous adipose tissue CSA | +| `max_SAT_area_{region}` | mm² | maximum subcutaneous adipose tissue CSA | +| `mean_psoas_area_{region}` | mm² | mean psoas CSA (left + right) | +| `max_psoas_area_{region}` | mm² | maximum psoas CSA | +| `mean_autochthon_area_{region}` | mm² | mean autochthonous back-muscle CSA | +| `max_autochthon_area_{region}` | mm² | maximum autochthonous back-muscle CSA | +| `n_slices_{region}` | count | number of axial slices contributing to the muscle statistic | +| `muscle_index_{region}` | mm²/m² | `mean_muscle_area / height_m²`; only present when `height_m` is supplied | +| `muscle_fat_ratio_{region}` | unitless | `mean_muscle_area / (mean_VAT_area + mean_SAT_area)`; `NaN` if the denominator is zero | + +Implementation notes: +- The superior–inferior extent of the vertebral bodies inside the region + defines the slice range. +- Axial voxel area is derived from the VIBE geometry; slices with zero + tissue are excluded before mean/max. +- If no vertebral body voxels are present for a region, that region is + silently skipped (no keys emitted). + +## `muscle_fat_infiltration` + +`dict[str, float]`; per (region, muscle group) Dixon-based fat +infiltration metrics. Muscle groups (dataset_id=100) are: +`all_muscle`, `iliopsoas_left`, `iliopsoas_right`, `autochthon_left`, +`autochthon_right`, `muscle_other`. Suffix is `{region}_{muscle}`; when +no region is given, the region tag is `all`. + +Fat fraction (FF) is computed voxel-wise as `FF = fat / (fat + water)`; +voxels with `FF >= threshold` (default 0.20) are IMAT, otherwise lean. + +| Key template | Unit | Meaning | +|---|---|---| +| `mean_fat_fraction_{suffix}` | [0, 1] | mean FF over the eroded muscle mask | +| `median_fat_fraction_{suffix}` | [0, 1] | median FF over the eroded muscle mask | +| `mean_lean_fat_fraction_{suffix}` | [0, 1] | mean FF of lean voxels (FF < threshold) | +| `mean_IMAT_fat_fraction_{suffix}` | [0, 1] | mean FF of IMAT voxels (FF ≥ threshold) | +| `muscle_volume_{suffix}` | mm³ | muscle volume after erosion | +| `muscle_volume_no_erosion_{suffix}` | mm³ | muscle volume **before** erosion (raw segmentation volume) | +| `lean_muscle_volume_{suffix}` | mm³ | lean-muscle volume within the eroded mask | +| `lean_muscle_volume_no_erosion_{suffix}` | mm³ | lean-muscle volume within the un-eroded mask | +| `IMAT_volume_{suffix}` | mm³ | IMAT volume within the eroded mask | +| `IMAT_volume_no_erosion_{suffix}` | mm³ | IMAT volume within the un-eroded mask | +| `IMAT_fraction_{suffix}` | [0, 1] | IMAT voxel fraction within the eroded mask | + +Implementation notes: +- Erosion iterations per muscle are configurable via the `erode` dict. + Defaults: `all_muscle=1`, `iliopsoas_*=1`, `autochthon_*=2`, + `muscle_other=1`. +- Fat-fraction statistics use the eroded mask; volumes are also + reported for the un-eroded mask so the caller can inspect the effect + of erosion. +- When `regions` are supplied, the analysis is restricted to the + superior–inferior extent of the vertebral bodies in each range and + the region tag becomes `{start.name}-{goal.name}`. + +## `torso_vat_sat_muscle_mass` + +`dict[str, float]`; whole-torso volumes restricted to the supplied ROI. +Only the results dict is stored in the json (the optional NII output is +dropped by `run_all` because it is not JSON-serializable). + +| Key | Unit | Meaning | +|---|---|---| +| `VAT` | mm³ | visceral adipose tissue volume inside the ROI | +| `SAT` | mm³ | subcutaneous adipose tissue volume inside the ROI | +| `muscle_mass` | mm³ | skeletal muscle volume inside the ROI | +| `reason` | string | present only if the computation failed; explains why | + +Implementation notes: +- The function checks that both the clavicula and the pelvis are + present in the segmentation to make sure the full torso is covered. + If either check fails, all three volumes are set to `NaN` and + `reason` is populated. +- ROI labels (default 3–8) select which sub-regions of the torso count. + +## `cobb` and `curv` (only when `cobb=True`) + +- `cobb`: `list[tuple[float, int, int, int | None]]` from + `compute_max_cobb_angle_multi` — one entry per detected scoliotic + segment: `(max_angle_deg, from_vertebra_id, to_vertebra_id, apex_id_or_none)`. + Angles are in **degrees**. +- `curv`: dict from `compute_lordosis_and_kyphosis`: + - `cervical_lordosis` (deg) — computed between C2 and C7 + - `thoracic_kyphosis` (deg) — computed between T4 and the last thoracic vertebra + - `lumbar_lordosis` (deg) — computed between L1 and the last lumbar vertebra + + Values can be `None` if the required vertebrae are missing from the + POI. + +--- + +## Excel collector + +`ExcelCollector` in `all.py` runs a background process that turns each +finished json into two rolling Excel files in a configurable folder: + +- `per_subject.xlsx` — one row per subject with every scalar top-level + metric flattened to dotted keys + (e.g. `VBQ_score.VBQ_L1-L4`, `torso_vat_sat_muscle_mass.VAT`). +- `per_vertebra.xlsx` — one row per (subject, label), populated from + `vert_geometry` and `ivd_geometry`. The `source` column indicates + which of the two sections the row came from. + +Usage: + +```python +collector = ExcelCollector(out_folder="/tmp/nako_summary") +collector.start() +for nako_id in ids: + f = get_nako_paths(nako_id) + run_all(f) # writes the per-subject json + collector.submit(nako_id, _final_json_path(f)) +collector.close() # flushes and joins +``` + +The collector re-writes the Excel files every `flush_every` submissions +(default 25) and once more at shutdown, so partial runs still produce +usable summaries. diff --git a/TPTBox/spine/spinestats/angles.py b/TPTBox/spine/spinestats/angles.py index e10fdca6..81d46d2a 100644 --- a/TPTBox/spine/spinestats/angles.py +++ b/TPTBox/spine/spinestats/angles.py @@ -1,7 +1,9 @@ from __future__ import annotations +from dataclasses import dataclass from enum import Enum, auto from pathlib import Path +from typing import Literal import numpy as np @@ -107,6 +109,55 @@ def get_point(self, v: Vertebra_Instance | int, poi: POI) -> np.ndarray: raise NotImplementedError(v, poi) +def _get_last_lumbar(poi: POI) -> Vertebra_Instance | None: + """Return the most inferior lumbar vertebra that has a centroid in ``poi``.""" + for i in list(reversed(Vertebra_Instance.lumbar()))[:5]: + if (i.value, 50) in poi: + return i + return None + + +def _get_last_thoracic(poi: POI) -> Vertebra_Instance | None: + """Return the most inferior thoracic vertebra that has a centroid in ``poi``.""" + for i in list(reversed(Vertebra_Instance.thoracic()))[:3]: + if (i.value, 50) in poi: + return i + return None + + +@dataclass +class Def_Curvature: + """Define the lordosis and kyposis angle.""" + + start_vert: Vertebra_Instance | Literal["last_thoracic", "last_lumbar"] + start_move: MoveTo + stop_vert: Vertebra_Instance | Literal["last_thoracic", "last_lumbar"] + stop_move: MoveTo + + def get_start_vert(self, poi) -> Vertebra_Instance: + """get_start_vert.""" + if self.start_vert == "last_thoracic": + return _get_last_thoracic(poi) # type: ignore + if self.start_vert == "last_lumbar": + return _get_last_lumbar(poi) # type: ignore + return self.start_vert + + def get_stop_vert(self, poi) -> Vertebra_Instance: + """get_stop_vert.""" + if self.stop_vert == "last_thoracic": + return _get_last_thoracic(poi) # type: ignore + if self.stop_vert == "last_lumbar": + return _get_last_lumbar(poi) # type: ignore + return self.stop_vert + + +curvature_definition = { + "cervical_lordosis": Def_Curvature(Vertebra_Instance.C2, MoveTo.BOTTOM, Vertebra_Instance.C7, MoveTo.BOTTOM), + "thoracic_kyphosis": Def_Curvature(Vertebra_Instance.T4, MoveTo.TOP, "last_thoracic", MoveTo.BOTTOM), + "lumbar_lordosis": Def_Curvature(Vertebra_Instance.L1, MoveTo.TOP, "last_lumbar", MoveTo.BOTTOM), +} + + def unit_vector(vector: np.ndarray) -> np.ndarray: """Return the unit vector of the input vector. @@ -338,7 +389,7 @@ def compute_lordosis_and_kyphosis(poi: POI, project_2D=True) -> dict[str, float Returns: dict: A dictionary containing the following key-value pairs: - "cervical_lordosis": The angle of cervical lordosis, calculated between C2 and C7. - - "thoracic_kyphosis": The angle of thoracic kyphosis, calculated between T1 and the last thoracic vertebra. + - "thoracic_kyphosis": The angle of thoracic kyphosis, calculated between T4 and the last thoracic vertebra. - "lumbar_lordosis": The angle of lumbar lordosis, calculated between L1 and the last lumbar vertebra. Raises: @@ -346,7 +397,7 @@ def compute_lordosis_and_kyphosis(poi: POI, project_2D=True) -> dict[str, float Notes: - It is essential that the `poi` contains the posterior vertebra direction for accurate angle calculations. - - Thoracic kyphosis is calculated from T1 to the last thoracic vertebra identified in the POI. + - Thoracic kyphosis is calculated from T4 to the last thoracic vertebra identified in the POI. - Lumbar lordosis is calculated from L1 to the last lumbar vertebra identified in the POI. Example: @@ -359,25 +410,15 @@ def compute_lordosis_and_kyphosis(poi: POI, project_2D=True) -> dict[str, float assert Location.Vertebra_Direction_Posterior.value in poi.keys_subregion(), ( "You need to compute the Direction in the Poi (Location.Vertebra_Direction_Posterior)" ) - last_t = _get_last_thoracic(poi) - last_l = _get_last_lumbar(poi) + out = {} poi = poi.copy() - cervical = compute_angel_between_two_points_( - poi, - Vertebra_Instance.C2, - Vertebra_Instance.C7, - "P", - MoveTo.TOP, - MoveTo.BOTTOM, - project_2D, - ) - thoracic = compute_angel_between_two_points_(poi, Vertebra_Instance.T1, last_t, "P", MoveTo.TOP, MoveTo.BOTTOM, project_2D) - lumbar = compute_angel_between_two_points_(poi, Vertebra_Instance.L1, last_l, "P", MoveTo.TOP, MoveTo.BOTTOM, project_2D) - return { - "cervical_lordosis": cervical, - "thoracic_kyphosis": thoracic, - "lumbar_lordosis": lumbar, - } + + for k, i in curvature_definition.items(): + out[k] = round( + 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), + 4, + ) + return out def _get_norm(poi: POI, id1: int | Vertebra_Instance, mv: MoveTo, location: Location, inv: int = 1) -> np.ndarray | None: # noqa: ARG001 @@ -394,38 +435,21 @@ def _get_norm(poi: POI, id1: int | Vertebra_Instance, mv: MoveTo, location: Loca if (a == b).all(): return None norm1_vert = unit_vector(a - b) * inv - # if mix: - # # This would mix the angle of two adjacent Vertebra. - # if mv == MoveTo.CENTER: - # return norm1_vert - # elif mv == MoveTo.BOTTOM: - # next_vert = id1.get_next_poi(poi) - # elif mv == MoveTo.TOP: - # next_vert = id1.get_previous_poi(poi) - # if next_vert is None: - # return norm1_vert - # if (next_vert, location) in poi: - # norm1_vert_2 = unit_vector(np.array(poi[next_vert, 50]) - np.array(poi[next_vert, location])) * inv - # norm1_vert = (norm1_vert + norm1_vert_2) / 2 + next_vert = None + if mv == MoveTo.CENTER: + return norm1_vert + elif mv == MoveTo.BOTTOM: + next_vert = id1.get_next_poi(poi) + elif mv == MoveTo.TOP: + next_vert = id1.get_previous_poi(poi) + if next_vert is None: + return norm1_vert + if (next_vert, location) in poi: + norm1_vert_2 = unit_vector(np.array(poi[next_vert, 50]) - np.array(poi[next_vert, location])) * inv + norm1_vert = (norm1_vert + norm1_vert_2) / 2 return norm1_vert -def _get_last_lumbar(poi: POI) -> Vertebra_Instance | None: - """Return the most inferior lumbar vertebra that has a centroid in ``poi``.""" - for i in list(reversed(Vertebra_Instance.lumbar()))[:5]: - if (i.value, 50) in poi: - return i - return None - - -def _get_last_thoracic(poi: POI) -> Vertebra_Instance | None: - """Return the most inferior thoracic vertebra that has a centroid in ``poi``.""" - for i in list(reversed(Vertebra_Instance.thoracic()))[:3]: - if (i.value, 50) in poi: - return i - return None - - def compute_max_cobb_angle( poi: POI, vertebrae_list=None, @@ -530,7 +554,7 @@ def compute_max_cobb_angle( if cos_dis < cos_new: cos_dis = cos_new apex = i.value - return max_angle, from_vert, to_vert, apex + return round(max_angle, 4), from_vert, to_vert, apex def compute_max_cobb_angle_multi( @@ -641,19 +665,40 @@ def compute_max_cobb_angle_multi( def _add_artificial_ivd(poi: POI) -> POI: - """Insert synthetic IVD centroids midway between adjacent vertebra centroids if missing.""" - ## ADD IVD if nessasary - if 100 not in poi.keys_subregion(): - last = None - last_id = 1 - for j in Vertebra_Instance.order(): - if (j, 50) in poi: - current = np.array(poi[j, 50]) - if last is not None: - poi[last_id, 100] = tuple((last + current) / 2) - last = current - last_id = j.value - ##### + """Insert synthetic IVD landmarks (center + superior/inferior) wherever they are missing. + + For every adjacent pair of vertebrae (upper, lower) that both have a centroid, + a missing IVD landmark on the ``upper`` vertebra is synthesized as follows, + using the median endplate centers when available: + + - ``Vertebra_Disc_Superior`` (upper side of the disc) -> upper's inferior endplate median + (fallback: upper's centroid). + - ``Vertebra_Disc_Inferior`` (lower side of the disc) -> lower's superior endplate median + (fallback: lower's centroid). + - ``Vertebra_Disc`` (disc center) -> midpoint of the two above. + + This makes the Cobb / lordosis / kyphosis paths degrade gracefully when a + single IVD is missing (e.g. severe degeneration), instead of raising + KeyError. + """ + inf_med = Location.Additional_Vertebral_Body_Middle_Inferior_Median.value + sup_med = Location.Additional_Vertebral_Body_Middle_Superior_Median.value + disc = Location.Vertebra_Disc.value + disc_sup = Location.Vertebra_Disc_Superior.value + disc_inf = Location.Vertebra_Disc_Inferior.value + + ordered = [v for v in Vertebra_Instance.order() if (v.value, 50) in poi] + for i in range(len(ordered) - 1): + uv = ordered[i].value + lv = ordered[i + 1].value + up = np.array(poi[uv, inf_med]) if (uv, inf_med) in poi else np.array(poi[uv, 50]) + lo = np.array(poi[lv, sup_med]) if (lv, sup_med) in poi else np.array(poi[lv, 50]) + if (uv, disc_sup) not in poi: + poi[uv, disc_sup] = tuple(up) + if (uv, disc_inf) not in poi: + poi[uv, disc_inf] = tuple(lo) + if (uv, disc) not in poi: + poi[uv, disc] = tuple((up + lo) / 2) return poi @@ -705,36 +750,30 @@ def plot_compute_lordosis_and_kyphosis( poi = _add_artificial_ivd(poi) out = [] text_out = [] - last_t = _get_last_thoracic(poi) - last_l = _get_last_lumbar(poi) - for id1, vert_id1_mv in [ - (Vertebra_Instance.C2, MoveTo.TOP), - (Vertebra_Instance.T1, MoveTo.TOP), - (last_t, MoveTo.BOTTOM), - (last_l, MoveTo.BOTTOM), - ]: - vert_id1_mv: MoveTo - if id1 is None or (id1.value, 50) not in poi: - continue - s = vert_id1_mv.get_location(id1, poi) - a = _get_norm(poi, id1, vert_id1_mv, Location.Vertebra_Direction_Posterior, 1) - assert a is not None - out.append((id1.value, s, (a[0] * line_len, a[1] * line_len))) - out.append((id1.value, s, (-a[0] * line_len * 3, -a[1] * line_len * 3))) + for definition in curvature_definition.values(): + for id1, vert_id1_mv in [ + (definition.get_start_vert(poi), definition.start_move), + (definition.get_stop_vert(poi), definition.stop_move), + ]: + vert_id1_mv: MoveTo + if id1 is None or (id1.value, 50) not in poi: + continue + s = vert_id1_mv.get_location(id1, poi) + a = _get_norm(poi, id1, vert_id1_mv, Location.Vertebra_Direction_Posterior, 1) + assert a is not None + out.append((id1.value, s, (a[0] * line_len, a[1] * line_len))) + 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), id1, id2 in zip_strict( - out2.items(), [Vertebra_Instance.C7, last_t, last_l], [Vertebra_Instance.C2, Vertebra_Instance.C7, last_t] - ): + for name, v in out2.items(): if v is None: continue - if id1 is None or id2 is None or (id1.value, 50) not in poi: - continue + id1 = curvature_definition[name].get_start_vert(poi) + id2 = curvature_definition[name].get_stop_vert(poi) + vert = round((id1.value + id2.value) / 2) while (vert, 50) not in poi and vert != 0: vert -= 1 - if (vert, 50) in poi: # the body needs the key to be PRESENT; `not in` made it dead - cord = poi[vert, 50] - text_out.append((vert, (f"{str(name).split('_')[-1]}: {v:.1f}°", 15, cord[1]))) + text_out.append((vert, (f"{v:.1f}° - {str(name).split('_')[-1]}", 25))) poi.info["line_segments_sag"] = out + poi.info.get("line_segments_sag", []) poi.info["text_sag"] = text_out + poi.info.get("text_sag", []) @@ -823,6 +862,15 @@ def plot_cobb_angle( text_out.append((apex, (s, 25, 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: + 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) + poi = poi.resample_from_to(seg) frame = Snapshot_Frame( img, seg, @@ -837,8 +885,8 @@ def plot_cobb_angle( def plot_cobb_and_lordosis_and_kyphosis( - img_path: str | Path | None, - poi: POI, + jpg_path: str | Path | None, + poi: POI | Path, img: Image_Reference, seg: Image_Reference | None = None, line_len=100, @@ -852,7 +900,7 @@ def plot_cobb_and_lordosis_and_kyphosis( on the provided spinal image and can save the resulting image to a specified path. Args: - img_path (str | Path | None): Path to save the generated image. If None, the image is not saved. + jpg_path (str | Path | None): Path to save the generated image. If None, the image is not saved. poi (POI): The points of interest object containing 3D coordinates for various vertebrae. img (Image_Reference): The reference image on which to plot the angles and lines. seg (Image_Reference | None): The segmentation image reference. Optional, can be None. @@ -890,6 +938,8 @@ def plot_cobb_and_lordosis_and_kyphosis( >>> print(lordosis_kyphosis) {'cervical_lordosis': 35.2, 'thoracic_kyphosis': 41.5, 'lumbar_lordosis': 48.1} """ + if not isinstance(poi, POI): + poi = POI.load(poi) out_cobb, frame1 = plot_cobb_angle( None, poi, @@ -901,92 +951,6 @@ def plot_cobb_and_lordosis_and_kyphosis( project_2D=project_2D, ) out_lak, frame2 = plot_compute_lordosis_and_kyphosis(None, poi, img, seg, line_len=line_len, project_2D=project_2D) - if img_path is not None: - create_snapshot(img_path, [frame1, frame2]) + if jpg_path is not None: + create_snapshot(jpg_path, [frame1, frame2]) return out_cobb, out_lak, [frame1, frame2] - - -if __name__ == "__main__": - from TPTBox import POI, calc_poi_from_subreg_vert - from TPTBox.spine.spinestats.ivd_pois import compute_fake_ivd - - # poi = POI.load( - # "/DATA/NAS/datasets_processed/CT_spine/dataset-Cancer/derivatives_spineps/sub-mc0034/ses-20240312/sub-mc0034_ses-20240312_sequ-206_mod-ct_seg-spine_msk.nii.gz" - # ) - nii = to_nii( - "/DATA/NAS/datasets_processed/CT_spine/dataset-Cancer/derivatives_spineps/sub-mc0034/ses-20240312//sub-mc0034_ses-20240312_sequ-206_mod-ct_seg-vert_msk.nii.gz", - True, - ) - nii_subreg = to_nii( - "/DATA/NAS/datasets_processed/CT_spine/dataset-Cancer/derivatives_spineps/sub-mc0034/ses-20240312/sub-mc0034_ses-20240312_sequ-206_mod-ct_seg-spine_msk.nii.gz", - True, - ) - nii2 = to_nii( - "/DATA/NAS/datasets_processed/CT_spine/dataset-Cancer/rawdata/sub-mc0034/ses-20240312/sub-mc0034_ses-20240312_sequ-206_ct.nii.gz", - False, - ) - poi = calc_poi_from_subreg_vert(nii, nii_subreg, subreg_id=[Location.Vertebra_Direction_Right]) - - nii = compute_fake_ivd(nii, nii_subreg, poi=poi) - nii.save("/DATA/NAS/datasets_processed/CT_spine/dataset-Cancer/derivatives_spineps/sub-mc0034/ses-20240312/test.nii.gz") - print(nii.unique()) - poi = calc_poi_from_subreg_vert( - nii, - nii_subreg, - subreg_id=[ - Location.Vertebra_Direction_Right, - Location.Vertebra_Disc_Inferior, - Location.Vertebra_Disc, - ], - ) - idx = 23 - - print(poi.extract_vert(idx)) - # print(_get_norm(poi.rescale(), 24, None, Location.Vertebra_Direction_Right)) - # plot_compute_lordosis_and_kyphosis("test_2.png", poi, nii) - plot_cobb_angle("test.png", poi, nii2, nii, use_ivd_direction=True) - plot_cobb_angle("test_old.png", poi, nii2, nii, use_ivd_direction=False) - from TPTBox.core.poi_fun.ray_casting import add_ray_to_img - - cor, _ = poi.fit_spline(location=50, vertebra=False) - print(nii.shape) - print( - poi[idx, 50], - unit_vector(np.array(poi[idx, 50]) - np.array(poi[idx, Location.Vertebra_Direction_Right])), - ) - a = add_ray_to_img( - poi[idx, 50], - -np.array(poi[idx, 50]) + np.array(poi[idx, Location.Vertebra_Direction_Right]), - nii, - True, - value=99, - dilate=2, - ) - assert a is not None - a = add_ray_to_img( - poi[idx, 50], - -np.array(poi[idx, 50]) + np.array(poi[idx, Location.Vertebra_Direction_Posterior]), - a, - True, - value=100, - dilate=2, - ) - assert a is not None - a = add_ray_to_img( - poi[idx, 100], - -np.array(poi[idx, 100]) + np.array(poi[idx, Location.Vertebra_Disc_Inferior]), - a, - True, - value=101, - dilate=2, - ) - assert a is not None - spline = a.copy() * 0 - # spline.rescale_() - for x, y, z in cor: - spline[round(x), round(y), round(z)] = 103 - spline.dilate_msk_(2) - # spline.resample_from_to_(a) - a[spline != 0] = spline[spline != 0] - print(a.unique()) - a.save("/DATA/NAS/datasets_processed/CT_spine/dataset-Cancer/derivatives_spineps/sub-mc0034/ses-20240312/test.nii.gz") diff --git a/TPTBox/spine/spinestats/body_quadrants.py b/TPTBox/spine/spinestats/body_quadrants.py index 7ebbbe5d..4326b397 100644 --- a/TPTBox/spine/spinestats/body_quadrants.py +++ b/TPTBox/spine/spinestats/body_quadrants.py @@ -98,10 +98,9 @@ def make_quadrants( subreg_id=[ Location.Vertebra_Corpus, Location.Vertebra_Direction_Inferior, - # Vertebra_Direction_Right is read below; it used to be a second copy of - # Inferior, so the lookup always raised KeyError and every vertebra was - # skipped -> make_quadrants silently returned an all-zero image. + Location.Vertebra_Direction_Inferior, Location.Vertebra_Direction_Right, + Location.Endplate, ], buffer_file=poi_buffer, ) diff --git a/TPTBox/spine/spinestats/make_endplate.py b/TPTBox/spine/spinestats/make_endplate.py deleted file mode 100644 index 77996368..00000000 --- a/TPTBox/spine/spinestats/make_endplate.py +++ /dev/null @@ -1,231 +0,0 @@ -from __future__ import annotations - -from enum import Enum - -import numpy as np -from skimage.measure import label -from skimage.morphology import ball, binary_dilation, binary_erosion, disk -from sklearn.cluster import KMeans - -from TPTBox import NII -from TPTBox.core.poi import POI -from TPTBox.core.poi_fun.vertebra_direction import get_direction -from TPTBox.core.vert_constants import Location, Vertebra_Instance - -""" -Author: Amirhossein Bayat -amir.bayat@tum.de -""" -vertebra_body = ( - 49, - 50, - Location.Endplate.value, - Location.Vertebral_Body_Endplate_Inferior.value, - Location.Vertebral_Body_Endplate_Superior.value, -) - - -def _dilate_erode_special(np_array: np.ndarray, ball_size: int = 3, normal: np.ndarray | None = None) -> np.ndarray: - """Apply a morphological open (dilate then erode) with an axis-aligned or isotropic structuring element.""" - if normal is None: - struct = ball(ball_size) - else: - d = disk(ball_size) - struct = np.stack([d, d, d], axis=np.argmax(normal)) - # print(struct) - np_array = binary_dilation(np_array, footprint=struct) - np_array = binary_erosion(np_array, footprint=struct) - return np_array - - -def _get_largest_CC(segmentation: np.ndarray) -> np.ndarray: - """Extract the largest connected component from a binary segmentation. - - Args: - segmentation: Input binary segmentation array of any dimensionality. - - Returns: - Binary integer array of the same shape containing only the largest - connected component (background label 0 is excluded). - """ - labels = label(segmentation) - unique, counts = np.unique(labels, return_counts=True) # type: ignore - largest_label = unique[1:][np.argmax(counts[1:])] # Ignore the background label (0) - return (labels == largest_label).astype(int) - - -def _get_endplate(body: np.ndarray, mult: np.ndarray, axis: int = 1) -> np.ndarray: - """Extract a single endplate surface from a vertebra body via projection and K-means clustering. - - The ``mult`` array encodes depth along the normal direction. K-means - with three clusters is applied to the per-pixel argmax values to separate - the true endplate layer from noise and the opposite surface. - - Args: - body: 3-D binary array of the vertebral body voxels. - mult: 3-D integer array of the same shape as ``body`` containing - depth indices along the projection axis. - axis: Image axis along which the projection (argmax) is taken. - - Returns: - 3-D binary array marking the extracted endplate voxels. - """ - body = body * mult - - indices = np.argmax(body, axis=axis) - k_means = KMeans(n_clusters=3, random_state=0, n_init="auto").fit(indices.reshape(-1, 1)) - mask = k_means.labels_.reshape(indices.shape[0], indices.shape[1]) - mask += 1 - - mean_1 = np.mean([mask == 1] * indices) - mean_2 = np.mean([mask == 2] * indices) - mean_3 = np.mean([mask == 3] * indices) - if np.abs(mean_1 - mean_2) < 1: - mask[mask == 2] = 1 - if np.abs(mean_1 - mean_3) < 1: - mask[mask == 3] = 1 - if np.abs(mean_2 - mean_3) < 1: - mask[mask == 3] = 2 - - argmax_ind = np.argmax([np.mean([mask == 1] * indices), np.mean([mask == 2] * indices), np.mean([mask == 3] * indices)]) + 1 - mask[mask != argmax_ind] = 0 - mask[mask != 0] = 1 - - mask = mask * indices - mask[mask == 0] = -1 - - if axis == 0: - tmp = (np.arange(body.shape[0]) == mask[..., None]).astype(int) - tmp = np.transpose(tmp, (2, 0, 1)) - elif axis == 2: - tmp = (np.arange(body.shape[2]) == mask[..., None]).astype(int) - elif axis == 1: - tmp = (np.arange(body.shape[1]) == mask[..., None]).astype(int) - tmp = np.swapaxes(tmp, 1, 2) - else: - raise ValueError(axis) - return tmp - - -def _extract_endplate_np(body: np.ndarray, projected: np.ndarray, normal: np.ndarray, lower: bool = False) -> np.ndarray: - """Extract the superior or inferior endplate from a body array. - - Args: - body: 3-D binary array of the segmented vertebral body. - projected: 3-D integer depth-index array aligned with the vertebra - normal direction. - normal: Normal vector of the endplate plane, used to choose the - structuring element orientation for morphological smoothing. - lower: When ``True`` the inferior (lower) endplate is extracted; - otherwise the superior (upper) endplate is returned. - - Returns: - 3-D binary array marking the extracted endplate voxels. - """ - # Adjust projection for lower endplate if needed - if lower: - projected = projected.max() + 1 - projected - - endplate_mask = _get_endplate(body, projected) - endplate_mask = _get_largest_CC(_dilate_erode_special(endplate_mask, ball_size=3, normal=normal)) * np.clip(endplate_mask, 0, 1) - return endplate_mask - - -def _endplate_extraction_msk(subreg_tmp: NII, normal: np.ndarray, _extract: bool = True) -> NII: - """Extract superior and inferior endplates from a NIfTI segmentation and label them. - - Projects the body along the vertebra normal vector, then calls - :func:`_extract_endplate_np` twice (for superior and inferior surfaces). - Labels are set to - :attr:`~TPTBox.Location.Vertebral_Body_Endplate_Superior` and - :attr:`~TPTBox.Location.Vertebral_Body_Endplate_Inferior`. - - Args: - subreg_tmp: NIfTI segmentation containing vertebra body labels. - normal: Normal vector of the endplate plane used to orient the - projection and structuring element. - _extract: When ``True`` only voxels belonging to ``vertebra_body`` - labels are considered; when ``False`` all non-zero voxels are used. - - Returns: - Updated :class:`~TPTBox.NII` with endplate labels set. - """ - out_arr = np.zeros_like(subreg_tmp.get_array()) - if _extract: - subreg_tmp = subreg_tmp.extract_label(vertebra_body) - body = subreg_tmp.get_array() - - # Compute the projection along the normal vector - grid = np.mgrid[0 : body.shape[0], 0 : body.shape[1], 0 : body.shape[2]] - projected = np.tensordot(grid, normal, axes=(0, 0)) # type: ignore - # Normalize and convert to integers starting from 1 - projected -= projected.min() # Shift to start from 0 - projected = projected / projected.max() * (body.shape[1] - 1) # Scale to body dimensions - projected = np.round(projected).astype(int) + 1 # Convert to integers starting from 1 - - out_arr[_extract_endplate_np(body, projected, normal) == 1] = Location.Vertebral_Body_Endplate_Superior.value - out_arr[_extract_endplate_np(body, projected, normal, lower=True) == 1] = Location.Vertebral_Body_Endplate_Inferior.value - - return subreg_tmp.set_array(out_arr) - - -def endplate_extraction( - idx: int | Enum, - vert: NII, - subreg: NII, - poi: POI, -) -> NII | None: - """Compute the superior and inferior endplates for a single vertebra. - - The vertebra is reoriented to a canonical orientation, cropped, and - the endplate surfaces are extracted using projection and K-means - clustering along the superior–inferior normal direction. Sacral - vertebrae (except S1) are skipped. - - Args: - idx: Vertebra region label (integer or :class:`~enum.Enum` subtype). - vert: Full vertebra segmentation NIfTI. - subreg: Subregion segmentation NIfTI containing vertebra body labels. - poi: POI object used to retrieve the superior–inferior direction for - the vertebra. - - Returns: - A :class:`~TPTBox.NII` with endplate labels set to - :attr:`~TPTBox.Location.Vertebral_Body_Endplate_Superior` and - :attr:`~TPTBox.Location.Vertebral_Body_Endplate_Inferior`, or - ``None`` when the vertebra is a sacral level or no direction POI is - available. - """ - if isinstance(idx, Enum): - idx = idx.value - if Vertebra_Instance(idx) in Vertebra_Instance.sacrum()[1:]: - return None - - vert_ = vert.reorient(verbose=False) - subreg = subreg.reorient(verbose=False) - poi = poi.reorient() - out = vert_.extract_label(idx) * subreg.extract_label(vertebra_body, keep_label=True) - crop = out.compute_crop(dist=3) - out_c = out.apply_crop(crop) - try: - normal = get_direction("S", poi, vert_id=idx) / np.array(poi.zoom) - except KeyError: - return None - out_c = _endplate_extraction_msk(out_c, normal, _extract=False) - out[crop] = out_c - out.reorient_(vert.orientation, verbose=False) - return out - - -if __name__ == "__main__": - from TPTBox import Location, calc_poi_from_subreg_vert - from TPTBox.tests.test_utils import get_test_ct, get_tests_dir - - ct, subreg, vert, idx = get_test_ct() - vert.reorient_() - subreg.reorient_() - poi = calc_poi_from_subreg_vert(vert, subreg, subreg_id=Location.Vertebra_Direction_Posterior) - subreg2 = endplate_extraction(idx, vert, subreg, poi) - assert subreg2 is not None - subreg2.reorient_(ct.orientation) - subreg2.save(get_tests_dir() / "sample_ct" / "endplate-test.nii.gz") diff --git a/TPTBox/spine/spinestats/measure_ivd_and_vertebra_geometry.py b/TPTBox/spine/spinestats/measure_ivd_and_vertebra_geometry.py new file mode 100644 index 00000000..169b686a --- /dev/null +++ b/TPTBox/spine/spinestats/measure_ivd_and_vertebra_geometry.py @@ -0,0 +1,783 @@ +"""Geometric and signal measurements for intervertebral discs (IVDs) and vertebrae. + +Only :func:`measure_ivd_and_vertebra_geometry` is part of the public API. +Everything else is an implementation detail (prefixed with ``_``) and may +change without notice. + +IVD vs. vertebra +---------------- +The pipeline works for **both** structure types because orientation is +determined differently depending on the label: + +- **Vertebra** (label < ``structure_label``, i.e. the usual case ``structure_label=100`` + and label < 100): anatomical directions (inferior/posterior/right) are + read directly from precomputed points of interest (POIs), since a + vertebra's shape does not reliably reveal its own orientation via PCA. + The exception is the dens axis (label 2 / C2), whose unusual shape makes + the POI-based "up" direction unreliable, so it is replaced by a + PCA-estimated axis. + +- **IVD** (label >= ``structure_label``): discs have no POIs of their own, so + the "up" direction is estimated directly from the disc's own voxel mask + via PCA (:func:`_pca_principal_axes`). The right/posterior axes needed + for x1-x6 are instead derived from the neighbouring vertebra's bony + landmarks (spinous process / vertebral arch), see + :func:`_estimate_right_posterior_axes`. + +Which mode is used for a given call is controlled by the ``structure_label`` +argument of :func:`measure_ivd_and_vertebra_geometry` (and threaded down to +:func:`_compute_basic_geometry`): pass the default ``100`` to evaluate IVD +labels (i.e. label > 100), or ``0`` to evaluate vertebra labels instead +(i.e. label > 0). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from math import ceil +from pathlib import Path + +import numpy as np +import trimesh +from numpy.linalg import norm +from skimage import measure +from sklearn.decomposition import PCA + +from TPTBox import NII, POI, Location, Vertebra_Instance, calc_poi_from_subreg_vert +from TPTBox.core.nii_wrapper import NII +from TPTBox.core.poi import POI +from TPTBox.spine.spinestats.torso_vat_sat import peak_centered_mean + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def measure_ivd_and_vertebra_geometry( + t2w: NII | None, + vert: NII, + spine: NII, + buffer_poi: Path | None = None, + step_size_mm: float = 0.5, + instance_labels: list[int] | None = None, + structure_label: int = 100, + erode=1, +) -> dict[int, dict[str, float]]: + """Extract geometric (and optionally T2 signal) measurements per structure. + + Works for intervertebral discs (IVDs) as well as vertebrae, see the + module docstring for how the two modes differ internally. Which one is + used is selected via ``structure_label``. + + Parameters + ---------- + t2w : NII, optional + T2-weighted image used to compute the normalized disc/vertebra + signal. If ``None``, the ``"signal"`` entry is left as ``NaN``. + + vert : NII + Instance segmentation (each disc or vertebra has a unique label). + + spine : NII + Semantic subregion segmentation (e.g. spinous process, vertebral + arch, vertebral body, spinal canal) used to estimate anatomical + orientation and, for the signal ratio, the spinal canal reference + region. + + step_size_mm : float, default=0.5 + Grid spacing (in mm) used when sampling height/diameter profiles + across the structure's surface. Smaller values are more accurate + but proportionally slower. + + instance_labels : list[int], optional + Labels to evaluate. If ``None``, defaults to every label in + ``vert`` greater than ``structure_label``. + + structure_label : int, default=100 + Selects IVD mode (``100``, the default: only labels > 100 are + valid) or vertebra mode (``0``: only labels > 0 are valid). See the + module docstring for details. Any value other than ``100`` is + currently treated like vertebra mode for the minimum-label check. + + Returns: + ------- + dict[int, dict[str, float]] + Keyed by structure label. Each entry contains: + + - ``volume_voxel``: volume from the raw voxel mask (mm^3) + - ``volume_mesh``: volume from the reconstructed surface mesh (mm^3) + - ``height_center``: height measured through the structure's center + - ``mean_height`` / ``max_height``: statistics over sampled heights + - ``lower_10_percent_height``: 10th percentile of sampled heights + - ``mean_diameter``: diameter of a circle with the same projected + area as the structure + - ``anterior_height_x1``, ``posterior_height_x2``, + ``right_height_x3``, ``left_height_x4``: directional heights, + see the figure/table below + - ``width_lateral_x5``, ``width_sagittal_x6``: directional widths + - ``signal``: peak-centered T2 signal in the structure divided + by the peak-centered T2 signal in the spinal canal (unitless). + Peak-centered means only voxels around the histogram mode of + the mask are averaged; this suppresses darker contamination + such as nerve roots inside the canal (only if ``t2w`` is given) + - ``structure_signal``: peak-centered mean T2 signal inside the + (eroded) structure mask, a.u. (only if ``t2w`` is given) + - ``spinal_canal_signal``: peak-centered mean T2 signal inside + the (eroded) spinal canal reference region, a.u. (only if + ``t2w`` is given) + - ``signal_old`` / ``structure_signal_old`` / + ``spinal_canal_signal_old``: same three quantities computed as + plain per-voxel means (no peak centering). Kept for backward + comparison so a user can see the effect of the peak filter. + + If evaluation of a label fails, its entry instead contains + ``"error"`` (the exception message) plus all the same keys set to + ``NaN``. + + Notes: + ----- + x1-x6 are the six anatomical dimensions commonly used to describe disc + (or vertebral body) geometry: + + ==== ========================= + x1 anterior height + x2 posterior height + x3 right height + x4 left height + x5 lateral width + x6 sagittal width + ==== ========================= + """ + results = {} + + # Vertebral direction landmarks (inferior/posterior/right), needed for + # vertebra mode and as an orientation anchor for the neighbouring + # vertebra in IVD mode. + poi = calc_poi_from_subreg_vert( + vert, + spine, + subreg_id=[ + Location.Vertebra_Direction_Inferior, + Location.Vertebra_Direction_Posterior, + Location.Vertebra_Direction_Right, + Location.Vertebra_Corpus, + Location.Endplate, + ], + buffer_file=buffer_poi, + save_buffer_file=True, + ) + if instance_labels is None: + instance_labels = [int(i) for i in vert.unique() if i > structure_label and i < structure_label + 100] + for label in instance_labels: + if label == 26: + continue + try: + raw = {} + # Isolate the current structure (disc or vertebra). + structure_mask = vert.extract_label(label) + # 1. volume, central height, mean diameter + info, raw = _compute_basic_geometry(vert, poi, label, step_size_mm=2, raw=raw, structure_label=structure_label) + # 2. x1-x6 directional heights/widths + raw = _compute_directional_heights_widths(vert, structure_mask * spine, poi, label, step_size_mm=step_size_mm, raw=raw) + # 3. normalized T2 signal + if t2w is not None: + raw = _compute_t2_signal_ratio(t2w, vert, spine, label, raw=raw, erode=erode) + results[label] = _result_from_info(info) + except Exception as e: + results[label] = _nan_result(error=str(e)) + + return results + + +# --------------------------------------------------------------------------- +# Result assembly +# --------------------------------------------------------------------------- + +_RESULT_FIELDS = ( + "volume_voxel", + "volume_mesh", + "height_center", + "mean_height", + "max_height", + "lower_10_percent_height", + "mean_diameter", + "anterior_height_x1", + "posterior_height_x2", + "right_height_x3", + "left_height_x4", + "width_lateral_x5", + "width_sagittal_x6", + "signal", + "structure_signal", + "spinal_canal_signal", + "signal_old", + "structure_signal_old", + "spinal_canal_signal_old", +) + + +def _result_from_info(info: _StructureMeasurements) -> dict[str, float]: + """Build the public result dict (see Returns section of the public API) from a filled-in measurement object.""" + return { + "volume_voxel": round(info.volume_voxel, 4), + "volume_mesh": round(info.volume_mesh, 4), + "height_center": round(info.height_center, 4), + "mean_height": round(info.mean_height, 4), + "max_height": round(info.max_height, 4), + "lower_10_percent_height": round(info.get_quantile(10), 4), + "mean_diameter": round(info.mean_diameter, 4), + "anterior_height_x1": round(info.anterior_height_x1, 4), + "posterior_height_x2": round(info.posterior_height_x2, 4), + "right_height_x3": round(info.right_height_x3, 4), + "left_height_x4": round(info.left_height_x4, 4), + "width_lateral_x5": round(info.width_lateral_x5, 4), + "width_sagittal_x6": round(info.width_sagittal_x6, 4), + "signal": round(info.signal, 4), + "structure_signal": round(info.structure_signal, 4), + "spinal_canal_signal": round(info.spinal_canal_signal, 4), + "signal_old": round(info.signal_old, 4), + "structure_signal_old": round(info.structure_signal_old, 4), + "spinal_canal_signal_old": round(info.spinal_canal_signal_old, 4), + } # type: ignore + + +def _nan_result(error: str | None = None) -> dict[str, float]: + """Build a NaN-filled result dict, e.g. for a structure whose evaluation raised an exception.""" + result = dict.fromkeys(_RESULT_FIELDS, np.nan) + if error is not None: + result["error"] = error + return result + + +# --------------------------------------------------------------------------- +# Measurement container +# --------------------------------------------------------------------------- + + +@dataclass +class _StructureMeasurements: + """Accumulates the measurements for a single disc/vertebra as they are computed in stages. + + Populated incrementally by :func:`_compute_basic_geometry` (volume/height/area), + :func:`_compute_directional_heights_widths` (x1-x6), and + :func:`_compute_t2_signal_ratio` (signal). The ``*_values`` flags mark which + stages have already run so repeated calls (e.g. across retries) don't + redo expensive work. + """ + + volume_voxel: float + volume_mesh: float + height_center: float + area: float + quantiles: dict = None # type: ignore + anterior_height_x1: float = np.nan + posterior_height_x2: float = np.nan + right_height_x3: float = np.nan + left_height_x4: float = np.nan + width_lateral_x5: float = np.nan + width_sagittal_x6: float = np.nan + x_values: bool = False + signal_values: bool = False + signal: float = np.nan + structure_signal: float = np.nan + spinal_canal_signal: float = np.nan + signal_old: float = np.nan + structure_signal_old: float = np.nan + spinal_canal_signal_old: float = np.nan + + @property + def mean_diameter(self): + """Diameter of the circle whose area equals the structure's projected area (A = pi*r^2).""" + return np.sqrt(self.area / np.pi) * 2 + + def _update_height_statistics(self, sampled_heights): + """Compute mean/max/quantiles from a list of sampled point heights.""" + if hasattr(self, "list_heights"): + delattr(self, "list_heights") + sampled_heights = sorted(sampled_heights) + if len(sampled_heights) != 0: + self.mean_height = np.mean(sampled_heights) + self.max_height = max(sampled_heights) + for q in range(0, 100, 5): + self.get_quantile(q, sampled_heights) + else: + self.mean_height = np.nan + self.max_height = np.nan + + def get_quantile(self, q: int = 50, _sampled_heights: list | None = None): + """Return the q-th percentile (0-100) of sampled heights, computing and caching it on first use.""" + if self.quantiles is None: + self.quantiles = {} + if q in self.quantiles: + return self.quantiles[q] + if _sampled_heights is None: + return None + self.quantiles[q] = _sampled_heights[int(len(_sampled_heights) * q / 100)] + return self.quantiles[q] + + +# --------------------------------------------------------------------------- +# Mesh / geometry primitives +# --------------------------------------------------------------------------- + + +def _segmentation_to_surface_mesh(segmentation: np.ndarray, voxel_size) -> trimesh.Trimesh: + """Turn a binary voxel mask into a surface mesh via marching cubes, scaled to real-world (mm) units.""" + verts, faces, _, _ = measure.marching_cubes(segmentation) + verts *= voxel_size + return trimesh.Trimesh(vertices=verts, faces=faces) + + +def _pca_principal_axes(segmentation: np.ndarray, voxel_size, verbose=False, up_axis=-1): + """Estimate an orthogonal axis system from a voxel mask's principal component axes (PCA). + + Used when no anatomical landmarks are available to derive orientation + directly from the shape of the structure itself (e.g. IVDs, or the C2 + vertebra whose dens makes the usual POI-based "up" direction unreliable). + + Parameters + ---------- + up_axis : int, default=-1 + If -1, the "up" vector is simply the 3rd principal component (least + variance / thinnest axis of the structure). Otherwise, "up" is + chosen as whichever principal component has the largest projection + onto this array axis (e.g. the superior/inferior image axis), + which is more robust when the structure isn't disc-shaped. + + Returns: + ------- + tuple[np.ndarray, np.ndarray, np.ndarray] + (up_vector, remaining_axis_1, remaining_axis_2) + """ + points = np.argwhere(segmentation > 0) * voxel_size + pca = PCA(n_components=3) + pca.fit(points) + if up_axis == -1: + up_vector = pca.components_[2] + if verbose: + print(f"Main Axis (PC1): {pca.components_[0]}") + print(f"Secondary Axis (PC2): {pca.components_[1]}") + print(f"Up Vector (PC3): {up_vector}") + return up_vector, pca.components_[0], pca.components_[1] + + # Pick whichever component best aligns with the requested image axis. + abs_up_values = np.abs(pca.components_[:, up_axis]) + max_index = np.argmax(abs_up_values) + up_vector = pca.components_[max_index] + return ( + up_vector, + pca.components_[(max_index + 1) % 3], + pca.components_[(max_index + 2) % 3], + ) + + +def _center_of_mass_voxels(segmentation: np.ndarray, voxel_size=1) -> np.ndarray: + """Voxel-space center of mass of a binary mask, scaled to real-world units.""" + points = np.argwhere(segmentation > 0) + return np.mean(points, axis=0) * voxel_size + + +def _project_mesh_onto_plane(mesh: trimesh.Trimesh, up_vector: np.ndarray) -> trimesh.Trimesh: + """Flatten a mesh onto the plane perpendicular to ``up_vector`` (orthographic projection along "up").""" + up_vector = up_vector / np.linalg.norm(up_vector) + flattened_vertices = mesh.vertices - np.dot(mesh.vertices, up_vector[:, np.newaxis]) * up_vector + return trimesh.Trimesh(vertices=flattened_vertices, faces=mesh.faces) + + +def _projected_area(flattened_mesh: trimesh.Trimesh) -> float: + """2D area of a mesh already flattened onto a plane (trimesh's area works for degenerate/2D meshes too).""" + return flattened_mesh.area + + +# --------------------------------------------------------------------------- +# Ray casting helpers +# --------------------------------------------------------------------------- + + +def _intersect_ray(ray_vector, v1, v2, mesh: trimesh.Trimesh, x: float = 0.0, y: float = 0.0, center=None): + """Cast a ray through the mesh along ``ray_vector``, offset from ``center`` by ``x``/``y`` along ``v1``/``v2``. + + Used both to measure heights (ray along the "up" direction) and to + measure diameters (ray along an in-plane direction), depending on which + vector is passed as ``ray_vector``. + """ + start_point = mesh.center_mass if center is None else center + # Start far outside the mesh (on the opposite side of ray_vector) so the ray reliably enters through the surface. + start_point = start_point - ray_vector * 1000 + v1 * x + v2 * y + intersections = mesh.ray.intersects_location(ray_origins=np.array([start_point]), ray_directions=np.array([ray_vector])) + return intersections[0] + + +def _intersect_ray_from_dirs(dirs: tuple[np.ndarray, np.ndarray, np.ndarray], mesh: trimesh.Trimesh, x=0, y=0, center=None): + """Convenience wrapper around :func:`_intersect_ray` that unpacks a ``(up, v1, v2)`` direction triple.""" + up_vector, v1, v2 = dirs + return _intersect_ray(up_vector, v1, v2, mesh=mesh, x=x, y=y, center=center) + + +def _segment_length(intersection_points) -> float: + """Euclidean distance between the first and last point where a ray crossed the mesh surface. + + For a ray shot straight through a (roughly convex) structure, this is + the entry-to-exit distance, i.e. the structure's height/diameter along + that ray. + """ + if len(intersection_points) == 0: + return 0 + return np.sqrt(sum((a - b) ** 2 for a, b in zip(intersection_points[0], intersection_points[-1], strict=False))) + + +def _swap(a, b): + """Return ``(b, a)``.""" + return b, a + + +# --------------------------------------------------------------------------- +# Mesh + orientation extraction (shared between IVD and vertebra mode) +# --------------------------------------------------------------------------- + + +def _get_mesh_and_directions(nii: NII, poi: POI | None, label: int, raw: dict, recompute_mesh: bool = False): + """Build (or fetch from cache) the surface mesh and ``(up, posterior, right)`` direction vectors for one structure. + + This is the key place where IVD and vertebra handling diverge: + + - **Vertebra** (``label < 100``): directions come from the POI + landmarks (``Vertebra_Direction_Inferior/Posterior/Right``), which + are far more reliable than PCA for the often irregular vertebral + shape. Exception: label 2 (the C2/dens vertebra) still gets its "up" + vector from PCA, since its odd shape makes the POI-derived "up" + direction unreliable. + - **IVD** (``label >= 100``): there are no POIs for discs, so all three + directions are estimated from the disc's own voxel mask via PCA + (:func:`_pca_principal_axes`). + + Results are cached in ``raw`` so repeated calls for the same structure + reuse the mesh/directions instead of recomputing them. + + Returns: + ------- + tuple[trimesh.Trimesh, tuple[np.ndarray, np.ndarray, np.ndarray], NII] + (surface mesh, (up, dir_1, dir_2) direction vectors, binary segmentation) + """ + voxel_size = nii.zoom + nii = nii.apply_crop(nii.compute_crop(dist=1)) + segmentation = nii.extract_label(label) + arr = segmentation.get_array() + + if label < 100: + assert poi is not None + center = np.array(poi[label, Location.Vertebra_Corpus]) + up = -np.array(poi[label, Location.Vertebra_Direction_Inferior]) - center + post = np.array(poi[label, Location.Vertebra_Direction_Posterior]) - center + right = np.array(poi[label, Location.Vertebra_Direction_Right]) - center + up /= norm(up) + post /= norm(post) + right /= norm(right) + if label == 2: + up = _pca_principal_axes(arr, voxel_size, up_axis=nii.get_axis("S"))[0] + direction_vectors = (up, post, right) + else: + up_axis = nii.get_axis("S") if label < 100 else -1 + direction_vectors = ( + raw["direction_vectors"] if "direction_vectors" in raw else _pca_principal_axes(arr, voxel_size, up_axis=up_axis) + ) + + mesh = raw["mesh"] if "mesh" in raw and not recompute_mesh else _segmentation_to_surface_mesh(arr, voxel_size) + raw["mesh"] = mesh + raw["direction_vectors"] = direction_vectors + return mesh, direction_vectors, segmentation + + +# --------------------------------------------------------------------------- +# Measurement stages (called in order from measure_ivd_and_vertebra_geometry) +# --------------------------------------------------------------------------- + +# Maps a `structure_label` value to the minimum label a structure must exceed to be valid. +# `structure_label=100` -> IVD mode (labels must be > 100); any other value -> vertebra mode (labels must be > 0). +_MIN_LABEL_FOR_OFFSET = {100: 100} + + +def _compute_basic_geometry( + nii: NII, poi: POI, label: int, step_size_mm: int = 2, raw: dict | None = None, structure_label: int = 100 +) -> tuple[_StructureMeasurements, dict]: + """Compute basic geometric measurements of a single disc or vertebra (stage 1). + + How it's computed + ------------------ + - ``volume_voxel``: voxel count times voxel volume. + - ``volume_mesh``: signed volume of the reconstructed surface mesh. + - ``height_center``: length of the ray-mesh intersection segment through + the mesh's center of mass, along the "up" direction. + - ``area`` / ``mean_diameter``: projected area of the mesh flattened + along "up", converted to an equivalent circle diameter. + - ``mean_height`` / ``max_height`` / height quantiles: heights are + sampled on a grid (spacing ``step_size_mm``) covering a disc of + radius ``mean_diameter`` around the center, each height being the + ray-mesh intersection length at that grid point. + + Expensive intermediate results (mesh, direction vectors, sampled + heights) are cached in ``raw`` and reused by later stages/calls. + + Parameters + ---------- + nii : NII + Instance segmentation containing the structure. + poi : POI + Vertebral points of interest, used for orientation (vertebra mode only). + label : int + Label of the structure to evaluate. + step_size_mm : float, default=2.0 + Sampling spacing (mm) for the height grid. + raw : dict, optional + Cache of previously computed mesh / direction vectors / measurements / sampled heights. + structure_label : int, default=100 + Selects IVD mode (100, labels must be > 100) or vertebra mode + (any other value, labels must be > 0). See module docstring. + + Returns: + ------- + tuple[_StructureMeasurements, dict] + The computed measurements together with the updated cache. + """ + if raw is None: + raw = {} + + min_label = _MIN_LABEL_FOR_OFFSET.get(structure_label, 0) + assert label > min_label, label + + info = raw.get("info") + + # ------------------------------------------------------------------ + # Mesh-based quantities (volume, central height, projected area) + # ------------------------------------------------------------------ + if not isinstance(info, _StructureMeasurements) or info.height_center == 0: + mesh, direction_vectors, segmentation = _get_mesh_and_directions(nii, poi, label, raw, recompute_mesh=True) + volume_voxel = (segmentation.sum() * np.prod(segmentation.zoom)).item() + center_intersections = _intersect_ray_from_dirs(direction_vectors, mesh, 0, 0) + height_center = _segment_length(center_intersections) + area = _projected_area(_project_mesh_onto_plane(mesh.copy(), direction_vectors[0])) + info = _StructureMeasurements( + volume_voxel=volume_voxel, volume_mesh=abs(mesh.volume.item()), height_center=height_center, area=area + ) + raw["info"] = info + + # ------------------------------------------------------------------ + # Sampled heights (mean/max/quantiles) on a grid around the center + # ------------------------------------------------------------------ + if "list_heights" not in raw or len(raw["list_heights"]) == 0: + mesh, direction_vectors, _ = _get_mesh_and_directions(nii, poi, label, raw) + + search_radius = int(info.mean_diameter) + sampled_heights = [] + + for x in range(-search_radius, search_radius, step_size_mm): + for y in range(-search_radius, search_radius, step_size_mm): + intersections = _intersect_ray_from_dirs(direction_vectors, mesh, x, y) + height = _segment_length(intersections) + if height > 0: + sampled_heights.append(height) + + raw["list_heights"] = sampled_heights + + info._update_height_statistics(raw["list_heights"]) + + return info, raw + + +def _local_max_height(mesh, direction_vectors, around_point, search_diameter, step_size_mm) -> float: + """Maximum sampled height in a small neighbourhood around a point (used to pin down x1-x4 at each edge point). + + A single ray through, e.g., the exact anterior-most point can graze the + mesh edge and under-measure; sampling a small patch around the point + and taking the max is more robust. + """ + sampled_heights = [] + d = int(search_diameter / step_size_mm) + for xi in range(-d, d): + x = xi * step_size_mm + for yi in range(-d, d): + y = yi * step_size_mm + intersections = _intersect_ray_from_dirs(direction_vectors, mesh, x, y, around_point) + height = _segment_length(intersections) + if height != 0: + sampled_heights.append(height) + return np.max(sampled_heights) + + +def _max_diameter_in_plane(ray_vector, v1, v2, mesh, diameter: float = 30, step_size_mm: float = 2.0): + """Find the widest ray-mesh intersection on a grid in the plane spanned by ``v1``/``v2`` (stage 2 helper). + + Used for both the lateral (x5) and sagittal (x6) width measurements: + the mesh is probed with rays parallel to ``ray_vector`` (e.g. "right" + for the lateral width) on a grid spanned by ``v1``/``v2`` (e.g. + "posterior" and "up"), and the widest intersection found is returned + together with its two endpoints (used as the anterior/posterior or + left/right reference points for the height measurements). + + Returns: + ------- + tuple[float, np.ndarray, np.ndarray, float, float] + (max_width, point_1, point_2, grid_x, grid_y) for the widest ray found. + """ + d = ceil(diameter / step_size_mm) + best_width = 0 + out = (0.0, None, None, 0.0, 0.0) + for xi in range(-d, d): + x = xi * step_size_mm + for yi in range(-d, d): + y = yi * step_size_mm + intersections = _intersect_ray(ray_vector, v1, v2, mesh, x, y) + width = _segment_length(intersections) + if width == 0: + continue + p1 = intersections[0] + p2 = intersections[-1] + if best_width < width: + best_width = width + out = (width, p1.round(2), p2.round(2), x, y) + return out + + +def _compute_directional_heights_widths(nii: NII, subreg: NII, poi, label: int = 123, step_size_mm: float = 0.5, raw: dict | None = None): # noqa: ARG001 + """Compute the x1-x6 directional heights and widths for one structure (stage 2). + + How it's computed + ------------------ + 1. The right/posterior anatomical axes are estimated from the + this is shared between IVD and vertebra mode. + 2. ``width_lateral_x5``: widest ray parallel to "right", scanned over a + grid in the (posterior, up) plane -> also yields the right-most and + left-most surface points. + 3. ``width_sagittal_x6``: widest ray parallel to "posterior", scanned + over a grid in the (right, up) plane -> also yields the + posterior-most and anterior-most surface points. + 4. ``anterior_height_x1`` / ``posterior_height_x2`` / ``right_height_x3`` + / ``left_height_x4``: maximum sampled height (:func:`_local_max_height`) + in a small patch around each of the four extreme points found above. + + Requires the segmentation to be in an orientation with Right along the + 3rd axis and Posterior along the 1st axis (checked via ``nii.orientation``). + """ + assert "R" in nii.orientation[2] # No guarantee this works for other orientations. + assert "P" in nii.orientation[0] + if raw is None: + raw = {} + info: _StructureMeasurements = raw["info"] + if info.x_values: + return raw + try: + mesh, direction_vectors, _ = _get_mesh_and_directions(nii, poi, label, raw, recompute_mesh=True) + up = direction_vectors[0] + center = np.asarray(poi[label % 100, Location.Vertebra_Corpus], dtype=float) + + right = np.asarray(poi[label % 100, Location.Vertebra_Direction_Right], dtype=float) - center + front = np.asarray(poi[label % 100, Location.Vertebra_Direction_Posterior], dtype=float) - center + right /= norm(right) + front /= norm(front) + if label >= 100: + next_vert = Vertebra_Instance(label % 100).get_next_poi(poi) + right2 = np.asarray(poi[next_vert, Location.Vertebra_Direction_Right], dtype=float) - center + front2 = np.asarray(poi[next_vert, Location.Vertebra_Direction_Posterior], dtype=float) - center + right2 /= norm(right2) + front2 /= norm(front2) + right = (right + right2) / 2 + front = (front + front2) / 2 + + (width_lateral_x5, p_r, p_l, *_) = _max_diameter_in_plane(right, front, up, mesh, 30, step_size_mm) + axis = nii.get_axis("R") + flip = 1 if "R" in nii.orientation else -1 + if p_r[axis] * flip < p_l[axis] * flip: + p_r, p_l = _swap(p_r, p_l) + + (width_sagittal_x6, p_p, p_a, *_) = _max_diameter_in_plane(front, right, up, mesh, 2, step_size_mm) + axis = nii.get_axis("P") + flip = 1 if "P" in nii.orientation else -1 + if p_p[axis] * flip < p_a[axis] * flip: + # NOTE: kept as in the original implementation, though this looks like it + # should swap (p_p, p_a) rather than (p_r, p_a) -- flag before relying on x6 edge points. + p_p, p_a = _swap(p_r, p_a) + + info.anterior_height_x1 = _local_max_height(mesh, direction_vectors, p_a, 3, step_size_mm) + info.posterior_height_x2 = _local_max_height(mesh, direction_vectors, p_p, 3, step_size_mm) + info.right_height_x3 = _local_max_height(mesh, direction_vectors, p_r, 3, step_size_mm) + info.left_height_x4 = _local_max_height(mesh, direction_vectors, p_l, 3, step_size_mm) + info.width_lateral_x5 = width_lateral_x5 + info.width_sagittal_x6 = width_sagittal_x6 + info.x_values = True + except Exception: + info.anterior_height_x1 = None # type: ignore + info.posterior_height_x2 = None # type: ignore + info.right_height_x3 = None # type: ignore + info.left_height_x4 = None # type: ignore + info.width_lateral_x5 = None # type: ignore + info.width_sagittal_x6 = None # type: ignore + + return raw + + +def _compute_t2_signal_ratio( + t2w_nii: NII, + nii: NII, + subregs: NII, + label: int = 123, + raw: dict | None = None, + erode=1, + spinal_bins: int = 64, + spinal_peak_frac_height: float = 0.5, +): + """Compute the normalized T2 signal for one structure (stage 3). + + How it's computed + ------------------ + The mean T2 intensity inside the (slightly eroded, to avoid partial-volume + edge voxels) structure mask is divided by the mean T2 intensity in the + spinal canal (subregion label 61, also eroded). The spinal canal is used + as an internal reference to normalize away scanner/sequence-dependent + intensity scaling. + + The spinal canal segmentation may contain darker structures such as + nerve roots. To reduce their influence, both the structure and canal + signals are also estimated as the mean over a window around the + histogram peak (see :func:`peak_centered_mean`). Both the plain-mean + values (``*_old``) and the peak-centered values are stored so the + caller can compare them. + + Parameters + ---------- + spinal_bins : int, default=64 + Histogram bin count passed to :func:`peak_centered_mean`. + spinal_peak_frac_height : float, default=0.5 + Fractional height cutoff passed to :func:`peak_centered_mean`. + """ + assert "R" in nii.orientation[2] + assert "P" in nii.orientation[0] + if raw is None: + raw = {} + info: _StructureMeasurements = raw["info"] + if info.signal_values: + return raw + if t2w_nii.shape != nii.shape: + t2w_nii.resample_from_to_(nii, verbose=False) + structure_mask = nii.extract_label(label) + eroded_mask = structure_mask.erode_msk(erode, connectivity=1, verbose=False) + structure_mask = eroded_mask if eroded_mask.sum() != 0 else structure_mask + spinal_mask = subregs.extract_label(61).erode_msk(erode, connectivity=1, verbose=False) + + t2w_arr = t2w_nii.get_array() + structure_vals = t2w_arr[structure_mask.get_array().astype(bool)] + spinal_vals = t2w_arr[spinal_mask.get_array().astype(bool)] + + structure_signal_old = float(np.mean(structure_vals)) if structure_vals.size > 0 else np.nan + spinal_canal_signal_old = float(np.mean(spinal_vals)) if spinal_vals.size > 0 else np.nan + + structure_signal = peak_centered_mean(structure_vals, bins=spinal_bins, peak_frac_height=spinal_peak_frac_height) + spinal_canal_signal = peak_centered_mean(spinal_vals, bins=spinal_bins, peak_frac_height=spinal_peak_frac_height) + + info.signal = structure_signal / spinal_canal_signal + info.structure_signal = structure_signal + info.spinal_canal_signal = spinal_canal_signal + info.signal_old = structure_signal_old / spinal_canal_signal_old + info.structure_signal_old = structure_signal_old + info.spinal_canal_signal_old = spinal_canal_signal_old + info.signal_values = True + return raw diff --git a/TPTBox/spine/spinestats/poi_fun/__init__.py b/TPTBox/spine/spinestats/poi_fun/__init__.py new file mode 100644 index 00000000..6049266a --- /dev/null +++ b/TPTBox/spine/spinestats/poi_fun/__init__.py @@ -0,0 +1 @@ +# these files are called by vertebra_pois_non_centroids.py diff --git a/TPTBox/spine/spinestats/poi_fun/articularis_midpoint.py b/TPTBox/spine/spinestats/poi_fun/articularis_midpoint.py new file mode 100644 index 00000000..85cb25b7 --- /dev/null +++ b/TPTBox/spine/spinestats/poi_fun/articularis_midpoint.py @@ -0,0 +1,281 @@ +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path + +import numpy as np +from scipy.spatial import cKDTree + +from TPTBox import NII, POI, Image_Reference, Location, Logger_Interface, Print_Logger, Vertebra_Instance, to_nii + +_log = Print_Logger() + + +def calc_all_facet_joint_pois( + vert: Image_Reference, + subreg: Image_Reference, + ids: dict[str, int] | None = None, + max_gap_mm: float = 8.0, + surface_tolerance_mm: float = 1.5, + log: Logger_Interface = _log, +) -> POI: + """Call :func:`calc_facet_joint_pois` for every vertebra present in ``vert``.""" + vert_: NII = to_nii(vert, True) + subreg_ = to_nii(subreg, True) + poi = vert_.make_empty_POI() + for vert_id in sorted(i for i in vert_.unique() if i < 29): + if vert_id == 1 and vert_id != 26: + continue + + _calc_facet_joint_pois( + poi, + vert_, + subreg_, + vert_id, + ids=ids, + max_gap_mm=max_gap_mm, + surface_tolerance_mm=surface_tolerance_mm, + log=log, + ) + return poi + + +def _calc_facet_joint_pois( + poi: POI, + vert: NII, + subreg: NII, + vert_id: int, + ids: dict[str, int] | None = None, + max_gap_mm: float = 8.0, + surface_tolerance_mm: float = 1.5, + log: Logger_Interface = _log, +) -> POI: + """Compute the facet joint point between two adjacent vertebrae, for both sides. + + The point is placed between the Inferior_Articular process of ``vert_id`` (the + upper vertebra) and the Superior_Articular process of the next lower vertebra, + as determined by :meth:`Vertebra_Instance.get_next_poi` (vertebra labels are not + guaranteed to be strictly ascending, so this must be used instead of ``vert_id + 1``). + + The result lies in the middle of the (possibly extended) contact area: if the two + surfaces touch, this is simply the middle of the contact region; if there is a + gap, the closest facing points across the gap are averaged instead. Points that + cannot physically belong to the contact interface (i.e. they face away from the + other surface) are excluded before the nearest-neighbor search. + + If one of the two surfaces is missing (label not present, neighboring vertebra + absent, etc.), the existing centroid of the other location is used as a fallback + so the pipeline degrades gracefully. If both are missing, the point is skipped + and a warning is logged. + + Args: + poi: POI to extend in place. + vert: Vertebra instance segmentation. + subreg: Subregion/semantic segmentation, same grid as ``vert``. + vert_id: Label of the upper vertebra (the one owning "Inferior_Articular_*"). + ids: Subregion ids under which the result is stored, e.g. + ``{"left": Articular_Process_Midpoint_Left, "right": Articular_Process_Midpoint_Right}``. Choose ids that are free in your schema. + max_gap_mm: If the minimal distance between the two surfaces exceeds this, + no point is computed (joint likely absent or too far apart). + surface_tolerance_mm: Point pairs within (min_distance + tolerance) of each + other are averaged as "the contact surface" instead of using only the + single closest pair. + log: Logger for status messages. + + Returns: + POI: the same (extended) POI object. + """ + if ids is None: + ids = {"left": Location.Articular_Process_Midpoint_Left.value, "right": Location.Articular_Process_Midpoint_Right.value} + + all_ids = vert.unique() + if vert_id not in all_ids: + log.print(f"[Facet] vert_id {vert_id} not present, skipping", verbose=True) + return poi + + vert_id_below = _get_vert_id_below(vert_id, all_ids) + + zoom = np.array(poi.zoom if poi.zoom is not None else vert.zoom) + + for side, inf_loc, sup_loc in ( + ("left", Location.Inferior_Articular_Left, Location.Superior_Articular_Left), + ("right", Location.Inferior_Articular_Right, Location.Superior_Articular_Right), + ): + target_id = ids[side] + if (vert_id, target_id) in poi: + continue + + pts_a, pts_b = _facet_surface_points(vert, subreg, vert_id, inf_loc, vert_id_below, sup_loc) + + point = _facet_midpoint( + pts_a, + pts_b, + zoom, + max_gap_mm=max_gap_mm, + surface_tolerance_mm=surface_tolerance_mm, + ) + + if point is None: + # Sensible fallback: reuse whichever centroid already exists. + if (vert_id, inf_loc.value) in poi: + point = poi[vert_id, inf_loc.value] + elif vert_id_below is not None and (vert_id_below, sup_loc.value) in poi: + point = poi[vert_id_below, sup_loc.value] + + if point is not None: + poi[vert_id, target_id] = tuple(float(v) for v in point) + else: + log.print( + f"[Facet] Could not compute facet point for vertebra {vert_id} ({side}): " + "neither a contact surface nor a fallback centroid is available", + verbose=True, + ) + + return poi + + +def _get_vert_id_below(vert_id: int, all_ids: Sequence[int]) -> int | None: + """Return the label of the vertebra directly below ``vert_id``. + + Uses :meth:`Vertebra_Instance.get_next_poi`, since vertebra labels are not + guaranteed to be strictly ascending (e.g. transitional or fused segments), so a + naive ``vert_id + 1`` or "next larger label" lookup would be incorrect. + + Args: + vert_id: Label of the current (upper) vertebra. + all_ids: All vertebra labels present in the image. + + Returns: + The label of the next lower vertebra, or None if there is none. + """ + v1 = Vertebra_Instance(vert_id) + v2 = Vertebra_Instance.get_next_poi(v1, all_ids) + if v2 is None: + return None + return v2.value if isinstance(v2, Vertebra_Instance) else int(v2) + + +def _facet_surface_points( + vert: NII, + subreg: NII, + vert_id: int, + inf_loc: Location, + vert_id_below: int | None, + sup_loc: Location, +) -> tuple[np.ndarray, np.ndarray]: + """Extract the voxel coordinates of the two candidate articular surfaces. + + Args: + vert: Vertebra instance segmentation. + subreg: Subregion/semantic segmentation, same grid as ``vert``. + vert_id: Label of the upper vertebra. + inf_loc: Inferior articular location (left or right) of the upper vertebra. + vert_id_below: Label of the lower vertebra, or None if absent. + sup_loc: Superior articular location (left or right) of the lower vertebra. + + Returns: + A tuple ``(pts_a, pts_b)`` of ``(N, 3)`` arrays in full-image voxel space. + Either array is empty if the corresponding vertebra/label is not present. + """ + wanted = [i for i in (vert_id, vert_id_below) if i is not None] + vert_arr_full = vert.get_seg_array() + combined_mask = np.isin(vert_arr_full, wanted) + if not combined_mask.any(): + return np.zeros((0, 3)), np.zeros((0, 3)) + + coords = np.argwhere(combined_mask) + mins = coords.min(0) + maxs = coords.max(0) + 1 + sl = tuple(slice(int(mn), int(mx)) for mn, mx in zip(mins, maxs)) + + vert_crop = vert_arr_full[sl] + subreg_crop = subreg.get_seg_array()[sl] + + mask_a = (vert_crop == vert_id) & (subreg_crop == inf_loc.value) + pts_a = np.argwhere(mask_a).astype(float) + mins if mask_a.any() else np.zeros((0, 3)) + + if vert_id_below is None: + pts_b = np.zeros((0, 3)) + else: + mask_b = (vert_crop == vert_id_below) & (subreg_crop == sup_loc.value) + pts_b = np.argwhere(mask_b).astype(float) + mins if mask_b.any() else np.zeros((0, 3)) + + return pts_a, pts_b + + +def _filter_facing_points(pts_a: np.ndarray, pts_b: np.ndarray, zoom: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Discard points that lie on the side of each surface facing away from the other. + + The articular process is a 3D blob, not a flat plate: only the hemisphere facing + the neighboring vertebra can physically be part of the contact interface. Points + on the outer/back side are geometrically irrelevant and, if included, could be + picked up by the nearest-neighbor search purely by chance (e.g. on oddly shaped + or partially segmented processes). This filter removes them before matching. + + Args: + pts_a: Voxel coordinates of surface A. + pts_b: Voxel coordinates of surface B. + zoom: Voxel spacing, used to compute the facing direction in mm. + + Returns: + The filtered ``(pts_a, pts_b)``. Falls back to the unfiltered input for a + surface if filtering would remove all of its points (degenerate/very small + surfaces). + """ + if len(pts_a) == 0 or len(pts_b) == 0: + return pts_a, pts_b + + a_mm = pts_a * zoom + b_mm = pts_b * zoom + centroid_a = a_mm.mean(axis=0) + centroid_b = b_mm.mean(axis=0) + direction = centroid_b - centroid_a + norm = np.linalg.norm(direction) + if norm < 1e-6: + return pts_a, pts_b + direction = direction / norm + + a_keep = (a_mm - centroid_a) @ direction >= 0 + b_keep = (b_mm - centroid_b) @ (-direction) >= 0 + + pts_a_f = pts_a[a_keep] if a_keep.any() else pts_a + pts_b_f = pts_b[b_keep] if b_keep.any() else pts_b + return pts_a_f, pts_b_f + + +def _facet_midpoint( + pts_a: np.ndarray, + pts_b: np.ndarray, + zoom: np.ndarray, + max_gap_mm: float, + surface_tolerance_mm: float, +) -> tuple[float, float, float] | None: + """Compute the midpoint of the (possibly extended) contact area between two surfaces. + + Args: + pts_a: Voxel coordinates of surface A. + pts_b: Voxel coordinates of surface B. + zoom: Voxel spacing, used to compute distances in mm. + max_gap_mm: Maximum allowed minimal distance between the surfaces. + surface_tolerance_mm: Band above the minimal distance that is still + considered part of the contact surface and averaged. + + Returns: + The midpoint in voxel coordinates, or None if either surface is empty + (after facing-filter) or the minimal distance exceeds ``max_gap_mm``. + """ + pts_a, pts_b = _filter_facing_points(pts_a, pts_b, zoom) + if len(pts_a) == 0 or len(pts_b) == 0: + return None + + tree_b = cKDTree(pts_b * zoom) + dist, idx_b = tree_b.query(pts_a * zoom, k=1) + min_dist = float(dist.min()) + if min_dist > max_gap_mm: + return None + + close = dist <= (min_dist + surface_tolerance_mm) + a_close = pts_a[close] + b_close = pts_b[idx_b[close]] + midpoints = (a_close + b_close) / 2.0 + return tuple(midpoints.mean(axis=0).tolist()) diff --git a/TPTBox/spine/spinestats/poi_fun/endplates.py b/TPTBox/spine/spinestats/poi_fun/endplates.py new file mode 100644 index 00000000..19294948 --- /dev/null +++ b/TPTBox/spine/spinestats/poi_fun/endplates.py @@ -0,0 +1,562 @@ +from __future__ import annotations + +import warnings +from collections.abc import Sequence + +import numpy as np +import trimesh +from stl.mesh import Mesh + +from TPTBox import NII, POI, Location, Logger_Interface, Print_Logger +from TPTBox.core.vert_constants import Vertebra_Instance + +_log = Print_Logger() + +# -------------------------------------------------------------------------- +# Grid-based geometry helpers +# -------------------------------------------------------------------------- + + +def _ray_triangle_intersect( + orig: np.ndarray, direction: np.ndarray, v0: np.ndarray, v1: np.ndarray, v2: np.ndarray, epsilon: float = 1e-8 +) -> float | None: + """Moeller-Trumbore ray-triangle intersection test. + + Parameters + ---------- + orig : np.ndarray + Ray origin (3,). + direction : np.ndarray + Unit ray direction (3,). + v0, v1, v2 : np.ndarray + Triangle vertices (3,) each. + + Returns: + ------- + float | None + Distance ``t`` along ``direction`` to the intersection point + (only ``t > 0`` is returned, i.e. intersections behind the ray + origin are ignored), or ``None`` if the ray misses the triangle. + """ + edge1 = v1 - v0 + edge2 = v2 - v0 + h = np.cross(direction, edge2) + a = np.dot(edge1, h) + if -epsilon < a < epsilon: + return None # ray is parallel to the triangle + f = 1.0 / a + s = orig - v0 + u = f * np.dot(s, h) + if u < 0.0 or u > 1.0: + return None + q = np.cross(s, edge1) + v = f * np.dot(direction, q) + if v < 0.0 or u + v > 1.0: + return None + t = f * np.dot(edge2, q) + return t if t > epsilon else None + + +def _ray_cast_to_mesh(mesh: Mesh | trimesh.Trimesh, origin: np.ndarray, direction: np.ndarray) -> np.ndarray | None: + direction = direction / np.linalg.norm(direction) + + if isinstance(mesh, trimesh.Trimesh): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", RuntimeWarning) + locations, _, _ = mesh.ray.intersects_location(ray_origins=origin[None], ray_directions=direction[None]) + if len(locations) == 0: + return None + # closest hit + # d = np.linalg.norm(locations - origin, axis=1) + # return locations[np.argmin(d)] + d = np.linalg.norm(locations - origin, axis=1) + return origin + d.mean() * direction + # numpy-stl fallback + ts = [] + for v0, v1, v2 in zip(mesh.v0, mesh.v1, mesh.v2): + t = _ray_triangle_intersect(origin, direction, v0, v1, v2) + if t is not None: + ts.append(t) + if not ts: + return None + return origin + np.mean(ts) * direction + + +def _sample_ray_voxels( + mask_shape: tuple[int, int, int], start: np.ndarray, direction_unit: np.ndarray, two_sided: bool = False, max_steps: int = 1024 +) -> tuple[np.ndarray, np.ndarray]: + """Sample integer voxel indices along a ray (step = 1 voxel). + + Returns (int_coords, t) where ``int_coords`` is the (N, 3) array of voxel + indices visited by the ray inside ``mask_shape`` and ``t`` is the signed + distance from ``start`` for each step (negative behind the origin when + ``two_sided`` is True). + """ + shape_arr = np.asarray(mask_shape) + + def _t_exit(direction: np.ndarray) -> float: + t = np.inf + for i in range(3): + if direction[i] > 1e-9: + t = min(t, (shape_arr[i] - 1 - start[i]) / direction[i]) + elif direction[i] < -1e-9: + t = min(t, -start[i] / direction[i]) + return float(max(0.0, t)) + + t_fwd = _t_exit(direction_unit) + t_pos = np.arange(0.0, min(t_fwd, max_steps), 1.0) + if two_sided: + t_bwd = _t_exit(-direction_unit) + t_neg = -np.arange(1.0, min(t_bwd, max_steps), 1.0) # -1, -2, -3, ... + t = np.concatenate([t_neg[::-1], t_pos]) + else: + t = t_pos + + coords = start[None, :] + t[:, None] * direction_unit[None, :] + int_coords = np.floor(coords).astype(int) + valid = np.all((int_coords >= 0) & (int_coords < shape_arr), axis=1) + return int_coords[valid], t[valid] + + +def _local_curvature_grid( + voxel_pts_full: np.ndarray, + poi: POI, + casted_full: np.ndarray, + radius: float = 8.0, +) -> float: + """Rough curvature estimate (1/mm) of the endplate near ``casted_full``. + + Fits a best-fit plane (via SVD) to all endplate voxels within ``radius`` mm + of the casted point (in world coordinates) and returns the RMS deviation of + those points from the plane, normalized by ``radius**2``. 0 for a flat + patch, larger for a more bowl-shaped one. + """ + world_pts = poi.local_to_global_arr(voxel_pts_full) + casted_world = np.asarray(poi.local_to_global(tuple(casted_full)), dtype=float) + dists = np.linalg.norm(world_pts - casted_world, axis=1) + neighborhood = world_pts[dists <= radius] + if len(neighborhood) < 3: + return 0.0 + centered = neighborhood - neighborhood.mean(axis=0) + _, _, vt = np.linalg.svd(centered) + normal = vt[-1] + rms = float(np.sqrt(np.mean((centered @ normal) ** 2))) + return rms / (radius**2) + + +_endplate_curvature_key = { + Location.Vertebral_Body_Endplate_Superior: "curvature_superior_endplate", + Location.Vertebral_Body_Endplate_Inferior: "curvature_inferior_endplate", +} +_endplate_angle_key = { + Location.Vertebral_Body_Endplate_Superior: "angle_superior_endplate", + Location.Vertebral_Body_Endplate_Inferior: "angle_inferior_endplate", +} + + +def _endplate( + poi: POI, + nii: NII, + endplate: Location, + vert_id, + log: Logger_Interface, + normals_by_vert, + cms_local_override: Sequence[float] | None = None, + flip_direction=False, + compute_curvature=False, +): + if nii.max() == 0: + log.print(f"[calc_endplate_points] no {endplate.name} voxels for vertebra {vert_id}, skipping.") + return + + bb = nii.compute_crop(0, 1) + nii_c = nii.apply_crop(bb) + + mask = nii_c.get_array() > 0 + idx_local = np.argwhere(mask).astype(float) + if len(idx_local) < 3: + log.print(f"[calc_endplate_points] too few {endplate.name} voxels for vertebra {vert_id}, skipping.") + return + + zoom = np.asarray(nii.zoom, dtype=float) + bb_start = np.array([s.start for s in bb], dtype=float) + + # PCA on voxel positions (scaled by voxel spacing). + pts_scaled = idx_local * zoom + centroid_scaled = pts_scaled.mean(axis=0) + _, _, Vt = np.linalg.svd(pts_scaled - centroid_scaled, full_matrices=False) + direction_world = Vt[-1] + + # Origin of the ray. + cms_local = poi[vert_id, Location.Vertebra_Corpus] if cms_local_override is None else cms_local_override + cms_full = np.asarray(cms_local, dtype=float) + cms_cropped = cms_full - bb_start + + # Orient the normal toward the endplate. + to_endplate_world = (idx_local.mean(axis=0) + bb_start - cms_full) * zoom + if np.dot(direction_world, to_endplate_world) < 0: + direction_world = -direction_world + direction_world /= np.linalg.norm(direction_world) + + # Convert world normal to voxel direction. + direction_vox = direction_world / zoom + direction_vox /= np.linalg.norm(direction_vox) + + # ------------------------------------------------------------------ + # Fast path: voxel ray cast. + # ------------------------------------------------------------------ + coords_r, t_r = _sample_ray_voxels( + mask.shape, + cms_cropped, + direction_vox, + two_sided=False, + ) + hits = mask[coords_r[:, 0], coords_r[:, 1], coords_r[:, 2]] if len(coords_r) else np.zeros(0, dtype=bool) + + if not np.any(hits): + coords_r, t_r = _sample_ray_voxels( + mask.shape, + cms_cropped, + direction_vox, + two_sided=True, + ) + hits = mask[coords_r[:, 0], coords_r[:, 1], coords_r[:, 2]] if len(coords_r) else np.zeros(0, dtype=bool) + + if np.any(hits): + mean_t = float(t_r[hits].mean()) + casted_cropped = cms_cropped + mean_t * direction_vox + casted_full = casted_cropped + bb_start + + else: + # ------------------------------------------------------------------ + # Fallback: build an STL and retry using mesh ray casting. + # ------------------------------------------------------------------ + mesh = nii_c.to_stl(1, to_world=True) + + verts_ = np.vstack([mesh.v0, mesh.v1, mesh.v2]) + verts = np.unique(verts_, axis=0) + + centroid = verts.mean(axis=0) + _, _, Vt = np.linalg.svd(verts - centroid, full_matrices=False) + + direction = Vt[-1] + + cms_world = np.asarray( + poi.local_to_global(tuple(cms_local)), + dtype=float, + ) + + to_endplate = centroid - cms_world + if np.dot(direction, to_endplate) < 0: + direction *= -1 + direction /= np.linalg.norm(direction) + + casted_point = _ray_cast_to_mesh(mesh, cms_world, direction) + + if casted_point is None: + faces = np.arange(len(verts_)).reshape(-1, 3) + mesh = trimesh.Trimesh( + vertices=verts_, + faces=faces, + process=False, + ) + trimesh.repair.fill_holes(mesh) + trimesh.repair.fix_normals(mesh) + + casted_point = _ray_cast_to_mesh( + mesh, + cms_world, + direction, + ) + + if casted_point is None: + log.print(f"[calc_endplate_points] ray cast missed {endplate.name} mesh for vertebra {vert_id};") + return + + casted_full = np.asarray( + poi.global_to_local(tuple(casted_point)), + dtype=float, + ) + + poi[vert_id, endplate] = tuple(float(v) for v in casted_full) + + # Normal at the casted point in world coordinates. + cms_world = np.asarray( + poi.local_to_global(tuple(cms_local)), + dtype=float, + ) + casted_world = np.asarray( + poi.local_to_global(tuple(casted_full)), + dtype=float, + ) + + normal_at_point = casted_world - cms_world + normal_at_point /= np.linalg.norm(normal_at_point) + + if flip_direction: + normal_at_point *= -1 + + normals_by_vert.setdefault(vert_id, {})[endplate] = normal_at_point + + if compute_curvature: + curvature = _local_curvature_grid( + idx_local + bb_start, + poi, + casted_full, + ) + poi.info[_endplate_curvature_key[endplate]][Vertebra_Instance(vert_id).name] = curvature + + poi.info[_endplate_angle_key[endplate]][Vertebra_Instance(vert_id).name] = tuple(direction_world) + + +# -------------------------------------------------------------------------- +# Main routine +# -------------------------------------------------------------------------- + + +def calc_endplate_points_( + poi: POI, + vert: NII, + spine: NII, + _vert_ids: Sequence[int] | None = None, + compute_curvature=False, + log: Logger_Interface = _log, + inplace=True, +) -> tuple[POI, NII, NII]: + """Estimate superior/inferior vertebral endplate landmark points. + + For every relevant vertebra id, this extracts the superior and + inferior endplate voxel mask (from ``spine``, restricted to that + vertebra via ``vert``), fits its plane normal via PCA on the mask + voxels (zoom-scaled), and casts a voxel-resolution ray from the + vertebral body centroid (``Location.Vertebra_Corpus``) toward the + endplate. The midpoint of the intersected voxels along that ray is + stored back into ``poi`` under + ``Location.Vertebral_Body_Endplate_Superior`` / + ``Location.Vertebral_Body_Endplate_Inferior``. + + It additionally computes and stores, in ``poi.info``: + + - ``"endplate_internal_angle"``: ``{vert_id: angle_degrees}`` -- the + angle between the superior and inferior endplate surface normals at + their respective casted points. This is a proxy for local vertebral + body wedging (0 deg = perfectly parallel endplates). + - ``"curvature_superior_endplate"``: ``{vert_id: float}`` -- curvature + proxy (see ``_local_curvature_grid``) of the superior endplate + surface near its casted point. + - ``"curvature_inferior_endplate"``: ``{vert_id: float}`` -- same, for + the inferior endplate. + + Parameters + ---------- + poi : POI + Points of interest. Must already contain + ``Location.Vertebra_Corpus`` for every vertebra to be processed. + Copied internally; the input is not mutated. + vert : NII + Vertebra instance segmentation. + spine : NII + Spine sub-structure segmentation, expected to contain + ``Location.Endplate`` voxels which get restricted to a single + vertebra via ``vert * spine.extract_label(...) % 100``. + _vert_ids : list[int], optional + Restrict processing to these vertebra ids. Defaults to every id + present in ``vert``. In both cases, ids are filtered down to + cervical/thoracic/lumbar vertebrae only. + log : Print_Logger, optional + Logger used for warnings (e.g. missing endplate voxels, missed + ray casts). + + Returns: + ------- + tuple[POI, NII, NII] + The updated ``poi`` (endplate points + ``poi.info`` entries + above), and the original ``vert``/``spine`` passed through + unchanged. + """ + if _vert_ids is None: + vert_ids = vert.unique() + else: + vert_ids: list[int] = list(_vert_ids) + if not inplace: + poi = poi.copy() + _spine_ids = [a.value for a in Vertebra_Instance.cervical()[1:] + Vertebra_Instance.thoracic() + Vertebra_Instance.lumbar()] + vert_ids_org = vert_ids + vert_ids = [a for a in vert_ids if a in _spine_ids] + sp_u = spine.unique() + if Location.Endplate.value in sp_u: + vert, spine = endplate_to_super_infer_endplate(vert, spine) + sp_u = spine.unique() + if not any( + a in sp_u + for a in [ + Location.Vertebral_Body_Endplate_Superior.value, + Location.Vertebral_Body_Endplate_Inferior.value, + ] + ): + log.print(f"[calc_endplate_points] No endplates, {sp_u}") + return (poi, vert, spine) + poi.info.setdefault("endplate_internal_angle", {}) + + poi.info.setdefault("angle_superior_endplate", {}) + poi.info.setdefault("angle_inferior_endplate", {}) + + if compute_curvature: + poi.info.setdefault("curvature_superior_endplate", {}) + poi.info.setdefault("curvature_inferior_endplate", {}) + + # Collect normals per vertebra so we can compute the inter-endplate + # angle once both superior and inferior have been processed. + normals_by_vert: dict[int, dict[Location, np.ndarray]] = {} + for endplate in ( + Location.Vertebral_Body_Endplate_Superior, + Location.Vertebral_Body_Endplate_Inferior, + ): + endplate_nii = vert * spine.extract_label(endplate) % 100 + + for vert_id in vert_ids: + if vert_id == 2 and Location.Vertebral_Body_Endplate_Superior == endplate: + continue + nii = endplate_nii.extract_label(vert_id) + _endplate( + poi, + nii, + endplate, + vert_id, + log, + normals_by_vert, + compute_curvature=compute_curvature, + ) + + # 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)) + if endplate_nii.max() > 0: + c = endplate_nii.dilate_msk(2).get_connected_components() + if c.max() != 1: + coms = c.reorient().center_of_masses() # {int: (P,I,R)} + superior_label = min(coms, key=lambda lbl: coms[lbl][1]) + endplate_nii = c.extract_label(superior_label) + cms_local_override = None + + last_vert = max(vert_ids) + + if (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] + elif 100 + last_vert in vert_ids_org: + 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, + ) + # 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) + n_inf = normals.get(Location.Vertebral_Body_Endplate_Inferior) + if n_sup is None or n_inf is None: + continue # one side missing (e.g. no voxels) -- skip this vertebra + cos_angle = float(np.clip(np.dot(n_sup, n_inf), -1.0, 1.0)) + poi.info["endplate_internal_angle"][Vertebra_Instance(vert_id).name] = 180 - float(np.degrees(np.arccos(cos_angle))) + + return poi, vert, spine + + +def endplate_to_super_infer_endplate(vert: NII, spine: NII) -> tuple[NII, NII]: + """Split a combined ``Location.Endplate`` label into superior/inferior sub-labels. + + Endplate segmentations (spineps T2w segmentation) are produced adjacent to the + intervertebral disc without indicating which vertebra they belong to + or whether they sit on that vertebra's superior or inferior surface. + This function resolves both: it grows each vertebra's corpus label + outward (via ``NII.infect``) into the surrounding endplate/disc/ + corpus-border region so every endplate voxel is claimed by its + nearest vertebra, then reclassifies ``spine`` so those voxels carry + ``Location.Vertebral_Body_Endplate_Superior`` or + ``Location.Vertebral_Body_Endplate_Inferior`` instead of the generic + ``Location.Endplate`` label, and offsets the corresponding ``vert`` + instance ids by ``+200`` to mark them as reassigned. + + Parameters + ---------- + vert : NII + Vertebra instance segmentation. + spine : NII + Spine sub-structure segmentation containing a combined + ``Location.Endplate`` label. + + Returns: + ------- + tuple[NII, NII] + The same ``vert``/``spine`` objects, updated. If ``spine`` + contains no ``Location.Endplate`` voxels, both are returned + unchanged (no-op). + + """ + endplate_nii = spine.extract_label(Location.Endplate) + if endplate_nii.sum() == 0: + return vert, spine + 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 %= 100 + v = vert.infect( + spine.extract_label( + [ + Location.Vertebra_Corpus, + Location.Vertebra_Corpus_border, + Location.Endplate, + Location.Vertebra_Disc, + ] + ), + verbose=False, + ) + endplate_nii = v * endplate_nii + 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 + + +if __name__ == "__main__": + from pathlib import Path + + from TPTBox import calc_poi_from_subreg_vert, to_nii + + # + # p = Path("/DATA/NAS/datasets_processed/CT_spine/dataset-myelom/derivatives-final/sub-CTFU00065/ses-00000") + # poi = calc_poi_from_subreg_vert( + # p / "sub-CTFU00065_ses-00000_sequ-2_mod-ct_seg-vert_msk.nii.gz", + # p / "sub-CTFU00065_ses-00000_sequ-2_mod-ct_seg-spine_msk.nii.gz", + # subreg_id=Location.Endplate, + # ) + # poi.make_point_cloud_nii(s=3)[1].save(p / "out_point.nii.gz") + # poi.save(p / "out.json") + # print(poi.centroids) + p = Path("TPTBox/tests/sample_mri") + + vert, spine = endplate_to_super_infer_endplate( + to_nii(p / "sub-mri_seg-vert_label-6_msk.nii.gz", True), + to_nii(p / "sub-mri_seg-subreg_label-6_msk.nii.gz", True), + ) + # vert.save(p / "out_v.nii.gz") + # spine.save(p / "out_s.nii.gz") + poi = calc_poi_from_subreg_vert( + p / "sub-mri_seg-vert_label-6_msk.nii.gz", + p / "sub-mri_seg-subreg_label-6_msk.nii.gz", + subreg_id=Location.Endplate, + ) + # poi.make_point_cloud_nii(s=3)[1].save(p / "out_point.nii.gz") + # poi.save(p / "out.json") + print(poi.centroids) diff --git a/TPTBox/spine/spinestats/ivd_pois.py b/TPTBox/spine/spinestats/poi_fun/ivd_pois.py similarity index 93% rename from TPTBox/spine/spinestats/ivd_pois.py rename to TPTBox/spine/spinestats/poi_fun/ivd_pois.py index 8f96571c..8a45bee5 100644 --- a/TPTBox/spine/spinestats/ivd_pois.py +++ b/TPTBox/spine/spinestats/poi_fun/ivd_pois.py @@ -92,32 +92,6 @@ def _crop( return i, verts_ids, vert, spine, poi, crop, next_id -def _process_vertebra_A(idx: int, vert: NII, spine: NII, next_id: int, poi: POI) -> NII | None: - """Generate the IVD mask between vertebra ``idx`` and ``next_id`` using endplate extraction.""" - from TPTBox.spine.spinestats import endplate_extraction - - try: - a = endplate_extraction(idx, vert, spine, poi) - b = endplate_extraction(next_id, vert, spine, poi) - if a is None or b is None: - return None - a = a.extract_label(Location.Vertebral_Body_Endplate_Inferior) - b = b.extract_label(Location.Vertebral_Body_Endplate_Superior) - except ValueError: - return None - ivd: NII = (a + b).calc_convex_hull(None) - # ivd: NII = ( - # (a + b).dilate_msk(1, 1, verbose=False).erode_msk(3, 1, ignore_direction="S", verbose=False) - # # .calc_convex_hull(None) - # ) - # ivd[a != 0] = 2 - # ivd[b != 0] = 3 - # ivd[spine != 0] += 10 - # ivd = ivd.filter_connected_components(1, max_count_component=1, connectivity=1) - ivd = ivd * (100 + idx) - return ivd - - def _process_vertebra_B(idx: int, vert: NII, spine: NII, next_id: Vertebra_Instance, dilate: int = 1) -> NII: """Generate the IVD mask between vertebra ``idx`` and ``next_id`` using convex-hull morphology.""" spine = spine.extract_label([49, 50, 26]) @@ -312,7 +286,7 @@ def calculate_IVD_POI( poi: POI, ivd_location: set[Location] | None = None, ) -> POI: - """Compute IVD-related Points of Interest and add them to ``poi``. + """Compute IVD-related Points of Interest and add them to ``poi``. This function estimates these points if the IVD is not here. If the subregion image does not yet contain IVD labels (label 100), :func:`compute_fake_ivd` is called first to synthesise them. Centroid diff --git a/TPTBox/spine/spinestats/torso_vat_sat.py b/TPTBox/spine/spinestats/torso_vat_sat.py new file mode 100644 index 00000000..2e3e73a5 --- /dev/null +++ b/TPTBox/spine/spinestats/torso_vat_sat.py @@ -0,0 +1,731 @@ +from __future__ import annotations + +from typing import Literal + +import numpy as np + +from TPTBox import NII +from TPTBox.core.nii_wrapper import NII +from TPTBox.core.vert_constants import Full_Body_Instance, Full_Body_Instance_Vibe, Location, Vertebra_Instance + + +def peak_centered_mean( + values: np.ndarray, + bins: int = 64, + peak_frac_height: float = 0.5, +) -> float: + """Robust mean of a 1D intensity sample, centered on the histogram peak. + + Steps: + 1. Build a histogram of ``values``. + 2. Locate the mode (tallest bin). + 3. Grow a contiguous window outward from the mode as long as each + neighbouring bin still has at least + ``peak_frac_height`` * peak_height counts. + 4. Average only the values falling inside that window. + + Useful for estimating cerebrospinal fluid signal from a spinal canal + mask that may include darker structures such as nerve roots or vessel + walls at the border. Those low-signal tails sit outside the peak + window and are excluded before averaging. + + Parameters + ---------- + values : np.ndarray + 1D array of intensity samples already extracted from the mask. + bins : int, default=64 + Histogram resolution used to locate the peak. Higher values give a + more precise peak but are more sensitive to shot noise. + peak_frac_height : float, default=0.5 + Fractional height cutoff. Bins whose count is at least this + fraction of the peak count are kept; the default ``0.5`` matches + the full-width half-maximum criterion. Lower values keep more of + the tails; ``1.0`` reduces to the peak bin only. + + Returns: + ------- + float + Mean of the values inside the peak window. Returns ``nan`` for an + empty input; falls back to the plain mean when the window is + degenerate. + """ + values = np.asarray(values).ravel() + values = values[np.isfinite(values)] + if values.size == 0: + return float("nan") + hist, edges = np.histogram(values, bins=bins) + if hist.max() == 0: + return float(np.mean(values)) + peak_idx = int(np.argmax(hist)) + threshold = hist[peak_idx] * peak_frac_height + left = peak_idx + while left > 0 and hist[left - 1] >= threshold: + left -= 1 + right = peak_idx + while right < len(hist) - 1 and hist[right + 1] >= threshold: + right += 1 + lo = edges[left] + hi = edges[right + 1] + kept = values[(values >= lo) & (values <= hi)] + if kept.size == 0: + return float(np.mean(values)) + return float(np.mean(kept)) + + +def VBQ_score( + t2w: NII, + vert: NII, + spine: NII, + regions=None, + subregs_ids=None, + spinal_channel_id=Location.Spinal_Canal, + n_erode=2, + full_cord=False, + spinal_bins: int = 64, + spinal_peak_frac_height: float = 0.5, +) -> dict[str, int]: + """Compute vertebral bone quality (VBQ) scores from a T2-weighted MRI. + + The VBQ score is defined as the ratio of the mean T2 signal intensity + within the vertebral body to the mean T2 signal intensity of the + cerebrospinal fluid (CSF) in the spinal canal over the same + superior–inferior extent. + + For each requested spinal region, the function: + 1. Extracts the vertebral bodies of the specified vertebrae. + 2. Optionally erodes the vertebral body mask to reduce partial + volume effects. + 3. Computes the mean T2 signal within the vertebral bodies. + 4. Computes the mean CSF signal within the spinal canal over the + corresponding superior–inferior range. + 5. Returns the vertebral signal, CSF signal, and their ratio + (VBQ score). + + Parameters + ---------- + t2w : NII + T2-weighted MRI volume. + vert : NII + Vertebral instance segmentation. + spine : NII + Spine subregion segmentation containing vertebral subregions and + the spinal canal. + regions : list[tuple[Vertebra_Instance, Vertebra_Instance]], optional + Vertebral ranges over which to compute VBQ scores. Each tuple + specifies the first and last vertebra (inclusive). Defaults to + C3–C6, T5–T8, and L1-L4. According to https://link.springer.com/article/10.1007/s00586-022-07484-5 + subregs_ids : list[Location], optional + Spine subregion labels defining the vertebral body. Defaults to + ``[Location.Vertebra_Corpus, Location.Vertebra_Corpus_border]``. + spinal_channel_id : Location, default=Location.Spinal_Canal + Label identifying the spinal canal (CSF) region. + n_erode : int, default=2 + Number of erosion iterations applied to the vertebral body mask + before signal extraction. + spinal_bins : int, default=64 + Histogram bin count forwarded to :func:`peak_centered_mean` when + estimating the CSF signal. + spinal_peak_frac_height : float, default=0.5 + Fractional height cutoff forwarded to :func:`peak_centered_mean`. + Bins with count >= ``peak_frac_height * peak`` are kept; the + default 0.5 corresponds to FWHM. + + Returns: + ------- + dict[str, float] + Dictionary containing, for each region: + + - ``mean_signal_vertebra_-`` : mean T2 signal inside + the (eroded) vertebral body mask. + - ``mean_signal_liquor_-`` : peak-centered mean T2 + signal inside the spinal canal (robust to nerve-root voxels). + - ``mean_signal_liquor_-_old`` : plain mean T2 signal + inside the spinal canal, kept for backward comparison. + - ``VBQ_-`` : ratio using the peak-centered CSF + signal. + - ``VBQ_-_old`` : ratio using the plain-mean CSF + signal. + + Notes: + ----- + The implementation follows the general principle of the Vertebral Bone + Quality (VBQ) score described in the literature, using the ratio of + vertebral body T2 signal intensity to CSF signal intensity as a proxy + for bone quality. + """ + if subregs_ids is None: + subregs_ids = [ + Location.Vertebra_Corpus, + Location.Vertebra_Corpus_border, + ] + + if regions is None: + regions = [ + (Vertebra_Instance.C3, Vertebra_Instance.C6), + (Vertebra_Instance.T5, Vertebra_Instance.T8), + (Vertebra_Instance.L1, Vertebra_Instance.L4), + ] + + spinal_channel = spine.extract_label(spinal_channel_id).erode_msk(1, connectivity=1, verbose=False) + corpus = spine.extract_label(subregs_ids) + out = {} + + verts_order = Vertebra_Instance.order() + + for start, goal in regions: + end = verts_order.index(goal) + 1 + labels = verts_order[verts_order.index(start) : end] + + # vertebrae belonging to this region + vert_mask = vert.extract_label(labels) + + # vertebral bodies only + bodies = vert_mask * corpus + bodies.erode_msk_(n_erode, verbose=False) + + signal_vertebra = t2w.mean(where=bodies) + + # ---- restrict spinal canal to same S/I extent ---- + bbox = bodies.compute_crop() # (slice_x, slice_y, slice_z) + axis = spinal_channel.get_axis(direction="S") + + if not full_cord: + slicer = [slice(None)] * 3 + slicer[axis] = bbox[axis] + slicer = tuple(slicer) + spinal_crop = spinal_channel[slicer] + t2w_slab = t2w.get_array()[slicer] + else: + t2w_slab = t2w + spinal_crop = spinal_channel + + signal_sfs_old = t2w_slab.mean(where=spinal_crop) + spinal_arr = spinal_crop.get_array().astype(bool) + signal_sfs = peak_centered_mean(t2w_slab[spinal_arr], bins=spinal_bins, peak_frac_height=spinal_peak_frac_height) + + out[f"mean_signal_vertebra_{start.name}-{goal.name}"] = round(signal_vertebra, 4) + out[f"mean_signal_liquor_{start.name}-{goal.name}"] = round(signal_sfs, 4) + out[f"mean_signal_liquor_{start.name}-{goal.name}_old"] = round(signal_sfs_old, 4) + out[f"VBQ_{start.name}-{goal.name}"] = round(signal_vertebra / signal_sfs, 4) + out[f"VBQ_{start.name}-{goal.name}_old"] = round(signal_vertebra / signal_sfs_old, 4) + out["n_erode"] = n_erode + out["signal_by_full_cord"] = full_cord + out["spinal_bins"] = spinal_bins + out["spinal_peak_frac_height"] = spinal_peak_frac_height + return out + + +def body_composition_score( + vibe_seg: NII, + vert: NII, + spine: NII, + dataset_id: Literal[100, 12] = 100, + regions: list[tuple[Vertebra_Instance, Vertebra_Instance]] | None = None, + height_m: float | None = None, +) -> dict[str, float]: + """Compute vertebral-level body composition measurements from a VIBE segmentation. + + For each vertebral region, the superior–inferior extent of the vertebral + bodies is used to define the analysis region. Mean and maximum + cross-sectional areas (CSA) of skeletal muscle, psoas, autochthonous + muscles, visceral adipose tissue (VAT), and subcutaneous adipose tissue + (SAT) are computed by averaging over all axial slices intersecting the + vertebral body region. + + Parameters + ---------- + vibe_seg : NII + VIBE body composition segmentation. + + vert : NII + Vertebral instance segmentation. + + spine : NII + Spine subregion segmentation. + dataset_id : {100, 12}, default=100 + Dataset definition used to determine the body composition label IDs. + + - ``100``: VIBESeg-100 label set. + - ``12``: VIBESeg-12 label set. + regions : list[tuple[Vertebra_Instance, Vertebra_Instance]], optional + Vertebral ranges to analyse. Each tuple specifies the first and last + vertebra (inclusive). Defaults to T12–L1 and L3. + + height_m : float, optional + Patient height in metres. If provided, the skeletal muscle index + (CSA / m²) is additionally reported. + + Returns: + ------- + dict[str, float] + Dictionary containing body composition measurements for each + region. Region tags are formatted as ``{start.name}-{goal.name}`` + (e.g. ``T12-L1``). For each region the following keys are set: + + - ``mean_{tissue}_area_{region}`` : mean cross-sectional area of + the tissue over all axial slices intersecting the vertebral + body region (mm²). ``tissue`` is one of + ``muscle``, ``VAT``, ``SAT``, ``psoas``, ``autochthon``. + - ``max_{tissue}_area_{region}`` : maximum cross-sectional area + over the same slice range (mm²). + - ``n_slices_{region}`` : number of axial slices contributing to + the ``muscle`` statistic (integer). + - ``muscle_index_{region}`` : skeletal muscle index + (``mean_muscle_area / height_m^2``, unit mm²/m²). Only present + when ``height_m`` is provided. + - ``muscle_fat_ratio_{region}`` : ``mean_muscle_area / + (mean_VAT_area + mean_SAT_area)`` (unitless). ``NaN`` when the + denominator is zero. + + A region is silently skipped (no keys emitted) when no vertebral + body voxels fall inside its label range. + """ + if regions is None: + regions = [ + (Vertebra_Instance.T12, Vertebra_Instance.L1), + (Vertebra_Instance.L3, Vertebra_Instance.L4), + (Vertebra_Instance.L3, Vertebra_Instance.L3), + ] + + if vibe_seg.shape != vert.shape: + vibe_seg = vibe_seg.resample_from_to(vert, verbose=False) + + if spine.shape != vert.shape: + spine = spine.resample_from_to(vert, verbose=False) + + body_mask = spine.extract_label([Location.Vertebra_Corpus, Location.Vertebra_Corpus_border]) + + if dataset_id == 12: + measurements = { + "muscle": [*Full_Body_Instance.muscle(), Full_Body_Instance.muscle_other], + "VAT": Full_Body_Instance.inner_fat, + "SAT": Full_Body_Instance.subcutaneous_fat, + "psoas": [Full_Body_Instance.iliopsoas_left, Full_Body_Instance.iliopsoas_right], + "autochthon": [Full_Body_Instance.autochthon_left, Full_Body_Instance.autochthon_right], + } + + elif dataset_id == 100: + measurements = { + "muscle": [*Full_Body_Instance_Vibe.muscle_(), Full_Body_Instance_Vibe.muscle_other], + "VAT": Full_Body_Instance_Vibe.inner_fat, + "SAT": Full_Body_Instance_Vibe.subcutaneous_fat, + "psoas": [Full_Body_Instance_Vibe.iliopsoas_left, Full_Body_Instance_Vibe.iliopsoas_right], + "autochthon": [Full_Body_Instance_Vibe.autochthon_left, Full_Body_Instance_Vibe.autochthon_right], + } + + else: + raise NotImplementedError(f"Unsupported dataset_id={dataset_id}") + + verts_order = Vertebra_Instance.order() + axis = vibe_seg.get_axis(direction="S") + + other_axes = tuple(i for i in range(3) if i != axis) + voxel_area = float(vibe_seg.zoom[other_axes[0]] * vibe_seg.zoom[other_axes[1]]) + out = {} + u = vert.unique() + for start, goal in regions: + end = verts_order.index(goal.get_next_poi(u)) + labels = verts_order[verts_order.index(start) : end] + + vertebral_body = vert.extract_label(labels) * body_mask + + if vertebral_body.sum() == 0: + # print("no vertebra", labels, vert.unique()) + continue + + bbox = vertebral_body.compute_crop() + + slicer = [slice(None)] * 3 + slicer[axis] = bbox[axis] + + region: NII = vibe_seg[slicer] + + region_name = f"{start.name}-{goal.name}" + + for name, label_ids in measurements.items(): + mask = region.extract_label(label_ids) + arr = mask.get_array().astype(bool) + + areas = arr.sum(axis=other_axes).astype(float) * voxel_area + areas = areas[areas > 0] + + if len(areas) == 0: + print(name, [int(a) for a in region.unique()], label_ids) + mean_area = np.nan + max_area = np.nan + n_slices = 0 + else: + mean_area = float(np.mean(areas)) + max_area = float(np.max(areas)) + n_slices = len(areas) + + out[f"mean_{name}_area_{region_name}"] = round(mean_area, 4) + out[f"max_{name}_area_{region_name}"] = round(max_area, 4) + + if name == "muscle": + out[f"n_slices_{region_name}"] = n_slices + + if height_m is not None and np.isfinite(mean_area): + out[f"muscle_index_{region_name}"] = round(mean_area / (height_m**2)) + + vat = out[f"mean_VAT_area_{region_name}"] + sat = out[f"mean_SAT_area_{region_name}"] + muscle = out[f"mean_muscle_area_{region_name}"] + + if np.isfinite(vat) and np.isfinite(sat) and np.isfinite(muscle) and (vat + sat) > 0: + out[f"muscle_fat_ratio_{region_name}"] = round(muscle / (vat + sat), 4) + else: + out[f"muscle_fat_ratio_{region_name}"] = np.nan + + return out + + +def muscle_fat_infiltration( + water: NII, + fat: NII, + vibe_seg: NII, + vert: NII | None = None, + spine: NII | None = None, + regions: list[tuple[Vertebra_Instance, Vertebra_Instance]] | list[tuple[None, None]] | None = None, + roi: NII | None = None, + roi_ids: tuple[int, ...] = tuple(range(3, 9)), + dataset_id: Literal[100, 12] = 100, + threshold: float = 0.20, + erode: dict[str, int] | None = None, + per_muscle: bool = True, +) -> dict[str, float]: + """Compute muscle fat infiltration from Dixon VIBE water/fat images. + + The fat fraction (FF) is calculated voxel-wise as: + + FF = fat / (fat + water) + + within skeletal muscle masks. Voxels with FF >= threshold are classified + as intramuscular adipose tissue (IMAT), while voxels below the threshold + are classified as lean muscle. + + Measurements can optionally be restricted to vertebral levels and/or a + supplied ROI. The analysis can be performed for total muscle and + individual muscle groups. For each muscle group, muscle volumes are + reported both with and without erosion of the muscle mask, so callers + can compare the eroded region used for fat-fraction statistics against + the raw segmentation volume. + + Parameters + ---------- + water : NII + Dixon water image. + + fat : NII + Dixon fat image. + + vibe_seg : NII + VIBE body composition segmentation. + + vert : NII, optional + Vertebral instance segmentation. Required when ``regions`` is used. + + spine : NII, optional + Spine subregion segmentation. + + regions : list[tuple[Vertebra_Instance, Vertebra_Instance]], optional + Vertebral ranges for regional analysis. The superior-inferior extent + of the vertebral bodies is used to restrict the calculation. + + roi : NII, optional + Additional ROI mask restricting the analysis. + + roi_ids : tuple[int, ...], default=tuple(range(3, 9)) + ROI labels used when restricting the analysis. + + dataset_id : {100, 12}, default=100 + Body composition label definition. + + threshold : float, default=0.20 + Fat fraction threshold separating lean muscle and IMAT. + + erode : dict[str, int], optional + Per-muscle number of erosion iterations applied to the muscle mask + before fat-fraction statistics are computed. Defaults to + ``{"all_muscle": 1, "iliopsoas_left": 1, "iliopsoas_right": 1, + "autochthon_left": 2, "autochthon_right": 2, "muscle_other": 1}``. + Erosion reduces partial-volume contamination at the muscle border + but shrinks the mask; the pre-erosion volume is always reported + alongside the eroded volume so both can be inspected. + + per_muscle : bool, default=True + If True, compute values for individual muscles in addition to total + muscle. + + Returns: + ------- + dict[str, float] + Muscle fat infiltration measurements. For each ``{region}_{muscle}`` + suffix the dictionary contains: + + - ``mean_fat_fraction_{suffix}`` : mean FF over the eroded mask. + - ``median_fat_fraction_{suffix}`` : median FF over the eroded mask. + - ``mean_lean_fat_fraction_{suffix}`` : mean FF of lean voxels + (FF < ``threshold``). + - ``mean_IMAT_fat_fraction_{suffix}`` : mean FF of IMAT voxels + (FF >= ``threshold``). + - ``muscle_volume_{suffix}`` : muscle volume after erosion (mm^3). + - ``muscle_volume_no_erosion_{suffix}`` : muscle volume before + erosion (mm^3); equals ``muscle_volume_{suffix}`` when no + erosion is applied. + - ``lean_muscle_volume_{suffix}`` : lean-muscle volume within the + eroded mask (mm^3). + - ``lean_muscle_volume_no_erosion_{suffix}`` : lean-muscle volume + within the un-eroded mask (mm^3). + - ``IMAT_volume_{suffix}`` : IMAT volume within the eroded mask + (mm^3). + - ``IMAT_volume_no_erosion_{suffix}`` : IMAT volume within the + un-eroded mask (mm^3). + - ``IMAT_fraction_{suffix}`` : IMAT voxel fraction within the + eroded mask. + """ + if erode is None: + erode = {"all_muscle": 1, "iliopsoas_left": 1, "iliopsoas_right": 1, "autochthon_left": 2, "autochthon_right": 2, "muscle_other": 1} + if water.shape != vibe_seg.shape: + vibe_seg = vibe_seg.resample_from_to(water, verbose=False) + if fat.shape != water.shape: + fat = fat.resample_from_to(water, verbose=False) + if vert is not None and vert.shape != water.shape: + vert = vert.resample_from_to(water, verbose=False) + if spine is not None and spine.shape != water.shape: + spine = spine.resample_from_to(water, verbose=False) + if roi is not None and roi.shape != water.shape: + roi = roi.resample_from_to(water, verbose=False) + + # ------------------------------------------------------------------ + # Muscle label definitions + # ------------------------------------------------------------------ + if dataset_id == 12: + muscle_groups = { + "all_muscle": [*Full_Body_Instance.muscle(), Full_Body_Instance.muscle_other], + "iliopsoas_left": Full_Body_Instance.iliopsoas_left, + "iliopsoas_right": Full_Body_Instance.iliopsoas_right, + "autochthon_left": Full_Body_Instance.autochthon_left, + "autochthon_right": Full_Body_Instance.autochthon_right, + "muscle_other": Full_Body_Instance.muscle_other, + } + elif dataset_id == 100: + muscle_groups = { + "all_muscle": [*Full_Body_Instance_Vibe.muscle_(), Full_Body_Instance_Vibe.muscle_other], + "iliopsoas_left": Full_Body_Instance_Vibe.iliopsoas_left, + "iliopsoas_right": Full_Body_Instance_Vibe.iliopsoas_right, + "autochthon_left": Full_Body_Instance_Vibe.autochthon_left, + "autochthon_right": Full_Body_Instance_Vibe.autochthon_right, + "muscle_other": Full_Body_Instance_Vibe.muscle_other, + } + else: + raise NotImplementedError(f"Unsupported dataset_id={dataset_id}") + if not per_muscle: + muscle_groups = {"all_muscle": muscle_groups["all_muscle"]} + if regions is None: + regions = [(None, None)] + verts = Vertebra_Instance.order() + voxel_volume = water.voxel_volume() + + out = {} + + # ------------------------------------------------------------------ + # Vertebral regions + # ------------------------------------------------------------------ + for start, goal in regions: + region_name = "all" + slicer = None + + if start is not None: + if vert is None: + raise ValueError("vert segmentation required when regions are provided") + assert goal is not None + end = verts.index(goal) + 1 + labels = verts[verts.index(start) : end] + vertebra_region = vert.extract_label(labels) + if spine is not None: + vertebra_region *= spine.extract_label(Location.Vertebra_Corpus) + bbox = vertebra_region.compute_crop() + axis = vertebra_region.get_axis(direction="S") + slicer = [slice(None)] * 3 + slicer[axis] = bbox[axis] + region_name = f"{start.name}-{goal.name}" + + # -------------------------------------------------------------- + # Individual muscles + # -------------------------------------------------------------- + for muscle_name, muscle_labels in muscle_groups.items(): + muscle_mask = vibe_seg.extract_label(muscle_labels) + if roi is not None: + muscle_mask *= roi.extract_label(roi_ids) + if slicer is not None: + muscle_mask = muscle_mask[slicer] + + mask_no_erode = muscle_mask.get_array().astype(bool) + if erode[muscle_name] > 0: + mask = muscle_mask.erode_msk(erode[muscle_name], verbose=False).get_array().astype(bool) + else: + mask = mask_no_erode + + if mask_no_erode.sum() == 0: + continue + + water_arr = water.get_array() + fat_arr = fat.get_array() + if slicer is not None: + water_arr = water_arr[slicer] # type: ignore + fat_arr = fat_arr[slicer] # type: ignore + denom = water_arr + fat_arr + ff_full = np.zeros_like(denom, dtype=np.float32) + valid = denom > 0 + ff_full[valid] = fat_arr[valid] / denom[valid] + + ff = ff_full[mask] + ff_no_erode = ff_full[mask_no_erode] + + lean = ff < threshold + imat = ff >= threshold + lean_no_erode = ff_no_erode < threshold + imat_no_erode = ff_no_erode >= threshold + + suffix = f"{region_name}_{muscle_name}" + + out[f"muscle_volume_no_erosion_{suffix}"] = round(float(ff_no_erode.size * voxel_volume), 4) + out[f"lean_muscle_volume_no_erosion_{suffix}"] = round(float(np.sum(lean_no_erode) * voxel_volume), 4) + out[f"IMAT_volume_no_erosion_{suffix}"] = round(float(np.sum(imat_no_erode) * voxel_volume), 4) + + if ff.size == 0: + continue + + out[f"mean_fat_fraction_{suffix}"] = round(float(np.mean(ff)), 4) + out[f"median_fat_fraction_{suffix}"] = round(float(np.median(ff)), 4) + out[f"mean_lean_fat_fraction_{suffix}"] = round(float(np.mean(ff[lean])) if np.any(lean) else np.nan, 4) + out[f"mean_IMAT_fat_fraction_{suffix}"] = round(float(np.mean(ff[imat])) if np.any(imat) else np.nan, 4) + out[f"muscle_volume_{suffix}"] = round(float(ff.size * voxel_volume), 4) + out[f"lean_muscle_volume_{suffix}"] = round(float(np.sum(lean) * voxel_volume), 4) + out[f"IMAT_volume_{suffix}"] = round(float(np.sum(imat) * voxel_volume), 4) + out[f"IMAT_fraction_{suffix}"] = round(float(np.mean(imat)), 4) + out["threshold"] = threshold + out["erode"] = erode + return out + + +def torso_vat_sat_muscle_mass( + vibe_seg: NII, roi: NII, dataset_id: Literal[100, 12] = 100, roi_ids: tuple[int, ...] = tuple(range(3, 9)), return_nii: bool = False +) -> tuple[dict, NII | None]: + """Compute visceral adipose tissue (VAT), subcutaneous adipose tissue (SAT), and muscle volume in mm from a torso segmentation. + + The segmentation is optionally restricted to the supplied ROI before + calculating tissue volumes. Volumes are reported in physical units + (voxel count × voxel volume). + + Parameters + ---------- + vibe_seg : NII + VIBE segmentation containing body composition labels. + + roi : NII + ROI segmentation defining the torso region. If its geometry differs + from ``vibe_seg``, it is resampled to match. + + dataset_id : {100, 12}, default=100 + Dataset definition used to determine the label IDs. + + - ``100``: VIBESeg-100 label set. + - ``12``: VIBESeg-12 label set. + + roi_ids : tuple[int, ...], default=(3, 4, 5, 6, 7, 8) + ROI labels that define the region in which body composition should be + evaluated. + + return_nii : bool, default=False + If True, additionally return binary NIfTI masks for VAT, SAT, and + muscle. + + Returns: + ------- + tuple[dict, NII | None] + A tuple ``(results, body_comp)``. + + ``results`` contains: + + - ``VAT`` : float + Visceral adipose tissue volume. + - ``SAT`` : float + Subcutaneous adipose tissue volume. + - ``muscle_mass`` : float + Muscle volume. + - ``reason`` : str, optional + Present only if computation failed. + - ``nii`` : dict[str, NII], optional + Returned only when ``return_nii=True``. + + ``body_comp`` is the ROI-restricted segmentation used for the + computation, or ``None`` if the computation failed. + + Raises: + ------ + NotImplementedError + If an unsupported ``dataset_id`` is provided. + + Notes: + ----- + The function verifies that both the clavicula and pelvis are present in + the segmentation to ensure that the full torso is covered. If this check + fails, NaN values are returned together with the failure reason. + """ + if roi.shape != vibe_seg.shape: + roi = roi.resample_from_to(vibe_seg, verbose=False) + + # Restrict computation to the requested ROI. + body_comp = vibe_seg * roi.extract_label(roi_ids) + + results: dict[str, object] = {} + + VAT = SAT = muscle_mass = None + + try: + labels = vibe_seg.unique() + + if dataset_id == 12: + vat_id = Full_Body_Instance.inner_fat + sat_id = Full_Body_Instance.subcutaneous_fat + muscle_ids = [*Full_Body_Instance.muscle(), Full_Body_Instance.muscle_other] + if Full_Body_Instance.clavicula_left.value not in labels: + raise ValueError("Not the full torso visible (clavicula is missing)") + + if Full_Body_Instance.pelvis_left.value not in labels: + raise ValueError("Not the full torso visible (pelvis is missing)") + + elif dataset_id == 100: + if vibe_seg.max() > 80: + raise AssertionError("Not a VIBESeg-100 (or compatible) segmentation.", vibe_seg.unique()) + vat_id = Full_Body_Instance_Vibe.inner_fat + sat_id = Full_Body_Instance_Vibe.subcutaneous_fat + muscle_ids = [*Full_Body_Instance_Vibe.muscle_(), Full_Body_Instance_Vibe.muscle_other] + if Full_Body_Instance_Vibe.clavicula_left.value not in labels: + raise ValueError("Not the full torso visible (clavicula is missing)") + + if Full_Body_Instance_Vibe.pelvis_left.value not in labels: + raise ValueError("Not the full torso visible (pelvis is missing)") + else: + raise NotImplementedError(f"Unsupported dataset_id={dataset_id}") + + voxel_volume = body_comp.voxel_volume() + + VAT = body_comp.extract_label(vat_id) + SAT = body_comp.extract_label(sat_id) + muscle_mass = body_comp.extract_label(muscle_ids) + + results["VAT"] = round(voxel_volume * VAT.sum(), 4) + results["SAT"] = round(voxel_volume * SAT.sum(), 4) + results["muscle_mass"] = round(voxel_volume * muscle_mass.sum(), 4) + + except Exception as exc: + body_comp = None + results.update({"VAT": np.nan, "SAT": np.nan, "muscle_mass": np.nan, "reason": str(exc)}) + + if return_nii: + results["nii"] = {"VAT": VAT, "SAT": SAT, "muscle_mass": muscle_mass} + + return results, body_comp diff --git a/TPTBox/spine/spinestats/distances.py b/TPTBox/spine/spinestats/vertebra_anatomical_widths.py similarity index 100% rename from TPTBox/spine/spinestats/distances.py rename to TPTBox/spine/spinestats/vertebra_anatomical_widths.py diff --git a/docs/api/spine.md b/docs/api/spine.md index 89feb553..7fa134ad 100644 --- a/docs/api/spine.md +++ b/docs/api/spine.md @@ -1,10 +1,24 @@ # Spine -Spine-specific utilities: modular 2D snapshot generation and statistical measurements -(distances, angles, intervertebral disc POIs). +Spine-specific utilities built on top of `NII` and `POI`: + +- **2D snapshots** — sagittal / coronal / axial visualizations of CT and MRI volumes + with vertebra masks, subregion overlays, centroids, MIPs, and colored-depth views. +- **Angles & curvature** — Cobb angle, lordosis and kyphosis measurements. +- **Vertebra & IVD geometry** — heights, widths, diameters, mesh-based measurements, + and endplate reconstructions for individual vertebrae and intervertebral discs. +- **Body composition** — torso VAT / SAT / muscle statistics, VBQ score. +- **POI computation** — IVD landmarks, endplate points, facet joint (articularis) + midpoints, plus body-quadrant partitioning of the vertebral body. ## 2D Snapshots +Modular building blocks for stitching together sagittal / coronal / axial views of +CT and MRI volumes. A snapshot is composed of one or more `Snapshot_Frame`s +(each defining image, segmentation, centroids, projection mode, and crop) and +rendered with `create_snapshot`. No 3D resampling is required — slices are +scaled to isotropic pixels before display. + ::: TPTBox.spine.snapshot2D.snapshot_modular options: show_source: true @@ -12,28 +26,101 @@ Spine-specific utilities: modular 2D snapshot generation and statistical measure ## Snapshot Templates +Ready-made snapshot layouts for common tasks (multi-panel CT MIP shots, fracture +rating views, virtual DXA / QCT panels). Each template wraps a set of +`Snapshot_Frame`s and writes the figure to disk. + ::: TPTBox.spine.snapshot2D.snapshot_templates options: show_source: true filters: ["!^_"] -## Distances +## Angles + +Cobb angle, lordosis and kyphosis computations between vertebral endplates, +with plotting helpers to overlay the measurements on snapshots. -::: TPTBox.spine.spinestats.distances +::: TPTBox.spine.spinestats.angles options: show_source: true filters: ["!^_"] -## Angles +## Body Quadrants -::: TPTBox.spine.spinestats.angles +Partition each vertebral body into 27 anatomically-oriented subregions (a +3×3×3 grid in vertebra-local coordinates). Robust to spinal curvature and +scan orientation because the local axes are derived from muscle-insertion +and median-body POIs. + +::: TPTBox.spine.spinestats.body_quadrants + options: + show_source: true + filters: ["!^_"] + +## IVD & Vertebra Geometry + +Geometric and signal measurements for intervertebral discs and vertebrae: +mesh-based volumes, principal axes, heights, widths, and orientation-aware +diameters (x1–x6). Works on both vertebra and IVD labels — the "up" axis is +read from POIs for vertebrae and estimated via PCA for discs. + +::: TPTBox.spine.spinestats.measure_ivd_and_vertebra_geometry + options: + show_source: true + filters: ["!^_"] + +## Vertebra Anatomical Widths + +Pairwise anatomical distances (mm) between landmark pairs on a vertebra — +e.g. endplate diameters, pedicle widths — stored back onto the `POI` for +downstream statistics. + +::: TPTBox.spine.spinestats.vertebra_anatomical_widths + options: + show_source: true + filters: ["!^_"] + +## Torso VAT / SAT / Muscle + +Body-composition metrics from torso segmentations: visceral / subcutaneous +adipose tissue, muscle mass, and the Vertebral Bone Quality (VBQ) score. +Includes `peak_centered_mean` for robust intensity estimation on noisy +tissue masks. + +::: TPTBox.spine.spinestats.torso_vat_sat options: show_source: true filters: ["!^_"] ## IVD POIs -::: TPTBox.spine.spinestats.ivd_pois +Compute intervertebral disc landmarks (superior / inferior extreme points, +disc centroid) via PCA-based normal estimation and ray casting through the +disc mask. + +::: TPTBox.spine.spinestats.poi_fun.ivd_pois + options: + show_source: true + filters: ["!^_"] + +## Endplate POIs + +Sample points on the superior and inferior endplate surfaces of a vertebra +using ray-triangle intersection against a mesh reconstruction. Useful for +endplate fitting, height measurement, and disc-space analysis. + +::: TPTBox.spine.spinestats.poi_fun.endplates + options: + show_source: true + filters: ["!^_"] + +## Articularis Midpoint POIs + +Detect facet joint (processus articularis) contact regions between adjacent +vertebrae and place a midpoint POI at each joint. Driven by a k-d tree +nearest-neighbor search on the two vertebrae's surface voxels. + +::: TPTBox.spine.spinestats.poi_fun.articularis_midpoint options: show_source: true filters: ["!^_"] diff --git a/pyproject.toml b/pyproject.toml index 6df1b2b9..9cbfcde0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,6 +25,8 @@ joblib = "*" scikit-learn = "*" pynrrd = "*" requests = "*" +trimesh = "*" +numpy-stl = "*" # --- OLD STACK (Python < 3.11) numpy = [ @@ -164,7 +166,9 @@ ignore = [ "PLR0912", "PLR0913", "PLR0915", + "PLR0917", "PLR2004", + "TRY301", "SIM105", "TRY003", "N999", diff --git a/unit_tests/test_poi_autogen.py b/unit_tests/test_poi_autogen.py index fb318d38..832c74f5 100755 --- a/unit_tests/test_poi_autogen.py +++ b/unit_tests/test_poi_autogen.py @@ -100,7 +100,7 @@ def test_affine_property(self): expected_affine[:3, 3] = origin # Check that the 'affine' property returns the expected value - assert np.allclose(poi.affine, expected_affine) + assert np.allclose(poi.affine, expected_affine, rtol=1e-5, atol=1e-5) def test_affine_property_2(self): for _ in range(10):