Skip to content

NNunet softmax #161

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
robert-graf wants to merge 3 commits into
base: main
Choose a base branch
from
Open

NNunet softmax #161

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
05d4d92
readd probabilities as an optional output for nnUnet
Nov 9, 2026
f3f9516
bug fix for uint8
Nov 9, 2026
81072b1
add POI info if missing
Nov 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions TPTBox/core/poi.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
TRIPLE,
ZOOMS,
Abstract_lvl,
Any,
Location,
Sentinel,
Vertebra_Instance,
Expand Down Expand Up @@ -1171,6 +1172,9 @@ def calc_poi_from_subreg_vert(
if loc in subreg_id_int:
log.print("Calc centroid from subregion id", Location(loc).name, verbose=verbose)
vert_arr = vert_msk.get_seg_array()
# Only widen when needed: the offsets (up to v + 100) must fit into the dtype, otherwise NumPy 2 raises OverflowError
if vert_arr.dtype.kind in "ui" and np.iinfo(vert_arr.dtype).max < v + 100:
vert_arr = vert_arr.astype(np.int32)
subreg_arr = subreg_msk.get_seg_array()
# IVD / Endplates Superior / Endplate Inferior
vert_arr[subreg_arr != loc] = 0
Expand All @@ -1186,6 +1190,12 @@ def calc_poi_from_subreg_vert(
inplace=True,
)
subreg_id_int.remove(loc)
# calc_centroids* build new POIs from plain int ids and therefore drop the level info of ``extend_to``.
# This function works on vertebra instances + Location sub-regions, so (re-)set it here once.
if extend_to.level_one_info is Any:
extend_to.level_one_info = Vertebra_Instance
if extend_to.level_two_info is Any:
extend_to.level_two_info = Location
# Step 5 call non_centroid_pois
# Prepare mask to binary mask
print("step 5", subreg_id_int) if _print_phases else None
Expand Down
28 changes: 24 additions & 4 deletions TPTBox/segmentation/VibeSeg/inference_nnunet.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ def run_inference_on_file(
gpu=None,
keep_size: bool = False,
fill_holes: bool = False,
logits: bool = False, # deprecated
logits: bool = False,
mapping=None,
crop: bool = False,
max_folds=None,
Expand Down Expand Up @@ -144,8 +144,11 @@ def run_inference_on_file(
original image size.
fill_holes: If ``True``, fill holes in the segmentation mask after
inference.
logits: Deprecated; Was commented out for less GPU waste.
If ``True``, also return the raw softmax logits array.
logits: If ``True``, also return the softmax probabilities as a ``float32`` array
``(num_classes, X, Y, Z)`` on the grid of ``input_nii[0]`` (linearly resampled
like the segmentation; channel ``c`` is nnUNet label ``c``, *before* ``mapping``).
Costs ``num_classes`` x volume of RAM. ``None`` in the returned tuple when the
run fell back to chunked/patched inference, which cannot return probabilities.
mapping: Optional label remapping dict applied to the segmentation after
inference.
crop: If ``True``, crop the input to its foreground bounding box before
Expand Down Expand Up @@ -191,7 +194,7 @@ def run_inference_on_file(
Returns:
A tuple ``(seg_nii, softmax_logits)`` where ``seg_nii`` is the
segmentation (as a :class:`~TPTBox.NII` or file path) and
``softmax_logits`` is the raw logit array or ``None``.
``softmax_logits`` is the softmax probability array or ``None``.
"""
if model_path is None:
auto_download = True
Expand Down Expand Up @@ -421,6 +424,23 @@ def _defaults_fit_gpu(memory: float, gpu: int | None, safety_factor: float = 2.0

if padd != 0:
seg_nii = seg_nii[padd:-padd, padd:-padd, padd:-padd]
if softmax_logits is not None:
softmax_logits = softmax_logits[:, padd:-padd, padd:-padd, padd:-padd]
if logits and softmax_logits is None:
logger.print("Probabilities requested but unavailable for this (chunked / low-memory) run.", Log_Type.WARNING)
if softmax_logits is not None and not keep_size:
# Same grid change as the segmentation, but linear so the maps stay smooth and within [0, 1]
# NII.set_array casts floats to int32 while *self* is a seg, so the template must not be one
grid = seg_nii.copy()
grid.seg = False
softmax_logits = np.stack(
[
grid.set_array(np.ascontiguousarray(c, dtype=np.float32))
.resample_from_to(og_nii, mode="nearest", order=1, verbose=False, out_dtype=np.float32)
.get_array()
for c in softmax_logits
]
)

if mapping is not None:
seg_nii.map_labels_(mapping)
Expand Down
85 changes: 53 additions & 32 deletions TPTBox/segmentation/nnUnet_utils/export_prediction.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
# method for deep learning-based biomedical image segmentation. Nature methods, 18(2), 203-211.
from __future__ import annotations

from typing import Literal, overload

import numpy as np
import torch
from acvl_utils.cropping_and_padding.bounding_boxes import bounding_box_to_slice
Expand Down Expand Up @@ -162,6 +164,30 @@ def convert_probabilities_to_segmentation(
return segmentation


@overload
def convert_predicted_logits_to_segmentation_with_correct_shape(
predicted_logits: torch.Tensor | np.ndarray,
plans_manager: PlansManager,
configuration_manager: ConfigurationManager,
label_manager: LabelManager,
properties_dict: dict,
return_probabilities: Literal[False] = False,
num_threads_torch: int = 8,
device=None,
logger=...,
) -> np.ndarray: ...
@overload
def convert_predicted_logits_to_segmentation_with_correct_shape(
predicted_logits: torch.Tensor | np.ndarray,
plans_manager: PlansManager,
configuration_manager: ConfigurationManager,
label_manager: LabelManager,
properties_dict: dict,
return_probabilities: Literal[True],
num_threads_torch: int = 8,
device=None,
logger=...,
) -> tuple[np.ndarray, np.ndarray]: ...
def convert_predicted_logits_to_segmentation_with_correct_shape(
predicted_logits: torch.Tensor | np.ndarray,
plans_manager: PlansManager,
Expand All @@ -172,7 +198,7 @@ def convert_predicted_logits_to_segmentation_with_correct_shape(
num_threads_torch: int = 8,
device=None,
logger=logger,
) -> np.ndarray:
) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
"""Revert all preprocessing steps and return a segmentation in the original image space.

Performs (in reverse order): resampling from network output spacing back to
Expand All @@ -190,23 +216,31 @@ def convert_predicted_logits_to_segmentation_with_correct_shape(
properties_dict: Case properties dict as updated by
:meth:`DefaultPreprocessor.run_case_npy` (contains cropping bbox,
pre-crop shape, and original spacing).
return_probabilities: Reserved for future use. Raises
:class:`NotImplementedError` if ``True``.
return_probabilities: If ``True``, also return the softmax probabilities
(``float32``, ``(num_classes, X, Y, Z)``) in the same original image
space as the segmentation.
num_threads_torch: Number of threads used by PyTorch during resampling.
device: Torch device on which the argmax runs; ``None`` uses the tensor's current device.
logger: Logger used for progress messages from the argmax GPU-fallback path.

Returns:
Integer segmentation array with dtype ``np.uint8`` or ``np.uint16``
in the original image space (before any preprocessing).

Raises:
NotImplementedError: If ``return_probabilities`` is ``True``.
in the original image space (before any preprocessing). With
``return_probabilities`` a tuple ``(segmentation, probabilities)``.
"""
if return_probabilities:
raise NotImplementedError()
old_threads = torch.get_num_threads()
torch.set_num_threads(num_threads_torch)
# run_inference_on_file rescales the input to the model spacing itself, so nnUNet's own
# preprocessing/export resampling should be a no-op (same shape before and after).
logits_shape = tuple(int(i) for i in predicted_logits.shape[1:])
target_shape = tuple(int(i) for i in properties_dict["shape_after_cropping_and_before_resampling"])
if logits_shape != target_shape:
logger.on_warning(
f"nnUNet resamples the prediction {logits_shape} -> {target_shape} "
f"(model spacing {list(configuration_manager.spacing)}, input spacing {list(properties_dict['spacing'])}). "
"The input was expected to be pre-rescaled to the model spacing; this adds an interpolation step "
"(check `resolution_range`/`spacing` in the model's dataset.json / inference_config.json)."
)
# resample to original shape
current_spacing = (
configuration_manager.spacing
Expand All @@ -222,6 +256,14 @@ def convert_predicted_logits_to_segmentation_with_correct_shape(
# segmentation: np.ndarray = label_manager.convert_probabilities_to_segmentation(predicted_logits) # type: ignore
segmentation: np.ndarray = convert_probabilities_to_segmentation(label_manager, predicted_logits, device, logger=logger)
segmentation = segmentation.astype(np.uint8 if len(label_manager.foreground_labels) < 255 else np.uint16)
probabilities: np.ndarray | None = None
if return_probabilities:
probs = predicted_logits.cpu().numpy() if isinstance(predicted_logits, torch.Tensor) else np.asarray(predicted_logits)
probabilities = np.zeros((probs.shape[0], *properties_dict["shape_before_cropping"]), dtype=np.float32)
probabilities[(slice(None), *bounding_box_to_slice(properties_dict["bbox_used_for_cropping"]))] = probs
del probs
# revert transpose (channel axis stays first)
probabilities = probabilities.transpose([0] + [i + 1 for i in plans_manager.transpose_backward])
del predicted_logits
# put segmentation in bbox (revert cropping)
segmentation_reverted_cropping = np.zeros(
Expand All @@ -235,30 +277,9 @@ def convert_predicted_logits_to_segmentation_with_correct_shape(
# revert transpose
segmentation = segmentation.transpose(plans_manager.transpose_backward)
logger.print(segmentation.shape)
# if return_probabilities:
# raise NotImplementedError()
# # revert cropping
# try:
# predicted_probabilities = predicted_logits # noqa: F821
# predicted_probabilities = label_manager.revert_cropping_on_probabilities( # type: ignore
# predicted_probabilities,
# properties_dict["bbox_used_for_cropping"],
# properties_dict["shape_before_cropping"],
# )
# except Exception:
# predicted_probabilities = label_manager.revert_cropping( # type: ignore
# predicted_probabilities.numpy() if isinstance(predicted_probabilities, torch.Tensor) else predicted_probabilities,
# properties_dict["bbox_used_for_cropping"],
# properties_dict["shape_before_cropping"],
# )
# if isinstance(predicted_probabilities, torch.Tensor):
# predicted_probabilities = predicted_probabilities.cpu().numpy()
# # revert transpose
# predicted_probabilities = predicted_probabilities.transpose([0] + [i + 1 for i in plans_manager.transpose_backward])
# torch.set_num_threads(old_threads)
# return segmentation, predicted_probabilities
# else:
torch.set_num_threads(old_threads)
if probabilities is not None:
return segmentation, probabilities
return segmentation


Expand Down
27 changes: 19 additions & 8 deletions TPTBox/segmentation/nnUnet_utils/inference_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -301,21 +301,18 @@ def run_inference(
predictor: Loaded ``nnUNetPredictor`` as returned by :func:`load_inf_model`.
reorient_PIR: If True, reorient each input image to PIR orientation before
passing it to the model.
logits: If True, return raw softmax logits. Currently not implemented and
will raise ``NotImplementedError``.
logits: If True, also return the softmax probabilities as a ``float32`` array
``(num_classes, X, Y, Z)`` on the grid (shape, orientation) of the first input.
logger: Logger used for progress and error output.
verbose: Unused; reserved for future logging support.

Raises:
NotImplementedError: If *logits* is True.
AssertionError: If *input_nii* is not a supported type, or if the output
shape does not match the input shape.

Returns:
A tuple of (segmentation_NII, uncertainty_map_or_None, softmax_logits_or_None).
"""
if logits:
raise NotImplementedError("logits=True")
if isinstance(input_nii, str):
assert input_nii.endswith(".nii.gz"), f"input file is not a .nii.gz! Got {input_nii}"
input_nii = NII.load(input_nii, seg=False)
Expand All @@ -339,13 +336,27 @@ def run_inference(
logger.on_fail("could not stack images; shapes=", [a.shape for a in img_arrs])
raise
props = {"spacing": i.zoom[::-1]} # PIR
out = predictor.predict_single_npy_array(img, props, save_or_return_probabilities=False, logger=logger)
segmentation: np.ndarray = out # type: ignore
softmax_logits = None
probabilities: np.ndarray | None = None
if logits:
segmentation, probabilities = predictor.predict_single_npy_array(img, props, save_or_return_probabilities=True, logger=logger)
else:
segmentation = predictor.predict_single_npy_array(img, props, save_or_return_probabilities=False, logger=logger)
segmentation = np.transpose(segmentation, axes=segmentation.ndim - 1 - np.arange(segmentation.ndim))
assert segmentation.shape == input_nii[0].shape, (segmentation.shape, input_nii[0].shape)
seg_nii = input_nii[0].set_array(segmentation.astype(np.uint8), seg=True)
seg_nii.reorient_(orientation, verbose=False)

softmax_logits = None
if probabilities is not None:
# nnUNet axis order (C, z, y, x) -> NII order (C, x, y, z), then the same reorientation as the segmentation
probabilities = np.transpose(probabilities, axes=[0, *range(probabilities.ndim - 1, 0, -1)])
assert probabilities.shape[1:] == input_nii[0].shape, (probabilities.shape, input_nii[0].shape)
channels = []
for c in probabilities:
channel = input_nii[0].set_array(np.ascontiguousarray(c, dtype=np.float32), seg=False)
channel.reorient_(orientation, verbose=False)
channels.append(channel.get_array())
softmax_logits = np.stack(channels).astype(np.float32, copy=False)
return seg_nii, None, softmax_logits


Expand Down
53 changes: 37 additions & 16 deletions TPTBox/segmentation/nnUnet_utils/predictor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from math import ceil, floor
from queue import Queue
from threading import Thread
from typing import Literal, overload

import numpy as np
import torch
Expand Down Expand Up @@ -356,24 +357,31 @@ def mapp(d: dict):
self.network.eval() # type: ignore
self.loaded_networks = [self.network]

@overload
def predict_single_npy_array(
self, input_image: np.ndarray, image_properties: dict, save_or_return_probabilities: Literal[False] = False, logger=...
) -> np.ndarray: ...
@overload
def predict_single_npy_array(
self, input_image: np.ndarray, image_properties: dict, save_or_return_probabilities: Literal[True], logger=...
) -> tuple[np.ndarray, np.ndarray]: ...
def predict_single_npy_array(
self, input_image: np.ndarray, image_properties: dict, save_or_return_probabilities: bool = False, logger=logger
) -> np.ndarray:
) -> np.ndarray | tuple[np.ndarray, np.ndarray]:
"""Run full inference on a single numpy image array.

Args:
input_image: Image array with shape ``(C, X, Y, Z)`` where ``C`` is
the number of input channels.
image_properties: Dictionary that must contain a ``'spacing'`` key
with the voxel spacing in mm (list/tuple of floats).
save_or_return_probabilities: If ``True``, also return softmax
probabilities in addition to the label map. Currently raises
:class:`NotImplementedError` inside the conversion step.
save_or_return_probabilities: If ``True``, also return the softmax
probabilities (``float32``, ``(C, ...)``) in addition to the label map.
logger: Logger used for progress and error output.

Returns:
The predicted segmentation array (or a tuple with probabilities if
``save_or_return_probabilities`` is ``True``).
The predicted segmentation array, or ``(segmentation, probabilities)`` if
``save_or_return_probabilities`` is ``True``.
"""
segmentation_previous_stage: np.ndarray = None # type: ignore # Was previously a parameter

Expand Down Expand Up @@ -402,16 +410,29 @@ def predict_single_npy_array(
import time

t = time.time()
ret = convert_predicted_logits_to_segmentation_with_correct_shape(
predicted_logits,
self.plans_manager,
self.configuration_manager,
self.label_manager,
dct["data_properites"],
return_probabilities=save_or_return_probabilities,
device=self.device,
logger=logger,
)
ret: np.ndarray | tuple[np.ndarray, np.ndarray]
if save_or_return_probabilities:
ret = convert_predicted_logits_to_segmentation_with_correct_shape(
predicted_logits,
self.plans_manager,
self.configuration_manager,
self.label_manager,
dct["data_properites"],
return_probabilities=True,
device=self.device,
logger=logger,
)
else:
ret = convert_predicted_logits_to_segmentation_with_correct_shape(
predicted_logits,
self.plans_manager,
self.configuration_manager,
self.label_manager,
dct["data_properites"],
return_probabilities=False,
device=self.device,
logger=logger,
)
print("convert_predicted_logits_to_segmentation_with_correct_shape; Took", time.time() - t, " seconds")

return ret
Expand Down