From 2b0fdb00f8b6d1b1edf9c5d041d0192bf23d5ebe Mon Sep 17 00:00:00 2001 From: Brijesh Murugesan Date: Fri, 28 Aug 2026 17:20:15 -0400 Subject: [PATCH] Add optional rib-anchored vertebral identity stage Vertebra predictions can delineate bone well while assigning level names that are displaced, so that a run of vertebrae carries the identity of a neighbouring level. This adds an optional stage that adjudicates level identity from costovertebral rib attachment geometry before the configured vertebrae_engine performs its usual shape cleanup. The cue is external: rib attachment position is a physical measurement of a structure outside the vertebral column, rather than an inference from the vertebral predictions themselves, their ordering, spacing, connected components or internal consistency. Architecture: - new config key vertebrae_identity (default 'none'), independent of vertebrae_engine rather than a value of it, so the selected engine still runs afterwards - utils/vertebrae_rib_identity_engine.py holds the scientific core - utils/vertebrae_rib_identity.py is the ShapeKit adapter: input discovery, grid validation, label-space conversion, logging, fallbacks Inputs are the case CT (via the existing ct_file_name / ct_root keys) and a precomputed TotalSegmentator 'total' volume (via new rib_file_name / rib_root keys). Only the geometry of the rib masks is read; TotalSegmentator vertebral labels are not consulted. TotalSegmentator is not invoked and is not added as a dependency. Conservative by construction: a missing CT, a missing rib volume, an incompatible voxel grid, fewer than four usable rib pairs, or no level exceeding the correction threshold all leave the masks unchanged, log the reason, and let the batch continue. The stage never raises into a worker. vertebrae_identity: rib_anchor requires vertebrae_engine: shapekit_pro. The combination with the default engine is rejected at startup, before any case is processed, because that engine may reassign contiguous cranio-caudal identities during cleanup and undo the correction. No existing ShapeKit algorithm was modified. Adds tests/test_vertebrae_rib_identity.py (64 tests, synthetic phantoms, no imaging data required) and docs/vertebrae_rib_identity.md. No new runtime dependencies. --- config.yaml | 28 +- docs/config.md | 30 ++ docs/vertebrae_rib_identity.md | 244 +++++++++ main.py | 35 ++ tests/test_vertebrae_rib_identity.py | 663 +++++++++++++++++++++++++ utils/vertebrae_rib_identity.py | 366 ++++++++++++++ utils/vertebrae_rib_identity_engine.py | 417 ++++++++++++++++ 7 files changed, 1782 insertions(+), 1 deletion(-) create mode 100644 docs/vertebrae_rib_identity.md create mode 100644 tests/test_vertebrae_rib_identity.py create mode 100644 utils/vertebrae_rib_identity.py create mode 100644 utils/vertebrae_rib_identity_engine.py diff --git a/config.yaml b/config.yaml index cd435fa..bd6b669 100644 --- a/config.yaml +++ b/config.yaml @@ -103,4 +103,30 @@ vertebrae_engine: shapekit # CT lookup for shapekit_pro: first /, then # // when ct_root is set. ct_file_name: ct.nii.gz -# ct_root: /path/to/original/ct/cases \ No newline at end of file +# ct_root: /path/to/original/ct/cases + +# optional vertebral identity stage, applied BEFORE vertebrae_engine: +# none - disabled (default; behaviour is unchanged) +# rib_anchor - adjudicate which level is which from costovertebral rib +# attachment geometry, then hand the masks to vertebrae_engine +# for the usual shape cleanup. It reads an external anatomical +# cue - the ribs - rather than inferring level identity from +# the vertebral predictions alone. Requires the case CT (found +# via ct_file_name / ct_root above) and a precomputed +# TotalSegmentator 'total' volume; only the rib masks in that +# volume are read. When either input is missing or sits on a +# different voxel grid, the stage logs the reason and leaves +# the masks untouched. +# Requires vertebrae_engine: shapekit_pro. Pairing it with +# the default engine is rejected at startup, because that +# engine may reassign contiguous cranio-caudal identities +# and undo the correction. +vertebrae_identity: none + +# rib volume lookup for rib_anchor: first /, then +# //, then /.nii.gz. +rib_file_name: total.nii.gz +# rib_root: /path/to/totalsegmentator/cases + +# optional directory for one small JSON QA record per case +# rib_qa_dir: /path/to/qa diff --git a/docs/config.md b/docs/config.md index 5a18483..0b29f87 100644 --- a/docs/config.md +++ b/docs/config.md @@ -75,3 +75,33 @@ 8. `ct_file_name` / `ct_root`: how `shapekit_pro` finds the CT. The engine first looks for `/`; when `ct_root` is set it also tries `//`. + +9. `vertebrae_identity`: an optional identity stage applied **before** + `vertebrae_engine`. `none` (default) disables it and leaves ShapeKit's + behaviour unchanged. `rib_anchor` adjudicates which vertebral level is which + from costovertebral rib attachment geometry, then hands the masks to + `vertebrae_engine` for the usual shape cleanup. + + It is independent of `vertebrae_engine` rather than a value of it: this stage + performs no shape cleanup, so the configured engine still runs afterwards. + + > [!IMPORTANT] + > `rib_anchor` requires `vertebrae_engine: shapekit_pro`. Combining it with + > the default engine is rejected at startup with a non-zero exit, because + > that engine may reassign contiguous cranio-caudal identities and undo the + > correction. + + It requires the case CT (located with `ct_file_name` / `ct_root`) and a + precomputed TotalSegmentator `total` volume, both on the prediction's voxel + grid. Only the rib masks in that volume are read. When either input is + missing or sits on a different grid, the stage logs the reason and leaves the + masks untouched. See [vertebrae_rib_identity.md](vertebrae_rib_identity.md). + +10. `rib_file_name` / `rib_root`: how `rib_anchor` finds the rib volume. It + looks for `/` first, then + `//`, then `/.nii.gz`. + +11. `rib_qa_dir`: optional directory for one small JSON QA record per case, + holding the ribs used, the ribs rejected for left/right disagreement, the + solved per-level displacements, the corrected levels and the voxel counts. + Omitted by default. diff --git a/docs/vertebrae_rib_identity.md b/docs/vertebrae_rib_identity.md new file mode 100644 index 0000000..624e45b --- /dev/null +++ b/docs/vertebrae_rib_identity.md @@ -0,0 +1,244 @@ +

Rib-Anchored Vertebral Identity

+ +An optional stage that adjudicates **which vertebral level is which** before the +configured `vertebrae_engine` performs its usual shape cleanup. It is disabled by +default; with the default configuration ShapeKit behaves exactly as it did before +this stage existed. + +--- + +## 1. What it does + +A vertebra segmentation can delineate bone well while assigning level *names* that +are displaced, so that a run of vertebrae carries the identity of a neighbouring +level. The bone is present and its shape is plausible, but the labels are wrong. + +This stage estimates each level's superior position from **costovertebral rib +attachment geometry** and relabels only those levels whose current position +disagrees with that estimate by more than a fixed threshold. + +The cue it uses is external: it reads the ribs, rather than deriving level identity +from the vertebral predictions themselves, their ordering, spacing, connected +components or internal consistency. + +The anatomy it encodes: + +| rib | articulation | predicted attachment | +|---|---|---| +| 1 | single facet on the T1 body | centroid of T1 | +| 2–9 | facets on T(N−1) and TN | midpoint of those two centroids | +| 10–12 | single facet on their own body | centroid of TN | + +Left and right ribs are measured independently. A pair whose two sides disagree by +more than the canonical tolerance is discarded rather than trusted. The remaining +pairs form a least-squares system, together with a smoothness term, a soft tie to +the model's own centroids, and a strongly weighted lumbar anchor. Levels displaced +beyond +the threshold are relabelled at boundaries placed on CT-detected disc planes. + +## 2. Why it is optional + +It requires two inputs ShapeKit does not otherwise need — the case CT and a +precomputed rib segmentation. Mask-only inputs cannot supply them, so the stage is +off unless it is explicitly enabled and its inputs are present. + +It is also deliberately conservative. It relabels only where the rib evidence +requires it, and leaves a case untouched when the evidence is absent, insufficient, +or in agreement with the existing labels. + +## 3. Required inputs + +| Input | Where it comes from | Requirement | +|---|---|---| +| Vertebra masks | the usual `segmentations/` folder | at least 3 non-empty | +| Case CT | `ct_file_name` / `ct_root` (shared with `shapekit_pro`) | same voxel grid as the prediction | +| Rib volume | `rib_file_name` / `rib_root` | TotalSegmentator `total` multilabel, same voxel grid | + +"Same voxel grid" means identical array shape **and** a matching affine. Volumes +that do not already share the grid are rejected; this stage performs no resampling +and no reorientation. + +Only the geometry of the rib masks (label ids 92–115) is read. TotalSegmentator's +vertebral labels are not consulted and play no part in any decision. + +**ShapeKit does not run TotalSegmentator.** The rib volume is a precomputed input, +generated once per dataset outside this pipeline. Keeping it external avoids adding +a GPU inference dependency to a CPU post-processor, and keeps the rib label +convention pinned to a file you control rather than to an installed version. + +## 4. Configuration + +```yaml +# optional identity stage, applied BEFORE vertebrae_engine +vertebrae_identity: rib_anchor # 'none' (default) disables the stage + +# CT lookup, shared with shapekit_pro +ct_file_name: ct.nii.gz +# ct_root: /path/to/original/ct/cases + +# rib volume lookup +rib_file_name: total.nii.gz +# rib_root: /path/to/totalsegmentator/cases + +# optional: one small JSON QA record per case +# rib_qa_dir: /path/to/qa +``` + +The rib volume is looked up in the same order the CT is: inside the case directory +first, then under an external root. + +1. `/` +2. `//` +3. `/.nii.gz` + +## 5. Interaction with `vertebrae_engine` + +The two settings are independent and compose: + +``` +vertebra prediction + │ + ▼ +vertebrae_identity (optional — decides which level is which) + │ + ▼ +vertebrae_engine (shapekit or shapekit_pro — cleans up shape) + │ + ▼ + ShapeKit output +``` + +`vertebrae_identity` is not a value of `vertebrae_engine` and does not replace it. +This stage performs no shape cleanup of its own, so the selected engine still runs +afterwards exactly as configured. + +## 6. Fallback behaviour + +Every condition below leaves the masks exactly as they arrived, logs the reason, +and lets processing continue. A case that cannot be adjudicated is never degraded, +and a batch run is never halted by one. + +| Condition | Behaviour | +|---|---| +| `vertebrae_identity: none` | stage does not run at all | +| Fewer than 3 vertebra masks | no correction | +| CT missing or unreadable | no correction | +| Rib volume missing or unreadable | no correction | +| CT or rib grid differs in shape or affine | no correction, no resampling | +| Spine centreline cannot be derived | no correction | +| Fewer than 4 usable rib pairs | no correction | +| A rib pair's two sides disagree beyond tolerance | that pair is excluded; the rest are used | +| No level exceeds the correction threshold | labels left unchanged | +| Any unexpected error | caught, logged, masks returned unchanged | + +## 7. Directory layout + +Rib volumes alongside each case: + +``` +INPUT (--input_folder) +└── BDMAP_00000031 + ├── ct.nii.gz + ├── total.nii.gz <- rib_file_name + └── segmentations + ├── vertebrae_L5.nii.gz + ... + └── vertebrae_C1.nii.gz +``` + +Or held separately, leaving the prediction folders untouched: + +``` +rib_root/ +├── BDMAP_00000006/total.nii.gz +└── BDMAP_00000031/total.nii.gz +``` + +## 8. Example + +```bash +python -W ignore main.py \ + --input_folder /path/to/predictions \ + --output_folder /path/to/output \ + --log_folder logs/rib_identity \ + --cpu_count 8 +``` + +With `vertebrae_identity: rib_anchor` set in `config.yaml`, each case logs one line: + +``` +[ShapeKit-RibIdentity] BDMAP_00000031: ribs_used=[1,...,10] rejected_lr=[11, 12] + corrected=['L1','T12','T11','T10','T9','T8'] moved=366151 untouched=0.7479 +``` + +and, when a CT or rib volume is absent: + +``` +[ShapeKit-RibIdentity] BDMAP_00000006: rib volume not found (...); identity + correction skipped +``` + +## 9. Resource notes + +The stage holds the CT and the rib volume in memory in addition to the masks +ShapeKit already loads. On a 0.7 mm whole-spine case those are roughly 0.7 GB each, +so budget workers accordingly with `--cpu_count`; high worker counts on +high-resolution cohorts are memory-bound rather than CPU-bound. Runtime is dominated +by connected-component and profile computations over the volume, not by the +least-squares solve, which has fourteen unknowns. + +## 10. Scope of validation + +The method was developed and validated on the two-case BodyMaps vertebrae warm-up +set. The tests in `tests/test_vertebrae_rib_identity.py` cover the label mapping, +the rib measurement and rejection rules, the costovertebral equations, the solver, +the correction threshold, and every fallback path, using synthetic phantoms that +require no imaging data. + +Behaviour on larger and more varied cohorts has not been measured. The stage is +default-off and conservative for that reason. + +## 11. Why `rib_anchor` requires `shapekit_pro` + +> [!IMPORTANT] +> `vertebrae_identity: rib_anchor` is only valid with +> `vertebrae_engine: shapekit_pro`. Pairing it with the default `shapekit` +> engine is rejected at startup, before any case is processed. + +The default `shapekit` vertebra engine may reassign contiguous cranio-caudal +identities during its own cleanup. When the identity stage has deliberately +changed which level a body belongs to, that reassignment can undo the identity +stage and restore the original labelling. The run would appear to succeed while +discarding the correction on every case, which is worse than not running it at +all — so the combination is refused rather than warned about. + +Measured on the synthetic shifted phantom, where the identity stage relabels +38,400 voxels: + +| engine after the identity stage | voxels still differing from the input | outcome | +|---|---|---| +| `shapekit` (default) | 0 | correction lost | +| `shapekit_pro` | 38,400 | correction retained | + +ShapeKit exits with status 2 and an explanatory message: + +``` +[ERROR] Incompatible configuration: vertebrae_identity: 'rib_anchor' cannot be +used with vertebrae_engine: 'shapekit'. +... +Resolve it in one of two ways: + - set vertebrae_engine: shapekit_pro, which preserves the corrected + identities; or + - set vertebrae_identity: none to disable rib-anchored correction and keep + the default engine. +``` + +No existing ShapeKit algorithm was modified to make the combination work. The +default engine's behaviour is left exactly as it is; only the unsupported +pairing is refused. + +This check is a *configuration* error and is distinct from the per-case +conditions in §6. A missing CT, a missing rib volume, too few usable ribs or an +incompatible grid are properties of one case: they are logged, that case is +skipped, and the batch continues. Only a configuration that would silently +discard the correction for every case stops the run. diff --git a/main.py b/main.py index b0ab4ef..929fec4 100644 --- a/main.py +++ b/main.py @@ -1,9 +1,15 @@ import argparse +import sys import multiprocessing from multiprocessing import cpu_count from utils.organs_postprocessing import * from utils.vertebrae_postprocessing import postprocessing_vertebrae from utils.vertebrae_pro import postprocessing_vertebrae_pro +from utils.vertebrae_rib_identity import (IDENTITY_RIB_ANCHOR, + IdentityConfigError, + postprocessing_vertebrae_rib_identity, + resolve_rib_path, + validate_identity_config) import logging import yaml import traceback @@ -36,6 +42,10 @@ vertebrae_engine = config.get('vertebrae_engine', 'shapekit') ct_file_name = config.get('ct_file_name', 'ct.nii.gz') ct_root = config.get('ct_root', None) +vertebrae_identity = config.get('vertebrae_identity', 'none') +rib_file_name = config.get('rib_file_name', 'total.nii.gz') +rib_root = config.get('rib_root', None) +rib_qa_dir = config.get('rib_qa_dir', None) ############################################################## @@ -110,6 +120,7 @@ def combine_segmentation_dict(segmentation_dict: dict, class_map: dict) -> np.nd def process_organs(segmentation_dict: dict, reference_img, combined_seg: np.array, target_organs: set, patient_id: str, logger: logging.Logger, ct_path: str = None, + rib_path: str = None, ): """ Apply organ-specific post-processing functions to the segmentation dict @@ -195,6 +206,18 @@ def process_organs(segmentation_dict: dict, reference_img, combined_seg: np.arra ) if 'vertebrae' in target_organs: + # optional identity adjudication, ahead of the shape-cleanup engine + if vertebrae_identity == IDENTITY_RIB_ANCHOR: + segmentation_dict, _ = postprocessing_vertebrae_rib_identity( + patient_id, + segmentation_dict, + reference_img, + ct_path, + rib_path, + logger=logger, + qa_dir=rib_qa_dir, + ) + if vertebrae_engine == 'shapekit_pro': segmentation_dict = postprocessing_vertebrae_pro( patient_id, @@ -246,6 +269,10 @@ def main(input_path, input_folder_name, output_path=None): if not os.path.exists(ct_path) and ct_root is not None: ct_path = os.path.join(ct_root, input_folder_name, ct_file_name) + # locate the precomputed rib volume for the optional identity stage + rib_path = resolve_rib_path(input_path, input_folder_name, + rib_file_name, rib_root) + postprocessed_segmentation_dict = process_organs( segmentation_dict, img, @@ -254,6 +281,7 @@ def main(input_path, input_folder_name, output_path=None): patient_id = patient_id, logger = logging, ct_path = ct_path, + rib_path = rib_path, ) save_folder_path = os.path.join(output_path, input_folder_name) @@ -352,6 +380,13 @@ def run_in_parallel(sub_folders, input_folder, output_folder, max_workers=4, tqd if __name__ == '__main__': + # validate the vertebrae configuration before any case is processed + try: + validate_identity_config(vertebrae_identity, vertebrae_engine) + except IdentityConfigError as config_error: + print(f"[ERROR] {config_error}") + sys.exit(2) + input_folder = args.input_folder output_folder= args.output_folder diff --git a/tests/test_vertebrae_rib_identity.py b/tests/test_vertebrae_rib_identity.py new file mode 100644 index 0000000..be8fda5 --- /dev/null +++ b/tests/test_vertebrae_rib_identity.py @@ -0,0 +1,663 @@ +"""Tests for the rib-anchored vertebral identity stage. + +Everything here runs on synthetic phantoms; no imaging data is required. The +phantom is a column of uniform blocks at a known pitch with rib slabs placed at +their anatomically correct costovertebral attachments, so the expected answer is +known in closed form. + +Run from the repository root: + + python -m pytest tests/test_vertebrae_rib_identity.py -v +""" + +import logging +import os +import sys + +import numpy as np +import pytest + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +from utils import vertebrae_rib_identity as adapter # noqa: E402 +from utils import vertebrae_rib_identity_engine as engine # noqa: E402 + + +# -------------------------------------------------------------------------- +# Phantom construction +# -------------------------------------------------------------------------- + +ZOOM = 2.0 # isotropic mm +PITCH_SLICES = 12 # 24 mm between level centres +BODY_SLICES = 8 # block height; the remaining 4 slices are the disc +N_LEVELS = 20 # engine ids 1..20, covering the 4..17 solve domain +Z0 = 10 # first block starts here +SHAPE = (64, 64, Z0 + N_LEVELS * PITCH_SLICES + 20) + + +def level_slices(k): + """Inclusive z-slice range occupied by the block for engine label k.""" + start = Z0 + (k - 1) * PITCH_SLICES + return start, start + BODY_SLICES - 1 + + +def level_centroid_mm(k): + """Centroid the engine will measure for a uniform block, in mm.""" + lo, hi = level_slices(k) + return (lo + hi) / 2.0 * ZOOM + + +def expected_attachment_mm(rib_number): + """Costovertebral attachment predicted by the canonical anatomy model.""" + if rib_number == 1 or rib_number >= 10: + return level_centroid_mm(engine.k_of_T(rib_number)) + above = level_centroid_mm(engine.k_of_T(rib_number - 1)) + below = level_centroid_mm(engine.k_of_T(rib_number)) + return (above + below) / 2.0 + + +def build_phantom(shift_levels=0, rib_lr_offset_mm=0.0, ribs=range(1, 13)): + """Takes: an optional whole-column label shift, an optional left/right rib + disagreement in mm, and which ribs to draw. + Does: builds a label volume, a matching CT, and a TotalSegmentator-style + rib volume on one grid. + Returns: (labels uint8, ct int16, rib volume int16, affine).""" + labels = np.zeros(SHAPE, dtype=np.uint8) + ct = np.full(SHAPE, -200, dtype=np.int16) + ribs_vol = np.zeros(SHAPE, dtype=np.int16) + + cx = cy = SHAPE[0] // 2 + half = 10 # 20 voxels = 40 mm across + + for k in range(1, N_LEVELS + 1): + lo, hi = level_slices(k) + stored = k + shift_levels + if not 1 <= stored <= engine.N_CLASSES: + continue + labels[cx - half:cx + half, cy - half:cy + half, lo:hi + 1] = stored + ct[cx - half:cx + half, cy - half:cy + half, lo:hi + 1] = 400 + + for rib_number in ribs: + z_mm = expected_attachment_mm(rib_number) + for side, sign, offset in (("left", 1, 0.0), + ("right", -1, rib_lr_offset_mm)): + z_idx = int(round((z_mm + offset) / ZOOM)) + if not 0 <= z_idx < SHAPE[2]: + continue + label_id = engine.TS_RIB_IDS[f"rib_{side}_{rib_number}"] + if sign > 0: + x0, x1 = cx + half, cx + half + 22 + else: + x0, x1 = cx - half - 22, cx - half + ribs_vol[x0:x1, cy - 3:cy + 3, z_idx:z_idx + 2] = label_id + + affine = np.diag([ZOOM, ZOOM, ZOOM, 1.0]) # already RAS + return labels, ct, ribs_vol, affine + + +def phantom_dict(labels): + """Turn a phantom label volume into a ShapeKit segmentation dict.""" + return {name: (labels == engine.CLASS_MAP_INV[name]).astype(np.uint8) + for name in adapter.VERTEBRA_NAMES + if np.any(labels == engine.CLASS_MAP_INV[name])} + + +# the engine does not expose a reverse map; build one for the tests +engine.CLASS_MAP_INV = {v: k for k, v in engine.CLASS_MAP.items()} + + +@pytest.fixture +def logger(): + log = logging.getLogger("rib_identity_test") + log.handlers = [] + log.addHandler(logging.NullHandler()) + log.setLevel(logging.DEBUG) + return log + + +class _RefImg: + """Minimal stand-in for the nibabel reference image the stage receives.""" + + def __init__(self, affine, zooms=(ZOOM, ZOOM, ZOOM)): + self.affine = affine + + class _H: + def get_zooms(self_inner): + return zooms + + self.header = _H() + + +# -------------------------------------------------------------------------- +# 1. k_of_T +# -------------------------------------------------------------------------- + +def test_k_of_T_endpoints_and_monotonicity(): + assert engine.k_of_T(12) == 6 # T12 + assert engine.k_of_T(1) == 17 # T1 + values = [engine.k_of_T(n) for n in range(1, 13)] + assert values == sorted(values, reverse=True) + assert len(set(values)) == 12 + + +def test_k_of_T_agrees_with_class_map(): + for n in range(1, 13): + assert engine.CLASS_MAP[engine.k_of_T(n)] == f"vertebrae_T{n}" + + +# -------------------------------------------------------------------------- +# 2. complete label mapping +# -------------------------------------------------------------------------- + +def test_label_mapping_endpoints(): + assert adapter.engine_id_to_shapekit_id(1) == 26 + assert engine.CLASS_MAP[1] == "vertebrae_L5" + assert adapter.SHAPEKIT_VERTEBRA_LABELS[26] == "vertebrae_L5" + + assert adapter.engine_id_to_shapekit_id(24) == 49 + assert engine.CLASS_MAP[24] == "vertebrae_C1" + assert adapter.SHAPEKIT_VERTEBRA_LABELS[49] == "vertebrae_C1" + + +def test_label_mapping_is_a_bijection_over_all_levels(): + for engine_id in range(1, 25): + shapekit_id = adapter.engine_id_to_shapekit_id(engine_id) + assert adapter.shapekit_id_to_engine_id(shapekit_id) == engine_id + assert (engine.CLASS_MAP[engine_id] + == adapter.SHAPEKIT_VERTEBRA_LABELS[shapekit_id]) + assert sorted(adapter.SHAPEKIT_VERTEBRA_LABELS) == list(range(26, 50)) + assert len(adapter.VERTEBRA_NAMES) == 24 + + +def test_label_mapping_matches_upstream_table(): + from utils.vertebrae_postprocessing import all_labels + assert {int(k): v for k, v in all_labels.items()} == \ + adapter.SHAPEKIT_VERTEBRA_LABELS + + +def test_label_mapping_matches_repository_config(): + import yaml + root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + with open(os.path.join(root, "config.yaml")) as handle: + class_map = yaml.safe_load(handle)["class_map"] + vertebrae = {int(k): v for k, v in class_map.items() + if str(v).startswith("vertebrae_")} + assert vertebrae == adapter.SHAPEKIT_VERTEBRA_LABELS + + +# -------------------------------------------------------------------------- +# 3 / 4. left-right rib measurement acceptance and rejection +# -------------------------------------------------------------------------- + +def _measure(rib_lr_offset_mm): + labels, _, ribs_vol, _ = build_phantom(rib_lr_offset_mm=rib_lr_offset_mm) + zooms = np.array([ZOOM, ZOOM, ZOOM]) + centre, _, _ = engine.spine_centreline(engine.debris_filtered(labels), zooms) + log = {} + found = engine.rib_attachments(ribs_vol, centre, zooms, log) + return found, log + + +def test_symmetric_ribs_are_accepted(): + found, log = _measure(0.0) + assert sorted(found) == list(range(1, 13)) + assert log["ribs_excluded_lr"] == [] + + +def test_left_right_disagreement_inside_tolerance_is_accepted(): + found, log = _measure(engine.RIB_LR_TOL_MM - 3.0) + assert sorted(found) == list(range(1, 13)) + assert log["ribs_excluded_lr"] == [] + + +def test_left_right_disagreement_beyond_tolerance_is_rejected(): + found, log = _measure(engine.RIB_LR_TOL_MM + 9.0) + assert found == {} + excluded = {entry["rib"] for entry in log["ribs_excluded_lr"]} + assert excluded == set(range(1, 13)) + for entry in log["ribs_excluded_lr"]: + assert entry["lr_mm"] > engine.RIB_LR_TOL_MM + + +def test_rib_tolerance_constant_is_canonical(): + assert engine.RIB_LR_TOL_MM == 15.0 + + +# -------------------------------------------------------------------------- +# 5. costovertebral rib equations +# -------------------------------------------------------------------------- + +def _row_for(rib_number, measurement=0.0): + """Extract the single rib row the solver builds for one rib.""" + ks = list(range(engine.K_LO, engine.K_HI + 1)) + centroid = {k: 0.0 for k in ks} + captured = {} + real_lstsq = np.linalg.lstsq + + def spy(A, b, rcond=None): + captured["A"] = A + captured["b"] = b + return real_lstsq(A, b, rcond=rcond) + + np.linalg.lstsq = spy + try: + engine.solve_centroids({rib_number: measurement}, centroid, + {k: 1.0 for k in ks}, {}) + finally: + np.linalg.lstsq = real_lstsq + # the rib row is the first row, weighted by W_RIB + return captured["A"][0] / engine.W_RIB, list( + range(engine.K_LO, engine.K_HI + 1)) + + +def test_rib_one_articulates_with_a_single_body(): + row, ks = _row_for(1) + assert row[ks.index(engine.k_of_T(1))] == pytest.approx(1.0) + assert np.count_nonzero(row) == 1 + + +@pytest.mark.parametrize("rib_number", list(range(2, 10))) +def test_ribs_two_to_nine_span_two_bodies(rib_number): + row, ks = _row_for(rib_number) + above = ks.index(engine.k_of_T(rib_number - 1)) + below = ks.index(engine.k_of_T(rib_number)) + assert row[above] == pytest.approx(0.5) + assert row[below] == pytest.approx(0.5) + assert np.count_nonzero(row) == 2 + + +@pytest.mark.parametrize("rib_number", [10, 11, 12]) +def test_ribs_ten_to_twelve_articulate_with_their_own_body(rib_number): + row, ks = _row_for(rib_number) + assert row[ks.index(engine.k_of_T(rib_number))] == pytest.approx(1.0) + assert np.count_nonzero(row) == 1 + + +# -------------------------------------------------------------------------- +# 6. least-squares construction and solution +# -------------------------------------------------------------------------- + +def _consistent_inputs(): + ks = list(range(engine.K_LO, engine.K_HI + 1)) + centroid = {k: level_centroid_mm(k) for k in ks} + ribs = {n: expected_attachment_mm(n) for n in range(1, 13)} + frac = {k: 1.0 for k in ks} + return ribs, centroid, frac, ks + + +def test_solve_reproduces_a_self_consistent_column(): + ribs, centroid, frac, ks = _consistent_inputs() + solved, delta = engine.solve_centroids(ribs, centroid, frac, {}) + for k in ks: + assert solved[k] == pytest.approx(centroid[k], abs=1e-6) + assert delta[k] == pytest.approx(0.0, abs=1e-6) + + +def test_pinned_lumbar_anchor_barely_moves(): + ribs, centroid, frac, _ = _consistent_inputs() + for n in ribs: + ribs[n] += 30.0 # push every rib superiorly + solved, delta = engine.solve_centroids(ribs, centroid, frac, {}) + assert abs(delta[engine.K_LO]) < 2.0 # W_PIN holds L2 in place + assert abs(delta[engine.K_HI]) > abs(delta[engine.K_LO]) + + +def test_solve_domain_is_l2_to_t1(): + assert (engine.K_LO, engine.K_HI) == (4, 17) + assert engine.CLASS_MAP[engine.K_LO] == "vertebrae_L2" + assert engine.CLASS_MAP[engine.K_HI] == "vertebrae_T1" + + +def test_solve_weights_are_canonical(): + assert (engine.W_RIB, engine.W_SMOOTH, engine.W_TIE, engine.W_PIN) == \ + (1.0, 0.9, 0.12, 50.0) + + +# -------------------------------------------------------------------------- +# 12. threshold semantics (controlled synthetic deltas, no patient data) +# -------------------------------------------------------------------------- + +def test_min_delta_constant_is_canonical(): + assert engine.MIN_DELTA_MM == 12.0 + + +def test_threshold_operator_in_source_is_strictly_greater(): + """Guard the exact gating expression against an accidental >= or epsilon.""" + source = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "utils", "vertebrae_rib_identity_engine.py") + with open(source, encoding="utf-8") as handle: + text = handle.read() + assert "abs(v) > MIN_DELTA_MM" in text + assert "abs(v) >= MIN_DELTA_MM" not in text + assert "MIN_DELTA_MM -" not in text and "MIN_DELTA_MM +" not in text + + +@pytest.mark.parametrize("delta_mm, corrected", [ + (0.0, False), + (11.9, False), + (12.0, False), # exactly on the threshold: NOT corrected + (-12.0, False), # sign must not change the decision + (12.000001, True), + (12.1, True), + (-12.1, True), + (30.0, True), +]) +def test_gate_decision_for_controlled_deltas(delta_mm, corrected, monkeypatch, + logger): + """Drive the real gate inside process() with a controlled delta.""" + labels, ct, ribs_vol, affine = build_phantom() + target = 10 # engine id inside the domain + + real_solve = engine.solve_centroids + + def controlled(ribs, centroid, frac, log): + solved, _ = real_solve(ribs, centroid, frac, log) + delta = {k: 0.0 for k in solved} + delta[target] = delta_mm + log["delta_mm"] = {engine.SHORT[k]: v for k, v in delta.items()} + return solved, delta + + monkeypatch.setattr(engine, "solve_centroids", controlled) + _, log = engine.process(labels, affine, (ZOOM, ZOOM, ZOOM), ct, ribs_vol) + + assert log.get("levels_corrected") is not None + was_corrected = engine.SHORT[target] in log["levels_corrected"] + assert was_corrected is corrected + + +# -------------------------------------------------------------------------- +# 11. a self-consistent column must be left completely alone +# -------------------------------------------------------------------------- + +def test_correct_column_is_not_modified(): + labels, ct, ribs_vol, affine = build_phantom() + out, log = engine.process(labels, affine, (ZOOM, ZOOM, ZOOM), ct, ribs_vol) + assert log["levels_corrected"] == [] + assert log["voxels_moved"] == 0 + assert np.array_equal(out[labels > 0], labels[labels > 0]) + + +def test_shifted_column_is_detected_and_corrected(): + labels, ct, ribs_vol, affine = build_phantom(shift_levels=1) + _, log = engine.process(labels, affine, (ZOOM, ZOOM, ZOOM), ct, ribs_vol) + assert log["levels_corrected"], "a one-level shift should be detected" + assert log["voxels_moved"] > 0 + + +# -------------------------------------------------------------------------- +# 10. too few usable ribs +# -------------------------------------------------------------------------- + +def test_too_few_ribs_makes_no_correction(): + labels, ct, ribs_vol, affine = build_phantom(shift_levels=1, + ribs=[4, 5]) + out, log = engine.process(labels, affine, (ZOOM, ZOOM, ZOOM), ct, ribs_vol) + assert "too few usable ribs" in log["status"] + assert np.array_equal(out, labels) + + +# -------------------------------------------------------------------------- +# 7 / 8 / 9. adapter fallbacks +# -------------------------------------------------------------------------- + +def _dict_and_ref(): + labels, ct, ribs_vol, affine = build_phantom() + return phantom_dict(labels), _RefImg(affine), labels, ct, ribs_vol, affine + + +def _save(tmp_path, name, array, affine): + import nibabel as nib + path = str(tmp_path / name) + nib.save(nib.Nifti1Image(array, affine), path) + return path + + +def test_missing_ct_is_a_safe_noop(tmp_path, logger): + seg, ref, labels, _, ribs_vol, affine = _dict_and_ref() + rib_path = _save(tmp_path, "total.nii.gz", ribs_vol, affine) + before = {k: v.copy() for k, v in seg.items()} + out, log = adapter.postprocessing_vertebrae_rib_identity( + "CASE", seg, ref, str(tmp_path / "absent_ct.nii.gz"), rib_path, logger) + assert "CT not found" in log["status"] + assert all(np.array_equal(out[k], before[k]) for k in before) + + +def test_missing_rib_volume_is_a_safe_noop(tmp_path, logger): + seg, ref, labels, ct, _, affine = _dict_and_ref() + ct_path = _save(tmp_path, "ct.nii.gz", ct, affine) + before = {k: v.copy() for k, v in seg.items()} + out, log = adapter.postprocessing_vertebrae_rib_identity( + "CASE", seg, ref, ct_path, str(tmp_path / "absent_ribs.nii.gz"), logger) + assert "rib volume not found" in log["status"] + assert all(np.array_equal(out[k], before[k]) for k in before) + + +def test_incompatible_grid_is_a_safe_noop(tmp_path, logger): + seg, ref, labels, ct, ribs_vol, affine = _dict_and_ref() + ct_path = _save(tmp_path, "ct.nii.gz", ct, affine) + small = ribs_vol[:32, :32, :32] + rib_path = _save(tmp_path, "total.nii.gz", small, affine) + before = {k: v.copy() for k, v in seg.items()} + out, log = adapter.postprocessing_vertebrae_rib_identity( + "CASE", seg, ref, ct_path, rib_path, logger) + assert "incompatible" in log["status"] + assert all(np.array_equal(out[k], before[k]) for k in before) + + +def test_mismatched_affine_is_a_safe_noop(tmp_path, logger): + seg, ref, labels, ct, ribs_vol, affine = _dict_and_ref() + ct_path = _save(tmp_path, "ct.nii.gz", ct, affine) + shifted = affine.copy() + shifted[0, 3] += 25.0 # same shape, different grid + rib_path = _save(tmp_path, "total.nii.gz", ribs_vol, shifted) + before = {k: v.copy() for k, v in seg.items()} + out, log = adapter.postprocessing_vertebrae_rib_identity( + "CASE", seg, ref, ct_path, rib_path, logger) + assert "incompatible" in log["status"] + assert all(np.array_equal(out[k], before[k]) for k in before) + + +def test_too_few_masks_is_a_safe_noop(tmp_path, logger): + seg, ref, labels, ct, ribs_vol, affine = _dict_and_ref() + ct_path = _save(tmp_path, "ct.nii.gz", ct, affine) + rib_path = _save(tmp_path, "total.nii.gz", ribs_vol, affine) + thin = {k: seg[k] for k in list(seg)[:2]} + out, log = adapter.postprocessing_vertebrae_rib_identity( + "CASE", thin, ref, ct_path, rib_path, logger) + assert "identity correction skipped" in log["status"] + assert out is thin + + +def test_adapter_never_raises_on_garbage(tmp_path, logger): + seg, ref, labels, ct, ribs_vol, affine = _dict_and_ref() + ct_path = _save(tmp_path, "ct.nii.gz", ct, affine) + rib_path = _save(tmp_path, "total.nii.gz", ribs_vol, affine) + + class Broken: + affine = "not-an-affine" + header = None + + out, log = adapter.postprocessing_vertebrae_rib_identity( + "CASE", seg, Broken(), ct_path, rib_path, logger) + assert "failed" in log["status"] or "skipped" in log["status"] + assert out is seg + + +# -------------------------------------------------------------------------- +# assembly / scatter round trip and overlap precedence +# -------------------------------------------------------------------------- + +def test_assemble_scatter_round_trip(): + labels, _, _, _ = build_phantom() + seg = phantom_dict(labels) + volume, present = adapter.assemble_label_volume(seg) + assert volume is not None and len(present) == N_LEVELS + assert np.array_equal(volume, labels) + rebuilt = adapter.scatter_label_volume(volume, dict(seg)) + again, _ = adapter.assemble_label_volume(rebuilt) + assert np.array_equal(again, labels) + + +def test_overlapping_masks_resolve_to_the_more_superior_level(): + shape = (4, 4, 4) + seg = {"vertebrae_L5": np.ones(shape, dtype=np.uint8), + "vertebrae_L4": np.ones(shape, dtype=np.uint8)} + volume, _ = adapter.assemble_label_volume(seg) + assert np.all(volume == adapter.name_to_engine_id("vertebrae_L4")) + + +def test_rib_path_resolution_prefers_the_case_directory(tmp_path): + case = tmp_path / "CASE" + case.mkdir() + (case / "total.nii.gz").write_bytes(b"x") + root = tmp_path / "ribs" + (root / "CASE").mkdir(parents=True) + (root / "CASE" / "total.nii.gz").write_bytes(b"x") + found = adapter.resolve_rib_path(str(case), "CASE", "total.nii.gz", str(root)) + assert found == os.path.join(str(case), "total.nii.gz") + + +def test_rib_path_resolution_falls_back_to_the_external_root(tmp_path): + case = tmp_path / "CASE" + case.mkdir() + root = tmp_path / "ribs" + root.mkdir() + (root / "CASE.nii.gz").write_bytes(b"x") + found = adapter.resolve_rib_path(str(case), "CASE", "total.nii.gz", str(root)) + assert found == os.path.join(str(root), "CASE.nii.gz") + + +def test_rib_path_resolution_returns_none_when_absent(tmp_path): + case = tmp_path / "CASE" + case.mkdir() + assert adapter.resolve_rib_path(str(case), "CASE", "total.nii.gz", None) is None + + +# -------------------------------------------------------------------------- +# Configuration validation: the rib_anchor / engine pairing +# +# This is a configuration error, not a per-case condition. It is checked once +# at startup and stops the run; the per-case fallbacks above must stay no-ops. +# -------------------------------------------------------------------------- + +def test_rib_anchor_with_default_engine_is_rejected(): + with pytest.raises(adapter.IdentityConfigError) as excinfo: + adapter.validate_identity_config("rib_anchor", "shapekit") + message = str(excinfo.value) + assert "cannot be used with" in message + assert "shapekit_pro" in message + assert "vertebrae_identity: none" in message + assert "reassign contiguous" in message + + +def test_rib_anchor_with_pro_engine_is_accepted(): + assert adapter.validate_identity_config("rib_anchor", "shapekit_pro") is None + + +@pytest.mark.parametrize("identity", ["none", None, ""]) +@pytest.mark.parametrize("engine_name", ["shapekit", "shapekit_pro"]) +def test_disabled_identity_accepts_any_engine(identity, engine_name): + assert adapter.validate_identity_config(identity, engine_name) is None + + +def test_unknown_identity_value_is_rejected(): + with pytest.raises(adapter.IdentityConfigError) as excinfo: + adapter.validate_identity_config("rib_anchr", "shapekit_pro") + assert "Unknown vertebrae_identity" in str(excinfo.value) + + +def test_unknown_engine_is_also_rejected_when_identity_is_on(): + """An engine we have not verified must not be assumed compatible.""" + with pytest.raises(adapter.IdentityConfigError): + adapter.validate_identity_config("rib_anchor", "some_future_engine") + + +def test_validation_runs_before_any_case_is_processed(): + """Guard the call site: validation must precede the processing call.""" + root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + with open(os.path.join(root, "main.py"), encoding="utf-8") as handle: + source = handle.read() + guard = source.index("if __name__ == '__main__':") + validate_at = source.index("validate_identity_config(", guard) + run_at = source.index("run_in_parallel(", guard) + assert validate_at < run_at, "config validation must run before processing" + assert "sys.exit(2)" in source[guard:run_at] + + +def test_missing_inputs_are_not_configuration_errors(tmp_path, logger): + """A case that lacks a CT or ribs must no-op, never raise IdentityConfigError.""" + seg, ref, labels, ct, ribs_vol, affine = _dict_and_ref() + rib_path = _save(tmp_path, "total.nii.gz", ribs_vol, affine) + ct_path = _save(tmp_path, "ct.nii.gz", ct, affine) + + for missing_ct, missing_rib in ((True, False), (False, True), (True, True)): + working = {k: v.copy() for k, v in seg.items()} + out, log = adapter.postprocessing_vertebrae_rib_identity( + "CASE", working, ref, + str(tmp_path / "absent.nii.gz") if missing_ct else ct_path, + str(tmp_path / "absent.nii.gz") if missing_rib else rib_path, + logger) + assert "skipped" in log["status"] + assert all(np.array_equal(out[k], seg[k]) for k in seg) + + +# -------------------------------------------------------------------------- +# End to end: the correction must survive the downstream engine +# -------------------------------------------------------------------------- + +def _run_stage_then(engine_callable, tmp_path): + """Identity stage, then a downstream engine; returns (before, after) volumes.""" + labels, ct, ribs_vol, affine = build_phantom(shift_levels=1) + seg = phantom_dict(labels) + ref = _RefImg(affine) + ct_path = _save(tmp_path, "ct.nii.gz", ct, affine) + rib_path = _save(tmp_path, "total.nii.gz", ribs_vol, affine) + + def volume(d): + out = np.zeros(labels.shape, dtype=np.uint8) + for name in adapter.VERTEBRA_NAMES: + mask = d.get(name) + if mask is not None: + out[mask > 0] = adapter.ENGINE_ID_BY_NAME[name] + return out + + before = volume(seg) + seg, log = adapter.postprocessing_vertebrae_rib_identity( + "CASE", seg, ref, ct_path, rib_path, logging.getLogger("e2e")) + corrected = volume(seg) + seg = engine_callable(seg, ref, ct_path) + return before, corrected, volume(seg), log + + +def test_shapekit_pro_retains_the_identity_correction(tmp_path): + """The supported pairing must keep the correction after downstream cleanup.""" + from utils.vertebrae_pro import postprocessing_vertebrae_pro + + def run_pro(seg, ref, ct_path): + return postprocessing_vertebrae_pro( + "CASE", seg, ref, ct_path, logging.getLogger("e2e")) + + before, corrected, after, log = _run_stage_then(run_pro, tmp_path) + + assert log["voxels_moved"] == 38400 + assert int((before != corrected).sum()) == 38400 + # the correction is still present once shapekit_pro has run + assert int((before != after).sum()) == 38400 + + +def test_default_engine_would_revert_it(tmp_path): + """Why the pairing is rejected: the default engine undoes the correction. + + This documents the measured behaviour that motivates + ``validate_identity_config``. It asserts the revert, so the day the default + engine stops reverting, this test fails and the restriction can be revisited. + """ + from utils.vertebrae_postprocessing import postprocessing_vertebrae + + def run_default(seg, ref, ct_path): + return postprocessing_vertebrae("CASE", seg, logger=logging.getLogger("e2e")) + + before, corrected, after, log = _run_stage_then(run_default, tmp_path) + + assert int((before != corrected).sum()) == 38400 # stage did its work + assert int((before != after).sum()) == 0 # engine undid all of it diff --git a/utils/vertebrae_rib_identity.py b/utils/vertebrae_rib_identity.py new file mode 100644 index 0000000..ef3df45 --- /dev/null +++ b/utils/vertebrae_rib_identity.py @@ -0,0 +1,366 @@ +"""ShapeKit adapter for rib-anchored vertebral identity correction. + +This module is the integration layer only. It performs no anatomical +reasoning: every scientific decision is made by ``vertebrae_rib_identity_engine``, +which is a byte-for-byte copy of the accepted warm-up implementation. The +adapter is responsible for + + * discovering the case CT and the precomputed rib volume, + * checking that they share the prediction's voxel grid, + * converting between ShapeKit's per-organ binary masks (label ids 26-49) + and the engine's combined label volume (ids 1-24), + * logging one case-level record, and + * failing safely, so a batch run is never interrupted by a case that + lacks the inputs this stage needs. + +The stage runs *before* the configured ``vertebrae_engine`` and is orthogonal +to it: it adjudicates which level is which, then hands the masks on for the +usual shape cleanup. It is disabled unless ``vertebrae_identity`` is set. + +The rib volume is a precomputed TotalSegmentator ``total`` multilabel output. +Nothing here runs TotalSegmentator, and only the geometry of its rib masks is +read - its vertebral labels are not consulted at any point. +""" + +import json +import os + +import nibabel as nib +import numpy as np + +from . import vertebrae_rib_identity_engine as engine + +LOG_PREFIX = "[ShapeKit-RibIdentity]" + +# Value of the ``vertebrae_identity`` config key that selects this stage. +IDENTITY_RIB_ANCHOR = "rib_anchor" + +#: ``vertebrae_engine`` values this stage can run in front of. The default +#: engine is excluded deliberately - see ``validate_identity_config``. +COMPATIBLE_ENGINES = ("shapekit_pro",) + + +INCOMPATIBLE_ENGINE_MESSAGE = """\ +Incompatible configuration: vertebrae_identity: '{identity}' cannot be used \ +with vertebrae_engine: '{engine}'. + +The rib-anchored identity stage changes which vertebral level a body is +assigned to. The default ShapeKit vertebra engine may then reassign contiguous +cranio-caudal identities during its own cleanup, which can undo the identity +stage and silently restore the original labelling. The run would appear to +succeed while discarding the correction on every case. + +Resolve it in one of two ways: + - set vertebrae_engine: shapekit_pro, which preserves the corrected + identities; or + - set vertebrae_identity: none to disable rib-anchored correction and keep + the default engine. + +See docs/vertebrae_rib_identity.md for details.""" + + +class IdentityConfigError(ValueError): + """Raised for a configuration that cannot produce a correct result. + + Distinct from the per-case conditions this stage tolerates: a missing CT, + a missing rib volume, too few usable ribs or an incompatible grid are + properties of one case, are handled by skipping that case, and must never + stop a batch. This exception is for a combination of settings that would + silently discard the identity correction for every case, which is worth + refusing before any work starts.""" + + +def validate_identity_config(vertebrae_identity, vertebrae_engine): + """Takes: the configured ``vertebrae_identity`` and ``vertebrae_engine``. + Does: checks the two settings can produce a meaningful result together. + Intended to be called once at startup, before any case is processed. + Returns: None when the configuration is usable. + Raises: IdentityConfigError when the identity stage is enabled in front of + an engine that would discard its output.""" + if vertebrae_identity in (None, "none", "None", ""): + return + if vertebrae_identity != IDENTITY_RIB_ANCHOR: + raise IdentityConfigError( + f"Unknown vertebrae_identity: {vertebrae_identity!r}. " + f"Valid values are 'none' (default) and '{IDENTITY_RIB_ANCHOR}'.") + if vertebrae_engine in COMPATIBLE_ENGINES: + return + raise IdentityConfigError( + INCOMPATIBLE_ENGINE_MESSAGE.format( + identity=IDENTITY_RIB_ANCHOR, engine=vertebrae_engine)) + + +# -------------------------------------------------------------------------- +# Label-space contract +# +# Two identity spaces exist and must never be confused: +# +# engine space 1 .. 24 L5 .. C1 (the accepted warm-up convention) +# ShapeKit space 26 .. 49 L5 .. C1 (config.yaml class_map) +# +# The two are related by a constant offset, but no arithmetic is used to +# convert between them anywhere in this integration. Both directions go +# through the explicit tables below, keyed by anatomical name, so a future +# change to either convention surfaces as a failed consistency check at +# import time rather than as silently mislabelled anatomy. +# -------------------------------------------------------------------------- + +#: ShapeKit combined-label id -> anatomical name, ordered L5 (26) to C1 (49). +SHAPEKIT_VERTEBRA_LABELS = { + 26: "vertebrae_L5", 27: "vertebrae_L4", 28: "vertebrae_L3", + 29: "vertebrae_L2", 30: "vertebrae_L1", 31: "vertebrae_T12", + 32: "vertebrae_T11", 33: "vertebrae_T10", 34: "vertebrae_T9", + 35: "vertebrae_T8", 36: "vertebrae_T7", 37: "vertebrae_T6", + 38: "vertebrae_T5", 39: "vertebrae_T4", 40: "vertebrae_T3", + 41: "vertebrae_T2", 42: "vertebrae_T1", 43: "vertebrae_C7", + 44: "vertebrae_C6", 45: "vertebrae_C5", 46: "vertebrae_C4", + 47: "vertebrae_C3", 48: "vertebrae_C2", 49: "vertebrae_C1", +} + +#: anatomical name -> engine label id (1 = L5 ... 24 = C1) +ENGINE_ID_BY_NAME = {name: k for k, name in engine.CLASS_MAP.items()} +#: anatomical name -> ShapeKit combined-label id +SHAPEKIT_ID_BY_NAME = {name: i for i, name in SHAPEKIT_VERTEBRA_LABELS.items()} + +#: vertebra names in inferior-to-superior order, L5 first. +VERTEBRA_NAMES = [engine.CLASS_MAP[k] for k in range(1, engine.N_CLASSES + 1)] + + +def _check_label_tables(): + """Fail loudly at import if the two identity spaces ever diverge.""" + if set(ENGINE_ID_BY_NAME) != set(SHAPEKIT_ID_BY_NAME): + raise RuntimeError( + "vertebra name sets differ between engine and ShapeKit label tables") + engine_order = sorted(ENGINE_ID_BY_NAME, key=ENGINE_ID_BY_NAME.get) + shapekit_order = sorted(SHAPEKIT_ID_BY_NAME, key=SHAPEKIT_ID_BY_NAME.get) + if engine_order != shapekit_order: + raise RuntimeError( + "engine and ShapeKit label tables disagree on vertebra ordering") + # Cross-check against the table the existing vertebrae module already + # carries, when that module can be imported (it pulls optional deps). + try: + from .vertebrae_postprocessing import all_labels as upstream + except Exception: # noqa: BLE001 - the check is advisory, not required + return + if {int(k): v for k, v in upstream.items()} != SHAPEKIT_VERTEBRA_LABELS: + raise RuntimeError( + "SHAPEKIT_VERTEBRA_LABELS disagrees with vertebrae_postprocessing." + "all_labels") + + +_check_label_tables() + + +def engine_id_to_shapekit_id(engine_id): + """Takes: an engine label id (1..24). + Returns: the corresponding ShapeKit combined-label id (26..49).""" + return SHAPEKIT_ID_BY_NAME[engine.CLASS_MAP[engine_id]] + + +def shapekit_id_to_engine_id(shapekit_id): + """Takes: a ShapeKit combined-label id (26..49). + Returns: the corresponding engine label id (1..24).""" + return ENGINE_ID_BY_NAME[SHAPEKIT_VERTEBRA_LABELS[shapekit_id]] + + +def name_to_engine_id(name): + """Takes: an anatomical name such as ``vertebrae_T9``. + Returns: the engine label id.""" + return ENGINE_ID_BY_NAME[name] + + +# -------------------------------------------------------------------------- +# Input discovery and validation +# -------------------------------------------------------------------------- + +def resolve_rib_path(input_path, patient_id, rib_file_name, rib_root=None): + """Takes: the case input directory, the case id, the configured rib file + name, and an optional external rib root. + Does: looks for the rib volume inside the case directory first, then under + ``//``, then for a flat + ``/.nii.gz``. This mirrors how ShapeKit already + locates the case CT. + Returns: the first path that exists, else None.""" + if not rib_file_name: + return None + candidates = [os.path.join(input_path, rib_file_name)] + if rib_root: + candidates.append(os.path.join(rib_root, patient_id, rib_file_name)) + candidates.append(os.path.join(rib_root, f"{patient_id}.nii.gz")) + for path in candidates: + if path and os.path.exists(path): + return path + return None + + +def _load_aligned(path, reference_img, shape, kind, patient_id, logger): + """Takes: a volume path, the case reference image, the mask array shape, a + short label for messages, the case id and a logger. + Does: loads the volume and requires it to sit on exactly the prediction's + grid - identical array shape and an affine matching the reference. No + resampling and no reorientation is performed: a volume that does not + already share the grid is rejected rather than silently transformed. + Returns: the array as int16, or None with the reason logged.""" + try: + img = nib.load(path) + except Exception as exc: # noqa: BLE001 - batch runs must not stall + logger.warning(f"{LOG_PREFIX} {patient_id}: {kind} unreadable " + f"({path}: {exc}); identity correction skipped") + return None + if tuple(img.shape) != tuple(shape): + logger.warning(f"{LOG_PREFIX} {patient_id}: {kind} grid {tuple(img.shape)} " + f"does not match prediction {tuple(shape)}; " + f"identity correction skipped") + return None + if not np.allclose(img.affine, reference_img.affine, atol=1e-4): + logger.warning(f"{LOG_PREFIX} {patient_id}: {kind} affine differs from the " + f"prediction affine; identity correction skipped " + f"(no resampling is performed)") + return None + try: + return np.asarray(img.dataobj).astype(np.int16) + except Exception as exc: # noqa: BLE001 + logger.warning(f"{LOG_PREFIX} {patient_id}: {kind} could not be read into " + f"memory ({exc}); identity correction skipped") + return None + + +# -------------------------------------------------------------------------- +# Label-volume assembly +# -------------------------------------------------------------------------- + +def assemble_label_volume(segmentation_dict): + """Takes: the ShapeKit segmentation dict. + Does: builds the engine's combined label volume from the individual binary + vertebra masks. Levels are written in ascending id order, so where two + masks claim the same voxel the more superior level wins - the same + precedence ShapeKit's own ``combine_segmentation_dict`` applies. + Returns: (uint8 label volume, list of names that were present), or + (None, []) when no vertebra mask is present.""" + present = [n for n in VERTEBRA_NAMES + if segmentation_dict.get(n) is not None + and np.any(segmentation_dict[n])] + if not present: + return None, [] + shape = segmentation_dict[present[0]].shape + volume = np.zeros(shape, dtype=np.uint8) + for name in VERTEBRA_NAMES: # ascending engine id == ascending ShapeKit id + mask = segmentation_dict.get(name) + if mask is None: + continue + volume[mask > 0] = ENGINE_ID_BY_NAME[name] + return volume, present + + +def scatter_label_volume(volume, segmentation_dict): + """Takes: an engine label volume and the ShapeKit segmentation dict. + Does: writes each level back as a binary mask, replacing what was there. + A level that was absent before and is still empty stays absent. + Returns: the segmentation dict.""" + for name in VERTEBRA_NAMES: + mask = (volume == ENGINE_ID_BY_NAME[name]).astype(np.uint8) + if mask.any() or segmentation_dict.get(name) is not None: + segmentation_dict[name] = mask + return segmentation_dict + + +# -------------------------------------------------------------------------- +# Stage entry point +# -------------------------------------------------------------------------- + +def _summarise(patient_id, log, logger): + """Emit exactly one case-level line; the detail lives in the QA record.""" + if "status" in log: + logger.info(f"{LOG_PREFIX} {patient_id}: {log['status']}") + return + churn = log.get("churn", {}) + logger.info( + f"{LOG_PREFIX} {patient_id}: ribs_used={log.get('ribs_used', [])} " + f"rejected_lr={[r['rib'] for r in log.get('ribs_excluded_lr', [])]} " + f"corrected={log.get('levels_corrected', [])} " + f"moved={log.get('voxels_moved', 0)} " + f"untouched={churn.get('fraction_untouched')}") + + +def _write_qa(qa_dir, patient_id, log, logger): + if not qa_dir: + return + try: + os.makedirs(qa_dir, exist_ok=True) + with open(os.path.join(qa_dir, f"{patient_id}.json"), "w") as handle: + json.dump(log, handle, indent=2) + except Exception as exc: # noqa: BLE001 - QA output is never load-bearing + logger.warning(f"{LOG_PREFIX} {patient_id}: could not write QA record ({exc})") + + +def postprocessing_vertebrae_rib_identity(patient_id, segmentation_dict, + reference_img, ct_path, rib_path, + logger, qa_dir=None): + """Takes: the case id, the ShapeKit segmentation dict, the case reference + image, the resolved CT path, the resolved rib-volume path, a logger, + and an optional directory for per-case QA records. + Does: runs rib-anchored identity correction on the vertebra masks. The CT + and the rib volume must both be present and share the prediction's + voxel grid; when either is missing or incompatible the masks are + returned untouched and the reason is logged. All anatomical decisions + are made by the vendored engine, unchanged. + Returns: (segmentation dict, per-case log dict). The dict is returned + unmodified on every fallback path, and this function does not raise.""" + log = {} + try: + volume, present = assemble_label_volume(segmentation_dict) + if volume is None or len(present) < 3: + log["status"] = (f"{len(present)} vertebra masks present; " + f"identity correction skipped") + _summarise(patient_id, log, logger) + return segmentation_dict, log + + shape = volume.shape + + if not ct_path or not os.path.exists(ct_path): + log["status"] = (f"CT not found ({ct_path}); " + f"identity correction skipped") + _summarise(patient_id, log, logger) + return segmentation_dict, log + + if not rib_path or not os.path.exists(rib_path): + log["status"] = (f"rib volume not found ({rib_path}); " + f"identity correction skipped") + _summarise(patient_id, log, logger) + return segmentation_dict, log + + ct = _load_aligned(ct_path, reference_img, shape, "CT", patient_id, logger) + if ct is None: + log["status"] = "CT incompatible; identity correction skipped" + return segmentation_dict, log + + ribs = _load_aligned(rib_path, reference_img, shape, "rib volume", + patient_id, logger) + if ribs is None: + log["status"] = "rib volume incompatible; identity correction skipped" + return segmentation_dict, log + + log["ct_path"] = ct_path + log["rib_path"] = rib_path + + corrected, engine_log = engine.process( + volume, + reference_img.affine, + reference_img.header.get_zooms()[:3], + ct, + ribs, + ) + log.update(engine_log) + + if corrected is not None: + segmentation_dict = scatter_label_volume(corrected, segmentation_dict) + + _summarise(patient_id, log, logger) + _write_qa(qa_dir, patient_id, log, logger) + return segmentation_dict, log + + except Exception as exc: # noqa: BLE001 - one bad case must not stop a batch + logger.error(f"{LOG_PREFIX} {patient_id}: identity stage failed ({exc}); " + f"masks left unchanged") + log["status"] = f"identity stage failed ({exc}); masks left unchanged" + return segmentation_dict, log diff --git a/utils/vertebrae_rib_identity_engine.py b/utils/vertebrae_rib_identity_engine.py new file mode 100644 index 0000000..98c7624 --- /dev/null +++ b/utils/vertebrae_rib_identity_engine.py @@ -0,0 +1,417 @@ +"""Rib-anchored vertebral identity correction - scientific core. + +PROVENANCE +---------- +This module is the scientific core of the accepted BodyMaps vertebrae warm-up +solution, vendored unchanged. The body below (from ``from __future__`` to the +end of ``process``) is a byte-for-byte copy of lines 59-422 of the canonical +source file ``postprocessing_vertebrae.py``: + + SHA-256 c98e3b7233c29e860c7afddc51ce50db0518b732a81e156edae13e19971b5e01 + +It is kept as a literal slice, rather than rewritten to fit ShapeKit's style, +so that a single ``diff`` can demonstrate the algorithm has not changed. + +For that reason the import block is also left exactly as it was, which means +``argparse``, ``json``, ``os``, ``copy`` and ``nibabel`` are imported but +unused here: they belonged to the original file's command-line wrapper, which +this module deliberately omits. They are kept so the slice stays verifiable +byte-for-byte; they can be trimmed once equivalence has been established. All ShapeKit integration - configuration, input discovery, +label-space conversion, logging and fallback handling - lives in the adapter +``vertebrae_rib_identity.py`` and never in this file. + +PURPOSE +------- +Vertebra segmentation models can produce a column whose bone is well delineated +but whose level names are displaced, so that a run of vertebrae carries the +identity of a neighbouring level. This module estimates each vertebral level's +superior position from costovertebral rib attachment geometry and relabels only +those levels whose position disagrees with that estimate by more than the +canonical threshold. + +The distinguishing property of this approach is the cue it uses: it reads an +external anatomical structure - the ribs - rather than deriving level identity +from the vertebral predictions themselves, their ordering, spacing, connected +components or internal consistency. + +EXTERNAL INPUTS +--------------- +Beyond the vertebra prediction it requires: + + * the case CT, on the same voxel grid as the prediction; and + * a precomputed TotalSegmentator ``total`` multilabel volume, on the same + grid, from which only the 24 rib masks (ids 92-115) are read. + +Only the *geometry* of the rib masks is consumed. TotalSegmentator's vertebral +labels are deliberately not read and play no part in any decision: the +prediction model's vertebra head was fine-tuned on those labels, so they cannot +serve as an independent check on it. Nothing in this module runs +TotalSegmentator; the rib volume is supplied as a precomputed input. + +Requires numpy, scipy, nibabel - all existing ShapeKit dependencies. +""" + +from __future__ import annotations + +import argparse +import copy +import json +import os +from typing import Dict + +import nibabel as nib +import numpy as np +from nibabel.orientations import (aff2axcodes, apply_orientation, axcodes2ornt, + ornt_transform) +from scipy import ndimage, signal + +# -------------------------------------------------------------------------------------- +# Label convention: index 1 = L5, rising superiorly to 24 = C1 +# (SuPreM dataset/dataloader_test.py, class_map_part_vertebrae) +# -------------------------------------------------------------------------------------- + +CLASS_MAP: Dict[int, str] = { + 1: "vertebrae_L5", 2: "vertebrae_L4", 3: "vertebrae_L3", 4: "vertebrae_L2", + 5: "vertebrae_L1", 6: "vertebrae_T12", 7: "vertebrae_T11", 8: "vertebrae_T10", + 9: "vertebrae_T9", 10: "vertebrae_T8", 11: "vertebrae_T7", 12: "vertebrae_T6", + 13: "vertebrae_T5", 14: "vertebrae_T4", 15: "vertebrae_T3", 16: "vertebrae_T2", + 17: "vertebrae_T1", 18: "vertebrae_C7", 19: "vertebrae_C6", 20: "vertebrae_C5", + 21: "vertebrae_C4", 22: "vertebrae_C3", 23: "vertebrae_C2", 24: "vertebrae_C1", +} +N_CLASSES = len(CLASS_MAP) +SHORT = {k: v.replace("vertebrae_", "") for k, v in CLASS_MAP.items()} +CONN26 = np.ones((3, 3, 3), dtype=bool) + +# ---- parameters (all physical units) -------------------------------------------------- +RIB_LR_TOL_MM = 15.0 # max left/right disagreement for a rib to be trusted +MIN_DELTA_MM = 12.0 # a level is corrected only if it must move further than this +SNAP_MM = 7.0 # snap a boundary onto a CT disc trough within this distance +DEBRIS_REL = 0.02 # foreground components below this fraction are far-field debris +BONE_HU = 175.0 +K_LO, K_HI = 4, 17 # solve over L2 .. T1; L2 is the pinned caudal anchor +W_RIB, W_SMOOTH, W_TIE, W_PIN = 1.0, 0.9, 0.12, 50.0 +TROUGH_PROMINENCE = 0.03 +TROUGH_MIN_SEP_MM = 8.0 + +# TotalSegmentator `total` label ids for the rib structures (v2 class map). +TS_RIB_IDS = { + "rib_left_1": 92, "rib_left_2": 93, "rib_left_3": 94, "rib_left_4": 95, + "rib_left_5": 96, "rib_left_6": 97, "rib_left_7": 98, "rib_left_8": 99, + "rib_left_9": 100, "rib_left_10": 101, "rib_left_11": 102, "rib_left_12": 103, + "rib_right_1": 104, "rib_right_2": 105, "rib_right_3": 106, "rib_right_4": 107, + "rib_right_5": 108, "rib_right_6": 109, "rib_right_7": 110, "rib_right_8": 111, + "rib_right_9": 112, "rib_right_10": 113, "rib_right_11": 114, "rib_right_12": 115, +} + + +def k_of_T(n: int) -> int: + """Our label index for thoracic vertebra TN (T12 -> 6, ... T1 -> 17).""" + return 6 + (12 - n) + + +# -------------------------------------------------------------------------------------- +# Orientation (lossless axis permutation, no resampling) +# -------------------------------------------------------------------------------------- + +def to_ras(arr, affine): + src = axcodes2ornt(aff2axcodes(affine)) + dst = axcodes2ornt(("R", "A", "S")) + xf = ornt_transform(src, dst) + return apply_orientation(arr, xf), xf + + +def from_ras(arr, affine): + src = axcodes2ornt(aff2axcodes(affine)) + dst = axcodes2ornt(("R", "A", "S")) + return apply_orientation(arr, ornt_transform(dst, src)) + + +def ras_zooms(zooms, xf): + out = np.zeros(3, dtype=float) + for i in range(3): + out[int(xf[i, 0])] = float(zooms[i]) + return out + + +# -------------------------------------------------------------------------------------- +# Spine geometry +# -------------------------------------------------------------------------------------- + +def spine_centreline(fg, z, smooth_mm=25.0): + """Per-slice robust centre of the spine, interpolated and smoothed along S.""" + ns = fg.shape[2] + cen = np.full((ns, 2), np.nan) + for s in range(ns): + sl = fg[:, :, s] + if sl.any(): + r, c = np.nonzero(sl) + cen[s] = (np.median(r), np.median(c)) + ok = ~np.isnan(cen[:, 0]) + idx = np.flatnonzero(ok) + if idx.size < 3: + return None, None, None + for j in range(2): + cen[:, j] = np.interp(np.arange(ns), idx, cen[ok, j]) + cen[:, j] = ndimage.gaussian_filter1d(cen[:, j], + sigma=max(smooth_mm / z[2], 1.0), + mode="nearest") + return cen, int(idx[0]), int(idx[-1]) + + +def body_radius_profile(fg, cen, z, lo, hi, frac=0.55, floor_mm=8.0, cap_mm=26.0): + """Column radius tracking the vertebral body; cervical bodies are far smaller.""" + ns = fg.shape[2] + rad = np.full(ns, np.nan) + rr, cc = np.meshgrid(np.arange(fg.shape[0]), np.arange(fg.shape[1]), indexing="ij") + for s in range(lo, hi + 1): + sl = fg[:, :, s] + if not sl.any(): + continue + d = np.hypot((rr[sl] - cen[s, 0]) * z[0], (cc[sl] - cen[s, 1]) * z[1]) + rad[s] = np.clip(np.percentile(d, 70) * frac, floor_mm, cap_mm) + ok = ~np.isnan(rad) + rad = np.interp(np.arange(ns), np.flatnonzero(ok), rad[ok]) + return ndimage.gaussian_filter1d(rad, sigma=max(30.0 / z[2], 1.0), mode="nearest") + + +def bone_profile(ct, cen, rad, z, lo, hi, bone_hu=BONE_HU): + """Fraction of the tracking column that is bone, per superior slice.""" + ns = ct.shape[2] + prof = np.zeros(ns) + rr, cc = np.meshgrid(np.arange(ct.shape[0]), np.arange(ct.shape[1]), indexing="ij") + for s in range(lo, hi + 1): + d2 = ((rr - cen[s, 0]) * z[0]) ** 2 + ((cc - cen[s, 1]) * z[1]) ** 2 + col = d2 < rad[s] ** 2 + n = col.sum() + if n: + prof[s] = np.count_nonzero(col & (ct[:, :, s] > bone_hu)) / n + return prof + + +def debris_filtered(labels): + """Foreground with far-field specks removed - used only to locate the centreline.""" + fg = labels > 0 + if not fg.any(): + return fg + cc, n = ndimage.label(fg, structure=CONN26) + sz = np.bincount(cc.ravel()) + sz[0] = 0 + return np.isin(cc, [i for i in range(1, n + 1) if sz[i] > DEBRIS_REL * sz.max()]) + + +def level_stats(labels, z): + """Centroid of each level's dominant component, and that component's share.""" + objs = ndimage.find_objects(labels.astype(np.int32), max_label=N_CLASSES) + centroid, frac = {}, {} + for k in range(1, N_CLASSES + 1): + sl = objs[k - 1] + if sl is None: + continue + sub = labels[sl] == k + if not sub.any(): + continue + ccl, _ = ndimage.label(sub, structure=CONN26) + s2 = np.bincount(ccl.ravel()) + s2[0] = 0 + main = ccl == int(s2.argmax()) + w = main.sum(axis=(0, 1)).astype(float) + nz = np.flatnonzero(w) + cum = np.cumsum(w[nz]) / w[nz].sum() + centroid[k] = float(np.interp(0.5, cum, nz + sl[2].start)) * z[2] + frac[k] = float(s2.max()) / float(sub.sum()) + return centroid, frac + + +# -------------------------------------------------------------------------------------- +# Rib evidence +# -------------------------------------------------------------------------------------- + +def rib_attachments(ts, cen, z, log): + """Superior coordinate where each rib pair meets the spine, left/right cross-checked.""" + def one(mask): + idx = np.nonzero(mask) + if idx[0].size < 50: + return None + rv, av, sv = idx + d = np.hypot((rv - cen[sv, 0]) * z[0], (av - cen[sv, 1]) * z[1]) + return float(np.median(sv[d <= np.percentile(d, 10)])) * z[2] + + ribs, excluded = {}, [] + for num in range(1, 13): + la = one(ts == TS_RIB_IDS[f"rib_left_{num}"]) + ra = one(ts == TS_RIB_IDS[f"rib_right_{num}"]) + if la is None or ra is None: + continue + if abs(la - ra) > RIB_LR_TOL_MM: + excluded.append({"rib": num, "lr_mm": round(abs(la - ra), 1)}) + continue + ribs[num] = float(np.mean([la, ra])) + log["ribs_used"] = sorted(ribs) + log["ribs_excluded_lr"] = excluded + return ribs + + +def solve_centroids(ribs, centroid, frac, log): + """Least squares: rib equations + smoothness + model tie + pinned lumbar anchor.""" + ks = list(range(K_LO, K_HI + 1)) + idx = {k: i for i, k in enumerate(ks)} + rows, rhs, wts = [], [], [] + + def add(row, val, w): + rows.append(row); rhs.append(val); wts.append(w) + + for num, meas in ribs.items(): + r = np.zeros(len(ks)) + if num == 1 or num >= 10: + k = k_of_T(num) + if k not in idx: + continue + r[idx[k]] = 1.0 + else: + ka, kb = k_of_T(num - 1), k_of_T(num) + if ka not in idx or kb not in idx: + continue + r[idx[ka]] = 0.5 + r[idx[kb]] = 0.5 + add(r, meas, W_RIB) + + for i in range(1, len(ks) - 1): + r = np.zeros(len(ks)) + r[i - 1], r[i], r[i + 1] = 1.0, -2.0, 1.0 + add(r, 0.0, W_SMOOTH) + + for k in ks: + if k not in centroid: + continue + r = np.zeros(len(ks)) + r[idx[k]] = 1.0 + add(r, centroid[k], W_TIE * max(frac.get(k, 0.0), 0.05)) + + if K_LO in centroid: + r = np.zeros(len(ks)) + r[idx[K_LO]] = 1.0 + add(r, centroid[K_LO], W_PIN) + + A = np.array(rows) * np.array(wts)[:, None] + b = np.array(rhs) * np.array(wts) + sol, *_ = np.linalg.lstsq(A, b, rcond=None) + solved = {k: float(sol[idx[k]]) for k in ks} + delta = {k: solved[k] - centroid[k] for k in ks if k in centroid} + log["solved_mm"] = {SHORT[k]: round(v, 1) for k, v in solved.items()} + log["delta_mm"] = {SHORT[k]: round(v, 1) for k, v in delta.items()} + return solved, delta + + +def disc_troughs(labels, ct, z): + """Candidate disc planes, in millimetres, from CT attenuation along the column.""" + fgc = debris_filtered(labels) + cen, lo, hi = spine_centreline(fgc, z) + if cen is None: + return np.array([]) + rad = body_radius_profile(fgc, cen, z, lo, hi) + prof = bone_profile(ct, cen, rad, z, lo, hi) + seg = ndimage.gaussian_filter1d(prof[lo:hi + 1], sigma=max(1.8 / z[2], 0.8), + mode="nearest") + win = int(max(60.0 / z[2], 5)) + env = ndimage.gaussian_filter1d( + ndimage.maximum_filter1d(seg, size=win, mode="nearest"), + sigma=max(win / 4, 1.0), mode="nearest") + norm = seg / np.maximum(env, 1e-6) + tr, _ = signal.find_peaks(-norm, distance=max(TROUGH_MIN_SEP_MM / z[2], 2), + prominence=TROUGH_PROMINENCE) + return (tr + lo) * z[2] + + +# -------------------------------------------------------------------------------------- +# Driver +# -------------------------------------------------------------------------------------- + +def process(labels_native, affine, zooms_native, ct_native, ts_native): + log: dict = {} + lab, xf = to_ras(labels_native, affine) + z = ras_zooms(zooms_native, xf) + ct = to_ras(ct_native, affine)[0] if ct_native is not None else None + before = lab.copy() + + if ts_native is None or ct is None: + log["status"] = "no rib volume or CT supplied - no correction applied" + return labels_native, log + + ts = to_ras(ts_native, affine)[0] + fgc = debris_filtered(lab) + cen, _, _ = spine_centreline(fgc, z) + if cen is None: + log["status"] = "spine centreline unavailable" + return labels_native, log + + centroid, frac = level_stats(lab, z) + ribs = rib_attachments(ts, cen, z, log) + if len(ribs) < 4: + log["status"] = "too few usable ribs - no correction applied" + return labels_native, log + + solved, delta = solve_centroids(ribs, centroid, frac, log) + corr = sorted(k for k, v in delta.items() if abs(v) > MIN_DELTA_MM) + log["levels_corrected"] = [SHORT[k] for k in corr] + + if corr: + lo_k = max(min(corr) - 1, 1) + hi_k = min(max(corr) + 1, N_CLASSES) + seq = [k for k in range(lo_k, hi_k + 1) if k in solved] + troughs = disc_troughs(lab, ct, z) + + bounds = [] + for i in range(len(seq) - 1): + mid = (solved[seq[i]] + solved[seq[i + 1]]) / 2.0 + if len(troughs): + j = int(np.argmin(np.abs(troughs - mid))) + if abs(troughs[j] - mid) <= SNAP_MM: + mid = float(troughs[j]) + bounds.append(mid) + log["boundaries_mm"] = [round(b, 1) for b in bounds] + + edges = [-1e9] + bounds + [1e9] + slab = np.zeros(lab.shape[2], dtype=np.int16) + for i, k in enumerate(seq): + a = int(np.ceil(max(edges[i], 0) / z[2])) + bb = int(np.floor(min(edges[i + 1], (lab.shape[2] - 1) * z[2]) / z[2])) + if bb >= a: + slab[a:bb + 1] = k + + in_band = np.isin(lab, list(set(seq))) + moved = 0 + for s in range(lab.shape[2]): + if slab[s] == 0: + continue + m = in_band[:, :, s] + if m.any(): + prev = lab[:, :, s][m] + lab[:, :, s][m] = slab[s] + moved += int((prev != slab[s]).sum()) + log["voxels_moved"] = moved + else: + log["voxels_moved"] = 0 + + # cavities inside a level, claiming only voxels nothing else owns + objs = ndimage.find_objects(lab.astype(np.int32), max_label=N_CLASSES) + added = 0 + for k in range(1, N_CLASSES + 1): + sl = objs[k - 1] + if sl is None: + continue + sub = lab[sl] == k + if not sub.any(): + continue + gain = ndimage.binary_fill_holes(sub) & ~sub & (lab[sl] == 0) + if gain.any(): + blk = lab[sl] + blk[gain] = k + added += int(gain.sum()) + log["hole_fill"] = added + + fgb, fga = before > 0, lab > 0 + kept = int((fgb & fga & (before == lab)).sum()) + log["churn"] = {"kept": kept, + "relabelled": int((fgb & fga & (before != lab)).sum()), + "fraction_untouched": round(kept / max(int(fgb.sum()), 1), 4)} + return from_ras(lab, affine).astype(np.uint8), log