Add anatomic-consistency vertebrae re-identification (stable replacement for size-based reallocation)#5
Open
farhatmasood wants to merge 1 commit into
Conversation
There was a problem hiding this comment.
Pull request overview
This PR introduces a new spacing-/orientation-aware vertebra re-identification algorithm and wires it into the existing vertebrae postprocessing flow when a reference image affine is available, while keeping the legacy heuristics as a fallback.
Changes:
- Added
utils/vertebrae_reidentification.pyimplementing anatomic-consistency vertebra level re-identification and non-destructive relabeling. - Updated
utils/vertebrae_postprocessing.pyto prefer the new re-identification path whenreference_imgis provided. - Passed
reference_imgthrough inmain.pyand documented the new function indocs/functions.md.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| utils/vertebrae_reidentification.py | New core algorithm for spacing-aware vertebra level re-identification and post-cleanup (largest CC + hole fill + monotonicity gate). |
| utils/vertebrae_postprocessing.py | Integrates re-identification as the preferred path (affine available), retains legacy heuristics fallback. |
| main.py | Passes reference_img into vertebrae postprocessing to enable affine-aware behavior. |
| docs/functions.md | Documents the new vertebrae re-identification utility. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+181
to
+183
| for v in range(1, NUM_CLASSES + 1): # never let a level vanish | ||
| if (Vs == v).any() and not (out == v).any(): | ||
| out[Vs == v] = v |
Comment on lines
6
to
10
| from tqdm import tqdm | ||
| from scipy.ndimage import generate_binary_structure | ||
| from scipy.ndimage import label, binary_dilation, binary_erosion | ||
| from skimage.measure import label, regionprops | ||
| from scipy import ndimage |
| from .vertebrae_reidentification import reidentify_vertebrae_dict | ||
|
|
||
|
|
||
| #### @jliu452 postprocessing codes for the vertabreas part |
|
|
||
| ## Vertebrae re-identification (anatomic-consistency) | ||
|
|
||
| `vertebrae_reidentification.reidentify_vertebrae_dict(segmentation_dict, reference_img, logger=None, patient_id="")` |
Adds utils/vertebrae_reidentification.py: a spacing-aware, non-destructive re-identification of vertebra levels that provides a stable replacement for the unfinished size-based reallocate_based_on_size step. The dominant error in TotalSegmentator/SuPreM-style vertebra predictions is level MIS-IDENTIFICATION in the mid/thoraco-lumbar spine: the network segments the bone well but assigns the wrong level, so per-level Dice collapses in the middle while the anchored ends stay high. Method: - isolate vertebral-body cores via a physical distance transform and grow them back over all bone -> ordered instances; - anchor the sequence by offset-voting (the two ends must agree - the anatomic consistency cycle); fall back to the model labels if unreliable; - keep the model's own (Dice-optimal) mask on every already-correct vertebra and rebuild only the mis-identified span (plus a small superior buffer whose model boundaries the compression corrupts) with a watershed of the shifted cores; - keep the largest component per level, fill holes, and accept only if the result is strictly ordered along the spine. Wiring: postprocessing_vertebrae() now takes reference_img and uses the re-identification when the affine is available (legacy heuristics kept as fallback); main.py passes the reference image through. Docs updated.
farhatmasood
force-pushed
the
feat/vertebrae-reidentification
branch
from
1c47079 to
9fc949e
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds
utils/vertebrae_reidentification.py, a spacing-aware, non-destructive re-identification of vertebra levels. It provides a stable replacement for the currently-unfinished size-basedreallocate_based_on_sizestep (marked "TODO Not finished").Motivation
The dominant error in TotalSegmentator / SuPreM-style vertebra predictions is level mis-identification in the mid / thoraco-lumbar spine: the network segments the bone well but assigns the wrong level (e.g. it under-segments one vertebra and shifts every level above it by one). A per-level Dice therefore collapses in the middle (T11/T9 can drop to ~1%) while the anchored ends (lumbar, cervical) stay high — the signature of a one-level shift. Since Dice punishes deleted true-positive bone, the fix must be re-labelling, not deletion.
Method
model_label − rank); a strong agreement means the count is reliable (the two anchored ends must agree — the anatomic consistency cycle). If unreliable, the case is left as the model.Results
On a SuPreM SwinUNETR demo case with a severe mid-thoracic shift, iterating this design lifted the mean vertebra Dice from 73.9% → 89.5% → ~90.5%, recovering the catastrophic levels (e.g. T11 1%→88%, T9 1.5%→88%) while preserving the already-correct ones (L3, T2, T3 at 96–98%). The method matches the standalone reference implementation bit-for-bit.
Changes
utils/vertebrae_reidentification.py(self-contained; deps already inrequirements.txt: numpy, scipy, nibabel, cc3d with a scipy fallback).utils/vertebrae_postprocessing.py:postprocessing_vertebrae()gains areference_imgarg and uses the re-identification when the affine is available; the legacy heuristics are kept as a fallback.main.py: passes the reference image through.docs/functions.md: documents the new function.Notes
Enable via
vertebraeinconfig.yamltarget_organs(no other config change). Fully backward-compatible: with noreference_img, the legacy path runs unchanged.