Rib segmentation + low-RAM nnU-Net inference

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

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())

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}"

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.

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.
 
@@ -191,7 +209,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:
@@ -230,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
@@ -282,20 +321,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:
@@ -326,12 +360,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

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

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