Add optional rib-anchored vertebral identity stage (vertebrae_identity)

[brijeshm2026-hash]

Summary

Adds an optional, default-off stage that adjudicates which vertebral level is which before the configured vertebrae_engine runs. It is not a new engine and performs no shape cleanup of its own; the selected engine still runs afterwards.

Enabled with a new config key:

vertebrae_identity: rib_anchor    # 'none' (default) leaves ShapeKit unchanged

Problem

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 foreground voxels are present and locally plausible, but a block of levels is named one position off. Per-class Dice then collapses across that block while the anchored ends stay high.

This error can evade checks that operate only on the vertebral prediction. A column where every level is named one position low is internally self-consistent: ordering is monotonic, spacing is regular, components are clean, and voxel counts are unremarkable. Resolving it needs evidence from outside the vertebral labelling task.

Why ribs

Ribs supply that evidence. The thoracic spine typically corresponds to twelve numbered rib pairs, the lowest of which is borne by T12, and where a rib meets the spine is a physical measurement of a structure that is not part of the vertebral column.

The distinguishing property of this approach is the cue it uses: it reads an external anatomical structure rather than inferring vertebral identity solely from the vertebral predictions, their ordering, spacing, connected components, or internal anatomical consistency. This is a statement about the input it consumes, not a claim about the effectiveness of any other approach.

Costovertebral anatomy, as encoded:

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 and a pair whose sides disagree by more than 15 mm is discarded rather than trusted. The remaining pairs form a weighted least-squares system with a second-difference smoothness term, a soft tie to the model's own centroids weighted by how clean each level looks, and a strongly weighted lumbar anchor. The anchor is a heavily weighted row in the least-squares system rather than an exact equality constraint, so the lumbar end is held firmly but not fixed absolutely. Only levels whose solved position differs by more than 12 mm are relabelled, at boundaries placed on CT-detected disc troughs. A scan whose geometry already satisfies the rib constraints is left untouched.

Architecture

vertebra prediction
        │
        ▼
vertebrae_identity: rib_anchor      ← optional, default 'none'
        │
        ▼
vertebrae_engine: shapekit_pro      ← unchanged, still runs
        │
        ▼
   ShapeKit output

vertebrae_identity is deliberately a separate key rather than another vertebrae_engine value. The stage does no shape cleanup, so selecting it as the engine would disable the cleanup ShapeKit already performs. Keeping the two orthogonal means the engine dispatch is untouched and the stage composes with whichever engine is configured.

Two new modules, following the existing adapter/engine split:

• utils/vertebrae_rib_identity_engine.py — the scientific core.
• utils/vertebrae_rib_identity.py — the ShapeKit adapter: input discovery, grid validation, label-space conversion, logging, QA record, and fallback handling. It contains no anatomical reasoning.

Configuration

# optional identity stage, applied BEFORE vertebrae_engine
vertebrae_identity: rib_anchor
vertebrae_engine: shapekit_pro       # required with rib_anchor, see below

# CT lookup — reuses the existing keys, no new convention
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 the same way the CT already is: <input_case>/<rib_file_name>, then <rib_root>/<case_id>/<rib_file_name>, then <rib_root>/<case_id>.nii.gz.

Required inputs

Input	Source	Requirement
Vertebra masks	the usual segmentations/ folder	at least 3 non-empty
Case CT	ct_file_name / ct_root	identical array shape and matching affine
Rib volume	rib_file_name / rib_root	TotalSegmentator total multilabel, same grid

Both extra inputs are precomputed. ShapeKit does not run TotalSegmentator, and TotalSegmentator is not added to requirements.txt. Keeping it external avoids introducing a GPU inference dependency into a CPU post-processor, and pins the rib label convention to a file the operator controls rather than to an installed version.

TotalSegmentator's vertebral labels are not used. Only the geometry of the 24 rib masks (ids 92–115) is read. The vertebral labels in that volume are never consulted and play no part in any decision.

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 not degraded, and a batch is not halted by one.

Condition	Behaviour
vertebrae_identity: none	stage does not run
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, and no resampling is attempted
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 excluded, the rest are used
No level exceeds the correction threshold	labels left unchanged
Any unexpected error	caught, logged, masks returned unchanged

The stage never raises into a Pool worker.

Compatibility requirement

vertebrae_identity: rib_anchor requires vertebrae_engine: shapekit_pro. Combining it with the default shapekit engine is rejected at startup, before any case is processed, with a non-zero exit and an explanatory message.

The reason is measurable. The default 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 restore the original labelling. On a synthetic shifted phantom where the identity stage relabels 38,400 voxels:

engine after the identity stage	voxels still differing from the input
shapekit (default)	0 — correction not retained
shapekit_pro	38,400 — correction retained

A run that appeared to succeed while discarding the correction on every case would be worse than not running the stage at all, so the combination is refused rather than warned about. No existing ShapeKit algorithm was modified to accommodate this — spine_adjacent_pairs and the rest of the default engine are untouched.

This is a configuration check and is separate from the per-case conditions above. Missing inputs remain safe no-ops that let the batch continue; only a combination of settings that would discard the correction for every case stops the run.

Testing

tests/test_vertebrae_rib_identity.py — 64 tests, all passing. Everything runs on synthetic phantoms; no imaging data is required. The phantom is a column of uniform blocks at a known pitch with rib slabs at their exact costovertebral attachments, so the expected answer is known in closed form and faults can be injected deliberately.

Coverage:

• thoracic label mapping (k_of_T) and the complete 24-level label mapping, cross-checked against both vertebrae_postprocessing.all_labels and config.yaml;
• rib attachment extraction, left/right acceptance and rejection at the tolerance boundary;
• the costovertebral equations for rib 1, ribs 2–9, and ribs 10–12;
• least-squares construction and solution, including the strongly weighted lumbar anchor;
• correction-threshold semantics driven by controlled synthetic displacements, including that a displacement of exactly the threshold is not corrected;
• every fallback path: missing CT, missing ribs, shape mismatch, affine mismatch, too few masks, too few ribs, and unexpected errors;
• configuration validation: the unsupported pairing is rejected, the supported pairing is accepted, the disabled path accepts any engine, and missing inputs are not treated as configuration errors;
• end-to-end retention of the correction through shapekit_pro.

Batch behaviour was exercised through main.py with multiprocessing.Pool over six cases covering every input combination. All six completed; the four with deliberately absent or mismatched inputs logged a reason and were left unchanged.

Regression against the reference implementation

The scientific core is vendored unchanged from the implementation the method was developed with: the module body is a byte-for-byte copy of that source, verifiable with a single diff, so the refactor into ShapeKit provably did not alter the algorithm.

Equivalence was then checked end to end. The reference implementation and the ShapeKit path were run on identical inputs, and the identity-corrected labels were compared voxel by voxel before any engine touched them:

case	total voxels	differing voxels	array hashes
BDMAP_00000006	102,498,304	0	identical
BDMAP_00000031	365,428,736	0	identical
total	467,927,040	0	—

Per-case reports match as well: the same ribs used, the same ribs excluded with the same left/right disagreements, the same levels corrected, and the same voxel counts. Comparison is on the anatomical arrays rather than the NIfTI files, since headers and gzip framing legitimately differ between writers.

Per-level displacements were also compared as raw IEEE-754 bit patterns between default and single-threaded BLAS execution and were bit-identical, so the correction set does not depend on threading configuration on the tested platform.

Deployment notes

• No new dependencies. The stage uses numpy, scipy and nibabel, all already required. requirements.txt is unmodified.
• Default off. With the shipped config, and with any config predating this change, behaviour is unchanged.
• Memory. The stage holds the CT and the rib volume in addition to the masks ShapeKit already loads — roughly 0.7 GB each on a 0.7 mm whole-spine case. Peak for the stage on such a case was measured at about 3.8 GB. Worker counts on high-resolution cohorts should be budgeted accordingly with --cpu_count.
• CPU only, single-threaded numpy/scipy. Runtime is dominated by connected-component and profile computations, not by the least-squares solve, which has fourteen unknowns.

Known limitations

• The method was developed and evaluated on the two-case BodyMaps vertebrae warm-up set, where it raised average DSC from 74.5% to 90.5%. One of those two cases required no correction. Behaviour on larger and more varied cohorts has not been measured, and no generalisation beyond the evaluated cases is claimed.
• This is not a clinically validated tool.
• The stage requires a CT and a precomputed rib volume. On mask-only inputs it no-ops on every case.
• Voxel grids must already match. No resampling or reorientation is performed; a case whose CT or rib volume sits on a different grid is skipped rather than transformed.
• Compatibility is an explicit allow-list containing only shapekit_pro. Any future engine is rejected until it has been verified to preserve corrected identities and added to the list.
• The reverting behaviour of the default engine was characterised on a synthetic phantom, not on patient data. The restriction errs on the conservative side.
• Rib label ids are pinned to the TotalSegmentator v2 total class map (92–115); a future renumbering would require the table to be updated.
• The vendored core retains a few imports that belonged to the original command-line wrapper and are unused here. They are kept deliberately so the module stays verifiable byte-for-byte against its source, and can be trimmed in a follow-up.

Files

Added

• utils/vertebrae_rib_identity_engine.py — scientific core, vendored unchanged
• utils/vertebrae_rib_identity.py — ShapeKit adapter
• tests/test_vertebrae_rib_identity.py — 64 tests
• docs/vertebrae_rib_identity.md — method, inputs, configuration, fallbacks, engine requirement

Modified

• main.py — import the stage and its validator, read four config keys, resolve the rib path beside the existing CT path, run the stage ahead of the engine branch, validate the configuration at startup. Additive only; no existing line changed.
• config.yaml — the new keys with an explanatory comment block
• docs/config.md — entries for the new keys

No other file is touched.