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