From 9fc949e1c276e2e4efd8b7b8a57722c30bb2b06c Mon Sep 17 00:00:00 2001 From: farhatmasood Date: Wed, 8 Jul 2026 23:04:11 +0500 Subject: [PATCH] Add anatomic-consistency vertebrae re-identification 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. --- docs/functions.md | 34 +++- main.py | 1 + utils/vertebrae_postprocessing.py | 24 ++- utils/vertebrae_reidentification.py | 285 ++++++++++++++++++++++++++++ 4 files changed, 337 insertions(+), 7 deletions(-) create mode 100644 utils/vertebrae_reidentification.py diff --git a/docs/functions.md b/docs/functions.md index 4b88fc1..9fbfba2 100644 --- a/docs/functions.md +++ b/docs/functions.md @@ -155,4 +155,36 @@ organ_mask = check_organ_location(segmentation_dict, organ_mask, 'organ_name', -**Applicable to:** `bladder`, `femurs`, `prostate` \ No newline at end of file +**Applicable to:** `bladder`, `femurs`, `prostate` +--- + +## Vertebrae re-identification (anatomic-consistency) + +`vertebrae_reidentification.reidentify_vertebrae_dict(segmentation_dict, reference_img, logger=None, patient_id="")` + +A spacing-aware, **stable** replacement for the size-based +`reallocate_based_on_size` step. It targets the dominant vertebra error — +level **mis-identification** in the mid / thoraco-lumbar spine (the network +segments the bone but assigns the wrong level, collapsing per-level Dice in the +middle while the anchored ends stay high). + +**Method:** isolate vertebral-body cores with a physical distance transform and +grow them back over all bone (ordered instances); anchor the sequence by +offset-voting (the two anchored ends must agree — the *anatomic consistency +cycle*); keep the model's own mask on every already-correct vertebra and rebuild +only the mis-identified span (plus a small superior buffer) with a watershed of +the shifted cores; keep the largest component per level and fill holes. The +result is accepted only if every level is strictly ordered along the spine, +otherwise the model labels are kept unchanged (non-destructive). + +**Parameters:** +- `segmentation_dict` (`dict`): `{vertebra_name: np.ndarray mask}`. +- `reference_img` (`nib.Nifti1Image`): supplies the affine (voxel spacing / axes). +- `logger`, `patient_id`: optional logging. + +**Returns:** the `segmentation_dict` with corrected vertebra masks. + +**Enable:** add `vertebrae` to `target_organs` in `config.yaml`; `main.py` passes +the reference image through automatically. + +**Applicable to:** all 24 vertebra levels (`vertebrae_L5` … `vertebrae_C1`). diff --git a/main.py b/main.py index a09b67a..aec250d 100644 --- a/main.py +++ b/main.py @@ -194,6 +194,7 @@ def process_organs(segmentation_dict: dict, reference_img, combined_seg: np.arra patient_id, segmentation_dict, logger=logger, + reference_img=reference_img, # affine -> spacing/axes for re-identification ) return segmentation_dict diff --git a/utils/vertebrae_postprocessing.py b/utils/vertebrae_postprocessing.py index dd4f041..98274cd 100644 --- a/utils/vertebrae_postprocessing.py +++ b/utils/vertebrae_postprocessing.py @@ -9,10 +9,13 @@ from skimage.measure import label, regionprops from scipy import ndimage from .utils import remove_small_components, fill_holes +from .vertebrae_reidentification import reidentify_vertebrae_dict #### @jliu452 postprocessing codes for the vertabreas part -#### TODO Not finished warning +#### Anatomic-consistency re-identification added by @farhatmasood (see +#### vertebrae_reidentification.py) — a stable replacement for the size-based +#### reallocate_based_on_size step. # the general mapping @@ -426,16 +429,25 @@ def supress_non_largest_components(img, default_val = 0): return img_mod -def postprocessing_vertebrae(patiend_id:str, segmentation_dict: dict, logger): +def postprocessing_vertebrae(patiend_id:str, segmentation_dict: dict, logger, + reference_img=None): """ Post-processing for vertebrae labels. - Steps: - 1. Reallocate label IDs based on size (e.g. largest → most important label) - 2. Suppress all non-largest connected components - 3. Fill holes within vertebrae volumes + Preferred path (``reference_img`` supplied): anatomic-consistency + **re-identification** (`vertebrae_reidentification.reidentify_vertebrae_dict`) + — a spacing-aware, stable replacement for the size-based reallocation. It + keeps the model's Dice-optimal mask on every already-correct vertebra and + rebuilds only the mis-identified span, then keeps the largest component per + level and fills holes. + + Legacy path (no affine available): the earlier size/adjacency heuristics. """ + if reference_img is not None: + return reidentify_vertebrae_dict( + segmentation_dict, reference_img, logger=logger, patient_id=patiend_id) + # --------- legacy path (kept for backward compatibility) ---------- # TODO WARNING fixing .... reallocate_based_on_size() vertebrae_segmentations = np.zeros_like(next(iter(segmentation_dict.values())), dtype=np.uint8) diff --git a/utils/vertebrae_reidentification.py b/utils/vertebrae_reidentification.py new file mode 100644 index 0000000..35b1893 --- /dev/null +++ b/utils/vertebrae_reidentification.py @@ -0,0 +1,285 @@ +# -*- coding: utf-8 -*- +""" +vertebrae_reidentification.py +============================= + +Anatomic-consistency **re-identification** of vertebra labels for ShapeKit. + +This module provides a stable replacement for the (currently unfinished) +size-based `reallocate_based_on_size` step in `vertebrae_postprocessing.py`. +The dominant error of TotalSegmentator / SuPreM-style vertebra predictions is +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 while the anchored ends stay high. + +Because Dice punishes deleted true-positive bone, the method is **non- +destructive**: it re-labels rather than deletes, and — crucially — it keeps the +model's own (Dice-optimal) mask on every vertebra that is already correct, +rebuilding only the mis-identified span. + +Method (see the accompanying report for the derivation): + 1. Treat the union of all vertebra labels as spine bone; isolate vertebral + **body cores** with a physical distance transform and grow them back over + all bone by a nearest-core watershed -> ordered vertebra instances. + 2. Anchor the sequence: each confident instance votes an offset + (model_label - rank); a strong agreement means the count is reliable + (the two anchored ends must agree -- the anatomic-consistency cycle), + otherwise the case is left as the model. + 3. Copy the model mask through on every correctly-labelled vertebra; rebuild + only the shifted bone (plus a small superior buffer whose model boundaries + the compression corrupts) with a watershed of the shifted cores. + 4. Keep the largest connected component per vertebra, fill interior holes, and + accept the result only if every level is strictly ordered along the spine. + +Works in world space via the affine, so it is orientation- and spacing-aware. + +Author: Rao Farhat Masood (contributed to ShapeKit) +""" + +import numpy as np +from collections import defaultdict +from scipy import ndimage + +try: + import cc3d + _HAS_CC3D = True +except Exception: # pragma: no cover + _HAS_CC3D = False + +# Vertebra order used internally: 1 = L5 (most inferior) ... 24 = C1 (most superior) +ORDER = ["vertebrae_L5", "vertebrae_L4", "vertebrae_L3", "vertebrae_L2", + "vertebrae_L1", "vertebrae_T12", "vertebrae_T11", "vertebrae_T10", + "vertebrae_T9", "vertebrae_T8", "vertebrae_T7", "vertebrae_T6", + "vertebrae_T5", "vertebrae_T4", "vertebrae_T3", "vertebrae_T2", + "vertebrae_T1", "vertebrae_C7", "vertebrae_C6", "vertebrae_C5", + "vertebrae_C4", "vertebrae_C3", "vertebrae_C2", "vertebrae_C1"] +NAME_TO_ID = {n: i + 1 for i, n in enumerate(ORDER)} +ID_TO_NAME = {i + 1: n for i, n in enumerate(ORDER)} +NUM_CLASSES = 24 + + +# --------------------------------------------------------------------------- # +# helpers +# --------------------------------------------------------------------------- # +def _connected_components(mask, connectivity=26): + if _HAS_CC3D: + return cc3d.connected_components(mask.astype(np.uint8), + connectivity=connectivity, return_N=True) + lab, n = ndimage.label(mask, structure=ndimage.generate_binary_structure(3, 3)) + return lab, n + + +def _bbox(mask): + if not mask.any(): + return None + c = np.array(np.nonzero(mask)) + lo, hi = c.min(1), c.max(1) + 1 + return tuple(slice(int(lo[d]), int(hi[d])) for d in range(3)) + + +def _voxel_spacing(affine): + return np.sqrt((affine[:3, :3] ** 2).sum(axis=0)) + + +# --------------------------------------------------------------------------- # +# core re-identification +# --------------------------------------------------------------------------- # +def reidentify_core(label_vol, affine, r_mm=5.0, min_core_ml=0.3, + conf=0.6, agree_thr=0.8, buffer_sup=2): + """Return the re-identified 1..24 label volume, or ``None`` if unreliable.""" + spacing = _voxel_spacing(affine) + voxel_ml = float(np.prod(spacing)) / 1000.0 + rot, trans = affine[:3, :3], affine[:3, 3] + + bone = ndimage.binary_fill_holes(label_vol > 0) + if not bone.any(): + return None + bb = _bbox(bone) + lo = np.array([bb[d].start for d in range(3)]) + B, Vs = bone[bb], label_vol[bb] + + edt = ndimage.distance_transform_edt(B, sampling=spacing) + clab, nc = _connected_components(edt > r_mm) + csz = np.bincount(clab.ravel()); csz[0] = 0 + keep = [c for c in range(1, nc + 1) if csz[c] >= min_core_ml / voxel_ml] + if len(keep) < 6: + return None + + def core_z(c): + return float(rot[2] @ (np.array(np.nonzero(clab == c)).mean(1) + lo) + trans[2]) + order = sorted(keep, key=core_z) + remap = np.zeros(nc + 1, int) + for rank, c in enumerate(order, 1): + remap[c] = rank + clab = remap[clab] + K = len(order) + + idx = ndimage.distance_transform_edt(clab == 0, return_indices=True, sampling=spacing)[1] + inst = np.where(B, clab[tuple(idx)], 0) + + maj, votes = {}, [] + for i in range(1, K + 1): + labs = Vs[inst == i]; labs = labs[labs > 0] + if labs.size == 0: + maj[i] = 0; continue + vals, cnts = np.unique(labs, return_counts=True) + j = int(np.argmax(cnts)); maj[i] = int(vals[j]) + if cnts[j] / cnts.sum() >= conf: + votes.append(int(vals[j]) - i) + if not votes: + return None + vv, vc = np.unique(votes, return_counts=True) + offset = int(vv[np.argmax(vc)]) + if vc.max() / len(votes) < agree_thr: + return None + + corr = {i: i + offset for i in range(1, K + 1)} + label_insts = defaultdict(list) + for i in range(1, K + 1): + if maj[i] > 0: + label_insts[maj[i]].append(i) + + wrong = {i for i in range(1, K + 1) + if maj[i] > 0 and 1 <= corr[i] <= NUM_CLASSES + and corr[i] != maj[i] and abs(corr[i] - maj[i]) <= 1} + watershed_insts = [] + if wrong: + wrong_majs = {maj[i] for i in wrong} + span = set(wrong) | {i for i in range(1, K + 1) if maj[i] in wrong_majs} + buf, j = [], max(span) + 1 + while j <= K and len(buf) < buffer_sup: + if 1 <= corr[j] <= NUM_CLASSES and abs(corr[j] - maj[j]) <= 1: + buf.append(j); j += 1 + else: + break + watershed_insts = sorted(span | set(buf)) + + ws_majs = {maj[i] for i in watershed_insts} + correct_M = set() + for M, insts_M in label_insts.items(): + if M in ws_majs: + continue + if (len(insts_M) == 1 and corr[insts_M[0]] == M) or \ + any(abs(corr[i] - M) > 1 for i in insts_M): + correct_M.add(M) + + out = np.zeros_like(Vs) + for M in correct_M: + out[Vs == M] = M + covered = np.isin(Vs, list(correct_M)) if correct_M else np.zeros(Vs.shape, bool) + shifted = B & ~covered + if watershed_insts and shifted.any(): + seed = np.isin(clab, watershed_insts) + sidx = ndimage.distance_transform_edt(~seed, return_indices=True, sampling=spacing)[1] + near = clab[tuple(sidx)] + lut = np.zeros(K + 1, int) + for i in watershed_insts: + lut[i] = corr[i] + out[shifted] = lut[near[shifted]] + + 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 + + full = np.zeros_like(label_vol) + full[bb] = out + return full + + +def keep_largest_component(label_vol): + """One connected component per vertebra (removes leak false-positives).""" + out = label_vol.copy() + for k in range(1, NUM_CLASSES + 1): + bb = _bbox(out == k) + if bb is None: + continue + cc, n = _connected_components(out[bb] == k) + if n <= 1: + continue + sizes = np.bincount(cc.ravel()); sizes[0] = 0 + big = int(sizes.argmax()) + drop = (cc > 0) & (cc != big) + if drop.any(): + region = out[bb]; region[drop] = 0; out[bb] = region + return out + + +def fill_interior_holes(label_vol): + out = label_vol.copy() + for k in range(1, NUM_CLASSES + 1): + bb = _bbox(out == k) + if bb is None: + continue + region = out[bb] + holes = ndimage.binary_fill_holes(region == k) & (region == 0) + if holes.any(): + region[holes] = k; out[bb] = region + return out + + +def _is_monotonic(label_vol, affine): + rot, trans = affine[:3, :3], affine[:3, 3] + zs = [float(rot[2] @ np.array(np.nonzero(label_vol == k)).mean(1) + trans[2]) + for k in range(1, NUM_CLASSES + 1) if (label_vol == k).any()] + return all(zs[i] < zs[i + 1] for i in range(len(zs) - 1)) + + +# --------------------------------------------------------------------------- # +# ShapeKit entry point +# --------------------------------------------------------------------------- # +def reidentify_vertebrae_dict(segmentation_dict, reference_img, logger=None, + patient_id=""): + """ + ShapeKit-facing wrapper. + + Parameters + ---------- + segmentation_dict : dict {vertebra_name: np.ndarray(bool/uint8 mask)} + reference_img : nibabel image whose affine gives voxel spacing / axes + logger : optional logging.Logger + + Returns the segmentation_dict with corrected vertebra masks. Falls back to + the input labels (unchanged) whenever re-identification is judged unreliable + or the result is not anatomically ordered. + """ + def _log(msg): + if logger is not None: + logger.info(msg) + + present = [n for n in ORDER if n in segmentation_dict + and np.any(segmentation_dict[n])] + if len(present) < 6: + _log(f"[Vertebrae] {patient_id}: <6 levels present, skipping re-identification.") + return segmentation_dict + + shape = np.asarray(segmentation_dict[present[0]]).shape + combined = np.zeros(shape, dtype=np.uint8) + for name in ORDER: + m = segmentation_dict.get(name) + if m is not None and np.any(m): + combined[np.asarray(m) > 0] = NAME_TO_ID[name] + + affine = reference_img.affine + reid = reidentify_core(combined, affine) + if reid is None: + _log(f"[Vertebrae] {patient_id}: re-identification unreliable -> kept model labels.") + result = combined + else: + result = reid + + result = keep_largest_component(result) + result = fill_interior_holes(result) + + if reid is not None and not _is_monotonic(result, affine): + _log(f"[Vertebrae] {patient_id}: non-monotonic result -> reverted to model labels.") + result = fill_interior_holes(keep_largest_component(combined)) + elif reid is not None: + _log(f"[Vertebrae] {patient_id}: re-identification applied.") + + out = dict(segmentation_dict) + for k, name in ID_TO_NAME.items(): + mask = (result == k).astype(np.uint8) + if mask.any() or name in segmentation_dict: + out[name] = mask + return out