diff --git a/TPTBox/core/internal/nii_help.py b/TPTBox/core/internal/nii_help.py index 644853f..f2cdb6a 100644 --- a/TPTBox/core/internal/nii_help.py +++ b/TPTBox/core/internal/nii_help.py @@ -18,6 +18,38 @@ from TPTBox.core.vert_constants import AFFINE, MODES, SHAPE, ZOOMS, Sentinel, _supported_img_files +# NIfTI-1 headers can encode most numpy dtypes, but a couple of common ones +# don't have a NIfTI datatype code (notably float16 and bool). The NII wrapper +# still wants to *carry* an array in one of those dtypes without crashing - +# e.g. an nnU-Net pre-processing step producing a float16 volume, or a +# boolean mask. We only upcast when the array actually has to enter a +# Nifti1Image (i.e. writing to disk or handing the underlying nibabel object +# out), so the caller-visible ``_arr.dtype`` stays whatever they set. +_NIFTI_UNSUPPORTED_DTYPE_UPCAST: dict[np.dtype, np.dtype] = { + np.dtype(np.float16): np.dtype(np.float32), + np.dtype(np.bool_): np.dtype(np.uint8), +} + + +def _nifti_safe_dtype(dtype: np.dtype | type) -> np.dtype: + """Return the closest nibabel-supported dtype for storing in a Nifti1 header. + + For a dtype that NIfTI-1 already accepts, this is the identity. For the + unsupported cases we upcast conservatively (``float16 → float32``, + ``bool → uint8``). The array itself is *not* touched – see + :func:`_arr_for_nifti1`. + """ + d = np.dtype(dtype) + return _NIFTI_UNSUPPORTED_DTYPE_UPCAST.get(d, d) + + +def _arr_for_nifti1(arr: np.ndarray) -> np.ndarray: + """Return *arr* (or a copy in a safe dtype) fit to be passed to Nifti1Image.""" + safe = _nifti_safe_dtype(arr.dtype) + if safe == arr.dtype: + return arr + return arr.astype(safe, copy=False) + def secure_save(func, *, file_types=tuple(_supported_img_files)) -> Callable: """Decorator that writes to a `.backup` file first and restores it if saving fails. @@ -172,6 +204,7 @@ def _resample_from_to( order: int = 3, mode: MODES = "nearest", align_corners: bool | Sentinel = Sentinel(), # noqa: B008 + out_dtype: np.dtype | type | str | None = None, ) -> tuple[np.ndarray, np.ndarray, object]: """Resample *from_img* into the voxel space defined by *to_img*. @@ -191,6 +224,19 @@ def _resample_from_to( ``order == 0``), voxel corners are aligned between source and target grids. When ``False``, voxel centres are aligned (standard nibabel/scipy behaviour). + out_dtype: Optional NumPy dtype (or dtype-like) requested for the + resampled array. When set, this is forwarded to ``scipy.ndimage`` + as its ``output=`` argument, so the cast happens in-place during + interpolation without an extra full-volume copy afterwards. + ``None`` (default) keeps the source dtype. + With ``order > 0`` and an integer target dtype the float→int cast + is a plain truncation (matches NumPy's cast rules); pass ``order=0`` + for label maps. ``scipy.ndimage.affine_transform`` only accepts + ``uint8/uint16/int16/int32/float32/float64`` for its ``output=`` + argument. Requesting an unsupported dtype (notably ``float16``) + transparently falls back to resampling into ``float32`` and casting + to the requested dtype in a single extra pass - still much cheaper + than staying in the source dtype (e.g. float64) throughout. Returns: A 3-tuple ``(data, affine, header)`` where *data* is the resampled @@ -254,5 +300,32 @@ def _resample_from_to( to_vox2from_vox = npl.inv(a_from_affine).dot(a_to_affine) rzs, trans = to_matvec(to_vox2from_vox) - data = scipy_img.affine_transform(from_img.get_array(), rzs, trans, to_shape, order=order, mode=mode, cval=from_img.get_c_val()) # type: ignore + # scipy.ndimage.affine_transform can only write into a small set of dtypes + # (u8/u16/i16/i32/f32/f64). For anything else - notably float16, which the + # nnU-Net inference path wants for memory - we resample into float32 first + # and cast in one pass at the end. + _scipy_supported = (np.uint8, np.uint16, np.int16, np.int32, np.float32, np.float64) + if out_dtype is None: + scipy_out = None + post_cast: np.dtype | None = None + else: + req = np.dtype(out_dtype) + if req.type in _scipy_supported: + scipy_out = req + post_cast = None + else: + scipy_out = np.dtype(np.float32) + post_cast = req + data = scipy_img.affine_transform( # type: ignore + from_img.get_array(), + rzs, + trans, + to_shape, + order=order, + mode=mode, + cval=from_img.get_c_val(), + output=scipy_out, + ) + if post_cast is not None: + data = data.astype(post_cast, copy=False) return data, to_affine, from_img.header diff --git a/TPTBox/core/nii_wrapper.py b/TPTBox/core/nii_wrapper.py index 833761c..f27341f 100755 --- a/TPTBox/core/nii_wrapper.py +++ b/TPTBox/core/nii_wrapper.py @@ -19,7 +19,12 @@ from TPTBox.core import bids_files from TPTBox.core.compat import zip_strict -from TPTBox.core.internal.nii_help import _resample_from_to, secure_save +from TPTBox.core.internal.nii_help import ( + _arr_for_nifti1, + _nifti_safe_dtype, + _resample_from_to, + secure_save, +) from TPTBox.core.nii_poi_abstract import Has_Grid from TPTBox.core.nii_wrapper_math import NII_Math from TPTBox.core.np_utils import ( @@ -158,13 +163,14 @@ def _check_if_nifty_is_lying_about_its_dtype(self: NII): stacklevel=3, ) - out_dtype = dtype - if dtype == np.float16: + # Delegate the "which dtypes are unsupported on disk?" question to the + # canonical mapping in nii_help so we only ever maintain it in one place. + out_dtype = _nifti_safe_dtype(dtype) + if out_dtype != dtype: warnings.warn( f"Loaded NIfTY: incorrect dtype detected: {dtype} is not supported", stacklevel=3, ) - out_dtype = np.float32 if "float" in dtype_s and not change_dtype: pass elif positive and change_dtype: @@ -200,6 +206,7 @@ def _check_if_nifty_is_lying_about_its_dtype(self: NII): Proxy = tuple[tuple[int, int, int], np.ndarray] suppress_dtype_change_printout_in_set_array = False + # fmt: off class NII(NII_Math): @@ -431,15 +438,20 @@ def nii(self) -> Nifti1Image: current array. Use ``nii_abstract`` if you want to avoid this reconstruction overhead. """ if self.__divergent: - self._nii = Nifti1Image(self._arr,self.affine,self.header) - if self.dtype == self._arr.dtype: #type: ignore - nii = Nifti1Image(self._arr,self.affine,self.header) + # NIfTI-1 cannot encode a handful of dtypes (float16, bool). Upcast + # only the *materialised* array/header; leave `self._arr` alone so + # callers who kept a reference still see the native dtype. + arr_for_nib = _arr_for_nifti1(self._arr) # type: ignore + safe_dtype = _nifti_safe_dtype(self._arr.dtype) # type: ignore + self._nii = Nifti1Image(arr_for_nib, self.affine, self.header) + if self.dtype == self._arr.dtype: # type: ignore + nii = Nifti1Image(arr_for_nib, self.affine, self.header) else: if not suppress_dtype_change_printout_in_set_array: log.print(f"'set_array' with different dtype: from {self.dtype} to {self._arr.dtype}",verbose=True) #type: ignore - nii2 = Nifti1Image(self._arr,self.affine,self.header) - nii2.set_data_dtype(self._arr.dtype) - nii = Nifti1Image(self._arr,nii2.affine,nii2.header) # type: ignore + nii2 = Nifti1Image(arr_for_nib, self.affine, self.header) + nii2.set_data_dtype(safe_dtype) + nii = Nifti1Image(arr_for_nib, nii2.affine, nii2.header) # type: ignore if all(a is None for a in self.header.get_slope_inter()): nii.header.set_slope_inter(1,self.get_c_val()) # type: ignore #if self.header is not None: @@ -473,7 +485,7 @@ def nii(self,nii:Nifti1Image|_unpacked_nii): header = header.copy() header.set_sform(aff, code='aligned') header.set_qform(aff, code='unknown') - header.set_data_dtype(arr.dtype) + header.set_data_dtype(_nifti_safe_dtype(arr.dtype)) rotation_zoom = aff[:n, :n] zoom = np.sqrt(np.sum(rotation_zoom * rotation_zoom, axis=0)) #print(aff.shape,arr.shape,zoom) @@ -481,7 +493,7 @@ def nii(self,nii:Nifti1Image|_unpacked_nii): self._header = header return else: - nii = Nifti1Image(arr,aff) + nii = Nifti1Image(_arr_for_nifti1(arr), aff) self.__unpacked = False self.__divergent = False self._nii = nii @@ -694,15 +706,13 @@ def set_array(self, arr: np.ndarray | Self, inplace=False, verbose: logging = Fa """ if hasattr(arr,"get_array"): arr = arr.get_array() # type: ignore - if arr.dtype == bool: - arr = arr.astype(np.uint8) - if arr.dtype == np.float16: - arr = arr.astype(np.float32) + # Segmentations must not be floats; force an integer dtype. + # bool/float16 are kept in memory here - only upcast for storage when + # we hand a Nifti1Image out (see `_nifti_safe_dtype`). if self.seg and np.issubdtype(arr.dtype, np.floating): arr = arr.astype(np.int32) - #if self.dtype == arr.dtype: #type: ignore nii:_unpacked_nii = (arr,self.affine,self.header.copy()) - self.header.set_data_dtype(arr.dtype) + self.header.set_data_dtype(_nifti_safe_dtype(arr.dtype)) #else: # if not suppress_dtype_change_printout_in_set_array: # log.print(f"'set_array' with different dtype: from {self.nii.dataobj.dtype} to {arr.dtype}",verbose=verbose) #type: ignore @@ -745,13 +755,16 @@ def set_dtype(self, dtype: type | Literal['smallest_int', 'smallest_uint'] = np. if self.__unpacked: self._unpack() sel._arr = sel._arr.astype(dtype) - sel.header.set_data_dtype(dtype) + # header may only encode nibabel-safe dtypes; the array itself + # keeps whatever the caller asked for. + sel.header.set_data_dtype(_nifti_safe_dtype(dtype)) else: - sel.nii.set_data_dtype(dtype) + sel.nii.set_data_dtype(_nifti_safe_dtype(dtype)) if sel.nii.get_data_dtype() != self.dtype: #type: ignore if arr is None: arr = self.get_array() - sel.nii = Nifti1Image(arr.astype(dtype,casting=casting,order=order),self.affine,self.header) + new_arr = arr.astype(dtype, casting=casting, order=order) + sel.nii = Nifti1Image(_arr_for_nifti1(new_arr), self.affine, self.header) return sel def set_dtype_(self, dtype: type | Literal['smallest_uint', 'smallest_int'] = np.float32, order: Literal["C", "F", "A", "K"] = 'K', casting: Literal["no", "equiv", "safe", "same_kind", "unsafe"] = "unsafe") -> Self: @@ -1077,7 +1090,7 @@ def apply_pad(self, padd: Sequence[tuple[int | None, int | None]] | int | None, return self.copy(nii) - def rescale_and_reorient(self, axcodes_to=None, voxel_spacing=(-1, -1, -1), verbose: logging = True, inplace=False, c_val: float | None = None, mode: MODES = 'nearest') -> Self: + def rescale_and_reorient(self, axcodes_to=None, voxel_spacing=(-1, -1, -1), verbose: logging = True, inplace=False, c_val: float | None = None, mode: MODES = 'nearest', out_dtype: np.dtype | type | str | None = None) -> Self: """Reorients and then rescales the image in a single step. Args: @@ -1089,6 +1102,9 @@ def rescale_and_reorient(self, axcodes_to=None, voxel_spacing=(-1, -1, -1), verb inplace: If True, modifies this NII in place. Defaults to False. c_val: Background fill value for resampling. Defaults to None. mode: Interpolation / boundary mode. Defaults to ``"nearest"``. + out_dtype: Forwarded to :meth:`rescale` – see there for semantics + and caveats. Only affects the rescale step; the preceding + reorient keeps the source dtype. Returns: The reoriented and rescaled NII. @@ -1100,11 +1116,11 @@ def rescale_and_reorient(self, axcodes_to=None, voxel_spacing=(-1, -1, -1), verb axcodes_to = nio.ornt2axcodes(ornt_img) else: curr = self.reorient(axcodes_to=axcodes_to, verbose=verbose, inplace=inplace) - return curr.rescale(voxel_spacing=voxel_spacing, verbose=verbose, inplace=inplace,c_val=c_val,mode=mode) + return curr.rescale(voxel_spacing=voxel_spacing, verbose=verbose, inplace=inplace,c_val=c_val,mode=mode, out_dtype=out_dtype) - def rescale_and_reorient_(self, axcodes_to=None, voxel_spacing=(-1, -1, -1), c_val: float | None = None, mode: MODES = 'nearest', verbose: logging = True) -> Self: + def rescale_and_reorient_(self, axcodes_to=None, voxel_spacing=(-1, -1, -1), c_val: float | None = None, mode: MODES = 'nearest', verbose: logging = True, out_dtype: np.dtype | type | str | None = None) -> Self: """In-place variant of `rescale_and_reorient`.""" - return self.rescale_and_reorient(axcodes_to=axcodes_to,voxel_spacing=voxel_spacing,c_val=c_val,mode=mode,verbose=verbose,inplace=True) + return self.rescale_and_reorient(axcodes_to=axcodes_to,voxel_spacing=voxel_spacing,c_val=c_val,mode=mode,verbose=verbose,inplace=True, out_dtype=out_dtype) def reorient_same_as(self, img_as: Nifti1Image | Self, verbose: logging = False, inplace=False) -> Self: """Reorients this image to match the orientation of another image. @@ -1123,7 +1139,7 @@ def reorient_same_as(self, img_as: Nifti1Image | Self, verbose: logging = False, def reorient_same_as_(self, img_as: Nifti1Image | Self, verbose: logging = False) -> Self: """In-place variant of `reorient_same_as`.""" return self.reorient_same_as(img_as=img_as,verbose=verbose,inplace=True) - def rescale(self, voxel_spacing:float|tuple[float,...]=(1, 1, 1), c_val:float|None=None, verbose:logging=False, inplace=False,mode:MODES='nearest',order: int |None = None,align_corners:bool=False,atol=0.001) -> Self: + def rescale(self, voxel_spacing:float|tuple[float,...]=(1, 1, 1), c_val:float|None=None, verbose:logging=False, inplace=False,mode:MODES='nearest',order: int |None = None,align_corners:bool=False,atol=0.001, out_dtype: np.dtype | type | str | None = None) -> Self: """Rescales the NIfTI image to a new voxel spacing. Args: @@ -1140,6 +1156,11 @@ def rescale(self, voxel_spacing:float|tuple[float,...]=(1, 1, 1), c_val:float|No None, which selects 0 for segmentations and 3 otherwise. align_corners (bool|default): If True or not set and seg==True. Aline corners for scaling. This prevents segmentation mask to shift in a direction. atol: absolute tolerance for skipping if already close in voxel_spacing + out_dtype (dtype-like | None, optional): Requested dtype of the resampled array. + Forwarded to scipy as ``output=``, so the cast happens in-place during interpolation + (no extra full-volume copy). ``None`` keeps the source dtype. With integer targets and + ``order > 0`` the float→int cast is a plain truncation - pass ``order=0`` for label maps. + Returns: NII: A new NII object with the resampled image data. """ @@ -1171,18 +1192,18 @@ def rescale(self, voxel_spacing:float|tuple[float,...]=(1, 1, 1), c_val:float|No new_shp = new_shp + shp[len(new_shp):] new_aff = _rescale_affine(aff, shp, voxel_spacing, new_shp) # type: ignore new_aff[:n, n] = nib.affines.apply_affine(aff, [0 for _ in range(n)])# type: ignore - new_img = _resample_from_to(self, (new_shp, new_aff,voxel_spacing), order=order, mode=mode,align_corners=align_corners) + new_img = _resample_from_to(self, (new_shp, new_aff,voxel_spacing), order=order, mode=mode,align_corners=align_corners, out_dtype=out_dtype) log.print(f"Image resampled from {zms} to voxel size {voxel_spacing}",verbose=verbose) if inplace: self.nii = new_img return self return self.copy(new_img) - def rescale_(self, voxel_spacing=(1, 1, 1), c_val: float | None = None, verbose: logging = False, mode: MODES = 'nearest') -> Self: + def rescale_(self, voxel_spacing=(1, 1, 1), c_val: float | None = None, verbose: logging = False, mode: MODES = 'nearest', out_dtype: np.dtype | type | str | None = None) -> Self: """In-place variant of `rescale`.""" - return self.rescale( voxel_spacing=voxel_spacing, c_val=c_val, verbose=verbose,mode=mode, inplace=True) + return self.rescale( voxel_spacing=voxel_spacing, c_val=c_val, verbose=verbose,mode=mode, inplace=True, out_dtype=out_dtype) - def resample_from_to(self, to_vox_map:Image_Reference|Has_Grid|tuple[SHAPE,AFFINE,ZOOMS], mode:MODES='nearest', order: int |None=None, c_val=None, inplace = False,verbose:logging=True,align_corners:bool=False) -> Self: + def resample_from_to(self, to_vox_map:Image_Reference|Has_Grid|tuple[SHAPE,AFFINE,ZOOMS], mode:MODES='nearest', order: int |None=None, c_val=None, inplace = False,verbose:logging=True,align_corners:bool=False, out_dtype: np.dtype | type | str | None = None) -> Self: """Self will be resampled in coordinate of given other image. Adheres to global space not to local pixel space. Args: @@ -1194,6 +1215,12 @@ def resample_from_to(self, to_vox_map:Image_Reference|Has_Grid|tuple[SHAPE,AFFIN align_corners (bool|default): If True or not set and seg==True. Aline corners for scaling. This prevents segmentation mask to shift in a direction. inplace (bool, optional): Defaults to False. verbose (logging, optional): If True, log resampling shortcuts (skip / reorient-only). Defaults to True. + out_dtype (dtype-like | None, optional): Requested dtype of the resampled array. + Forwarded to scipy as ``output=``, so the cast happens in-place during interpolation + (no extra full-volume copy). ``None`` keeps the source dtype. With integer targets and + ``order > 0`` the float→int cast is a plain truncation - pass ``order=0`` for label maps. + Only applies when the actual resample is executed (the skip / reorient-only / + pad-only shortcuts return the source dtype). Returns: NII: @@ -1245,7 +1272,7 @@ def resample_from_to(self, to_vox_map:Image_Reference|Has_Grid|tuple[SHAPE,AFFIN log.print(f"resample_from_to: {self} to {mapping}",verbose=verbose) if order is None: order = 0 if self.seg else 3 - nii = _resample_from_to(self, mapping,order=order, mode=mode,align_corners=align_corners) + nii = _resample_from_to(self, mapping,order=order, mode=mode,align_corners=align_corners, out_dtype=out_dtype) if inplace: @@ -1253,9 +1280,9 @@ def resample_from_to(self, to_vox_map:Image_Reference|Has_Grid|tuple[SHAPE,AFFIN return self else: return self.copy(nii) - def resample_from_to_(self, to_vox_map: Image_Reference | Has_Grid | tuple[SHAPE, AFFINE, ZOOMS], mode: MODES = 'nearest', c_val: float | None = None, verbose: logging = True, aline_corners=False) -> Self: + def resample_from_to_(self, to_vox_map: Image_Reference | Has_Grid | tuple[SHAPE, AFFINE, ZOOMS], mode: MODES = 'nearest', c_val: float | None = None, verbose: logging = True, aline_corners=False, out_dtype: np.dtype | type | str | None = None) -> Self: """In-place variant of `resample_from_to`.""" - return self.resample_from_to(to_vox_map,mode=mode,c_val=c_val,inplace=True,verbose=verbose,align_corners=aline_corners) + return self.resample_from_to(to_vox_map,mode=mode,c_val=c_val,inplace=True,verbose=verbose,align_corners=aline_corners, out_dtype=out_dtype) @property def is_empty(self) -> bool: @@ -2435,10 +2462,13 @@ def save(self, file: str | Path, make_parents=True, verbose: logging = True, dty # `save` is a query and must not mutate `self`. arr = arr.astype(_smallest_int_dtype(arr, unsigned=True)) - self.header.set_data_dtype(arr.dtype) - out = Nifti1Image(arr, self.affine,self.header)#,dtype=arr.dtype) + # NIfTI-1 has no float16 (or bool) datatype. Upcast just the on-disk + # copy; the caller-owned ``self._arr`` keeps its native dtype. + safe_arr = _arr_for_nifti1(arr) + self.header.set_data_dtype(_nifti_safe_dtype(safe_arr.dtype)) + out = Nifti1Image(safe_arr, self.affine, self.header) # ,dtype=arr.dtype) if dtype is not None: - out.set_data_dtype(dtype) + out.set_data_dtype(_nifti_safe_dtype(dtype)) if out.header["qform_code"] == 0: #NIFTI_XFORM_UNKNOWN Will cause an error for some rounding of the affine in ITKSnap ... # 1 means Scanner coordinate system # 2 means align (to something) coordinate system diff --git a/TPTBox/segmentation/VibeSeg/inference_nnunet.py b/TPTBox/segmentation/VibeSeg/inference_nnunet.py index d7a94ef..0ea9988 100644 --- a/TPTBox/segmentation/VibeSeg/inference_nnunet.py +++ b/TPTBox/segmentation/VibeSeg/inference_nnunet.py @@ -345,13 +345,15 @@ def run_inference_on_file( logger.print("orientation", orientation, f"from {input_nii[0].orientation}") if verbose else None input_nii = [i.reorient(orientation) for i in input_nii] + logger.print("squash to fit float16") if verbose else None + input_nii = [squash_so_it_fits_in_float16(i) for i in input_nii] + if zoom is not None: logger.print("rescale", f"{zoom=} from {input_nii[0].zoom}") if verbose else None - input_nii = [i.rescale_(zoom, mode=mode, verbose=True) for i in input_nii] + input_nii = [i.rescale_(zoom, mode=mode, verbose=True, out_dtype=np.float16) for i in input_nii] logger.print(input_nii) - logger.print("squash to float16") if verbose else None - - input_nii = [squash_so_it_fits_in_float16(i) for i in input_nii] + else: + input_nii = [i.set_dtype(np.float16) for i in input_nii] if crop: crop = input_nii[0].compute_crop(minimum=20)