Add optional rib-anchored vertebral identity stage (vertebrae_identity)

config.yaml

@@ -103,4 +103,30 @@ vertebrae_engine: shapekit
 # CT lookup for shapekit_pro: first <input_case>/<ct_file_name>, then
 # <ct_root>/<case_id>/<ct_file_name> when ct_root is set.
 ct_file_name: ct.nii.gz
-# ct_root: /path/to/original/ct/cases
+# 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 <input_case>/<rib_file_name>, then
+# <rib_root>/<case_id>/<rib_file_name>, then <rib_root>/<case_id>.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

docs/config.md

@@ -75,3 +75,33 @@
 8. `ct_file_name` / `ct_root`: how `shapekit_pro` finds the CT. The engine
    first looks for `<input_case>/<ct_file_name>`; when `ct_root` is set it
    also tries `<ct_root>/<case_id>/<ct_file_name>`.
+
+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 `<input_case>/<rib_file_name>` first, then
+    `<rib_root>/<case_id>/<rib_file_name>`, then `<rib_root>/<case_id>.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.

docs/vertebrae_rib_identity.md

@@ -0,0 +1,244 @@
+<h1 align="center">Rib-Anchored Vertebral Identity</h1>
+
+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. `<input_case>/<rib_file_name>`
+2. `<rib_root>/<case_id>/<rib_file_name>`
+3. `<rib_root>/<case_id>.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.

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