diff --git a/TPTBox/core/poi.py b/TPTBox/core/poi.py index 26013fb..294a9fe 100755 --- a/TPTBox/core/poi.py +++ b/TPTBox/core/poi.py @@ -33,6 +33,7 @@ TRIPLE, ZOOMS, Abstract_lvl, + Any, Location, Sentinel, Vertebra_Instance, @@ -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 @@ -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 diff --git a/TPTBox/segmentation/VibeSeg/inference_nnunet.py b/TPTBox/segmentation/VibeSeg/inference_nnunet.py index 0f0711b..8096954 100644 --- a/TPTBox/segmentation/VibeSeg/inference_nnunet.py +++ b/TPTBox/segmentation/VibeSeg/inference_nnunet.py @@ -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, @@ -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 @@ -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 @@ -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) diff --git a/TPTBox/segmentation/nnUnet_utils/export_prediction.py b/TPTBox/segmentation/nnUnet_utils/export_prediction.py index c5c7dfe..0b6feaf 100755 --- a/TPTBox/segmentation/nnUnet_utils/export_prediction.py +++ b/TPTBox/segmentation/nnUnet_utils/export_prediction.py @@ -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 @@ -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, @@ -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 @@ -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 @@ -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( @@ -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 diff --git a/TPTBox/segmentation/nnUnet_utils/inference_api.py b/TPTBox/segmentation/nnUnet_utils/inference_api.py index e79664e..cd8197b 100755 --- a/TPTBox/segmentation/nnUnet_utils/inference_api.py +++ b/TPTBox/segmentation/nnUnet_utils/inference_api.py @@ -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) @@ -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 diff --git a/TPTBox/segmentation/nnUnet_utils/predictor.py b/TPTBox/segmentation/nnUnet_utils/predictor.py index a04a2d6..eeb0ca0 100755 --- a/TPTBox/segmentation/nnUnet_utils/predictor.py +++ b/TPTBox/segmentation/nnUnet_utils/predictor.py @@ -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 @@ -356,9 +357,17 @@ 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: @@ -366,14 +375,13 @@ def predict_single_npy_array( 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 @@ -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