From d8f7a3b9a3bc402fdb0316db64881011fe21a4e2 Mon Sep 17 00:00:00 2001 From: robert Date: Wed, 19 Aug 2026 14:40:30 +0200 Subject: [PATCH 01/13] git add remove internal code --- gui/snapshot_reviewer/review_gui.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/gui/snapshot_reviewer/review_gui.py b/gui/snapshot_reviewer/review_gui.py index a1cb21ff..0b5e859a 100644 --- a/gui/snapshot_reviewer/review_gui.py +++ b/gui/snapshot_reviewer/review_gui.py @@ -1588,11 +1588,6 @@ def _on_overwrite(jpg_path: Path, path_str: str, fam: dict[str, list[BIDS_FILE]] for k1, k2, coord in p2.items(): p[k1, k2] = coord p.save(f) - for t in fam["msk_seg-treg"]: - logger.on_debug(t) - if t.parent == "derivatives-final-points": - logger.on_debug("unlink", t.file["nii.gz"]) - t.file["nii.gz"].unlink(missing_ok=True) dlg = SlicerLaunchDialog( p, From aa48da3201a1ce5adf12e6f60aa51b849e50837a Mon Sep 17 00:00:00 2001 From: robert Date: Wed, 19 Aug 2026 14:42:06 +0200 Subject: [PATCH 02/13] add Rib prediction and processing --- TPTBox/segmentation/__init__.py | 1 + TPTBox/segmentation/_rib/__init__.py | 6 + TPTBox/segmentation/_rib/_rib_assign.py | 290 ++++++++++++++++++ TPTBox/segmentation/_rib/add_ribs.py | 178 +++++++++++ .../nnUnet_utils/inference_api.py | 5 + 5 files changed, 480 insertions(+) create mode 100644 TPTBox/segmentation/_rib/__init__.py create mode 100644 TPTBox/segmentation/_rib/_rib_assign.py create mode 100644 TPTBox/segmentation/_rib/add_ribs.py diff --git a/TPTBox/segmentation/__init__.py b/TPTBox/segmentation/__init__.py index 0070a83f..545ec556 100644 --- a/TPTBox/segmentation/__init__.py +++ b/TPTBox/segmentation/__init__.py @@ -1,4 +1,5 @@ from __future__ import annotations +from TPTBox.segmentation._rib.add_ribs import add_ribs_to_vert_spine from TPTBox.segmentation.spineps import _run_spineps_all, get_outpaths_spineps, run_spineps from TPTBox.segmentation.VibeSeg.vibeseg import extract_vertebra_bodies_from_VibeSeg, run_inference_on_file, run_nnunet, run_vibeseg diff --git a/TPTBox/segmentation/_rib/__init__.py b/TPTBox/segmentation/_rib/__init__.py new file mode 100644 index 00000000..607ae07a --- /dev/null +++ b/TPTBox/segmentation/_rib/__init__.py @@ -0,0 +1,6 @@ +from __future__ import annotations + +from TPTBox.segmentation._rib._rib_assign import assign_ribs_to_vert_segmentation, split_touching_rib_ccs +from TPTBox.segmentation._rib.add_ribs import add_ribs_to_vert_spine + +__all__ = ["add_ribs_to_vert_spine", "assign_ribs_to_vert_segmentation", "split_touching_rib_ccs"] diff --git a/TPTBox/segmentation/_rib/_rib_assign.py b/TPTBox/segmentation/_rib/_rib_assign.py new file mode 100644 index 00000000..faab2fc5 --- /dev/null +++ b/TPTBox/segmentation/_rib/_rib_assign.py @@ -0,0 +1,290 @@ +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass +from enum import Enum + +import numpy as np + +from TPTBox import NII, POI, Location, No_Logger, Vertebra_Instance, calc_centroids +from TPTBox.core.vert_constants import Full_Body_Instance + +logger = No_Logger(prefix="RibAssignment") + + +@dataclass(frozen=True) +class RibCandidate: + cc_label: int + volume: float + z: float + x: float + + +def _thoracic_like_labels(labels: Iterable[int]) -> list[int]: + """Keep only vertebrae that can carry ribs and sort cranial -> caudal. + + The original implementation used a fragile hard-coded range check. + Here we simply keep the contiguous block that was previously allowed + (7..20 inclusive) and sort ascending = top to bottom. + """ + return [*sorted(v for v in labels if 7 <= int(v) <= 20), 28] + + +def _split_ccs_by_side( + rib_cc: NII, + cms_cc: POI, + cms_cc2: POI, + cms_vert: POI, # , ref_vert: int +) -> tuple[list[RibCandidate], list[RibCandidate]]: + right_axis = cms_vert.get_axis("R") + inf_axis = cms_vert.get_axis("I") + # ref_x = cms_vert[ref_vert, 50][right_axis] + + left: list[RibCandidate] = [] + right: list[RibCandidate] = [] + vols = rib_cc.volumes() + + for cc in cms_cc2.keys_region(): + if cc == 0 or (cc, 50) not in cms_cc: + continue + center = cms_cc[cc, 50] + + center2 = cms_cc2[cc, 50] if (cc, 50) in cms_cc2 else center # noqa: SIM401 + cand = RibCandidate( + cc_label=cc, + volume=vols.get(cc, 0), + z=center[inf_axis], + x=center2[right_axis], + ) + distances = cms_vert.calculate_distances_cord(center2) + min_key = min(distances, key=distances.get) # type: ignore + # x < vertebra center => patient right in RAS-like orientation + if cand.x < cms_vert[min_key][right_axis]: + right.append(cand) + else: + left.append(cand) + + # sort top -> bottom (smallest inferior coordinate first) + left.sort(key=lambda c: c.z) + right.sort(key=lambda c: c.z) + return left, right + + +def _vert_z_spacing(cms_vert: POI) -> tuple[int, float, list[tuple[int, float]]]: + """Return (inferior axis index, median vertebra Z-spacing in voxels, sorted [(vert_id, z)]).""" + inf_axis = cms_vert.get_axis("I") + vz = [(int(v), float(cms_vert[v, 50][inf_axis])) for v in cms_vert.keys_region() if (v, 50) in cms_vert] + vz.sort(key=lambda t: t[1]) + if len(vz) < 2: + return inf_axis, 0.0, vz + spacings = [abs(vz[i + 1][1] - vz[i][1]) for i in range(len(vz) - 1)] + return inf_axis, float(np.median(spacings)) if spacings else 0.0, vz + + +def _cc_extent_on_axis(arr: np.ndarray, cc_label: int, axis: int) -> float: + """Return the extent of a CC on a single axis.""" + coords = np.where(arr == cc_label)[axis] + if coords.size == 0: + return 0.0 + return float(coords.max() - coords.min()) + + +def _try_erosion_split( + cc_mask_template: NII, + binary_cc: np.ndarray, + erosion_pixels: int, + min_volume: int, +) -> list[np.ndarray] | None: + """Erode a single binary CC and re-run CC labelling. Return list of binary sub-CCs (infected back to the original mask) if the erode+CC produced 2+ components, else None.""" + if erosion_pixels <= 0: + return None + cc_nii = cc_mask_template.copy().set_array_(binary_cc.astype(np.uint8), verbose=False) + try: + eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False) + except Exception: + return None + eroded_cc = eroded.get_connected_components(connectivity=3) + sub_labels = [int(x) for x in eroded_cc.unique() if x != 0] + if len(sub_labels) < 2: + return None + infected = eroded_cc.infect(cc_nii, verbose=False) + infected_arr = infected.get_seg_array() + sub_masks = [] + for sub in sub_labels: + m = infected_arr == sub + if int(m.sum()) >= min_volume: + sub_masks.append(m) + return sub_masks if len(sub_masks) >= 2 else None + + +def split_touching_rib_ccs( + rib_cc: NII, + cms_vert: POI, + max_span_factor: float = 1.4, + erosion_pixels: int = 2, + min_volume: int = 100, + max_passes: int = 3, + verbose: bool = False, +) -> NII: + """Split rib connected components that likely fuse the ribs of adjacent vertebrae. + + A CC whose extent on the inferior axis exceeds ``max_span_factor`` × the + median vertebra Z-spacing is treated as merged ribs and processed by + erosion: erode the CC, re-run CC labelling, infect the new labels back + onto the original CC voxels. Resolves ribs bridged by a thin strip of + voxels. + + Runs up to ``max_passes`` sweeps so a newly split sub-CC that is still + oversized gets another chance. + """ + inf_axis, median_spacing, _vz_sorted = _vert_z_spacing(cms_vert) + if median_spacing <= 0: + return rib_cc + threshold = median_spacing * max_span_factor + + arr = rib_cc.get_seg_array() + next_label = int(arr.max()) + 1 if arr.size and arr.max() > 0 else 1 + + for _pass in range(max_passes): + changed = False + for cc_label in sorted({int(x) for x in np.unique(arr) if x != 0}): + extent = _cc_extent_on_axis(arr, cc_label, inf_axis) + if extent <= threshold: + continue + + binary_cc = arr == cc_label + sub_masks = _try_erosion_split(rib_cc, binary_cc, erosion_pixels, min_volume) + if sub_masks is None: + continue + + arr[binary_cc] = 0 + for i, m in enumerate(sub_masks): + new_lbl = cc_label if i == 0 else next_label + if i > 0: + next_label += 1 + arr[m] = new_lbl + if verbose: + logger.print(f"split CC {cc_label} via erosion -> {new_lbl} (voxels={int(m.sum())})") + changed = True + if not changed: + break + + rib_cc.set_array_(arr, verbose=False) + return rib_cc + + +def assign_ribs_to_vert_segmentation( + vert_seg: NII, + sem_seg: NII, + rib_seg: NII, + verbose: bool = False, + min_volume: int = 100, + no_7=False, + split_touching: bool = True, + max_span_factor: float = 1.4, + erosion_pixels: int = 2, + left_id: int = Full_Body_Instance.rib_left.value, + right_id: int = Full_Body_Instance.rib_right.value, + error_value=255, + add_error=True, +) -> tuple[NII, NII]: + """Assign rib connected components deterministically from top to bottom. + + Refactor goals: + - deterministic top->bottom assignment + - no dominance / repeated while-loop logic + - robust left/right split using centroids + - preserve original output contract + + Args: + split_touching: If True, run ``split_touching_rib_ccs`` on the rib CC + map before assignment so ribs that touch front/middle/back get + separated (erosion first, Z-band fallback). + max_span_factor: Threshold (× median vertebra spacing) above which a rib + CC is treated as merged and eligible for splitting. + erosion_pixels: Voxels to erode by when attempting the erosion split. + """ + rib_seg.assert_affine(other=vert_seg, verbose=verbose) + rib_seg.assert_affine(other=sem_seg, verbose=verbose) + + ori = vert_seg.orientation + vert_seg = vert_seg.reorient() + sem_seg = sem_seg.reorient() + rib_seg = rib_seg.reorient(verbose=verbose) + vert_pred = vert_seg.extract_label([Vertebra_Instance.C7, *Vertebra_Instance.thoracic(), Vertebra_Instance.L1]) + rib_seg = rib_seg.extract_label([left_id, right_id], keep_label=True) # type: ignore + logger.on_debug(f"{rib_seg.unique()=}") + # Remove rib voxels overlapping vertebrae + rib_seg[vert_pred != 0] = 0 + + rib_cc = rib_seg.filter_connected_components(None, min_volume=min_volume, keep_label=False) + cms_vert = calc_centroids(vert_seg) + if split_touching: + rib_cc = split_touching_rib_ccs( + rib_cc, cms_vert, max_span_factor=max_span_factor, erosion_pixels=erosion_pixels, min_volume=min_volume, verbose=verbose + ) + cms_cc = calc_centroids(rib_cc * vert_pred.calc_convex_hull(None).dilate_msk_euclid(5)) + cms_cc2 = calc_centroids(rib_cc) # .dilate_msk_euclid(5) + + vert_labels = _thoracic_like_labels(vert_seg.unique()) + if not vert_labels: + logger.on_warning("No rib-bearing vertebrae found") + return vert_seg.reorient_(ori), sem_seg.reorient_(ori) + + # Use the top-most available vertebra as side reference + left_ccs, right_ccs = _split_ccs_by_side(rib_cc, cms_cc, cms_cc2, cms_vert) + + rib_vert_map = {cc: error_value for cc in rib_cc.unique() if cc != 0} + rib_subreg_map = {cc: 0 for cc in rib_cc.unique() if cc != 0} + + # Deterministic cranial -> caudal assignment + # Smarter T1 heuristic: skip the first vertebra only if there is no + # plausible rib CC close to it on either side. This avoids the common + # off-by-one while preserving rare valid T1 rib assignments or partial FOV + if no_7 and 7 in vert_labels: + vert_labels.remove(7) + assign_vert_labels = vert_labels + assert isinstance(vert_labels, (tuple, list)), type(vert_labels) + + if vert_labels and ( + (vert_labels[0] != 28 and vert_labels[1] != 28) or ((vert_labels[0], 50) in vert_labels and (vert_labels[1], 50) in vert_labels) + ): + inf_axis = cms_vert.get_axis("I") + t1_z = cms_vert[vert_labels[0], 50][inf_axis] + # use median vertebral spacing as adaptive threshold + spacing = abs(cms_vert[vert_labels[1], 50][inf_axis] - t1_z) if len(vert_labels) > 1 else 30 + threshold = max(5, spacing * 0.6) + + nearest_left = abs(left_ccs[0].z - t1_z) if left_ccs else float("inf") + nearest_right = abs(right_ccs[0].z - t1_z) if right_ccs else float("inf") + has_t1_rib = min(nearest_left, nearest_right) < threshold + + if not has_t1_rib: + assign_vert_labels = vert_labels[1:] + + for i, raw_vid in enumerate(assign_vert_labels): + vid = 20 if raw_vid == 28 else raw_vid + if i < len(left_ccs): + cc = left_ccs[i].cc_label + rib_vert_map[cc] = Vertebra_Instance(vid).RIB + rib_subreg_map[cc] = Location.Rib_Right.value + + if i < len(right_ccs): + cc = right_ccs[i].cc_label + rib_vert_map[cc] = Vertebra_Instance(vid).RIB + rib_subreg_map[cc] = Location.Rib_Left.value + + # rib_sem = rib_cc.map_labels(rib_subreg_map, verbose=False) + rib_inst = rib_cc.map_labels(rib_vert_map, verbose=False) + logger.on_debug(f"{rib_inst.unique()=}") + # Merge rib assignments back into the original vert / sem segmentations + # without disturbing existing (non-rib) labels. Skip unmatched CCs (sentinel error_value=255). + matched = (rib_inst != 0) & (rib_inst != error_value) + vert_seg[matched] = rib_inst[matched] + if add_error: + vert_seg[rib_inst == 255] = 255 + sem_seg[rib_seg != 0] = rib_seg.map_labels({left_id: Location.Rib_Left.value, right_id: Location.Rib_Right.value})[rib_seg != 0] # type: ignore + undefined = sum(v == error_value for v in rib_vert_map.values()) + logger.print(f"Unmatched rib CCs: {undefined}") + + return vert_seg.reorient_(ori), sem_seg.reorient_(ori) diff --git a/TPTBox/segmentation/_rib/add_ribs.py b/TPTBox/segmentation/_rib/add_ribs.py new file mode 100644 index 00000000..efdd7e18 --- /dev/null +++ b/TPTBox/segmentation/_rib/add_ribs.py @@ -0,0 +1,178 @@ +"""Add rib labels to an existing vertebra instance + spine subregion segmentation. + +Public entry point: :func:`add_ribs_to_vert_spine`. + +Given a vertebra instance segmentation (`vert`) and a spine subregion +segmentation (`spine`), this module optionally runs VibeSeg (dataset 12) on +the source CT to obtain a whole-body rib segmentation, then delegates label +assignment to :func:`assign_ribs_to_vert_segmentation` in ``rib_chatgpt``, +which handles left/right splitting, cranial-caudal matching, and the +fallback that separates touching ribs. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +from TPTBox import BIDS_FILE, NII, Image_Reference, Location, No_Logger, to_nii +from TPTBox.segmentation._rib._rib_assign import assign_ribs_to_vert_segmentation + +logger = No_Logger(prefix="AddRibs") + +_VIBESEG_RIB_DATASET_ID = 12 + + +def _resolve_rib_seg_path( + rib_seg: Image_Reference | None, + ct: Image_Reference | None, + rib_seg_out: str | Path | None, + dataset: str | Path | None, + derivatives_folder: str, +) -> tuple[Image_Reference | None, Path | None]: + """Determine the rib segmentation input and the output path used if we need to run VibeSeg. + + Returns ``(rib_seg_ref, out_path)``. ``rib_seg_ref`` is the existing segmentation + to load (may be a path that does not exist yet — the caller is expected to write + to it after inference). ``out_path`` is where VibeSeg should save if it runs. + """ + if rib_seg is not None: + return rib_seg, None + + if rib_seg_out is not None: + p = Path(rib_seg_out) + return (p if p.exists() else None), p + + if isinstance(ct, BIDS_FILE): + out = ct.get_changed_path( + "nii.gz", + "msk", + parent=derivatives_folder, + info={"seg": f"VIBESeg-{_VIBESEG_RIB_DATASET_ID}"}, + dataset_path=dataset, + ) + return (out if out.exists() else None), out + + raise ValueError( + "rib_seg is None and no output path could be derived. Pass one of:\n" + " * rib_seg (existing rib segmentation), or\n" + " * rib_seg_out (explicit output path), or\n" + " * ct as a BIDS_FILE plus a `dataset` root so a BIDS path can be generated." + ) + + +def _ensure_rib_seg( + rib_seg: Image_Reference | None, + ct: Image_Reference | None, + rib_seg_out: str | Path | None, + dataset: str | Path | None, + derivatives_folder: str, + override: bool, + vibeseg_kwargs: dict[str, Any] | None, +) -> NII: + """Return a ready-to-use rib segmentation NII, running VibeSeg if needed.""" + resolved, out_path = _resolve_rib_seg_path(rib_seg, ct, rib_seg_out, dataset, derivatives_folder) + + if resolved is not None and not override: + return to_nii(resolved, seg=True) + + if ct is None: + raise ValueError("rib_seg is missing (or override=True) and no `ct` was given to run VibeSeg on.") + if out_path is None: + raise ValueError("Cannot run VibeSeg without an output path. Pass `rib_seg_out` or a BIDS_FILE ct + dataset.") + + # Import lazily so importing this module never forces the segmentation stack. + from TPTBox.segmentation import run_vibeseg + + out_path.parent.mkdir(parents=True, exist_ok=True) + kwargs = dict(vibeseg_kwargs or {}) + logger.print(f"Running VibeSeg (dataset {_VIBESEG_RIB_DATASET_ID}) -> {out_path}") + return run_vibeseg(ct, out_path, dataset_id=_VIBESEG_RIB_DATASET_ID, override=override, **kwargs) + + +def _save_if_path(nii: NII, target: Image_Reference | str | Path | None) -> None: + if target is None: + return + if isinstance(target, BIDS_FILE): + nii.save(target.file["nii.gz"]) + return + if isinstance(target, (str, Path)): + nii.save(Path(target)) + + +def add_ribs_to_vert_spine( + vert: Image_Reference, + spine: Image_Reference, + rib_seg: Image_Reference | None = None, + ct: Image_Reference | None = None, + dataset: str | Path | None = None, + *, + rib_seg_out: str | Path | None = None, + derivatives_folder: str = "derivatives-rib", + override: bool = False, + save: bool = False, + vibeseg_kwargs: dict[str, Any] | None = None, + verbose: bool = False, + vert_path_out: Path | None = None, + spine_path_out: Path | None = None, + **assign_kwargs: Any, +) -> tuple[NII, NII]: + """Merge rib labels into a vertebra + spine segmentation. + + If ``rib_seg`` is not given, VibeSeg (dataset 12) is run on ``ct`` and the + result is written to a path derived from a BIDS_FILE ct or given explicitly + via ``rib_seg_out``. Otherwise a ``dataset`` root must resolve a BIDS path. + + Args: + vert: Vertebra instance segmentation (path, NII, or BIDS_FILE). + spine: Spine subregion segmentation (path, NII, or BIDS_FILE). + rib_seg: Optional rib segmentation. If ``None``, VibeSeg is executed. + ct: Source CT, required when ``rib_seg`` is missing. + dataset: BIDS dataset root, used together with a ``BIDS_FILE`` ``ct`` to + derive the VibeSeg output path. + rib_seg_out: Explicit output path for the generated rib segmentation + (overrides BIDS derivation). + derivatives_folder: Sub-folder of ``dataset`` for the VibeSeg output. + override: If True, re-run VibeSeg even when its output already exists. + save: If True, write the merged results back to the ``vert`` and + ``spine`` locations (only when they are paths or BIDS_FILEs). + vibeseg_kwargs: Extra kwargs forwarded to :func:`run_vibeseg`. + verbose: Verbose logging for the assignment step. + **assign_kwargs: Extra kwargs forwarded to + :func:`assign_ribs_to_vert_segmentation` (e.g. ``split_touching``, + ``max_span_factor``, ``erosion_pixels``). + + Returns: + ``(vert_with_ribs, spine_with_ribs)`` — instance and subregion masks. + """ + vert_nii = to_nii(vert, seg=True) + spine_nii = to_nii(spine, seg=True) + + if spine_nii.extract_label(Location.Rib_Left).sum() != 0 or spine_nii.extract_label(Location.Rib_Right).sum() != 0: + logger.on_warning("ribs already present in spine segmentation — returning input unchanged") + return vert_nii, spine_nii + + rib_nii = _ensure_rib_seg( + rib_seg=rib_seg, + ct=ct, + rib_seg_out=rib_seg_out, + dataset=dataset, + derivatives_folder=derivatives_folder, + override=override, + vibeseg_kwargs=vibeseg_kwargs, + ) + rib_nii = rib_nii.resample_from_to(vert_nii) + + vert_out, spine_out = assign_ribs_to_vert_segmentation( + vert_nii, + spine_nii, + rib_nii, + verbose=verbose, + **assign_kwargs, + ) + + if save: + _save_if_path(vert_out, vert_path_out if vert_path_out is not None else vert) + _save_if_path(spine_out, spine_path_out if spine_path_out is not None else spine) + + return vert_out, spine_out diff --git a/TPTBox/segmentation/nnUnet_utils/inference_api.py b/TPTBox/segmentation/nnUnet_utils/inference_api.py index c423e092..27455d04 100755 --- a/TPTBox/segmentation/nnUnet_utils/inference_api.py +++ b/TPTBox/segmentation/nnUnet_utils/inference_api.py @@ -169,6 +169,8 @@ def _run_inference_patches(input_nii: list[NII], nnunet, _cpu_chunks, logger=log Should only be used if there is not enough RAM on the system. """ logger.on_debug("Run: _run_inference_patches, You should only run this if you have limited RAM.") + import gc + from TPTBox.segmentation.nnUnet_utils.predictor import empty_cache empty_cache(nnunet.device) @@ -198,6 +200,9 @@ def _run_inference_patches(input_nii: list[NII], nnunet, _cpu_chunks, logger=log sl[split_axis] = slice(crop_start, crop_end) seg_chunk = seg_chunk[tuple(sl)] seg_chunks.append(seg_chunk) + del chunk_inputs + gc.collect() + empty_cache(nnunet.device) seg_arr = np.concatenate([s.get_array() for s in seg_chunks], axis=split_axis) seg_nii = input_nii[0].copy() seg_nii.seg = True From e386ce107f3869c6c290b565bd5b45638f20ff0e Mon Sep 17 00:00:00 2001 From: ga84mun Date: Fri, 21 Aug 2026 08:39:53 +0000 Subject: [PATCH 03/13] fix overflow issue if uint8 is exactly 255 --- TPTBox/core/np_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/TPTBox/core/np_utils.py b/TPTBox/core/np_utils.py index 10a2ae74..c7200f48 100755 --- a/TPTBox/core/np_utils.py +++ b/TPTBox/core/np_utils.py @@ -663,7 +663,7 @@ def np_map_labels(arr: UINTARRAY, label_map: LABEL_MAP) -> np.ndarray: if len(k) == 0: return arr - max_value = max(arr.max(), *k, *v) + 1 + max_value = int(max(arr.max(), *k, *v)) + 1 # The lookup table must be able to hold every mapping target. Building it in the input # dtype silently wraps targets outside that range (uint8: 300 -> 44, -5 -> 251). @@ -740,6 +740,8 @@ def np_bbox_binary(img: np.ndarray, px_dist: int | Sequence[int] | np.ndarray = n = img.ndim shp = img.shape + if isinstance(px_dist, float): + px_dist = ceil(px_dist) if isinstance(px_dist, int): px_dist = np.ones(n, dtype=int) * px_dist # uint8 overflows for px_dist > 255 assert len(px_dist) == n, f"dimension mismatch, got img shape {shp} and px_dist {px_dist}" From a53d46959acaa44563856324aa38269c86efeda7 Mon Sep 17 00:00:00 2001 From: ga84mun Date: Fri, 21 Aug 2026 08:40:33 +0000 Subject: [PATCH 04/13] better splitting --- TPTBox/segmentation/_rib/_rib_assign.py | 234 +++++++++++++++--------- TPTBox/segmentation/_rib/add_ribs.py | 30 ++- 2 files changed, 179 insertions(+), 85 deletions(-) diff --git a/TPTBox/segmentation/_rib/_rib_assign.py b/TPTBox/segmentation/_rib/_rib_assign.py index faab2fc5..8037f353 100644 --- a/TPTBox/segmentation/_rib/_rib_assign.py +++ b/TPTBox/segmentation/_rib/_rib_assign.py @@ -1,10 +1,14 @@ from __future__ import annotations +import os from collections.abc import Iterable +from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass from enum import Enum +from functools import partial import numpy as np +from tqdm import tqdm from TPTBox import NII, POI, Location, No_Logger, Vertebra_Instance, calc_centroids from TPTBox.core.vert_constants import Full_Body_Instance @@ -30,32 +34,19 @@ def _thoracic_like_labels(labels: Iterable[int]) -> list[int]: return [*sorted(v for v in labels if 7 <= int(v) <= 20), 28] -def _split_ccs_by_side( - rib_cc: NII, - cms_cc: POI, - cms_cc2: POI, - cms_vert: POI, # , ref_vert: int -) -> tuple[list[RibCandidate], list[RibCandidate]]: +def _split_ccs_by_side(rib_cc: NII, cms_cc: POI, cms_cc2: POI, cms_vert: POI) -> tuple[list[RibCandidate], list[RibCandidate]]: right_axis = cms_vert.get_axis("R") inf_axis = cms_vert.get_axis("I") # ref_x = cms_vert[ref_vert, 50][right_axis] - left: list[RibCandidate] = [] right: list[RibCandidate] = [] vols = rib_cc.volumes() - for cc in cms_cc2.keys_region(): if cc == 0 or (cc, 50) not in cms_cc: continue center = cms_cc[cc, 50] - center2 = cms_cc2[cc, 50] if (cc, 50) in cms_cc2 else center # noqa: SIM401 - cand = RibCandidate( - cc_label=cc, - volume=vols.get(cc, 0), - z=center[inf_axis], - x=center2[right_axis], - ) + cand = RibCandidate(cc_label=cc, volume=vols.get(cc, 0), z=center[inf_axis], x=center2[right_axis]) distances = cms_vert.calculate_distances_cord(center2) min_key = min(distances, key=distances.get) # type: ignore # x < vertebra center => patient right in RAS-like orientation @@ -70,106 +61,183 @@ def _split_ccs_by_side( return left, right -def _vert_z_spacing(cms_vert: POI) -> tuple[int, float, list[tuple[int, float]]]: - """Return (inferior axis index, median vertebra Z-spacing in voxels, sorted [(vert_id, z)]).""" - inf_axis = cms_vert.get_axis("I") - vz = [(int(v), float(cms_vert[v, 50][inf_axis])) for v in cms_vert.keys_region() if (v, 50) in cms_vert] - vz.sort(key=lambda t: t[1]) - if len(vz) < 2: - return inf_axis, 0.0, vz - spacings = [abs(vz[i + 1][1] - vz[i][1]) for i in range(len(vz) - 1)] - return inf_axis, float(np.median(spacings)) if spacings else 0.0, vz +def _touching_surface(mask_a: np.ndarray, mask_b: np.ndarray, nii: NII, axis_weights=None) -> float: + """Return weighted number of voxel faces shared by two masks. + + Touching along ``up_down_axis`` is weighted by ``up_down_weight``. + All other axes have weight 1.0. + """ + if axis_weights is None: + axis_weights = {nii.get_axis("S"): 0.01, nii.get_axis("A"): 0.1, nii.get_axis("R"): 1} + surface = 0.0 + + for axis in range(mask_a.ndim): + sl1 = [slice(None)] * mask_a.ndim + sl2 = [slice(None)] * mask_a.ndim + + sl1[axis] = slice(None, -1) + sl2[axis] = slice(1, None) + touching = ( + np.logical_and(mask_a[tuple(sl1)], mask_b[tuple(sl2)]).sum() + np.logical_and(mask_b[tuple(sl1)], mask_a[tuple(sl2)]).sum() + ) + + weight = axis_weights.get(axis, 1) + surface += weight * touching -def _cc_extent_on_axis(arr: np.ndarray, cc_label: int, axis: int) -> float: - """Return the extent of a CC on a single axis.""" - coords = np.where(arr == cc_label)[axis] - if coords.size == 0: - return 0.0 - return float(coords.max() - coords.min()) + return surface def _try_erosion_split( - cc_mask_template: NII, - binary_cc: np.ndarray, - erosion_pixels: int, - min_volume: int, -) -> list[np.ndarray] | None: - """Erode a single binary CC and re-run CC labelling. Return list of binary sub-CCs (infected back to the original mask) if the erode+CC produced 2+ components, else None.""" + cc_label: int, binary_cc: NII, erosion_pixels: int, min_volume: int, _pass=0, verbose=True +) -> tuple[int, list[np.ndarray]] | None: + """Try to split one CC by erosion. + + Returns ``(cc_label, sub_masks)`` when successful, otherwise ``None``. + ``sub_masks`` are returned in the original image shape. + """ if erosion_pixels <= 0: return None - cc_nii = cc_mask_template.copy().set_array_(binary_cc.astype(np.uint8), verbose=False) + + # Crop to the CC's bounding box (+ padding for erosion/infection). + cc_nii = binary_cc + + crop = cc_nii.compute_crop(0, 2) + cc_nii = cc_nii.apply_crop(crop) try: - eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False) - except Exception: + if _pass == 0: + eroded = cc_nii.erode_msk_euclid(n_pixel=erosion_pixels, verbose=False) + elif _pass == 1: + eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False) + elif _pass == 2: + eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False, ignore_direction="R") + elif _pass == 3: + eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False, ignore_direction="A") + else: + eroded = cc_nii.erode_msk(n_pixel=erosion_pixels, verbose=False) + except Exception as e: + if verbose: + logger.on_fail(f"Erosion failed for CC {cc_label}: {e}") return None eroded_cc = eroded.get_connected_components(connectivity=3) sub_labels = [int(x) for x in eroded_cc.unique() if x != 0] + if len(sub_labels) < 2: return None infected = eroded_cc.infect(cc_nii, verbose=False) infected_arr = infected.get_seg_array() - sub_masks = [] + full_size = cc_nii.sum() + # remerge if to small + for sub in sub_labels.copy(): + count = int((infected_arr == sub).sum()) + + should_merge = count * cc_nii.voxel_volume() < min_volume or count / full_size < 0.2 + if not should_merge: + continue + # print("merge", count, full_size, count / full_size) + remaining = [x for x in sub_labels if x != sub] + if not remaining: + break + + sub_mask = infected_arr == sub + if len(remaining) == 1: + target = remaining[0] + # print("merge", target) + else: + target = max(remaining, key=lambda candidate: _touching_surface(sub_mask, infected_arr == candidate, infected)) + # print("merge of many", target) + infected_arr[infected_arr == sub] = target + sub_labels.remove(sub) + if len(sub_labels) < 2: + return None + # Restore masks to the original shape here, so callers do not need to + # track either the crop or the original CC label separately. + full_masks = [] for sub in sub_labels: m = infected_arr == sub - if int(m.sum()) >= min_volume: - sub_masks.append(m) - return sub_masks if len(sub_masks) >= 2 else None + full = np.zeros(binary_cc.shape, dtype=bool) + full[crop] = m + full_masks.append(full) + if len(full_masks) < 2: + return None + return cc_label, full_masks def split_touching_rib_ccs( rib_cc: NII, - cms_vert: POI, - max_span_factor: float = 1.4, - erosion_pixels: int = 2, - min_volume: int = 100, - max_passes: int = 3, + erosion_pixels: int = 4, + min_volume: int = 1000, + max_passes: int = 2, + vert_ids=None, verbose: bool = False, + num_workers: int = 1, ) -> NII: """Split rib connected components that likely fuse the ribs of adjacent vertebrae. - A CC whose extent on the inferior axis exceeds ``max_span_factor`` × the - median vertebra Z-spacing is treated as merged ribs and processed by - erosion: erode the CC, re-run CC labelling, infect the new labels back - onto the original CC voxels. Resolves ribs bridged by a thin strip of - voxels. - - Runs up to ``max_passes`` sweeps so a newly split sub-CC that is still - oversized gets another chance. + ``_try_erosion_split`` can optionally be evaluated in parallel for all + connected components in a pass. Label assignment and writes to ``arr`` + remain sequential to keep labels deterministic and avoid races. """ - inf_axis, median_spacing, _vz_sorted = _vert_z_spacing(cms_vert) - if median_spacing <= 0: - return rib_cc - threshold = median_spacing * max_span_factor + # arr = rib_cc.get_seg_array() - arr = rib_cc.get_seg_array() - next_label = int(arr.max()) + 1 if arr.size and arr.max() > 0 else 1 + if vert_ids is None: + vert_ids = [] + next_label = int(rib_cc.max()) + 1 if rib_cc.shape and rib_cc.max() > 0 else 1 + u = {int(x) for x in rib_cc.unique() if x != 0} + # ribs counte twice (left/right) + expected_number_of_ccs = len([a for a in Vertebra_Instance.thoracic() if a.value in vert_ids]) * 2 + if expected_number_of_ccs == 0: + expected_number_of_ccs = 24 for _pass in range(max_passes): - changed = False - for cc_label in sorted({int(x) for x in np.unique(arr) if x != 0}): - extent = _cc_extent_on_axis(arr, cc_label, inf_axis) - if extent <= threshold: - continue + print("Separate RIBs - Pass", _pass + 1, f"{expected_number_of_ccs=}") - binary_cc = arr == cc_label - sub_masks = _try_erosion_split(rib_cc, binary_cc, erosion_pixels, min_volume) + labels = sorted(u) + + binary_ccs = [(cc_label, rib_cc.extract_label(cc_label)) for cc_label in labels] + if num_workers == 1: + results = [] + for cc_label, binary_cc in tqdm(binary_ccs, total=len(binary_ccs), desc="Separate RIBs"): + result = _try_erosion_split(cc_label, binary_cc, erosion_pixels, min_volume, _pass) + if result is not None: + results.append(result) + else: + results = [] + + with ThreadPoolExecutor(max_workers=num_workers) as executor: + futures = [ + executor.submit(_try_erosion_split, cc_label, binary_cc, erosion_pixels, min_volume, _pass) + for cc_label, binary_cc in binary_ccs + ] + + # tqdm advances immediately whenever an individual future completes. + for future in tqdm(as_completed(futures), total=len(futures), desc="Separate RIBs"): + result = future.result() + if result is not None: + results.append(result) + + # Apply mutations sequentially so label assignment is deterministic. + for cc_label, sub_masks in results: if sub_masks is None: continue - arr[binary_cc] = 0 + binary_cc = rib_cc.extract_label(cc_label) + u.discard(cc_label) + rib_cc[binary_cc] = 0 + for i, m in enumerate(sub_masks): new_lbl = cc_label if i == 0 else next_label if i > 0: next_label += 1 - arr[m] = new_lbl + + rib_cc[m] = new_lbl + u.add(new_lbl) + if verbose: logger.print(f"split CC {cc_label} via erosion -> {new_lbl} (voxels={int(m.sum())})") - changed = True - if not changed: + if rib_cc.max() >= expected_number_of_ccs: break - rib_cc.set_array_(arr, verbose=False) return rib_cc @@ -181,7 +249,6 @@ def assign_ribs_to_vert_segmentation( min_volume: int = 100, no_7=False, split_touching: bool = True, - max_span_factor: float = 1.4, erosion_pixels: int = 2, left_id: int = Full_Body_Instance.rib_left.value, right_id: int = Full_Body_Instance.rib_right.value, @@ -208,6 +275,8 @@ def assign_ribs_to_vert_segmentation( rib_seg.assert_affine(other=sem_seg, verbose=verbose) ori = vert_seg.orientation + vert_ids = vert_seg.unique() + vert_seg = vert_seg.reorient() sem_seg = sem_seg.reorient() rib_seg = rib_seg.reorient(verbose=verbose) @@ -216,13 +285,10 @@ def assign_ribs_to_vert_segmentation( logger.on_debug(f"{rib_seg.unique()=}") # Remove rib voxels overlapping vertebrae rib_seg[vert_pred != 0] = 0 - rib_cc = rib_seg.filter_connected_components(None, min_volume=min_volume, keep_label=False) cms_vert = calc_centroids(vert_seg) if split_touching: - rib_cc = split_touching_rib_ccs( - rib_cc, cms_vert, max_span_factor=max_span_factor, erosion_pixels=erosion_pixels, min_volume=min_volume, verbose=verbose - ) + rib_cc = split_touching_rib_ccs(rib_cc, erosion_pixels=erosion_pixels, min_volume=min_volume, vert_ids=vert_ids, verbose=verbose) cms_cc = calc_centroids(rib_cc * vert_pred.calc_convex_hull(None).dilate_msk_euclid(5)) cms_cc2 = calc_centroids(rib_cc) # .dilate_msk_euclid(5) @@ -281,10 +347,14 @@ def assign_ribs_to_vert_segmentation( # without disturbing existing (non-rib) labels. Skip unmatched CCs (sentinel error_value=255). matched = (rib_inst != 0) & (rib_inst != error_value) vert_seg[matched] = rib_inst[matched] - if add_error: - vert_seg[rib_inst == 255] = 255 + sem_seg[rib_seg != 0] = rib_seg.map_labels({left_id: Location.Rib_Left.value, right_id: Location.Rib_Right.value})[rib_seg != 0] # type: ignore undefined = sum(v == error_value for v in rib_vert_map.values()) logger.print(f"Unmatched rib CCs: {undefined}") - + if add_error: # or split_touching: + vert_seg[rib_inst == error_value] = error_value + if np.any(vert_seg.get_seg_array() == error_value): + vert_seg2 = vert_seg.remove_labels(error_value).infect(vert_seg.extract_label(error_value), verbose=False) + vert_seg[vert_seg != vert_seg2] = vert_seg2[vert_seg != vert_seg2] + vert_seg[np.logical_and(rib_inst == error_value, vert_seg == 0)] = error_value return vert_seg.reorient_(ori), sem_seg.reorient_(ori) diff --git a/TPTBox/segmentation/_rib/add_ribs.py b/TPTBox/segmentation/_rib/add_ribs.py index efdd7e18..f8c0f197 100644 --- a/TPTBox/segmentation/_rib/add_ribs.py +++ b/TPTBox/segmentation/_rib/add_ribs.py @@ -17,6 +17,7 @@ from TPTBox import BIDS_FILE, NII, Image_Reference, Location, No_Logger, to_nii from TPTBox.segmentation._rib._rib_assign import assign_ribs_to_vert_segmentation +from TPTBox.segmentation.spineps import run_spineps logger = No_Logger(prefix="AddRibs") @@ -45,12 +46,22 @@ def _resolve_rib_seg_path( if isinstance(ct, BIDS_FILE): out = ct.get_changed_path( + "nii.gz", + "msk", + parent=derivatives_folder, + info={"seg": f"VIBESeg-{_VIBESEG_RIB_DATASET_ID}", "mod": ct.bids_format}, + dataset_path=dataset, + ) + # TODO remove + out_ = ct.get_changed_path( "nii.gz", "msk", parent=derivatives_folder, info={"seg": f"VIBESeg-{_VIBESEG_RIB_DATASET_ID}"}, dataset_path=dataset, ) + if out_.exists(): + out_.rename(out) return (out if out.exists() else None), out raise ValueError( @@ -101,10 +112,10 @@ def _save_if_path(nii: NII, target: Image_Reference | str | Path | None) -> None def add_ribs_to_vert_spine( - vert: Image_Reference, - spine: Image_Reference, + vert: Image_Reference | None, + spine: Image_Reference | None, rib_seg: Image_Reference | None = None, - ct: Image_Reference | None = None, + ct: str | Path | BIDS_FILE | None = None, dataset: str | Path | None = None, *, rib_seg_out: str | Path | None = None, @@ -145,6 +156,19 @@ def add_ribs_to_vert_spine( Returns: ``(vert_with_ribs, spine_with_ribs)`` — instance and subregion masks. """ + if vert is None or spine is None: + assert ct is not None, "Provide vert/spine or a ct." + out = run_spineps( + ct, + dataset, + "ct", + "ct_instance", + "ct_labeling", + derivatives_folder, + ignore_compatibility_issues=True, + ) + vert = out["out_vert"] + spine = out["out_spine"] vert_nii = to_nii(vert, seg=True) spine_nii = to_nii(spine, seg=True) From 179ee327525d0a73e7d4bfe7be80364e5c943d3b Mon Sep 17 00:00:00 2001 From: robert Date: Fri, 21 Aug 2026 11:22:34 +0200 Subject: [PATCH 05/13] enable running with low RAM pc by splitting the nnunet prediction. (for the 0.8 full ct model we need 250 GB of RAM for a full body image.) --- .../segmentation/VibeSeg/inference_nnunet.py | 42 +++++++-- .../nnUnet_utils/inference_api.py | 88 +++++++++++++++++-- 2 files changed, 118 insertions(+), 12 deletions(-) diff --git a/TPTBox/segmentation/VibeSeg/inference_nnunet.py b/TPTBox/segmentation/VibeSeg/inference_nnunet.py index edf2732b..920ec20e 100644 --- a/TPTBox/segmentation/VibeSeg/inference_nnunet.py +++ b/TPTBox/segmentation/VibeSeg/inference_nnunet.py @@ -191,7 +191,16 @@ def run_inference_on_file( logger.on_fail(f"{shape=} has only {min(shape)} slice in a dimension.") return None, None - from TPTBox.segmentation.nnUnet_utils.inference_api import _run_inference_patches, load_inf_model, run_inference + from math import ceil + + from TPTBox.segmentation.nnUnet_utils.inference_api import ( + _get_total_ram_mb, + _run_inference_patches, + compute_cpu_chunks_for_ram, + estimate_peak_ram_mb, + load_inf_model, + run_inference, + ) if isinstance(idx, int): if auto_download: @@ -326,12 +335,33 @@ def run_inference_on_file( p = (padd, padd) input_nii = [i.apply_pad([p, p, p], mode="reflect") for i in input_nii] if _cpu_chunks is None or _cpu_chunks <= 1: - try: - seg_nii, _, softmax_logits = run_inference(input_nii, nnunet, logits=logits, logger=logger) - except MemoryError: - logger.print_error() - seg_nii = _run_inference_patches(input_nii, nnunet, None, logger=logger) + num_classes = int(nnunet.label_manager.num_segmentation_heads) + total_ram_mb = _get_total_ram_mb() + target_ram_mb = total_ram_mb * 0.5 + est_full_mb = estimate_peak_ram_mb(input_nii[0].shape, num_classes, len(input_nii)) + if est_full_mb > target_ram_mb: + shape = input_nii[0].shape + split_axis = int(np.argmax(shape)) + patch_size = nnunet.configuration_manager.patch_size + overlap = ceil(patch_size[split_axis] * (1 - nnunet.tile_step_size)) + auto_chunks = compute_cpu_chunks_for_ram(shape, split_axis, num_classes, len(input_nii), overlap, target_ram_mb) + logger.print( + f"Estimated peak RAM ~{est_full_mb:.0f} MB exceeds 50% of RAM ({target_ram_mb:.0f} MB of {total_ram_mb:.0f} MB); " + f"switching to _cpu_chunks={auto_chunks}.", + Log_Type.WARNING, + ) + seg_nii = _run_inference_patches(input_nii, nnunet, auto_chunks, logger=logger) softmax_logits = None + else: + logger.print( + f"Estimated peak RAM ~{est_full_mb:.0f} MB fits within 50% of RAM ({target_ram_mb:.0f} MB of {total_ram_mb:.0f} MB)." + ) if verbose else None + try: + seg_nii, _, softmax_logits = run_inference(input_nii, nnunet, logits=logits, logger=logger) + except MemoryError: + logger.print_error() + seg_nii = _run_inference_patches(input_nii, nnunet, None, logger=logger) + softmax_logits = None else: seg_nii = _run_inference_patches(input_nii, nnunet, _cpu_chunks, logger=logger) softmax_logits = None diff --git a/TPTBox/segmentation/nnUnet_utils/inference_api.py b/TPTBox/segmentation/nnUnet_utils/inference_api.py index 27455d04..e79664e7 100755 --- a/TPTBox/segmentation/nnUnet_utils/inference_api.py +++ b/TPTBox/segmentation/nnUnet_utils/inference_api.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Sequence from math import ceil from pathlib import Path @@ -16,6 +17,69 @@ _interop = False +def _get_total_ram_mb() -> float: + """Return total system RAM in MB, or a conservative fallback if it cannot be determined.""" + try: + import psutil + + return psutil.virtual_memory().total / (1024 * 1024) + except ImportError: + pass + try: + import os + + return (os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")) / (1024 * 1024) + except (AttributeError, ValueError, OSError): + return 8000.0 + + +def estimate_peak_ram_mb(shape: Sequence[int], num_classes: int, num_channels: int = 1) -> float: + """Estimate peak CPU RAM usage during a single-pass nnU-Net inference in MB. + + Accounts for the dominant allocations held roughly simultaneously by + :func:`run_inference` and the underlying predictor: + + * ``predicted_logits`` : ``num_classes * prod(shape) * 2 B`` (float16) + * ``n_predictions`` : ``prod(shape) * 2 B`` (float16) + * Model input tensor + a transposed / stacked copy : + ``2 * num_channels * prod(shape) * 2 B`` (float16) + * Output segmentation : ``prod(shape) * 1 B`` (uint8) + + A 1.3x headroom multiplier is applied to cover temporary transpose + buffers, per-tile predictions, gaussian weights, and torch allocator + overhead. + """ + n_voxels = int(np.prod(shape)) + bytes_per_voxel = (num_classes * 2) + 2 + (num_channels * 2 * 2) + 1 + return n_voxels * bytes_per_voxel * 1.3 / (1024 * 1024) + + +def compute_cpu_chunks_for_ram( + shape: Sequence[int], + split_axis: int, + num_classes: int, + num_channels: int, + overlap: int, + target_mb: float, +) -> int: + """Return the smallest ``n_chunks`` such that a single chunk's peak RAM stays under ``target_mb``. + + Returns ``1`` when the whole volume already fits. + """ + length = int(shape[split_axis]) + if length <= 1 or target_mb <= 0: + return 1 + for n_chunks in range(1, length + 1): + chunk_length = length // n_chunks + if chunk_length == 0: + return length + chunk_shape = list(shape) + chunk_shape[split_axis] = chunk_length + 2 * overlap + if estimate_peak_ram_mb(chunk_shape, num_classes, num_channels) <= target_mb: + return n_chunks + return length + + # Adapted from https://github.com/MIC-DKFZ/nnUNet # Isensee, F., Jaeger, P. F., Kohl, S. A., Petersen, J., & Maier-Hein, K. H. (2021). nnU-Net: a self-configuring # method for deep learning-based biomedical image segmentation. Nature methods, 18(2), 203-211. @@ -163,10 +227,13 @@ def _split_ranges(length: int, n_chunks: int, overlap: int): return ranges -def _run_inference_patches(input_nii: list[NII], nnunet, _cpu_chunks, logger=logger): +def _run_inference_patches(input_nii: list[NII], nnunet, _cpu_chunks, ram_fraction: float = 0.5, logger=logger): """Split image into k _cpu_chunks along the largest dimension. - Should only be used if there is not enough RAM on the system. + Should only be used if there is not enough RAM on the system. If + ``_cpu_chunks`` is ``None`` the chunk count is chosen so a single chunk's + estimated peak RAM stays under ``ram_fraction`` (default 50%) of total + system RAM. """ logger.on_debug("Run: _run_inference_patches, You should only run this if you have limited RAM.") import gc @@ -176,12 +243,21 @@ def _run_inference_patches(input_nii: list[NII], nnunet, _cpu_chunks, logger=log empty_cache(nnunet.device) shape = input_nii[0].shape split_axis = int(np.argmax(shape)) - - if _cpu_chunks is None: - _cpu_chunks = shape[split_axis] // 250 patch_size = nnunet.configuration_manager.patch_size overlap = ceil(patch_size[split_axis] * (1 - nnunet.tile_step_size)) - logger.print(f"{overlap=}") + + if _cpu_chunks is None: + num_classes = int(nnunet.label_manager.num_segmentation_heads) + total_mb = _get_total_ram_mb() + target_mb = total_mb * ram_fraction + _cpu_chunks = compute_cpu_chunks_for_ram(shape, split_axis, num_classes, len(input_nii), overlap, target_mb) + est_full_mb = estimate_peak_ram_mb(shape, num_classes, len(input_nii)) + logger.print( + f"Auto _cpu_chunks={_cpu_chunks} " + f"(full-run peak ~{est_full_mb:.0f} MB, target {target_mb:.0f} MB = {ram_fraction * 100:.0f}% of {total_mb:.0f} MB RAM)" + ) + _cpu_chunks = max(2, int(_cpu_chunks)) + logger.print(f"{overlap=}", f"chunks={_cpu_chunks}") ranges = _split_ranges( shape[split_axis], _cpu_chunks, From 9ab0a99c71e9486b1d4fc17ff6d1f4d592e5cc27 Mon Sep 17 00:00:00 2001 From: ga84mun Date: Fri, 21 Aug 2026 09:40:10 +0000 Subject: [PATCH 06/13] add option to not shortcut --- TPTBox/segmentation/_rib/_rib_assign.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/TPTBox/segmentation/_rib/_rib_assign.py b/TPTBox/segmentation/_rib/_rib_assign.py index 8037f353..3b4f5319 100644 --- a/TPTBox/segmentation/_rib/_rib_assign.py +++ b/TPTBox/segmentation/_rib/_rib_assign.py @@ -171,6 +171,7 @@ def split_touching_rib_ccs( vert_ids=None, verbose: bool = False, num_workers: int = 1, + short_cut=True, ) -> NII: """Split rib connected components that likely fuse the ribs of adjacent vertebrae. @@ -235,7 +236,7 @@ def split_touching_rib_ccs( if verbose: logger.print(f"split CC {cc_label} via erosion -> {new_lbl} (voxels={int(m.sum())})") - if rib_cc.max() >= expected_number_of_ccs: + if short_cut and rib_cc.max() >= expected_number_of_ccs: break return rib_cc @@ -254,6 +255,7 @@ def assign_ribs_to_vert_segmentation( right_id: int = Full_Body_Instance.rib_right.value, error_value=255, add_error=True, + short_cut=True, ) -> tuple[NII, NII]: """Assign rib connected components deterministically from top to bottom. @@ -288,7 +290,9 @@ def assign_ribs_to_vert_segmentation( rib_cc = rib_seg.filter_connected_components(None, min_volume=min_volume, keep_label=False) cms_vert = calc_centroids(vert_seg) if split_touching: - rib_cc = split_touching_rib_ccs(rib_cc, erosion_pixels=erosion_pixels, min_volume=min_volume, vert_ids=vert_ids, verbose=verbose) + rib_cc = split_touching_rib_ccs( + rib_cc, erosion_pixels=erosion_pixels, min_volume=min_volume, vert_ids=vert_ids, short_cut=short_cut, verbose=verbose + ) cms_cc = calc_centroids(rib_cc * vert_pred.calc_convex_hull(None).dilate_msk_euclid(5)) cms_cc2 = calc_centroids(rib_cc) # .dilate_msk_euclid(5) From 73462381c5dd2e196a3f90139c2483091a48d041 Mon Sep 17 00:00:00 2001 From: robert Date: Fri, 21 Aug 2026 12:07:57 +0200 Subject: [PATCH 07/13] move folder and update documentaiton --- TPTBox/segmentation/README.md | 31 +++++++- TPTBox/segmentation/__init__.py | 2 +- TPTBox/segmentation/_rib/__init__.py | 6 -- TPTBox/segmentation/rib/__init__.py | 5 ++ .../segmentation/{_rib => rib}/_rib_assign.py | 76 ++++++++++++++++--- TPTBox/segmentation/{_rib => rib}/add_ribs.py | 37 ++++++--- docs/api/segmentation.md | 11 +++ gui/snapshot_reviewer/review_gui.py | 2 +- 8 files changed, 140 insertions(+), 30 deletions(-) delete mode 100644 TPTBox/segmentation/_rib/__init__.py create mode 100644 TPTBox/segmentation/rib/__init__.py rename TPTBox/segmentation/{_rib => rib}/_rib_assign.py (79%) rename TPTBox/segmentation/{_rib => rib}/add_ribs.py (82%) diff --git a/TPTBox/segmentation/README.md b/TPTBox/segmentation/README.md index 0d7daea5..9edad320 100644 --- a/TPTBox/segmentation/README.md +++ b/TPTBox/segmentation/README.md @@ -13,6 +13,7 @@ from TPTBox.segmentation import ( run_nnunet, run_inference_on_file, extract_vertebra_bodies_from_VibeSeg, + add_ribs_to_vert_spine, ) ``` @@ -25,14 +26,16 @@ from TPTBox.segmentation import ( | `run_totalvibeseg(img_nii, ...)` | `VibeSeg/vibeseg.py` | Run TotalVibeSeg — extended label set | | `run_nnunet(img_nii, model_dir, ...)` | `VibeSeg/vibeseg.py` | Generic nnU-Net inference on a single NIfTI | | `run_inference_on_file(path, ...)` | `nnUnet_utils/inference_api.py` | Low-level nnU-Net inference on a file path | +| `add_ribs_to_vert_spine(vert, spine, ...)` | `rib/add_ribs.py` | Merge left/right rib labels into an existing vertebra + spine segmentation; optionally runs VibeSeg (dataset 12) on the source CT to obtain the raw rib mask | ## Dependencies | Pipeline | Requirement | |---|---| | SPINEPS | `pip install spineps` + model weights | -| VibeSeg / TotalVibeSeg | `pip install nnunetv2` + model weights (auto-downloaded on first run) | +| VibeSeg | `pip install nnunetv2` + model weights (auto-downloaded on first run) | | Generic nnU-Net | `pip install nnunetv2` + custom model directory | +| Rib assignment (`add_ribs_to_vert_spine`) | calls into VibeSeg/SPINEPS if the segmentation is missing. | All external tools are imported lazily — the core TPTBox package installs and imports cleanly without them. @@ -87,3 +90,29 @@ def main() -> None: if __name__ == "__main__": main() ``` + +## Adding ribs to an existing spine segmentation + +```python +from TPTBox import to_nii +from TPTBox.segmentation import add_ribs_to_vert_spine + +# Case 1: raw rib mask already exists +vert_out, spine_out = add_ribs_to_vert_spine( + vert="sub-01_seg-vert.nii.gz", + spine="sub-01_seg-spine.nii.gz", + rib_seg="sub-01_seg-VIBESeg-12.nii.gz", + save=True, # writes back to vert / spine paths +) + +# Case 2: no rib mask — VibeSeg dataset 12 is run on the CT +vert_out, spine_out = add_ribs_to_vert_spine( + vert="sub-01_seg-vert.nii.gz", + spine="sub-01_seg-spine.nii.gz", + ct="sub-01_ct.nii.gz", + rib_seg_out="sub-01_seg-VIBESeg-12.nii.gz", +) +``` + +Pass `split_touching=True` (default) so ribs of adjacent vertebrae that touch +front / middle / back get separated via erosion before assignment. diff --git a/TPTBox/segmentation/__init__.py b/TPTBox/segmentation/__init__.py index 545ec556..0dd197d8 100644 --- a/TPTBox/segmentation/__init__.py +++ b/TPTBox/segmentation/__init__.py @@ -1,5 +1,5 @@ from __future__ import annotations -from TPTBox.segmentation._rib.add_ribs import add_ribs_to_vert_spine +from TPTBox.segmentation.rib.add_ribs import add_ribs_to_vert_spine from TPTBox.segmentation.spineps import _run_spineps_all, get_outpaths_spineps, run_spineps from TPTBox.segmentation.VibeSeg.vibeseg import extract_vertebra_bodies_from_VibeSeg, run_inference_on_file, run_nnunet, run_vibeseg diff --git a/TPTBox/segmentation/_rib/__init__.py b/TPTBox/segmentation/_rib/__init__.py deleted file mode 100644 index 607ae07a..00000000 --- a/TPTBox/segmentation/_rib/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -from __future__ import annotations - -from TPTBox.segmentation._rib._rib_assign import assign_ribs_to_vert_segmentation, split_touching_rib_ccs -from TPTBox.segmentation._rib.add_ribs import add_ribs_to_vert_spine - -__all__ = ["add_ribs_to_vert_spine", "assign_ribs_to_vert_segmentation", "split_touching_rib_ccs"] diff --git a/TPTBox/segmentation/rib/__init__.py b/TPTBox/segmentation/rib/__init__.py new file mode 100644 index 00000000..0a9c4cbe --- /dev/null +++ b/TPTBox/segmentation/rib/__init__.py @@ -0,0 +1,5 @@ +from __future__ import annotations + +from TPTBox.segmentation.rib.add_ribs import add_ribs_to_vert_spine + +__all__ = ["add_ribs_to_vert_spine"] diff --git a/TPTBox/segmentation/_rib/_rib_assign.py b/TPTBox/segmentation/rib/_rib_assign.py similarity index 79% rename from TPTBox/segmentation/_rib/_rib_assign.py rename to TPTBox/segmentation/rib/_rib_assign.py index 3b4f5319..3908eb03 100644 --- a/TPTBox/segmentation/_rib/_rib_assign.py +++ b/TPTBox/segmentation/rib/_rib_assign.py @@ -1,3 +1,16 @@ +"""Rib-to-vertebra assignment on already-loaded segmentations (internal module). + +The functions here are implementation details of +:func:`TPTBox.segmentation.rib.add_ribs.add_ribs_to_vert_spine`, which is the +supported public entry point. They are not re-exported from the package. + +* ``assign_ribs_to_vert_segmentation`` — deterministic cranial→caudal mapping + of rib connected components to vertebrae T1..T12/L1, merging the result + back into the passed ``vert``/``spine`` NIIs. +* ``split_touching_rib_ccs`` — erosion-based pre-processing pass that + separates rib CCs which fuse the ribs of neighbouring vertebrae. +""" + from __future__ import annotations import os @@ -176,8 +189,28 @@ def split_touching_rib_ccs( """Split rib connected components that likely fuse the ribs of adjacent vertebrae. ``_try_erosion_split`` can optionally be evaluated in parallel for all - connected components in a pass. Label assignment and writes to ``arr`` + connected components in a pass. Label assignment and writes to ``rib_cc`` remain sequential to keep labels deterministic and avoid races. + + Args: + rib_cc: Connected-component-labelled rib mask, modified in place. + erosion_pixels: Voxels to erode per pass when attempting to split a CC. + min_volume: Minimum voxel count for a split sub-component to survive + the merge-back heuristic (values below are re-merged into the + largest touching neighbour). + max_passes: Number of erosion passes tried before giving up. Each pass + uses a slightly different erosion strategy (Euclidean, standard, + direction-restricted). + vert_ids: Optional list of vertebra labels present in the case, used to + estimate the expected number of rib CCs (thoracic vertebrae × 2). + verbose: Emit per-split log messages. + num_workers: Threads used to evaluate ``_try_erosion_split`` in + parallel. ``1`` runs sequentially. + short_cut: If True (default), stop as soon as ``rib_cc.max()`` reaches + the expected CC count. Set to False to always run every pass. + + Returns: + The (in-place-mutated) ``rib_cc`` NII. """ # arr = rib_cc.get_seg_array() @@ -259,19 +292,42 @@ def assign_ribs_to_vert_segmentation( ) -> tuple[NII, NII]: """Assign rib connected components deterministically from top to bottom. - Refactor goals: - - deterministic top->bottom assignment + Design: + + - deterministic top→bottom assignment - no dominance / repeated while-loop logic - - robust left/right split using centroids - - preserve original output contract + - robust left/right split using vertebra + rib centroids + - merges the results back into the passed ``vert_seg`` / ``sem_seg`` so + non-rib labels are preserved Args: - split_touching: If True, run ``split_touching_rib_ccs`` on the rib CC - map before assignment so ribs that touch front/middle/back get - separated (erosion first, Z-band fallback). - max_span_factor: Threshold (× median vertebra spacing) above which a rib - CC is treated as merged and eligible for splitting. + vert_seg: Vertebra instance segmentation. Modified in place: matched + rib CCs are labelled with ``Vertebra_Instance(vid).RIB`` and + unmatched CCs get ``error_value`` (when ``add_error=True``). + sem_seg: Spine subregion segmentation. Modified in place: rib voxels + get ``Location.Rib_Left`` / ``Location.Rib_Right``. + rib_seg: Raw rib segmentation containing ``left_id`` / ``right_id``. + verbose: Verbose logging. + min_volume: Minimum voxel volume for a connected component to be kept + (both for the initial CC filter and inside the touching-split). + no_7: If True, skip vertebra 7 (C7) even if it appears rib-bearing. + split_touching: If True, run :func:`split_touching_rib_ccs` on the rib + CC map before assignment so ribs that touch front/middle/back get + separated via erosion. erosion_pixels: Voxels to erode by when attempting the erosion split. + left_id: Label value in ``rib_seg`` that marks the left rib mask. + right_id: Label value in ``rib_seg`` that marks the right rib mask. + error_value: Sentinel label written into ``vert_seg`` for rib CCs that + could not be matched to a vertebra (default 255). + add_error: If True, propagate ``error_value`` into ``vert_seg`` and + grow neighbouring labels into the unmatched region; if False, drop + the unmatched CCs silently. + short_cut: Forwarded to :func:`split_touching_rib_ccs` — stop the + erosion loop early once the expected number of CCs is reached. + + Returns: + ``(vert_seg, sem_seg)`` — the mutated inputs, re-oriented back to the + original orientation of ``vert_seg``. """ rib_seg.assert_affine(other=vert_seg, verbose=verbose) rib_seg.assert_affine(other=sem_seg, verbose=verbose) diff --git a/TPTBox/segmentation/_rib/add_ribs.py b/TPTBox/segmentation/rib/add_ribs.py similarity index 82% rename from TPTBox/segmentation/_rib/add_ribs.py rename to TPTBox/segmentation/rib/add_ribs.py index f8c0f197..a4805f25 100644 --- a/TPTBox/segmentation/_rib/add_ribs.py +++ b/TPTBox/segmentation/rib/add_ribs.py @@ -5,9 +5,10 @@ Given a vertebra instance segmentation (`vert`) and a spine subregion segmentation (`spine`), this module optionally runs VibeSeg (dataset 12) on the source CT to obtain a whole-body rib segmentation, then delegates label -assignment to :func:`assign_ribs_to_vert_segmentation` in ``rib_chatgpt``, -which handles left/right splitting, cranial-caudal matching, and the -fallback that separates touching ribs. +assignment to the internal ``assign_ribs_to_vert_segmentation`` in +:mod:`TPTBox.segmentation.rib._rib_assign`, which handles left/right +splitting, cranial-caudal matching, and the erosion-based fallback that +separates touching ribs. """ from __future__ import annotations @@ -16,7 +17,7 @@ from typing import Any from TPTBox import BIDS_FILE, NII, Image_Reference, Location, No_Logger, to_nii -from TPTBox.segmentation._rib._rib_assign import assign_ribs_to_vert_segmentation +from TPTBox.segmentation.rib._rib_assign import assign_ribs_to_vert_segmentation from TPTBox.segmentation.spineps import run_spineps logger = No_Logger(prefix="AddRibs") @@ -130,15 +131,23 @@ def add_ribs_to_vert_spine( ) -> tuple[NII, NII]: """Merge rib labels into a vertebra + spine segmentation. - If ``rib_seg`` is not given, VibeSeg (dataset 12) is run on ``ct`` and the - result is written to a path derived from a BIDS_FILE ct or given explicitly - via ``rib_seg_out``. Otherwise a ``dataset`` root must resolve a BIDS path. + If ``vert`` and ``spine`` are both ``None``, SPINEPS is run on ``ct`` first + to produce them. If ``rib_seg`` is not given, VibeSeg (dataset 12) is run on + ``ct`` and the result is written to a path derived from a BIDS_FILE ct or + given explicitly via ``rib_seg_out``. Otherwise a ``dataset`` root must + resolve a BIDS path. + + If the input ``spine`` already contains ``Rib_Left``/``Rib_Right`` voxels, + the inputs are returned unchanged. Args: - vert: Vertebra instance segmentation (path, NII, or BIDS_FILE). - spine: Spine subregion segmentation (path, NII, or BIDS_FILE). + vert: Vertebra instance segmentation (path, NII, or BIDS_FILE). May be + ``None`` together with ``spine`` to trigger a SPINEPS run. + spine: Spine subregion segmentation (path, NII, or BIDS_FILE). May be + ``None`` together with ``vert``. rib_seg: Optional rib segmentation. If ``None``, VibeSeg is executed. - ct: Source CT, required when ``rib_seg`` is missing. + ct: Source CT, required when ``rib_seg`` is missing or when + ``vert``/``spine`` need to be generated. dataset: BIDS dataset root, used together with a ``BIDS_FILE`` ``ct`` to derive the VibeSeg output path. rib_seg_out: Explicit output path for the generated rib segmentation @@ -149,9 +158,15 @@ def add_ribs_to_vert_spine( ``spine`` locations (only when they are paths or BIDS_FILEs). vibeseg_kwargs: Extra kwargs forwarded to :func:`run_vibeseg`. verbose: Verbose logging for the assignment step. + vert_path_out: Optional explicit output path for the merged + vertebra segmentation. Only used when ``save=True``; falls back to + the ``vert`` location when omitted. + spine_path_out: Optional explicit output path for the merged spine + segmentation. Only used when ``save=True``; falls back to the + ``spine`` location when omitted. **assign_kwargs: Extra kwargs forwarded to :func:`assign_ribs_to_vert_segmentation` (e.g. ``split_touching``, - ``max_span_factor``, ``erosion_pixels``). + ``erosion_pixels``, ``min_volume``, ``no_7``, ``short_cut``). Returns: ``(vert_with_ribs, spine_with_ribs)`` — instance and subregion masks. diff --git a/docs/api/segmentation.md b/docs/api/segmentation.md index de6e0815..cd34a527 100644 --- a/docs/api/segmentation.md +++ b/docs/api/segmentation.md @@ -23,3 +23,14 @@ VibeSeg / nnU-Net (general deep learning inference). options: show_source: true filters: ["!^_"] + +## Rib assignment + +Adds left/right rib labels to an existing vertebra + spine segmentation. +When no raw rib mask is supplied, VibeSeg dataset 12 is invoked on the +source CT. + +::: TPTBox.segmentation.rib.add_ribs + options: + show_source: true + filters: ["!^_"] diff --git a/gui/snapshot_reviewer/review_gui.py b/gui/snapshot_reviewer/review_gui.py index 0b5e859a..df7a686a 100644 --- a/gui/snapshot_reviewer/review_gui.py +++ b/gui/snapshot_reviewer/review_gui.py @@ -1568,7 +1568,7 @@ def _open_in_slicer(self): search.append(self._parent) self._slicer_bgi = BIDS_Global_info(self.dataset_path, search) - def _on_overwrite(jpg_path: Path, path_str: str, fam: dict[str, list[BIDS_FILE]], viewed: list[BIDS_FILE]): + def _on_overwrite(jpg_path: Path, path_str: str, _fam: dict[str, list[BIDS_FILE]], viewed: list[BIDS_FILE]): """Default back-hook: log and update status bar.""" path = Path(path_str) fname = path.name From 51aaf152c45553ba0d250aa4791e6477b83a94a2 Mon Sep 17 00:00:00 2001 From: robert Date: Fri, 21 Aug 2026 12:12:28 +0200 Subject: [PATCH 08/13] remove temporary code --- TPTBox/segmentation/rib/add_ribs.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/TPTBox/segmentation/rib/add_ribs.py b/TPTBox/segmentation/rib/add_ribs.py index a4805f25..350d956f 100644 --- a/TPTBox/segmentation/rib/add_ribs.py +++ b/TPTBox/segmentation/rib/add_ribs.py @@ -53,16 +53,6 @@ def _resolve_rib_seg_path( info={"seg": f"VIBESeg-{_VIBESEG_RIB_DATASET_ID}", "mod": ct.bids_format}, dataset_path=dataset, ) - # TODO remove - out_ = ct.get_changed_path( - "nii.gz", - "msk", - parent=derivatives_folder, - info={"seg": f"VIBESeg-{_VIBESEG_RIB_DATASET_ID}"}, - dataset_path=dataset, - ) - if out_.exists(): - out_.rename(out) return (out if out.exists() else None), out raise ValueError( From edc0b4424ff365f00f6ed1693d29becbbb95cefc Mon Sep 17 00:00:00 2001 From: robert Date: Fri, 21 Aug 2026 23:20:11 +0200 Subject: [PATCH 09/13] re add missed ccs --- TPTBox/segmentation/rib/_rib_assign.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/TPTBox/segmentation/rib/_rib_assign.py b/TPTBox/segmentation/rib/_rib_assign.py index 3908eb03..426880d3 100644 --- a/TPTBox/segmentation/rib/_rib_assign.py +++ b/TPTBox/segmentation/rib/_rib_assign.py @@ -413,8 +413,12 @@ def assign_ribs_to_vert_segmentation( logger.print(f"Unmatched rib CCs: {undefined}") if add_error: # or split_touching: vert_seg[rib_inst == error_value] = error_value + if np.any(vert_seg.get_seg_array() == error_value): vert_seg2 = vert_seg.remove_labels(error_value).infect(vert_seg.extract_label(error_value), verbose=False) vert_seg[vert_seg != vert_seg2] = vert_seg2[vert_seg != vert_seg2] vert_seg[np.logical_and(rib_inst == error_value, vert_seg == 0)] = error_value + vert_seg[np.logical_and(vert_seg == 0, sem_seg.extract_label([Location.Rib_Left.value, Location.Rib_Right.value] == 1))] = ( + error_value + ) return vert_seg.reorient_(ori), sem_seg.reorient_(ori) From 72cc32415e66124bafa53e93a9565823ad905dd7 Mon Sep 17 00:00:00 2001 From: robert Date: Mon, 24 Aug 2026 13:35:01 +0200 Subject: [PATCH 10/13] reimplement fetching zoom from nnUNet config. --- .../segmentation/VibeSeg/inference_nnunet.py | 21 +++++++------------ 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/TPTBox/segmentation/VibeSeg/inference_nnunet.py b/TPTBox/segmentation/VibeSeg/inference_nnunet.py index 920ec20e..d86d23fd 100644 --- a/TPTBox/segmentation/VibeSeg/inference_nnunet.py +++ b/TPTBox/segmentation/VibeSeg/inference_nnunet.py @@ -291,20 +291,15 @@ def run_inference_on_file( zoom = ds_info.get("resolution_range", zoom) if zoom is None: + # nnUNet stores spacing in the transposed (internal) axis order used during + # training: internal_spacing[i] = original_spacing[transpose_forward[i]]. + # Invert with transpose_backward to recover spacing in the training-time + # numpy axis order, then reverse: run_inference() will reverse it again via + # zoom[::-1] before handing it to nnUNet, so the double reversal restores the + # exact plans order and prevents nnUNet from triggering a second resample. zoom_ = plans_info["configurations"]["3d_fullres"]["spacing"] - if all(zoom[0] == z for z in zoom_): - zoom = zoom_ - # order = plans_info["transpose_backward"] - ## order2 = plans_info["transpose_forward"] - # zoom = [zoom[order[0]], zoom[order[1]], zoom[order[2]]][::-1] - # orientation_ref = ("P", "I", "R") - # orientation_ref = [ - # orientation_ref[order[0]], - # orientation_ref[order[1]], - # orientation_ref[order[2]], - # ] # [::-1] - - # zoom_old = zoom_old[::-1] + transpose_backward = plans_info["transpose_backward"] + zoom = [zoom_[transpose_backward[i]] for i in range(len(zoom_))][::-1] zoom = [float(z) for z in zoom] except Exception: From 5c05f807e0beebbc5babbe1500f9a7e69ac2aaf4 Mon Sep 17 00:00:00 2001 From: ga84mun Date: Mon, 24 Aug 2026 13:23:38 +0000 Subject: [PATCH 11/13] fix dataset process --- TPTBox/core/internal/train_nnUnet/_prep_ds.py | 8 ++++---- .../core/internal/train_nnUnet/prepere_dataset.py | 15 +++------------ 2 files changed, 7 insertions(+), 16 deletions(-) diff --git a/TPTBox/core/internal/train_nnUnet/_prep_ds.py b/TPTBox/core/internal/train_nnUnet/_prep_ds.py index 637a6927..2328da62 100644 --- a/TPTBox/core/internal/train_nnUnet/_prep_ds.py +++ b/TPTBox/core/internal/train_nnUnet/_prep_ds.py @@ -181,20 +181,20 @@ def _add_file_async( return # except Exception: # seg.unlink(missing_ok=True) + if delete_brocken: if isinstance(img, Path): img = [img] for i in img: try: - to_nii(i, True).max() + to_nii(i).max() except Exception: - [Path(i).unlink(missing_ok=True) for i in img] - Path(seg).unlink(missing_ok=True) + Path(i).unlink(missing_ok=True) return try: to_nii(seg, True).max() except Exception: - [Path(i).unlink(missing_ok=True) for i in img] + # [Path(i).unlink(missing_ok=True) for i in img] Path(seg).unlink(missing_ok=True) return # load image diff --git a/TPTBox/core/internal/train_nnUnet/prepere_dataset.py b/TPTBox/core/internal/train_nnUnet/prepere_dataset.py index 19a86e14..c4dddf00 100644 --- a/TPTBox/core/internal/train_nnUnet/prepere_dataset.py +++ b/TPTBox/core/internal/train_nnUnet/prepere_dataset.py @@ -149,13 +149,13 @@ def _build_label_mapping( left_id = left.value if isinstance(left, Enum) else left right_id = right.value if isinstance(right, Enum) else right - if left_id not in mapping_forward: + if left_id not in mapping_forward and left_id not in labels_mapping.values(): raise ValueError(f"Mirror label {left_id} not present in raw_label_ids") - if right_id not in mapping_forward: + if right_id not in mapping_forward and right_id not in labels_mapping.values(): raise ValueError(f"Mirror label {right_id} not present in raw_label_ids") - mirror_out.append((mapping_forward[left_id], mapping_forward[right_id])) + mirror_out.append((mapping_forward.get(left_id, left_id), mapping_forward.get(right_id, right_id))) return (labels_mapping, mapping_forward, labels_mapping_return, mirror_out) @@ -208,18 +208,9 @@ def build_dataset(cfg: DatasetConfig) -> None: expected_labels = set(labels_mapping.values()) expected_labels.remove(0) - missing_mapping = labels_found - expected_labels - unused_mapping = expected_labels - labels_found - logger.on_text(f"Sample segmentation: {seg}") logger.on_text(f"Labels found : {sorted(labels_found)}") - if missing_mapping: - logger.on_fail(f"Labels present in segmentation but missing in mapping: {sorted(missing_mapping)}") - - if unused_mapping: - logger.on_ok(f"Unmapped labels {sorted(unused_mapping)}") - # Test remapping out = seg_nii.map_labels(mapping_forward) remapped_labels = sorted(out.unique()) From 3b9f709d58d835161b448b3dba256d073fd26021 Mon Sep 17 00:00:00 2001 From: robert Date: Mon, 24 Aug 2026 16:42:59 +0200 Subject: [PATCH 12/13] tune GPU memory patching --- .../segmentation/VibeSeg/inference_nnunet.py | 30 + TPTBox/segmentation/VibeSeg/vibeseg.py | 3 +- .../nnUnet_utils/estimate_nnunet_memory.py | 641 ++++++++++++++++++ TPTBox/segmentation/nnUnet_utils/predictor.py | 5 + 4 files changed, 678 insertions(+), 1 deletion(-) create mode 100644 TPTBox/segmentation/nnUnet_utils/estimate_nnunet_memory.py diff --git a/TPTBox/segmentation/VibeSeg/inference_nnunet.py b/TPTBox/segmentation/VibeSeg/inference_nnunet.py index d86d23fd..d7a94ef5 100644 --- a/TPTBox/segmentation/VibeSeg/inference_nnunet.py +++ b/TPTBox/segmentation/VibeSeg/inference_nnunet.py @@ -23,6 +23,24 @@ _model_cache: dict = {} +def _suggest_memory_estimation_script(idx, model_path: Path, reason: str, logger=logger) -> None: + """Point the user at ``estimate_nnunet_memory.py`` to fit ``memory_base``/``memory_factor``. + + Invoked when the model's ``dataset.json`` does not carry memory parameters + (fallback defaults are used) and when inference dies with a GPU OOM. + """ + script = Path(__file__).parent.parent / "nnUnet_utils/estimate_nnunet_memory.py" + dataset_arg = f"--dataset-id {idx} " if isinstance(idx, int) else "" + logger.on_warning( + f"{reason}\n" + f"To measure and set 'memory_base'/'memory_factor' for this model, run:\n" + f" python {script} --gpu 0 {dataset_arg}--model-path {model_path}\n" + "If have GPU memory issues, run this code above; The inference can than split correctly to still fit in GPU memory. " + f"It probes several input shapes and patches the model's dataset.json in place." + "The GPU should not be occupied when running this code, that interferes with the measurements. " + ) + + def get_ds_info(idx: int, _model_path: str | Path | None = None, exit_one_fail: bool = True, logger=logger) -> dict: """Load and return the ``dataset.json`` for the model with the given dataset index. @@ -239,10 +257,22 @@ def run_inference_on_file( if "labels" in ds_info2: ds_info["labels_mapping"] = ds_info2["labels"] + missing_mem_keys: list[str] = [] if memory_base is None: + if "memory_base" not in ds_info: + missing_mem_keys.append("memory_base") memory_base = float(ds_info.get("memory_base", 5000)) if memory_factor is None: + if "memory_factor" not in ds_info: + missing_mem_keys.append("memory_factor") memory_factor = float(ds_info.get("memory_factor", 160)) + if missing_mem_keys: + _suggest_memory_estimation_script( + idx, + model_path, + f"Memory parameter(s) {missing_mem_keys} not set in the model's dataset.json; falling back to defaults. {memory_base=}, {memory_factor=}", + logger=logger, + ) use_folds_arg = tuple(folds) if len(folds) != 5 else None # Include every setting that changes the loaded predictor so a cache hit is always equivalent diff --git a/TPTBox/segmentation/VibeSeg/vibeseg.py b/TPTBox/segmentation/VibeSeg/vibeseg.py index 4268ed74..6bec72a5 100644 --- a/TPTBox/segmentation/VibeSeg/vibeseg.py +++ b/TPTBox/segmentation/VibeSeg/vibeseg.py @@ -86,6 +86,7 @@ defaults = { 100: {"memory_base": 5500, "memory_factor": 25}, + 12: {"memory_base": 7000, "memory_factor": 200}, } @@ -121,7 +122,7 @@ def run_vibeseg( Returns: Segmentation ``NII`` saved at *out_seg*. """ - if dataset_id in defaults: + if dataset_id in defaults and model_path is None: for k, v in defaults[dataset_id].items(): if k not in args: args[k] = v diff --git a/TPTBox/segmentation/nnUnet_utils/estimate_nnunet_memory.py b/TPTBox/segmentation/nnUnet_utils/estimate_nnunet_memory.py new file mode 100644 index 00000000..e34fe498 --- /dev/null +++ b/TPTBox/segmentation/nnUnet_utils/estimate_nnunet_memory.py @@ -0,0 +1,641 @@ +"""estimate_vibeseg_memory.py +========================== +Measures peak GPU RAM consumed by run_vibeseg across a grid of synthetic +input shapes, then fits the three memory parameters used by TPTBox's +check_mem guard: + + check_mem passes when: + (n_voxels / 1e6 * memory_factor) + memory_base + < clamp(0.80 * gpu_total, lo=memory_base, hi=memory_max) + + So the fitted curve must be a CONSERVATIVE UPPER BOUND of actual usage, + not a mean — otherwise ~50 % of runs would be incorrectly skipped. + This script fits via quantile regression (default q=0.95) so the curve + sits above nearly all observations while remaining tight. + +Usage +----- + python estimate_vibeseg_memory.py [--gpu 0] [--model-path /path/to/nnUNet] + [--dataset-id 12] [--out-dir /tmp/vibeseg_probe] + [--shapes-csv shapes.csv] + +The script prints recommended values for memory_base, memory_factor, and +memory_max at the end, saves a CSV + PNG summary plot, and patches the +dataset.json in the nnUNet model folder. +""" + +import argparse +import csv +import gc +import json +import shutil +import subprocess +import sys +import time +from pathlib import Path + +import numpy as np + +from TPTBox import NII +from TPTBox.segmentation import run_vibeseg +from TPTBox.segmentation.VibeSeg.auto_download import download_weights + +try: + from tqdm import tqdm +except ImportError: + + def tqdm(it, **kwargs): + return it + + +# --------------------------------------------------------------------------- +# GPU helpers +# --------------------------------------------------------------------------- + + +def _nvidia_smi_query(field: str, gpu: int) -> float: + out = ( + subprocess.check_output( + ["nvidia-smi", f"--query-gpu={field}", "--format=csv,noheader,nounits", f"--id={gpu}"], + text=True, + ) + .strip() + .splitlines() + ) + return float(out[0]) + + +def gpu_used_mb(gpu: int) -> float: + return _nvidia_smi_query("memory.used", gpu) + + +def gpu_total_mb(gpu: int) -> float: + return _nvidia_smi_query("memory.total", gpu) + + +def wait_for_gpu_idle(gpu: int, stable_seconds: float = 2.0, poll_interval: float = 0.5) -> float: + """Poll until GPU memory is stable; return baseline usage in MB.""" + prev = gpu_used_mb(gpu) + stable_since = time.time() + while True: + time.sleep(poll_interval) + cur = gpu_used_mb(gpu) + if abs(cur - prev) < 50: + if time.time() - stable_since >= stable_seconds: + return cur + else: + stable_since = time.time() + prev = cur + + +class PeakPoller: + """Context manager: polls nvidia-smi in a thread, records peak usage.""" + + def __init__(self, gpu: int, poll_interval: float = 0.1): + self.gpu = gpu + self.poll_interval = poll_interval + self._peak = 0.0 + self._stop = False + self._thread = None + + def __enter__(self): + import threading + + self._stop = False + self._peak = gpu_used_mb(self.gpu) + + def _poll(): + while not self._stop: + try: + v = gpu_used_mb(self.gpu) + self._peak = max(self._peak, v) + except Exception: + pass + time.sleep(self.poll_interval) + + self._thread = threading.Thread(target=_poll, daemon=True) + self._thread.start() + return self + + def __exit__(self, *_): + self._stop = True + self._thread.join(timeout=2) + + @property + def peak_mb(self) -> float: + return self._peak + + +# --------------------------------------------------------------------------- +# Core probe +# --------------------------------------------------------------------------- + + +def derive_model_zoom(nnunet_path: Path, dataset_id: int) -> tuple[float, float, float] | None: + """Return the (X, Y, Z) target spacing the model rescales inputs to. + + Mirrors the resolution logic in ``inference_nnunet.run_inference_on_file``: + ``dataset.json['spacing']`` (reversed unless dataset 527) → ``resolution_range`` + → ``plans.json['configurations']['3d_fullres']['spacing']`` de-transposed. + Returns None if none of these are present so the caller can fall back. + """ + with open(nnunet_path / "dataset.json") as f: + ds_info = json.load(f) + plans_path = nnunet_path / "plans.json" + plans_info = json.loads(plans_path.read_text()) if plans_path.exists() else None + + zoom = ds_info.get("spacing") + if dataset_id not in [527] and zoom is not None: + zoom = zoom[::-1] + zoom = ds_info.get("resolution_range", zoom) + if zoom is None and plans_info is not None: + try: + zoom_ = plans_info["configurations"]["3d_fullres"]["spacing"] + transpose_backward = plans_info["transpose_backward"] + zoom = [zoom_[transpose_backward[i]] for i in range(len(zoom_))][::-1] + except (KeyError, IndexError): + zoom = None + if zoom is None: + return None + zoom = [float(z) for z in zoom] + assert len(zoom) == 3, zoom + return (zoom[0], zoom[1], zoom[2]) + + +def probe_shape( + shape: tuple, + gpu: int, + dataset_id: int, + model_path: str, + out_dir: Path, + voxel_size: float | tuple[float, float, float] = 0.8, +) -> dict: + """Run run_vibeseg on a synthetic volume, measure peak GPU RAM. + memory_max is set to 999 GB so the check_mem guard never fires here. + + ``voxel_size`` accepts either an isotropic scalar or a 3-tuple. Passing the + model's own zoom keeps run_vibeseg's internal rescale a no-op, so the shape + that reaches nnUNet matches the shape we probed. + """ + zoom_vec = (voxel_size, voxel_size, voxel_size) if isinstance(voxel_size, (int, float)) else tuple(voxel_size) + affine = np.diag([float(zoom_vec[0]), float(zoom_vec[1]), float(zoom_vec[2]), 1.0]) + nii = NII.from_numpy(np.random.rand(*shape).astype(np.float32), affine=affine) + out_path = str(out_dir / f"probe_{'x'.join(map(str, shape))}.nii.gz") + + baseline_mb = wait_for_gpu_idle(gpu) + + result = { + "shape": shape, + "n_voxels": int(np.prod(shape)), + "peak_mb": None, + "net_mb": None, + "elapsed_s": None, + "ok": False, + } + + try: + t0 = time.time() + with PeakPoller(gpu) as poller: + run_vibeseg( + nii, + out_path, + gpu=gpu, + dataset_id=dataset_id, + model_path=model_path, + memory_base=0, + memory_factor=0, + memory_max=999_000, # disable guard during probing + override=True, + fail_on_missing_memory=True, + ) + result["peak_mb"] = poller.peak_mb + result["net_mb"] = max(0.0, poller.peak_mb - baseline_mb) + result["elapsed_s"] = time.time() - t0 + result["ok"] = True + except Exception as exc: + print(f" [WARN] shape {shape} failed: {exc}", file=sys.stderr) + result["error"] = str(exc) + + gc.collect() + try: + import torch + + torch.cuda.empty_cache() + except ImportError: + pass + + return result + + +# --------------------------------------------------------------------------- +# Fitting – quantile regression (upper envelope, not mean) +# --------------------------------------------------------------------------- + + +def fit_parameters(records: list, quantile: float = 0.95) -> dict: + """Fit an upper-envelope line: + net_mb ≈ memory_base + n_voxels / 1e6 * memory_factor + + Uses quantile regression at `quantile` (default 0.95) so the predicted + curve lies above ~95 % of observations. This is intentional: the + check_mem guard must never *wrongly skip* a run that would have fit in + VRAM, so over-estimating slightly is safer than under-estimating. + + Falls back to OLS + residual-std shift if scipy is unavailable. + """ + ok = [r for r in records if r["ok"] and r["net_mb"] is not None] + if len(ok) < 2: + raise ValueError("Need at least 2 successful probes to fit parameters.") + + x = np.array([r["n_voxels"] / 1e6 for r in ok]) # M voxels + y = np.array([r["net_mb"] for r in ok]) + + try: + from scipy.optimize import linprog # quantile regression via LP + + # min q * sum(u) + (1-q) * sum(v) + # s.t. y - (a + b*x) = u - v, u,v >= 0 + # variables: [a, b, u_0..u_n, v_0..v_n] + n = len(x) + # Objective + c = np.zeros(2 + 2 * n) + c[2 : 2 + n] = quantile + c[2 + n :] = 1.0 - quantile + + # Equality: a + b*x_i + u_i - v_i = y_i + A_eq = np.zeros((n, 2 + 2 * n)) + A_eq[:, 0] = 1.0 + A_eq[:, 1] = x + A_eq[np.arange(n), 2 + np.arange(n)] = 1.0 + A_eq[np.arange(n), 2 + n + np.arange(n)] = -1.0 + b_eq = y + + bounds = [(None, None), (None, None)] + [(0, None)] * (2 * n) + res = linprog(c, A_eq=A_eq, b_eq=b_eq, bounds=bounds, method="highs") + memory_base = float(res.x[0]) + memory_factor = float(res.x[1]) + method_used = f"quantile regression (q={quantile})" + + except Exception: + # Fallback: OLS + push intercept up by (1-quantile) sigma + X = np.column_stack([np.ones(len(ok)), x]) + coeffs, *_ = np.linalg.lstsq(X, y, rcond=None) + memory_base = float(coeffs[0]) + memory_factor = float(coeffs[1]) + resid_std = float(np.std(y - X @ coeffs)) + from scipy.stats import norm + + memory_base += norm.ppf(quantile) * resid_std + method_used = f"OLS + {quantile:.0%}-quantile shift (scipy.optimize unavailable)" + + # OLS fit — used both for R² diagnostics and as a slope floor. + X2 = np.column_stack([np.ones(len(ok)), x]) + ols_coeffs, *_ = np.linalg.lstsq(X2, y, rcond=None) + ols_base, ols_slope = float(ols_coeffs[0]), float(ols_coeffs[1]) + y_pred_ols = X2 @ ols_coeffs + ss_res = np.sum((y - y_pred_ols) ** 2) + ss_tot = np.sum((y - y.mean()) ** 2) + r2 = 1.0 - ss_res / ss_tot if ss_tot > 0 else float("nan") + + # Guard against an outlier flattening the quantile fit: take the LARGER of + # the quantile-regression slope and the OLS slope. If OLS wins, re-fit the + # intercept at the same quantile so the line still sits above ~q of points. + if ols_slope > memory_factor: + memory_factor = ols_slope + memory_base = float(np.quantile(y - memory_factor * x, quantile)) + method_used += " + OLS slope floor (quantile slope was flattened)" + + # memory_max: largest observed net usage + headroom (caller applies factor) + memory_max = float(max(y)) + + return { + "memory_base": max(0.0, memory_base), + "memory_factor": max(0.0, memory_factor), + "memory_max": memory_max, + "r2": r2, + "method": method_used, + "n_points": len(ok), + "ols_base": ols_base, + "ols_slope": ols_slope, + } + + +# --------------------------------------------------------------------------- +# Default shape grid +# --------------------------------------------------------------------------- + +DEFAULT_SHAPES = [ + (96, 96, 96), + (160, 160, 160), + (192, 192, 192), + (256, 256, 128), + (256, 256, 256), + (320, 320, 160), + (320, 320, 320), + (400, 400, 200), + (400, 400, 400), + (512, 512, 200), + (512, 512, 400), + (512, 512, 512), +] + + +def load_shapes_csv(path: str) -> list: + shapes = [] + with open(path) as f: + for row in csv.DictReader(f): + shapes.append((int(row["d"]), int(row["h"]), int(row["w"]))) + return shapes + + +# --------------------------------------------------------------------------- +# Plotting +# --------------------------------------------------------------------------- + + +def make_plot(records: list, fit: dict, out_path: Path): + try: + import matplotlib.pyplot as plt + except ImportError: + print("[INFO] matplotlib not available – skipping plot.") + return + + ok = [r for r in records if r["ok"] and r["net_mb"] is not None] + if not ok: + return + + n_vox = np.array([r["n_voxels"] for r in ok]) + net_mb = np.array([r["net_mb"] for r in ok]) + x_M = n_vox / 1e6 # M voxels for plotting + + x_line = np.linspace(0, x_M.max(), 300) + y_line = fit["memory_base"] + x_line * fit["memory_factor"] # x already in M + + # OLS mean line for comparison + X = np.column_stack([np.ones(len(ok)), x_M]) + ols, *_ = np.linalg.lstsq(X, net_mb, rcond=None) + y_ols = ols[0] + x_line * ols[1] + + fig, axes = plt.subplots(1, 2, figsize=(13, 5)) + + ax = axes[0] + ax.scatter(x_M, net_mb, zorder=3, label="measured peak", color="steelblue") + ax.plot(x_line, y_ols, "k:", linewidth=1.2, label="OLS mean") + ax.plot(x_line, y_line, "r--", linewidth=1.8, label=f"upper envelope ({fit['method'].split('(')[1].rstrip(')')} )") + ax.set_xlabel("Volume size (M voxels)") + ax.set_ylabel("Net GPU RAM (MB)") + ax.set_title("GPU RAM vs volume size") + ax.legend(fontsize=8) + ax.grid(True, alpha=0.3) + + ax = axes[1] + y_pred = fit["memory_base"] + x_M * fit["memory_factor"] + resid = net_mb - y_pred + colors = ["tomato" if r > 0 else "steelblue" for r in resid] + ax.bar(range(len(ok)), resid, color=colors) + ax.axhline(0, color="k", linewidth=0.8) + ax.set_xlabel("Probe index (sorted by size)") + ax.set_ylabel("Residual (MB) [+ve = under-predicted]") + ax.set_title("Fit residuals (negative = safe headroom)") + ax.grid(True, alpha=0.3, axis="y") + + summary = ( + f"memory_base={fit['memory_base']:.0f} MB " + f"memory_factor={fit['memory_factor']:.2f} " + f"memory_max={fit['memory_max']:.0f} MB " + f"OLS R²={fit['r2']:.4f} n={fit['n_points']}" + ) + fig.suptitle("run_nnunet GPU memory parameter estimation", fontsize=13) + fig.text(0.5, -0.01, summary, ha="center", fontsize=9, bbox={"boxstyle": "round", "facecolor": "wheat", "alpha": 0.6}) + + plt.tight_layout() + plt.savefig(out_path, dpi=150, bbox_inches="tight") + print(f"[INFO] Plot saved → {out_path}") + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def parse_args(): + p = argparse.ArgumentParser(description="Estimate run_nnunet memory parameters from GPU measurements.") + p.add_argument("--gpu", type=int, default=0) + p.add_argument( + "--model-path", + default=None, + help="Base directory containing 'nnUNet_results'. If omitted, the TPTBox default weights path is used and missing models are auto-downloaded.", + ) + p.add_argument("--dataset-id", type=int, default=12) + p.add_argument("--out-dir", default="/tmp/run_nnunet_probe") + p.add_argument("--shapes-csv", default=None, help="CSV with columns d,h,w for custom shape grid") + p.add_argument( + "--voxel-size", + type=float, + default=None, + help="Isotropic voxel size for synthetic inputs. If omitted, the model's own " + "target spacing is derived from dataset.json/plans.json so run_vibeseg's " + "internal rescale is a no-op and the probed shape matches the shape reaching nnUNet.", + ) + p.add_argument("--results-csv", default="memory_probes.csv") + p.add_argument("--plot", default="memory_fit.png") + p.add_argument("--headroom", type=float, default=1.20, help="Multiplier applied to max observed net MB → memory_max (default 1.20)") + p.add_argument("--quantile", type=float, default=0.95, help="Quantile for upper-envelope fit (default 0.95)") + return p.parse_args() + + +def main(): + args = parse_args() + + # Resolve nnUNet model path (auto-download when using the default TPTBox path) + if args.model_path is None: + weights_dir = download_weights(args.dataset_id) + model_path = weights_dir.parent # /nnUNet_results + # When no --model-path is passed, run_vibeseg below must find the same weights, + # so we hand it the same base directory instead of the outdated CLI default. + args.model_path = str(model_path) + else: + model_path = Path(args.model_path) + if model_path.name != "nnUNet_results": + model_path = model_path / "nnUNet_results" + model_path.mkdir(parents=True, exist_ok=True) + # If the dataset isn't present under the supplied base, download into it. + if not any(model_path.glob(f"*{args.dataset_id:03}*")): + print(f"[INFO] Dataset {args.dataset_id:03} not found under {model_path}; downloading…") + download_weights(args.dataset_id, model_path=model_path) + assert model_path.exists(), model_path + + _key_ResEnc = "__nnUNet*ResEnc" + + def _resolve_nnunet_path(): + try: + return next(next(iter(model_path.glob(f"*{args.dataset_id:03}*"))).glob(f"*{_key_ResEnc}*")) + except StopIteration: + return next(next(iter(model_path.glob(f"*{args.dataset_id:03}*"))).glob("*__nnUNetPlans*")) + + try: + nnunet_path = _resolve_nnunet_path() + except StopIteration: + # Last-ditch: try one more download, then re-resolve. + print(f"[INFO] No nnUNet configuration found under {model_path}/Dataset{args.dataset_id:03}; retrying download…") + download_weights(args.dataset_id, model_path=model_path) + try: + nnunet_path = _resolve_nnunet_path() + except StopIteration as e: + raise RuntimeError(f"No nnUNet model found for dataset {args.dataset_id}") from e + + json_path = nnunet_path / "dataset.json" + assert json_path.exists(), json_path + + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + # Shape grid + if args.shapes_csv: + shapes = load_shapes_csv(args.shapes_csv) + print(f"[INFO] Loaded {len(shapes)} shapes from {args.shapes_csv}") + else: + shapes = DEFAULT_SHAPES + print(f"[INFO] Using default grid of {len(shapes)} shapes") + + total_gpu_mb = gpu_total_mb(args.gpu) + print(f"[INFO] GPU {args.gpu}: {total_gpu_mb:.0f} MB total VRAM") + print(f"[INFO] Fitting with quantile={args.quantile} headroom={args.headroom}") + + # Resolve voxel_size once: prefer the model's own target spacing so run_vibeseg's + # internal rescale is a no-op and the probed shape == the shape reaching nnUNet. + if args.voxel_size is None: + derived = derive_model_zoom(nnunet_path, args.dataset_id) + if derived is None: + voxel_size: float | tuple[float, float, float] = 0.8 + print("[WARN] Could not derive model zoom; falling back to isotropic 0.8 mm.") + else: + voxel_size = derived + print(f"[INFO] Derived model zoom from configs: {derived} mm") + else: + voxel_size = args.voxel_size + print(f"[INFO] Using explicit --voxel-size {voxel_size} mm (isotropic)") + print() + + shapes = sorted(shapes, key=lambda s: np.prod(s)) + + # ------------------------------------------------------------------ probes + records = [] + results_csv = nnunet_path / args.results_csv + + # Warm-up: the first inference pays for CUDA context init, kernel autotune, + # cuDNN benchmark, and one-time allocator growth. Discard it so the smallest + # shape isn't systematically inflated (that outlier drags the slope down). + print(f"[INFO] Warm-up pass on {shapes[0]} (discarded)") + _ = probe_shape( + shape=shapes[0], + gpu=args.gpu, + dataset_id=args.dataset_id, + model_path=args.model_path, + out_dir=out_dir, + voxel_size=voxel_size, + ) + + with open(results_csv, "w", newline="") as csvfile: + fieldnames = ["shape_d", "shape_h", "shape_w", "n_voxels", "peak_mb", "net_mb", "elapsed_s", "ok", "error"] + writer = csv.DictWriter(csvfile, fieldnames=fieldnames) + writer.writeheader() + + for shape in tqdm(shapes, desc="Probing shapes"): + n_vox = int(np.prod(shape)) + print(f"\n→ shape={shape} ({n_vox / 1e6:.2f} M voxels)") + + rec = probe_shape( + shape=shape, + gpu=args.gpu, + dataset_id=args.dataset_id, + model_path=args.model_path, + out_dir=out_dir, + voxel_size=voxel_size, + ) + records.append(rec) + + writer.writerow( + { + "shape_d": shape[0], + "shape_h": shape[1], + "shape_w": shape[2], + "n_voxels": n_vox, + "peak_mb": f"{rec['peak_mb']:.1f}" if rec["peak_mb"] else "", + "net_mb": f"{rec['net_mb']:.1f}" if rec["net_mb"] else "", + "elapsed_s": f"{rec['elapsed_s']:.1f}" if rec["elapsed_s"] else "", + "ok": rec["ok"], + "error": rec.get("error", ""), + } + ) + csvfile.flush() + + if rec["ok"]: + print(f" peak={rec['peak_mb']:.0f} MB net={rec['net_mb']:.0f} MB time={rec['elapsed_s']:.1f} s") + else: + print(f" FAILED: {rec.get('error', '?')}") + + print(f"\n[INFO] Raw results saved → {results_csv}") + + # ------------------------------------------------------------------ fit + try: + fit = fit_parameters(records, quantile=args.quantile) + except ValueError as e: + print(f"[ERROR] Fitting failed: {e}") + sys.exit(1) + + # Apply headroom to memory_max + fit["memory_max"] = fit["memory_max"] * args.headroom + + # ------------------------------------------------------------------ report + bar = "=" * 60 + print(f"\n{bar}") + print(" FITTED MEMORY PARAMETERS") + print(f" method: {fit['method']}") + print(bar) + print(f" memory_base = {fit['memory_base']:>10.0f} # MB (fixed overhead)") + print(f" memory_factor = {fit['memory_factor']:>10.2f} # n_voxels/1e6 * factor MB") + print(f" memory_max = {fit['memory_max']:>10.0f} # MB ({args.headroom:.0%} headroom on max observed)") + print(f" OLS R² = {fit['r2']:>10.4f} (diagnostic; fit targets upper envelope)") + print(bar) + print() + print(" Suggested call:") + print(" run_vibeseg(") + print(" nii, out,") + print(f" memory_base={fit['memory_base']:.0f},") + print(f" memory_factor={fit['memory_factor']:.2f},") + print(f" memory_max={fit['memory_max']:.0f},") + print(" )") + print(bar) + + # ------------------------------------------------------------------ plot + make_plot(records, fit, nnunet_path / args.plot) + + # --------------------------------------------------------- patch JSON + backup_path = json_path.with_suffix(".json.bak") + shutil.copy2(json_path, backup_path) + try: + with open(json_path) as f: + data = json.load(f) + data["memory_base"] = fit["memory_base"] + data["memory_factor"] = fit["memory_factor"] + tmp_path = json_path.with_suffix(".json.tmp") + with open(tmp_path, "w") as f: + json.dump(data, f, indent=4) + tmp_path.replace(json_path) + backup_path.unlink(missing_ok=True) + print(f"[INFO] dataset.json patched → {json_path}") + except Exception: + if backup_path.exists(): + shutil.copy2(backup_path, json_path) + raise + + +if __name__ == "__main__": + main() + # python estimate_vibeseg_memory.py \ + # --gpu 1 \ + # --model-path /DATA/NAS/FASTDATA/robert/nnUNet \ + # --dataset-id 12 \ + # --quantile 0.95 \ + # --headroom 1.20 diff --git a/TPTBox/segmentation/nnUnet_utils/predictor.py b/TPTBox/segmentation/nnUnet_utils/predictor.py index 9a90266d..610381a3 100755 --- a/TPTBox/segmentation/nnUnet_utils/predictor.py +++ b/TPTBox/segmentation/nnUnet_utils/predictor.py @@ -756,6 +756,11 @@ def _allocate(self, data: torch.Tensor, results_device, pbar: tqdm, gauss: bool device=results_device, ) except RuntimeError as e: + if self.fail_on_missing_memory: + # Probing / benchmarking mode: don't hide the OOM behind a slow CPU fallback, + # let the caller record the failure and move on to the next shape. + empty_cache(self.device) + raise try: n_predictions = None gaussian = 1 From db262ba72cf21ee24da5724d8987c97c65a928d5 Mon Sep 17 00:00:00 2001 From: robert Date: Mon, 24 Aug 2026 16:48:14 +0200 Subject: [PATCH 13/13] ruff --- .../nnUnet_utils/estimate_nnunet_memory.py | 44 ++++++++++--------- 1 file changed, 24 insertions(+), 20 deletions(-) diff --git a/TPTBox/segmentation/nnUnet_utils/estimate_nnunet_memory.py b/TPTBox/segmentation/nnUnet_utils/estimate_nnunet_memory.py index e34fe498..4ab8a4d0 100644 --- a/TPTBox/segmentation/nnUnet_utils/estimate_nnunet_memory.py +++ b/TPTBox/segmentation/nnUnet_utils/estimate_nnunet_memory.py @@ -1,5 +1,5 @@ -"""estimate_vibeseg_memory.py -========================== +"""Fit the nnUNet memory guard parameters from GPU probes. + Measures peak GPU RAM consumed by run_vibeseg across a grid of synthetic input shapes, then fits the three memory parameters used by TPTBox's check_mem guard: @@ -9,15 +9,15 @@ < clamp(0.80 * gpu_total, lo=memory_base, hi=memory_max) So the fitted curve must be a CONSERVATIVE UPPER BOUND of actual usage, - not a mean — otherwise ~50 % of runs would be incorrectly skipped. + not a mean, otherwise ~50 % of runs would be incorrectly skipped. This script fits via quantile regression (default q=0.95) so the curve sits above nearly all observations while remaining tight. Usage ----- - python estimate_vibeseg_memory.py [--gpu 0] [--model-path /path/to/nnUNet] - [--dataset-id 12] [--out-dir /tmp/vibeseg_probe] - [--shapes-csv shapes.csv] + python estimate_nnunet_memory.py [--gpu 0] [--model-path /path/to/nnUNet] + [--dataset-id 12] [--out-dir /tmp/nnunet_probe] + [--shapes-csv shapes.csv] The script prints recommended values for memory_base, memory_factor, and memory_max at the end, saves a CSV + PNG summary plot, and patches the @@ -44,7 +44,8 @@ from tqdm import tqdm except ImportError: - def tqdm(it, **kwargs): + def tqdm(it, **_kwargs): # type: ignore[no-redef] # noqa: ANN201 + """Fallback no-op progress bar when tqdm is not installed.""" return it @@ -66,10 +67,12 @@ def _nvidia_smi_query(field: str, gpu: int) -> float: def gpu_used_mb(gpu: int) -> float: + """Return currently used GPU memory (MB) as reported by nvidia-smi.""" return _nvidia_smi_query("memory.used", gpu) def gpu_total_mb(gpu: int) -> float: + """Return total GPU memory (MB) as reported by nvidia-smi.""" return _nvidia_smi_query("memory.total", gpu) @@ -123,6 +126,7 @@ def __exit__(self, *_): @property def peak_mb(self) -> float: + """Peak GPU memory (MB) observed while the poller was active.""" return self._peak @@ -170,9 +174,9 @@ def probe_shape( out_dir: Path, voxel_size: float | tuple[float, float, float] = 0.8, ) -> dict: - """Run run_vibeseg on a synthetic volume, measure peak GPU RAM. - memory_max is set to 999 GB so the check_mem guard never fires here. + """Run run_vibeseg on a synthetic volume and measure peak GPU RAM. + memory_max is set to 999 GB so the check_mem guard never fires here. ``voxel_size`` accepts either an isotropic scalar or a 3-tuple. Passing the model's own zoom keeps run_vibeseg's internal rescale a no-op, so the shape that reaches nnUNet matches the shape we probed. @@ -228,13 +232,12 @@ def probe_shape( # --------------------------------------------------------------------------- -# Fitting – quantile regression (upper envelope, not mean) +# Fitting - quantile regression (upper envelope, not mean) # --------------------------------------------------------------------------- def fit_parameters(records: list, quantile: float = 0.95) -> dict: - """Fit an upper-envelope line: - net_mb ≈ memory_base + n_voxels / 1e6 * memory_factor + """Fit an upper-envelope line ``net_mb ≈ memory_base + n_voxels / 1e6 * memory_factor``. Uses quantile regression at `quantile` (default 0.95) so the predicted curve lies above ~95 % of observations. This is intentional: the @@ -341,11 +344,9 @@ def fit_parameters(records: list, quantile: float = 0.95) -> dict: def load_shapes_csv(path: str) -> list: - shapes = [] + """Load a shape grid from a CSV with columns ``d,h,w``.""" with open(path) as f: - for row in csv.DictReader(f): - shapes.append((int(row["d"]), int(row["h"]), int(row["w"]))) - return shapes + return [(int(row["d"]), int(row["h"]), int(row["w"])) for row in csv.DictReader(f)] # --------------------------------------------------------------------------- @@ -353,11 +354,12 @@ def load_shapes_csv(path: str) -> list: # --------------------------------------------------------------------------- -def make_plot(records: list, fit: dict, out_path: Path): +def make_plot(records: list, fit: dict, out_path: Path) -> None: + """Render the measured/fitted memory curves and residuals to ``out_path``.""" try: import matplotlib.pyplot as plt except ImportError: - print("[INFO] matplotlib not available – skipping plot.") + print("[INFO] matplotlib not available - skipping plot.") return ok = [r for r in records if r["ok"] and r["net_mb"] is not None] @@ -418,7 +420,8 @@ def make_plot(records: list, fit: dict, out_path: Path): # --------------------------------------------------------------------------- -def parse_args(): +def parse_args() -> argparse.Namespace: + """Parse CLI arguments for the memory-estimation script.""" p = argparse.ArgumentParser(description="Estimate run_nnunet memory parameters from GPU measurements.") p.add_argument("--gpu", type=int, default=0) p.add_argument( @@ -444,7 +447,8 @@ def parse_args(): return p.parse_args() -def main(): +def main() -> None: + """Entry point: probe shapes, fit memory parameters, patch dataset.json.""" args = parse_args() # Resolve nnUNet model path (auto-download when using the default TPTBox path)