From edd6ecfc2c710473377fd3d5ad189b58bcd52d52 Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 5 Aug 2026 06:30:08 +0000 Subject: [PATCH 1/8] Fix core NII/np_utils methods broken on their default paths Four core operations failed or returned wrong results on ordinary inputs: - NII.rescale(scalar) built a generator, so the following len() raised TypeError. Every scalar-spacing call was broken. - NII.threshold() set arr[arr2>=t]=1 then arr[arr2<=t]=0, so the second write erased the == case; thresholding a binary mask at 1 returned an all-zero image. Rewritten as a single comparison, which also drops a redundant full-volume copy. - NII.normalize() divided by max_out instead of scaling into the target range and offset by min_out/max_out, so normalize(0,255) produced values in [0,1] and tripped its own assert (and max_out=0 raised ZeroDivisionError). Now scales properly, handles constant images, and uses np.isclose instead of exact float equality. - np_dilate_msk(mask=..., use_crop=True) cropped `mask` with the global crop but indexed it against the per-label crop, raising IndexError for any multi-label segmentation - i.e. NII.dilate_msk(mask=...) on its default path. - np_dilate_msk_euclid(labels=...) applied the label filter only when use_crop=False; the default path dilated every label. Both dilate helpers also stopped binarising the caller's mask array in place. Additionally corrects three early-return branches that had `inplace` inverted (nii_wrapper.py:811, 1127, 1140, 2700), so an in-place call no longer returns a different object and an out-of-place call no longer returns an alias of self. Co-Authored-By: Claude Opus 5 --- TPTBox/core/nii_wrapper.py | 10 +-- TPTBox/core/nii_wrapper_math.py | 24 ++--- TPTBox/core/np_utils.py | 15 ++-- unit_tests/test_regressions.py | 151 ++++++++++++++++++++++++++++++++ 4 files changed, 177 insertions(+), 23 deletions(-) create mode 100644 unit_tests/test_regressions.py diff --git a/TPTBox/core/nii_wrapper.py b/TPTBox/core/nii_wrapper.py index 3dbe1a1e..235a8161 100755 --- a/TPTBox/core/nii_wrapper.py +++ b/TPTBox/core/nii_wrapper.py @@ -808,7 +808,7 @@ def reorient(self:Self, axcodes_to: AX_CODES|str|None = ("P", "I", "R"), verbose new_img = arr, new_aff,self.header log.print("Image reoriented from", nio.ornt2axcodes(ornt_fr), "to", axcodes_to,verbose=verbose) else: - return self if not inplace else self.copy() + return self if inplace else self.copy() if inplace: self.nii = new_img return self @@ -1118,13 +1118,13 @@ def rescale(self, voxel_spacing:float|tuple[float,...]=(1, 1, 1), c_val:float|No NII: A new NII object with the resampled image data. """ if isinstance(voxel_spacing, (int,float)): - voxel_spacing =(voxel_spacing for _ in range(min(3,self.affine.shape[0]-1))) + voxel_spacing =tuple(voxel_spacing for _ in range(min(3,self.affine.shape[0]-1))) n = self.dims while n> len(voxel_spacing): voxel_spacing = (*voxel_spacing, -1) if all(a in (-1, b) for a,b in zip(voxel_spacing, self.zoom)): log.print(f"Image already resampled to voxel size {self.zoom}",verbose=verbose) - return self.copy() if inplace else self + return self if inplace else self.copy() c_val = self.get_c_val(c_val) # resample to new voxel spacing based on the current x-y-z-orientation @@ -1137,7 +1137,7 @@ def rescale(self, voxel_spacing:float|tuple[float,...]=(1, 1, 1), c_val:float|No voxel_spacing = tuple([v if v != -1 else z for v,z in zip_strict(voxel_spacing,zms)]) if np.isclose(voxel_spacing, self.zoom,atol=atol).all(): log.print(f"Image already resampled to voxel size {self.zoom}",verbose=verbose) - return self.copy() if inplace else self + return self if inplace else self.copy() # Calculate new shape new_shp = tuple(np.rint([shp[i] * zms[i] / voxel_spacing[i] for i in range(len(voxel_spacing))]).astype(int)) @@ -2697,7 +2697,7 @@ def extract_label(self,label:int|Enum|Sequence[int]|Sequence[Enum]|None, keep_la assert self.seg, "extracting a label only makes sense for a segmentation mask" if label is None: if keep_label: - return self.copy() if inplace else self + return self if inplace else self.copy() else: return self.clamp(0,1,inplace=inplace) seg_arr = self.get_seg_array() diff --git a/TPTBox/core/nii_wrapper_math.py b/TPTBox/core/nii_wrapper_math.py index c695ead6..08c54ea4 100755 --- a/TPTBox/core/nii_wrapper_math.py +++ b/TPTBox/core/nii_wrapper_math.py @@ -286,12 +286,19 @@ def normalize(self,min_out = 0, max_out = 1, quantile = 1., clamp_lower:float|No """ arr = self.get_array() max_v = np.quantile(arr[arr>0],q=quantile) - arr = self.clamp(clamp_lower,max_v,inplace=inplace) - arr -= arr.min() - min_out/max_out - arr /= arr.max() *max_out - assert arr.max() == max_out, f"{arr.max()} == {max_out}" - assert arr.min() == min_out - return self.set_array(arr.get_array(),inplace) + clamped = self.clamp(clamp_lower,max_v,inplace=inplace).get_array() + if not np.issubdtype(clamped.dtype, np.floating): + clamped = clamped.astype(np.float32) + mi = clamped.min() + ma = clamped.max() + if ma == mi: + # Degenerate (constant) image: nothing to spread out, map onto the lower bound. + clamped = np.full_like(clamped, min_out) + else: + clamped = (clamped - mi) / (ma - mi) * (max_out - min_out) + min_out + assert np.isclose(clamped.max(), max_out), f"{clamped.max()} == {max_out}" + assert np.isclose(clamped.min(), min_out), f"{clamped.min()} == {min_out}" + return self.set_array(clamped,inplace) def normalize_(self,min_out = 0, max_out = 1, quantile = 1., clamp_lower:float|None=None)->Self: """In-place variant of `normalize`.""" return self.normalize(min_out = min_out, max_out = max_out, quantile = quantile, clamp_lower=clamp_lower,inplace=True) @@ -413,10 +420,7 @@ def threshold(self,threshold=0.5, inplace=False)->Self: Returns: Self: Binarised segmentation instance. """ - arr = self.get_array() - arr2 = arr.copy() - arr[arr2>=threshold] = 1 - arr[arr2<=threshold] = 0 + arr = (self.get_array() >= threshold).astype(np.uint8) nii = self if inplace else self.copy() nii.seg = True nii:NII = nii.set_array(arr,inplace,verbose=False) diff --git a/TPTBox/core/np_utils.py b/TPTBox/core/np_utils.py index 03270ce8..0ee2abcd 100755 --- a/TPTBox/core/np_utils.py +++ b/TPTBox/core/np_utils.py @@ -468,14 +468,12 @@ def np_dilate_msk_euclid(arr: np.ndarray, n_pixel: int = 3, use_crop=True, label if use_crop: crop = np_bbox_binary(arr_bin, px_dist=1 + n_pixel, raise_error=False) - arrc = arr[crop] + # use the label-filtered array so `labels` is honoured here too, not just on the no-crop path + arrc = arr_bin[crop] else: - arrc = arr - if labels is not None: - arrc = arrc.copy() - arrc[np_isin(arr_bin, labels, invert=True)] = 0 + arrc = arr_bin if mask is not None: - mask[mask != 0] = 1 + mask = mask != 0 # do not mutate the caller's array if use_crop: mask = mask[crop] foreground = arrc > 0 @@ -544,7 +542,7 @@ def np_dilate_msk( arrc = arr if mask is not None: - mask[mask != 0] = 1 + mask = mask != 0 # do not mutate the caller's array if use_crop: mask = mask[crop] if ignore_axis is None: @@ -568,7 +566,8 @@ def np_dilate_msk( oc = out[lcrop] == 0 out[lcrop][oc] = msk_ibe_data[oc] * i if mask is not None: - out[lcrop][mask == 0] = 0 + # `mask` follows the global crop; index it with the per-label crop to match `out[lcrop]` + out[lcrop][mask[lcrop] == 0] = 0 else: out[out == 0] = msk_ibe_data[out == 0] * i if mask is not None: diff --git a/unit_tests/test_regressions.py b/unit_tests/test_regressions.py new file mode 100644 index 00000000..661f7f64 --- /dev/null +++ b/unit_tests/test_regressions.py @@ -0,0 +1,151 @@ +"""Regression tests for bugs found in the repository audit. + +Each test pins down behaviour that was previously broken; the comment on the class +names the symptom that used to occur. +""" + +from __future__ import annotations + +import unittest + +import numpy as np + +from TPTBox import NII +from TPTBox.core.np_utils import np_dilate_msk_euclid + + +def _make_nii(arr: np.ndarray, seg: bool = True, zoom=(1.0, 1.0, 1.0)) -> NII: + affine = np.diag([*zoom, 1.0]) + return NII.from_numpy(arr, affine=affine, seg=seg) + + +def _cube(shape=(10, 10, 10)) -> np.ndarray: + arr = np.zeros(shape, dtype=np.uint8) + arr[3:7, 3:7, 3:7] = 1 + return arr + + +class Test_Rescale_Scalar(unittest.TestCase): + """`rescale(1.5)` used to build a generator and raise TypeError on len().""" + + def test_scalar_voxel_spacing(self): + nii = _make_nii(_cube()) + out = nii.rescale(1.5) + self.assertEqual(len(out.shape), 3) + + def test_scalar_matches_tuple(self): + nii = _make_nii(_cube()) + self.assertEqual(nii.rescale(1.5).shape, nii.rescale((1.5, 1.5, 1.5)).shape) + + +class Test_Threshold(unittest.TestCase): + """`threshold` used to zero out voxels exactly equal to the threshold.""" + + def test_value_equal_to_threshold_is_kept(self): + nii = _make_nii(_cube()) + self.assertEqual(list(nii.threshold(1).unique()), [1]) + + def test_below_threshold_is_dropped(self): + arr = np.zeros((5, 5, 5), dtype=np.uint8) + arr[0, 0, 0] = 1 + arr[1, 1, 1] = 5 + out = _make_nii(arr).threshold(5).get_seg_array() + self.assertEqual(out[1, 1, 1], 1) + self.assertEqual(out[0, 0, 0], 0) + + +class Test_Normalize(unittest.TestCase): + """`normalize` divided by max_out instead of scaling into [min_out, max_out].""" + + def test_range_0_255(self): + nii = _make_nii(_cube(), seg=False) + out = nii.normalize(0, 255) + self.assertAlmostEqual(float(out.min()), 0.0, places=4) + self.assertAlmostEqual(float(out.max()), 255.0, places=3) + + def test_range_1_2(self): + nii = _make_nii(_cube(), seg=False) + out = nii.normalize(1, 2) + self.assertAlmostEqual(float(out.min()), 1.0, places=4) + self.assertAlmostEqual(float(out.max()), 2.0, places=4) + + def test_constant_image_does_not_raise(self): + arr = np.full((5, 5, 5), 3.0, dtype=np.float32) + _make_nii(arr, seg=False).normalize(0, 1) + + +class Test_Dilate_Msk_Mask(unittest.TestCase): + """`dilate_msk(mask=...)` raised IndexError for multi-label segmentations.""" + + @staticmethod + def _two_labels() -> NII: + arr = np.zeros((30, 30, 30), dtype=np.uint8) + arr[2:5, 2:5, 2:5] = 1 + arr[22:26, 22:26, 22:26] = 2 + return _make_nii(arr) + + def test_multilabel_with_mask_does_not_raise(self): + nii = self._two_labels() + full = _make_nii(np.ones((30, 30, 30), dtype=np.uint8)) + self.assertEqual(list(nii.dilate_msk(n_pixel=2, mask=full, verbose=False).unique()), [1, 2]) + + def test_mask_restricts_output(self): + nii = self._two_labels() + half = np.zeros((30, 30, 30), dtype=np.uint8) + half[:15] = 1 + out = nii.dilate_msk(n_pixel=2, mask=_make_nii(half), verbose=False).get_seg_array() + self.assertEqual(int((out[15:] != 0).sum()), 0) + + def test_caller_mask_is_not_mutated(self): + nii = self._two_labels() + mask_arr = np.zeros((30, 30, 30), dtype=np.uint8) + mask_arr[:15] = 7 + before = mask_arr.copy() + nii.dilate_msk(n_pixel=2, mask=_make_nii(mask_arr), verbose=False) + np.testing.assert_array_equal(mask_arr, before) + + +class Test_Dilate_Euclid_Labels(unittest.TestCase): + """`np_dilate_msk_euclid(labels=...)` ignored `labels` when use_crop=True.""" + + @staticmethod + def _adjacent() -> np.ndarray: + arr = np.zeros((12, 12, 12), dtype=np.uint8) + arr[5, 5, 5] = 1 + arr[5, 5, 7] = 2 + return arr + + def test_crop_and_nocrop_agree(self): + arr = self._adjacent() + with_crop = np_dilate_msk_euclid(arr.copy(), n_pixel=2, labels=[1], use_crop=True) + without_crop = np_dilate_msk_euclid(arr.copy(), n_pixel=2, labels=[1], use_crop=False) + np.testing.assert_array_equal(with_crop, without_crop) + + def test_unselected_label_is_not_dilated(self): + arr = self._adjacent() + out = np_dilate_msk_euclid(arr.copy(), n_pixel=2, labels=[1], use_crop=True) + self.assertEqual(int((out == 2).sum()), 0) + + +class Test_Inplace_Contract(unittest.TestCase): + """Early-return branches had `inplace` inverted, leaking/copying the wrong object.""" + + def test_rescale_inplace_returns_self(self): + nii = _make_nii(_cube()) + self.assertIs(nii.rescale_((1, 1, 1)), nii) + + def test_rescale_out_of_place_returns_copy(self): + nii = _make_nii(_cube()) + self.assertIsNot(nii.rescale((1, 1, 1)), nii) + + def test_reorient_out_of_place_returns_copy(self): + nii = _make_nii(_cube()) + self.assertIsNot(nii.reorient(nii.orientation), nii) + + def test_extract_label_none_keep_label_returns_copy(self): + nii = _make_nii(_cube()) + self.assertIsNot(nii.extract_label(None, keep_label=True), nii) + + +if __name__ == "__main__": + unittest.main() From 49f4dfc8f17206a316f572243052774167e7be2d Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 5 Aug 2026 06:33:29 +0000 Subject: [PATCH 2/8] Fix dtype selection so segmentations are not silently oversized or wrapped Three float->int guards used isinstance() against numpy scalar types: isinstance(self.dtype, np.floating) # a np.dtype object, never a scalar isinstance(arr, np.floating) # an ndarray, never a scalar Both are always False, so NII.__init__, NII.set_array and NII.save never downcast float segmentations. A float64 segmentation therefore stayed float64 for its whole lifetime - 8x the memory of the uint8 it should be, through every get_seg_array() copy and onto disk - and it defeated the unsigned-int fast paths in np_utils. Replaced with np.issubdtype(...). set_dtype("smallest_int"/"smallest_uint") chose the target type from arr.max() alone. Since the cast uses casting="unsafe", a negative minimum wrapped silently (-1 -> 255), and "smallest_uint" fell back to the *signed* np.int32. Dtype selection is now factored into _smallest_int_dtype(), which considers both bounds, asserts non-negativity for unsigned targets, and extends to 32/64-bit. np_map_labels built its lookup table with dtype=arr.dtype, so any mapping target outside the input dtype's range wrapped without warning - on a uint8 mask, 1 -> 300 produced 44 and 1 -> -5 produced 251. The table dtype is now derived from the mapping targets as well; in-range mappings keep the input dtype as before. NII.save no longer mutates self via set_dtype_ to achieve its cast, and no longer re-copies the whole volume afterwards. Co-Authored-By: Claude Opus 5 --- TPTBox/core/nii_wrapper.py | 56 ++++++++++++++++++----------- TPTBox/core/np_utils.py | 5 ++- unit_tests/test_regressions.py | 64 ++++++++++++++++++++++++++++++++++ 3 files changed, 103 insertions(+), 22 deletions(-) diff --git a/TPTBox/core/nii_wrapper.py b/TPTBox/core/nii_wrapper.py index 235a8161..8a4c8207 100755 --- a/TPTBox/core/nii_wrapper.py +++ b/TPTBox/core/nii_wrapper.py @@ -99,6 +99,33 @@ def formatwarning_tb(*args, **kwargs) -> str: _dtype_non_u = {"int8", "int16"} +def _smallest_int_dtype(arr: np.ndarray, unsigned: bool) -> type: + """Smallest integer dtype that can represent every value in ``arr``. + + Both bounds are considered: picking on ``max()`` alone silently wraps negative + values, because the casts here use ``casting="unsafe"``. + + Args: + arr: Array whose value range determines the dtype. + unsigned: If True, choose an unsigned type (requires ``arr.min() >= 0``). + + Returns: + The selected numpy dtype. + """ + mi = arr.min() + ma = arr.max() + if unsigned: + assert mi >= 0, f"an unsigned dtype requires non-negative values, but the minimum is {mi}" + for cand, limit in ((np.uint8, 256), (np.uint16, 65536), (np.uint32, 2**32)): + if ma < limit: + return cand + return np.uint64 + for cand, limit in ((np.int8, 128), (np.int16, 32768), (np.int32, 2**31)): + if ma < limit and mi >= -limit: + return cand + return np.int64 + + def _check_if_nifty_is_lying_about_its_dtype(self: NII): """Infers the correct dtype by inspecting the actual value range of the NIfTI dataobj.""" change_dtype = False @@ -240,7 +267,7 @@ def __init__(self, nii: Nifti1Image|_unpacked_nii, seg=False,c_val=None, desc:st self.set_description(desc) if seg: self._unpack() - if isinstance(self.dtype,np.floating): + if np.issubdtype(self.dtype,np.floating): self.set_dtype_("smallest_uint") @@ -672,7 +699,7 @@ def set_array(self, arr: np.ndarray | Self, inplace=False, verbose: logging = Fa arr = arr.astype(np.uint8) if arr.dtype == np.float16: arr = arr.astype(np.float32) - if self.seg and isinstance(arr, (np.floating, float)): + 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()) @@ -712,22 +739,8 @@ def set_dtype(self, dtype: type | Literal['smallest_int', 'smallest_uint'] = np. The NII with the new dtype (``self`` when ``inplace=True``, a new NII otherwise). """ sel = self if inplace else self.copy() - if dtype == "smallest_uint": - arr = self.get_array() - if arr.max()<256: - dtype = np.uint8 - elif arr.max()<65536: - dtype = np.uint16 - else: - dtype = np.int32 - elif dtype == "smallest_int": - arr = self.get_array() - if arr.max()<128: - dtype = np.int8 - elif arr.max()<32768: - dtype = np.int16 - else: - dtype = np.int32 + if dtype in ("smallest_uint", "smallest_int"): + dtype = _smallest_int_dtype(self.get_array(), unsigned=dtype == "smallest_uint") if self.__unpacked: self._unpack() sel._arr = sel._arr.astype(dtype) @@ -2361,9 +2374,10 @@ def save(self, file: str | Path, make_parents=True, verbose: logging = True, dty return self.save_nrrd(file,verbose=verbose) arr = self.get_array() if not self.seg else self.get_seg_array() - if isinstance(arr,np.floating) and self.seg: - self.set_dtype_("smallest_uint") - arr = self.get_array() if not self.seg else self.get_seg_array() + if self.seg and np.issubdtype(arr.dtype, np.floating): + # A segmentation must never be written out as float. Cast the local array: + # `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) diff --git a/TPTBox/core/np_utils.py b/TPTBox/core/np_utils.py index 0ee2abcd..b2f4d753 100755 --- a/TPTBox/core/np_utils.py +++ b/TPTBox/core/np_utils.py @@ -663,7 +663,10 @@ def np_map_labels(arr: UINTARRAY, label_map: LABEL_MAP) -> np.ndarray: max_value = max(arr.max(), *k, *v) + 1 - mapping_ar = np.arange(max_value, dtype=arr.dtype) + # 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). + lut_dtype = np.result_type(arr.dtype, np.min_scalar_type(int(v.max())), np.min_scalar_type(int(v.min()))) + mapping_ar = np.arange(max_value, dtype=lut_dtype) mapping_ar[k] = v return mapping_ar[arr] diff --git a/unit_tests/test_regressions.py b/unit_tests/test_regressions.py index 661f7f64..78d2e4cc 100644 --- a/unit_tests/test_regressions.py +++ b/unit_tests/test_regressions.py @@ -147,5 +147,69 @@ def test_extract_label_none_keep_label_returns_copy(self): self.assertIsNot(nii.extract_label(None, keep_label=True), nii) +class Test_Float_Seg_Downcast(unittest.TestCase): + """`isinstance(dtype, np.floating)` is never True, so float segs stayed float (8x memory).""" + + def test_float64_seg_is_downcast_on_construction(self): + arr = np.zeros((8, 8, 8), dtype=np.float64) + arr[2:5, 2:5, 2:5] = 1 + self.assertEqual(_make_nii(arr).dtype, np.uint8) + + def test_float_seg_keeps_its_labels(self): + arr = np.zeros((8, 8, 8), dtype=np.float32) + arr[2:5, 2:5, 2:5] = 7 + self.assertEqual(list(_make_nii(arr).unique()), [7]) + + def test_non_seg_float_is_left_alone(self): + arr = np.zeros((8, 8, 8), dtype=np.float32) + self.assertEqual(_make_nii(arr, seg=False).dtype, np.float32) + + +class Test_Map_Labels_Dtype(unittest.TestCase): + """np_map_labels built its lookup table in the input dtype, wrapping out-of-range targets.""" + + @staticmethod + def _arr() -> np.ndarray: + arr = np.zeros((8, 8, 8), dtype=np.uint8) + arr[2:5, 2:5, 2:5] = 1 + return arr + + def test_target_above_input_dtype_range(self): + from TPTBox.core.np_utils import np_map_labels + + self.assertIn(300, np_map_labels(self._arr(), {1: 300})) + + def test_negative_target(self): + from TPTBox.core.np_utils import np_map_labels + + self.assertIn(-5, np_map_labels(self._arr(), {1: -5})) + + def test_in_range_target_keeps_dtype(self): + from TPTBox.core.np_utils import np_map_labels + + self.assertEqual(np_map_labels(self._arr(), {1: 2}).dtype, np.uint8) + + +class Test_Smallest_Int_Dtype(unittest.TestCase): + """set_dtype('smallest_int') picked from max() only, wrapping negative values.""" + + def test_negative_range_uses_wide_enough_type(self): + arr = np.full((4, 4, 4), -200, dtype=np.int32) + arr[0, 0, 0] = 100 + nii = _make_nii(arr, seg=False) + self.assertEqual(nii.set_dtype("smallest_int").dtype, np.int16) + + def test_values_survive_the_cast(self): + arr = np.full((4, 4, 4), -200, dtype=np.int32) + arr[0, 0, 0] = 100 + out = _make_nii(arr, seg=False).set_dtype("smallest_int").get_array() + self.assertEqual(sorted(set(out.ravel().tolist())), [-200, 100]) + + def test_smallest_uint_rejects_negatives(self): + arr = np.full((4, 4, 4), -1, dtype=np.int32) + with self.assertRaises(AssertionError): + _make_nii(arr, seg=False).set_dtype("smallest_uint") + + if __name__ == "__main__": unittest.main() From b11343e876c8a0af5481b76530206bfd3f4aa806 Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 5 Aug 2026 06:37:00 +0000 Subject: [PATCH 3/8] Fix np_bbox_binary bounds, np_unique return types, and empty-mask border check np_bbox_binary clamped the stop index to the array shape and only then added 1, so a bounding box touching the far border produced stop == shape + 1. Slicing tolerates that, but every consumer computing stop - start (np_center_of_bbox_binary, NII.compute_crop, is_segmentation_in_border) saw a size one voxel too large. It also stored px_dist in a uint8 array, so any px_dist > 255 raised OverflowError under NumPy 2. Interior boxes are unaffected. np_unique / np_unique_withoutzero are annotated -> list[int] but returned three different scalar types depending on the input dtype and which of the four code paths was taken: Python int from bincount, np.int64 from the withoutzero fast path, and np.float32/np.int16 from the np.unique fallbacks. The numpy scalars are not JSON-serializable, so serializing label lists failed for non-uint inputs. All paths now go through .tolist(), which yields native scalars without truncating float values. is_segmentation_in_border() guarded on `slices is None`, but compute_crop(raise_error=False) returns full-extent slices for an empty mask and never None, so an empty segmentation was reported as touching the border. Now short-circuits on the existing NII.is_empty property. Co-Authored-By: Claude Opus 5 --- TPTBox/core/nii_wrapper.py | 5 +++ TPTBox/core/np_utils.py | 13 ++++--- unit_tests/test_regressions.py | 70 ++++++++++++++++++++++++++++++++++ 3 files changed, 83 insertions(+), 5 deletions(-) diff --git a/TPTBox/core/nii_wrapper.py b/TPTBox/core/nii_wrapper.py index 8a4c8207..9dc4ca49 100755 --- a/TPTBox/core/nii_wrapper.py +++ b/TPTBox/core/nii_wrapper.py @@ -2072,6 +2072,11 @@ def is_segmentation_in_border(self,minimum=0,voxel_tolerance: int = 2,use_mm: bo - bool: True if the segmentation is within the defined tolerance of the border, False otherwise. """ + # compute_crop(raise_error=False) returns full-extent slices for an empty mask rather + # than None, so an explicit emptiness check is needed - otherwise "nothing segmented" + # is reported as "touching the border". + if self.is_empty: + return False slices = self.compute_crop(minimum,dist=0,use_mm=use_mm,raise_error=False) if slices is None: return False diff --git a/TPTBox/core/np_utils.py b/TPTBox/core/np_utils.py index b2f4d753..c7dde01b 100755 --- a/TPTBox/core/np_utils.py +++ b/TPTBox/core/np_utils.py @@ -214,7 +214,9 @@ def old_np_unique(arr: np.ndarray) -> list[int]: return [idx for idx, i in enumerate(cc3dstatistics(arr)["voxel_counts"]) if i > 0] except Exception: pass - return list(np.unique(arr)) + # .tolist() yields native Python scalars; list() would leak numpy scalars, which are + # not JSON-serializable and differ from the ints the cc3d/bincount paths return. + return np.unique(arr).tolist() def np_unique(arr: np.ndarray) -> list[int]: @@ -256,8 +258,8 @@ def np_unique_withoutzero(arr: UINTARRAY) -> list[int]: return [] if max_val < 2**20: counts = np.bincount(arr.ravel()) - return list(np.where(counts[1:] > 0)[0] + 1) - return [i for i in np.unique(arr) if i != 0] + return (np.where(counts[1:] > 0)[0] + 1).tolist() + return [i for i in np.unique(arr).tolist() if i != 0] def old_np_unique_withoutzero(arr: UINTARRAY) -> list[int]: @@ -737,7 +739,7 @@ def np_bbox_binary(img: np.ndarray, px_dist: int | Sequence[int] | np.ndarray = n = img.ndim shp = img.shape if isinstance(px_dist, int): - px_dist = np.ones(n, dtype=np.uint8) * px_dist + 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}" bbox: list[float] = [] @@ -754,7 +756,8 @@ def np_bbox_binary(img: np.ndarray, px_dist: int | Sequence[int] | np.ndarray = out: tuple[slice, ...] = tuple( slice( max(bbox[i] - px_dist[i // 2], 0), - min(bbox[i + 1] + px_dist[i // 2], shp[i // 2]) + 1, + # clamp AFTER the +1, otherwise a bbox touching the far border yields stop == shape + 1 + min(bbox[i + 1] + px_dist[i // 2] + 1, shp[i // 2]), ) for i in range(0, len(bbox), 2) ) diff --git a/unit_tests/test_regressions.py b/unit_tests/test_regressions.py index 78d2e4cc..099b6ed2 100644 --- a/unit_tests/test_regressions.py +++ b/unit_tests/test_regressions.py @@ -211,5 +211,75 @@ def test_smallest_uint_rejects_negatives(self): _make_nii(arr, seg=False).set_dtype("smallest_uint") +class Test_Bbox_Binary(unittest.TestCase): + """np_bbox_binary clamped before adding 1 and stored px_dist as uint8.""" + + @staticmethod + def _touching_border() -> np.ndarray: + arr = np.zeros((20, 20, 20), dtype=np.uint8) + arr[:, 5, 5] = 1 + return arr + + def test_stop_never_exceeds_shape(self): + from TPTBox.core.np_utils import np_bbox_binary + + for sl, dim in zip(np_bbox_binary(self._touching_border(), px_dist=2), (20, 20, 20)): + self.assertLessEqual(sl.stop, dim) + + def test_large_px_dist_does_not_overflow(self): + from TPTBox.core.np_utils import np_bbox_binary + + self.assertEqual(np_bbox_binary(self._touching_border(), px_dist=300)[0], slice(0, 20)) + + def test_interior_bbox_is_unchanged(self): + from TPTBox.core.np_utils import np_bbox_binary + + arr = np.zeros((20, 20, 20), dtype=np.uint8) + arr[5:8, 5:8, 5:8] = 1 + self.assertEqual(np_bbox_binary(arr, px_dist=0)[0], slice(5, 8)) + + +class Test_Unique_Return_Types(unittest.TestCase): + """np_unique returned numpy scalars on the fallback paths, breaking json.dumps.""" + + def test_all_dtypes_return_native_scalars(self): + import json + + from TPTBox.core.np_utils import np_unique, np_unique_withoutzero + + for dtype in (np.uint8, np.int16, np.float32): + arr = np.zeros((8, 8, 8), dtype=dtype) + arr[0, 0, 0] = 3 + for values in (np_unique(arr), np_unique_withoutzero(arr)): + json.dumps(values) # would raise for numpy scalars + for value in values: + self.assertIn(type(value), (int, float), f"{dtype} produced {type(value)}") + + def test_values_are_correct(self): + from TPTBox.core.np_utils import np_unique, np_unique_withoutzero + + arr = np.zeros((8, 8, 8), dtype=np.int16) + arr[0, 0, 0] = 3 + self.assertEqual(np_unique(arr), [0, 3]) + self.assertEqual(np_unique_withoutzero(arr), [3]) + + +class Test_Segmentation_In_Border(unittest.TestCase): + """An empty mask was reported as touching the border.""" + + def test_empty_is_not_in_border(self): + self.assertFalse(_make_nii(np.zeros((10, 10, 10), dtype=np.uint8)).is_segmentation_in_border()) + + def test_touching_border_is_detected(self): + arr = np.zeros((10, 10, 10), dtype=np.uint8) + arr[0, 5, 5] = 1 + self.assertTrue(_make_nii(arr).is_segmentation_in_border()) + + def test_centred_is_not_in_border(self): + arr = np.zeros((20, 20, 20), dtype=np.uint8) + arr[9:11, 9:11, 9:11] = 1 + self.assertFalse(_make_nii(arr).is_segmentation_in_border()) + + if __name__ == "__main__": unittest.main() From ffc42f43010215d23afb37bc8d6132ef0e257654 Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 5 Aug 2026 06:41:09 +0000 Subject: [PATCH 4/8] Stop two caches from leaking state across the whole process sag_cor_curve_projection did: order = v_idx_order order += [i for i in range(256) if i not in v_idx_order] `order` aliases the module-level v_idx_order list (vert_constants.py), which is re-exported as TPTBox.v_idx_order, and `+=` extends a list in place. So rendering a single snapshot permanently grew the shared global from 105 to 256 entries, affecting anything else that reads it. `order` was never used afterwards, so both lines are dead - removed, along with the now-unused import. POI._vert_orientation_pir was a bare class attribute, i.e. one dict shared by every POI instance in the process. get_vert_direction_PIR compounded this by writing the cache onto a temporary extract_subregion() copy, so a lookup could return vertebra directions computed for a different subject processed earlier. It is now a per-instance dataclass field (default_factory=dict, excluded from repr/compare, so a copy still starts empty as documented), and the cache is written to the object the function was called with. Co-Authored-By: Claude Opus 5 --- TPTBox/core/poi.py | 5 ++- TPTBox/core/poi_fun/vertebra_direction.py | 3 +- TPTBox/spine/snapshot2D/snapshot_modular.py | 3 -- unit_tests/test_regressions.py | 41 +++++++++++++++++++++ 4 files changed, 47 insertions(+), 5 deletions(-) diff --git a/TPTBox/core/poi.py b/TPTBox/core/poi.py index 41cb0e08..d8ebf048 100755 --- a/TPTBox/core/poi.py +++ b/TPTBox/core/poi.py @@ -122,7 +122,10 @@ class POI(Abstract_POI, Has_Grid): # internal _rotation: ROTATION = field(init=False, default=None, repr=False, compare=False) # type: ignore _zoom: ZOOMS = field(init=False, default=(1, 1, 1), repr=False, compare=False) - _vert_orientation_pir = {} # Elusive; will not be saved; will not be copied. For Buffering results # noqa: RUF012 + # Elusive; will not be saved; will not be copied. For buffering results. + # Must be a per-instance field: as a bare class attribute it was shared by every POI + # object in the process, so cached directions leaked between subjects. + _vert_orientation_pir: dict = field(init=False, default_factory=dict, repr=False, compare=False) def _set_inplace(self, poi: Self) -> Self: """Copy all grid/affine attributes and centroids from ``poi`` into ``self``.""" diff --git a/TPTBox/core/poi_fun/vertebra_direction.py b/TPTBox/core/poi_fun/vertebra_direction.py index 9fc0ec21..ce473123 100644 --- a/TPTBox/core/poi_fun/vertebra_direction.py +++ b/TPTBox/core/poi_fun/vertebra_direction.py @@ -266,6 +266,7 @@ def get_vert_direction_PIR(poi: POI, vert_id: int, do_norm: bool = True, to_pir: """ if vert_id in poi._vert_orientation_pir and to_pir: return poi._vert_orientation_pir[vert_id] # Elusive buffer of iso/PIR directions. + cache_owner = poi # `poi` is rebound below; the cache belongs on the object we were called with poi = poi.extract_subregion( Location.Vertebra_Corpus, Location.Vertebra_Direction_Posterior, @@ -284,7 +285,7 @@ def n(x): right = np.array(poi[vert_id : Location.Vertebra_Direction_Right]) out = n(post - center), n(down - center), n(right - center) if to_pir: - poi._vert_orientation_pir[vert_id] = out + cache_owner._vert_orientation_pir[vert_id] = out return out diff --git a/TPTBox/spine/snapshot2D/snapshot_modular.py b/TPTBox/spine/snapshot2D/snapshot_modular.py index 7c31f066..b9e11dfb 100755 --- a/TPTBox/spine/snapshot2D/snapshot_modular.py +++ b/TPTBox/spine/snapshot2D/snapshot_modular.py @@ -28,7 +28,6 @@ to_nii, to_nii_optional, v_idx2name, - v_idx_order, ) from TPTBox.mesh3D.mesh_colors import _color_map_in_row, get_color_by_label @@ -162,8 +161,6 @@ def sag_cor_curve_projection( # Sagittal and coronal projections of a curved plane defined by centroids # Note: Will assume IPL orientation! # if x-direction (=S/I) is not fully incremental, a straight, not an interpolated plane will be returned - order = v_idx_order - order += [i for i in range(256) if i not in v_idx_order] # ctd_list.sorting_list = v_idx_order ctd_list.round_(3) diff --git a/unit_tests/test_regressions.py b/unit_tests/test_regressions.py index 099b6ed2..c24e9187 100644 --- a/unit_tests/test_regressions.py +++ b/unit_tests/test_regressions.py @@ -281,5 +281,46 @@ def test_centred_is_not_in_border(self): self.assertFalse(_make_nii(arr).is_segmentation_in_border()) +class Test_Global_Vert_Order(unittest.TestCase): + """sag_cor_curve_projection aliased the module-level v_idx_order list and extended it in place.""" + + def test_v_idx_order_matches_its_definition(self): + from TPTBox.core.vert_constants import v_idx2name, v_idx_order + + self.assertEqual(list(v_idx_order), list(v_idx2name.keys())) + + def test_snapshot_module_does_not_extend_it(self): + import TPTBox + from TPTBox.spine.snapshot2D import snapshot_modular # noqa: F401 + + self.assertEqual(len(TPTBox.v_idx_order), len(TPTBox.v_idx2name)) + + +class Test_POI_PIR_Cache(unittest.TestCase): + """_vert_orientation_pir was a class attribute, so it was shared by every POI in the process.""" + + @staticmethod + def _poi(value: float): + from TPTBox import POI + + return POI({1: {50: (value, value, value)}}, orientation=("R", "A", "S"), zoom=(1, 1, 1), shape=(10, 10, 10)) + + def test_cache_is_not_shared_between_instances(self): + first = self._poi(1.0) + second = self._poi(4.0) + first._vert_orientation_pir[99] = "subject-1" + self.assertNotIn(99, second._vert_orientation_pir) + + def test_own_cache_is_retained(self): + poi = self._poi(1.0) + poi._vert_orientation_pir[99] = "subject-1" + self.assertEqual(poi._vert_orientation_pir[99], "subject-1") + + def test_copy_starts_with_an_empty_cache(self): + poi = self._poi(1.0) + poi._vert_orientation_pir[99] = "subject-1" + self.assertNotIn(99, poi.copy()._vert_orientation_pir) + + if __name__ == "__main__": unittest.main() From 69436cb71d432e478d4b78e15243bca9c3b9d649 Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 5 Aug 2026 06:45:47 +0000 Subject: [PATCH 5/8] Fix runtime crashes on public API entry points - `from TPTBox import *` raised AttributeError: __all__ advertised load_poi but nothing imported it. Now re-exported from core.poi_fun.save_load. - nii[0:5] delegated to a key containing Ellipsis, which __getitem__ itself rejects two branches earlier, and dropped the `return`. It now pads the trailing dimensions with full slices. - extract_label("12") / remove_labels("12"): str is a Sequence, so the str branch was unreachable and the label was iterated character by character (IndexError). The str check now precedes the Sequence check. - ssim()/psnr() normalised via in-place `img_1 /= img_1.max()`, which raises UFuncTypeError on integer images. The neighbouring img_2 already used the out-of-place form; both now match. - sitk_utils.transform_centroid called the non-existent NII.get_empty_POI; the method is make_empty_POI (used correctly at eight other call sites). Also drops a duplicated assignment in the deformable branch. - stitching_tools.n4_bias passed dilate_msk_(mm=3), which is not a parameter of that method -> TypeError on every call. Now n_pixel=3. - inference_nnunet forwarded stacklevel= into the logger, which passes **qargs to print() -> TypeError whenever the input affine is the identity. - poi_global.to_other tested `isinstance(ref, Self)`; Self is a typing special form and isinstance() against it raises. Now checks POI_Global, and the method no longer falls off the end returning None. - poi_global.__init__ gated level_two_info on level_one_info (copy-paste), so passing only level_two_info silently dropped it. - calc_centroids took type(stage) after unwrapping the enum to .value, so level_one_info/level_two_info were always int and the saved POI header recorded "int" instead of the enum class. - BIDS auto_add_run_id did info["run"] += 1 on a value that must remain a decimal string for validate_entities(). Co-Authored-By: Claude Opus 5 --- TPTBox/__init__.py | 1 + TPTBox/core/bids_files.py | 7 +- TPTBox/core/nii_wrapper.py | 9 ++- TPTBox/core/nii_wrapper_math.py | 4 +- TPTBox/core/poi.py | 6 +- TPTBox/core/poi_fun/poi_global.py | 5 +- TPTBox/core/sitk_utils.py | 3 +- .../segmentation/VibeSeg/inference_nnunet.py | 1 - TPTBox/stitching/stitching_tools.py | 2 +- unit_tests/test_regressions.py | 76 +++++++++++++++++++ 10 files changed, 97 insertions(+), 17 deletions(-) diff --git a/TPTBox/__init__.py b/TPTBox/__init__.py index 70c66dc8..0b66dd2a 100755 --- a/TPTBox/__init__.py +++ b/TPTBox/__init__.py @@ -27,6 +27,7 @@ from TPTBox.core.poi import calc_poi_from_two_segs from TPTBox.core.poi import calc_poi_from_two_segs as calc_poi_labeled_buffered from TPTBox.core.poi_fun.poi_global import POI_Global +from TPTBox.core.poi_fun.save_load import load_poi from TPTBox.core.vert_constants import ZOOMS, Location, Vertebra_Instance, v_idx2name, v_idx_order, v_name2idx # Logger diff --git a/TPTBox/core/bids_files.py b/TPTBox/core/bids_files.py index f3849913..882445fa 100755 --- a/TPTBox/core/bids_files.py +++ b/TPTBox/core/bids_files.py @@ -1147,10 +1147,9 @@ def get_changed_path( # noqa: C901 return out_path if not out_path.exists(): return out_path - if "run" in info: - info["run"] += 1 - else: - info["run"] = 2 + # "run" is a decimal entity and must stay a string: validate_entities() calls + # .isdecimal() on it, which an int does not support. + info["run"] = str(int(info["run"]) + 1) if "run" in info else "2" def save_changed_path( self, diff --git a/TPTBox/core/nii_wrapper.py b/TPTBox/core/nii_wrapper.py index 9dc4ca49..6fd809fa 100755 --- a/TPTBox/core/nii_wrapper.py +++ b/TPTBox/core/nii_wrapper.py @@ -2620,7 +2620,8 @@ def __getitem__(self, key)-> Any: elif isinstance(key,np.ndarray): return self.get_array()[key] elif isinstance(key,slice): - self.__getitem__((key,Ellipsis,Ellipsis)) + # pad with full slices for the trailing dimensions; Ellipsis is rejected above + return self.__getitem__((key, *(slice(None) for _ in range(len(self.shape) - 1)))) else: raise TypeError("Invalid argument type:", type(key)) def __setitem__(self, key,value): @@ -2721,14 +2722,14 @@ def extract_label(self,label:int|Enum|Sequence[int]|Sequence[Enum]|None, keep_la return self.clamp(0,1,inplace=inplace) seg_arr = self.get_seg_array() + if isinstance(label,str): + label = int(label) # a str is also a Sequence, so this must come first if isinstance(label, Sequence): labels:int|list[int] = [idx.value if isinstance(idx,Enum) else idx for idx in label] assert 0 not in labels, 'Zero label does not make sense. This is the background' else: if isinstance(label,Enum): label = label.value - if isinstance(label,str): - label = int(label) assert label != 0, 'Zero label does not make sense. This is the background' labels = label @@ -2759,6 +2760,8 @@ def extract_label_(self, label: int | Enum | Sequence[int] | Sequence[Enum], kee def remove_labels(self,label:int|Enum|Sequence[int]|Sequence[Enum], inplace=False, verbose:logging=True, removed_to_label=0) -> Self: """If this NII is a segmentation you can single out one label.""" assert label != 0, 'Zero label does not make sens. This is the background' + if isinstance(label,str): + label = int(label) # a str is also a Sequence, so this must come first if not isinstance(label,Sequence): label = [label] # type: ignore flat: list[int] = [] diff --git a/TPTBox/core/nii_wrapper_math.py b/TPTBox/core/nii_wrapper_math.py index 08c54ea4..7d94b39e 100755 --- a/TPTBox/core/nii_wrapper_math.py +++ b/TPTBox/core/nii_wrapper_math.py @@ -456,7 +456,7 @@ def ssim(self, nii:NII_Proxy, min_v = 0)->float: """ img_1 = nii.get_array() - min_v img_2 = self.get_array() - min_v - img_1/= img_1.max() + img_1 = img_1/ img_1.max() # out-of-place: /= fails on integer arrays img_1[img_1<=0] = 0 img_2= img_2/ img_2.max() img_2[img_2<=0] = 0 @@ -480,7 +480,7 @@ def psnr(self,nii: NII_Proxy,min_v=0)->float: """ img_1 = nii.get_array() - min_v img_2 = self.get_array() - min_v - img_1/= img_1.max() + img_1 = img_1/ img_1.max() # out-of-place: /= fails on integer arrays img_1[img_1<=0] = 0 img_2= img_2/img_2.max() img_2[img_2<=0] = 0 diff --git a/TPTBox/core/poi.py b/TPTBox/core/poi.py index d8ebf048..45e06a35 100755 --- a/TPTBox/core/poi.py +++ b/TPTBox/core/poi.py @@ -1317,12 +1317,14 @@ def calc_centroids( - NaN values in the binary mask are ignored. """ args = {} + # Capture the enum class BEFORE unwrapping to .value, otherwise type() just reports int + # and the POI header records "int" instead of e.g. "Location". if isinstance(second_stage, Abstract_lvl): - second_stage = second_stage.value args["level_two_info"] = type(second_stage) + second_stage = second_stage.value if isinstance(first_stage, Abstract_lvl): - first_stage = first_stage.value args["level_one_info"] = type(first_stage) + first_stage = first_stage.value assert first_stage == -1 or second_stage == -1, "first or second dimension must be fixed." msk_nii = to_nii(msk, seg=True) msk_data = msk_nii.get_seg_array() diff --git a/TPTBox/core/poi_fun/poi_global.py b/TPTBox/core/poi_fun/poi_global.py index 0ff83048..0f798795 100755 --- a/TPTBox/core/poi_fun/poi_global.py +++ b/TPTBox/core/poi_fun/poi_global.py @@ -37,7 +37,7 @@ def __init__( args = {} if level_one_info is not None: args["level_one_info"] = level_one_info - if level_one_info is not None: + if level_two_info is not None: args["level_two_info"] = level_two_info self.itk_coords = itk_coords _format = FORMAT_GLOBAL @@ -124,8 +124,9 @@ def to_other_poi(self, ref: poi.POI | Self) -> poi.POI | Self | None: p = poi.POI.load(ref) if isinstance(ref, poi.POI): return self.to_other(p) - elif isinstance(ref, Self): + elif isinstance(ref, POI_Global): # `Self` is a typing form; isinstance() against it raises return self.to_cord_system(ref.itk_coords) + return p def to_global(self, itk_coords: bool | None = None) -> Self: """Return this object unchanged (already in global coordinates).""" diff --git a/TPTBox/core/sitk_utils.py b/TPTBox/core/sitk_utils.py index 1734d20d..aff2d36b 100755 --- a/TPTBox/core/sitk_utils.py +++ b/TPTBox/core/sitk_utils.py @@ -123,7 +123,6 @@ def transform_centroid(ctd: POI, transform: sitk.Transform, img_fixed: sitk.Imag for key, key2, (x, y, z) in ctd.items(): ctr_b = transform.TransformPoint((x, y, z)) out[key, key2] = ctr_b - out[key, key2] = ctr_b else: for key, key2, (x, y, z) in ctd.items(): ctr_b = img_moving.TransformContinuousIndexToPhysicalPoint((x, y, z)) @@ -131,7 +130,7 @@ def transform_centroid(ctd: POI, transform: sitk.Transform, img_fixed: sitk.Imag ctr_b = img_fixed.TransformPhysicalPointToContinuousIndex(ctr_b) out[key, key2] = ctr_b nii = sitk_to_nii(img_fixed, True) - return nii.get_empty_POI(out) + return nii.make_empty_POI(out) def get_sitk_metadata_from_ras_affine(affine: np.ndarray) -> tuple[tuple, tuple, tuple]: diff --git a/TPTBox/segmentation/VibeSeg/inference_nnunet.py b/TPTBox/segmentation/VibeSeg/inference_nnunet.py index 3d9cc85f..5232c728 100644 --- a/TPTBox/segmentation/VibeSeg/inference_nnunet.py +++ b/TPTBox/segmentation/VibeSeg/inference_nnunet.py @@ -481,7 +481,6 @@ def run_VibeSeg( if (in_niis[0].affine == np.eye(4)).all(): logger.on_warning( "Your affine matrix is the identity. Make sure that the spacing and orientation is correct. For NAKO VIBE it should be 1.40625 mm for R/L and A/P and 3 mm S/I. For UKBB R/L and A/P should be around 2.2 mm", - stacklevel=3, ) return run_inference_on_file( dataset_id, diff --git a/TPTBox/stitching/stitching_tools.py b/TPTBox/stitching/stitching_tools.py index 2e44fc49..3febd68c 100755 --- a/TPTBox/stitching/stitching_tools.py +++ b/TPTBox/stitching/stitching_tools.py @@ -166,7 +166,7 @@ def n4_bias( mask[mask != 0] = 1 mask_nii = nii.set_array(mask) mask_nii.seg = True - mask_nii.dilate_msk_(mm=3, verbose=False) + mask_nii.dilate_msk_(n_pixel=3, verbose=False) n4: NII = nii.n4_bias_field_correction(mask=from_nibabel(mask_nii.nii), spline_param=spline_param) if norm != -1: n4 *= norm / n4.max() diff --git a/unit_tests/test_regressions.py b/unit_tests/test_regressions.py index c24e9187..4f23aeb2 100644 --- a/unit_tests/test_regressions.py +++ b/unit_tests/test_regressions.py @@ -322,5 +322,81 @@ def test_copy_starts_with_an_empty_cache(self): self.assertNotIn(99, poi.copy()._vert_orientation_pir) +class Test_Public_API(unittest.TestCase): + """__all__ advertised load_poi, but nothing imported it, so `import *` raised.""" + + def test_load_poi_is_exported(self): + import TPTBox + + self.assertTrue(callable(TPTBox.load_poi)) + + def test_everything_in_all_is_importable(self): + import TPTBox + + missing = [name for name in TPTBox.__all__ if not hasattr(TPTBox, name)] + self.assertEqual(missing, []) + + +class Test_Getitem_Slice(unittest.TestCase): + """nii[0:5] delegated to a key containing Ellipsis, which the same method rejects.""" + + def test_single_slice(self): + nii = _make_nii(_cube()) + self.assertEqual(nii[0:5].shape, (5, 10, 10)) + + def test_full_slice_tuple_still_works(self): + nii = _make_nii(_cube()) + self.assertEqual(nii[2:5, 2:5, 2:5].shape, (3, 3, 3)) + + +class Test_Label_As_String(unittest.TestCase): + """str is a Sequence, so the str branch was unreachable and '12' iterated as characters.""" + + @staticmethod + def _seg() -> NII: + arr = np.zeros((10, 10, 10), dtype=np.uint8) + arr[1, 1, 1] = 12 + arr[2:5, 2:5, 2:5] = 3 + return _make_nii(arr) + + def test_extract_label_string_matches_int(self): + seg = self._seg() + self.assertEqual(int(seg.extract_label("12").sum()), int(seg.extract_label(12).sum())) + + def test_remove_labels_string(self): + self.assertEqual(list(self._seg().remove_labels("12", verbose=False).unique()), [3]) + + +class Test_Metrics_On_Integer_Images(unittest.TestCase): + """ssim/psnr used in-place `/=`, which fails on integer arrays.""" + + @staticmethod + def _int_nii() -> NII: + arr = np.zeros((10, 10, 10), dtype=np.int16) + arr[2:5, 2:5, 2:5] = 300 + return _make_nii(arr, seg=False) + + def test_ssim_identical_is_one(self): + nii = self._int_nii() + self.assertAlmostEqual(float(nii.ssim(nii)), 1.0, places=5) + + def test_psnr_runs_on_int16(self): + nii = self._int_nii() + self.assertTrue(np.isinf(nii.psnr(nii))) + + +class Test_Calc_Centroids_Level_Info(unittest.TestCase): + """type() was taken after unwrapping the enum to .value, so it was always int.""" + + def test_level_two_info_records_the_enum_class(self): + from TPTBox import Location, calc_centroids + + arr = np.zeros((10, 10, 10), dtype=np.uint8) + arr[2:5, 2:5, 2:5] = 1 + arr[6:9, 6:9, 6:9] = 2 + poi = calc_centroids(_make_nii(arr), second_stage=Location.Vertebra_Corpus) + self.assertIs(poi.level_two_info, Location) + + if __name__ == "__main__": unittest.main() From b2ee83b4edd6e7ba73a0275cd353ffeba167b5ae Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 5 Aug 2026 06:50:39 +0000 Subject: [PATCH 6/8] Fix conditions and mappings that silently produced wrong results None of these raised; they just computed the wrong thing. vert_constants: - Full_Body_Instance_Vibe.get_Full_Body_Instance_mapping() referenced cls.hip_left/hip_right, which do not exist on that enum (they are pelvis_left/pelvis_right, 50/51, matching the inverse map) -> the whole classmethod raised AttributeError. - The same dict literal had duplicate keys: lung_left twice, lung_right three times, channel twice. Only the last of each survived, so the shadowed entries were dead. Removed them, keeping the previously-winning target so behaviour is unchanged, with a note that one FBI lung label cannot address several Vibe lobes. - Abstract_lvl._get_id passed `cls` positionally into a classmethod, shifting the arguments and raising TypeError on every lookup - swallowed by a bare `except Exception`. Name resolution therefore never worked: Any._get_id("L1") fell through to int("L1"). Now resolves to 20. np_utils: - np_filter_connected_components compared a list[tuple] to an int, so that shortcut branch was unreachable. NOTE: fixing this to len(...) changes which components are preserved when the count equals largest_k_components; worth a domain review. - The background_threshold in the label-smoothing helper was applied to the argmax *indices* rather than the winning confidence, so it removed labels by index rather than by probability. Other: - vertebra_direction dropped the result of set_array().reorient().rescale_(), then read the array back off the un-reoriented object, writing the fill-back image in the wrong orientation and spacing. Mirrors the correct chained form in calc_center_spinal_cord. - SegmentationMesh discarded `int_arr.astype(np.uint16)`, so the float->int conversion it announces never happened. - snapshot_modular: cmap(color - 1 % LABEL_MAX % cmap.N) parses as color - 1, losing the wrap-around; parenthesised to match the correct site above it. - angles: `if (vert, 50) not in poi: cord = poi[vert, 50]` was inverted, so no lordosis/kyphosis label was ever emitted. - body_quadrants requested Vertebra_Direction_Inferior twice but reads Vertebra_Direction_Right; the KeyError was swallowed, so every vertebra was skipped and an all-zero image returned. - inference_nnunet built its label mapping with the raw string key instead of the parsed int. - predictor accumulated the chunk bounding-box max from self.min_s. - deepali_model gated source_seg on fixed_seg and target_seg on moving_seg. - point_registration.load_ assigned the dumped (moving, fixed) pair to (_img_fixed, _img_moving), so reloaded registrations mapped the wrong way. - save_mkr passed split_by_region=split_by_subregion, so with the documented defaults neither branch was taken and markers got random colours. - ray_casting negated z instead of flipping the half-space inequality. Co-Authored-By: Claude Opus 5 --- TPTBox/core/np_utils.py | 9 ++++++-- TPTBox/core/poi_fun/ray_casting.py | 4 +++- TPTBox/core/poi_fun/save_mkr.py | 2 +- TPTBox/core/poi_fun/vertebra_direction.py | 5 +++-- TPTBox/core/vert_constants.py | 22 +++++++++++-------- TPTBox/mesh3D/mesh.py | 2 +- TPTBox/registration/_deepali/deepali_model.py | 5 +++-- .../_ridged_points/point_registration.py | 6 +++-- .../segmentation/VibeSeg/inference_nnunet.py | 4 ++-- TPTBox/segmentation/nnUnet_utils/predictor.py | 2 +- TPTBox/spine/snapshot2D/snapshot_modular.py | 4 ++-- TPTBox/spine/spinestats/angles.py | 2 +- TPTBox/spine/spinestats/body_quadrants.py | 5 ++++- 13 files changed, 45 insertions(+), 27 deletions(-) diff --git a/TPTBox/core/np_utils.py b/TPTBox/core/np_utils.py index c7dde01b..c5f63f2a 100755 --- a/TPTBox/core/np_utils.py +++ b/TPTBox/core/np_utils.py @@ -1014,7 +1014,9 @@ def np_filter_connected_components( largest_k_components = min(largest_k_components, len(label_volume_pairs)) label_volume_pairs.sort(key=lambda x: x[1], reverse=True) - if len(labels) == 1 or label_volume_pairs == largest_k_components or largest_k_components_org is None or k_larges_global: + # `label_volume_pairs == largest_k_components` compared a list[tuple] to an int and was + # therefore always False, so this shortcut never fired when every component is kept. + if len(labels) == 1 or len(label_volume_pairs) == largest_k_components or largest_k_components_org is None or k_larges_global: preserve: list[int] = [x[0] for x in label_volume_pairs[:largest_k_components]] else: counter = dict.fromkeys(labels, 0) @@ -1294,7 +1296,10 @@ def np_smooth_gaussian_labelwise( seg_arr_s = seg_arr_smoothed.copy() if background_threshold is not None: - seg_arr_smoothed[seg_arr_smoothed < background_threshold] = len(sem_labels_plus_background) - 1 # background label + # Threshold the winning *confidence*, not the argmax index: seg_arr_smoothed holds + # label indices, so comparing it to a probability threshold zeroed out whichever + # labels happened to sort below it. + seg_arr_smoothed[arr_stack.max(axis=0) < background_threshold] = len(sem_labels_plus_background) - 1 # background label for idx, l in enumerate(sem_labels_plus_background): seg_arr_s[seg_arr_smoothed == idx] = l diff --git a/TPTBox/core/poi_fun/ray_casting.py b/TPTBox/core/poi_fun/ray_casting.py index 08c61401..1155a98f 100644 --- a/TPTBox/core/poi_fun/ray_casting.py +++ b/TPTBox/core/poi_fun/ray_casting.py @@ -669,5 +669,7 @@ def set_label_above_3_point_plane( plane_z = (-a * x - b * y - d) / c # Create the 3D array and set values above the plane to 0 - array[np.logical_and(mask, z * invert > plane_z)] = value + # Negating z is not the same as flipping the inequality; scale both sides so that + # invert=-1 really selects the opposite half-space. + array[np.logical_and(mask, invert * z > invert * plane_z)] = value return array diff --git a/TPTBox/core/poi_fun/save_mkr.py b/TPTBox/core/poi_fun/save_mkr.py index 17eec8d1..4aa21fe0 100644 --- a/TPTBox/core/poi_fun/save_mkr.py +++ b/TPTBox/core/poi_fun/save_mkr.py @@ -533,7 +533,7 @@ def _save_mrk( display=_get_display_dict( display, selectedColor=_get_markup_color( - {"color": color}, region, subregion, split_by_region=split_by_subregion, split_by_subregion=split_by_subregion + {"color": color}, region, subregion, split_by_region=split_by_region, split_by_subregion=split_by_subregion ), **addendum, ), diff --git a/TPTBox/core/poi_fun/vertebra_direction.py b/TPTBox/core/poi_fun/vertebra_direction.py index ce473123..112cf2cf 100644 --- a/TPTBox/core/poi_fun/vertebra_direction.py +++ b/TPTBox/core/poi_fun/vertebra_direction.py @@ -126,8 +126,9 @@ def calc_orientation_of_vertebra_PIR( cond = np.where(curr_slice != 0) x_slice[cond] = np.minimum(curr_slice[cond], x_slice[cond]) fill_back[i] = x_slice - subreg_sar.set_array(fill_back).reorient(poi.orientation).rescale_(poi.zoom) - arr = subreg_sar.get_array() + # set_array/reorient are out-of-place: the chained result must be captured, otherwise + # `arr` is still in (S,A,R) at iso spacing. Mirrors calc_center_spinal_cord below. + arr = subreg_sar.set_array(fill_back).reorient(poi.orientation).rescale_(poi.zoom).get_array() fill_back_nii.set_array_(arr) ret = calc_centroids(subreg_iso.set_array(out), second_stage=subreg_id, extend_to=poi_iso.copy(), inplace=True) diff --git a/TPTBox/core/vert_constants.py b/TPTBox/core/vert_constants.py index f8bbedee..4ee2afae 100755 --- a/TPTBox/core/vert_constants.py +++ b/TPTBox/core/vert_constants.py @@ -181,9 +181,11 @@ def _get_id(cls, s: str | int, no_raise=True) -> int: # noqa: ARG003 if n == self_name: continue try: - s = cl._get_id(cls, s, no_raise=False) - return s # type: ignore # noqa: TRY300 - except Exception: + # _get_id is a classmethod: passing `cls` positionally bound it to `s` + # and shifted `s` into `no_raise`, so every lookup raised TypeError and + # name resolution silently never worked. + return cl._get_id(s, no_raise=False) # type: ignore # noqa: TRY300 + except (KeyError, ValueError, AttributeError): pass return int(s) @@ -283,10 +285,11 @@ def get_Full_Body_Instance_mapping(cls) -> dict[int, int]: Full_Body_Instance.pancreas.value: cls.pancreas.value, # pancreas Full_Body_Instance.adrenal_gland_right.value: cls.adrenal_gland_right.value, # adrenal_gland_right Full_Body_Instance.adrenal_gland_left.value: cls.adrenal_gland_left.value, # adrenal_gland_left - Full_Body_Instance.lung_left.value: cls.lung_upper_lobe_left.value, # lung_upper_lobe_left + # Full_Body_Instance has one label per lung, while the Vibe scheme splits each + # lung into lobes. Only one lobe can be the target, so the lower lobe is used - + # this preserves the behaviour of the previous duplicate-key literal, where the + # last entry silently won. The shadowed upper/middle-lobe entries were dead. Full_Body_Instance.lung_left.value: cls.lung_lower_lobe_left.value, # lung_lower_lobe_left - Full_Body_Instance.lung_right.value: cls.lung_upper_lobe_right.value, # lung_upper_lobe_right - Full_Body_Instance.lung_right.value: cls.lung_middle_lobe_right.value, # lung_middle_lobe_right Full_Body_Instance.lung_right.value: cls.lung_lower_lobe_right.value, # lung_lower_lobe_right Full_Body_Instance.esophagus.value: cls.esophagus.value, # esophagus Full_Body_Instance.trachea.value: cls.trachea.value, # trachea @@ -323,9 +326,10 @@ def get_Full_Body_Instance_mapping(cls) -> dict[int, int]: Full_Body_Instance.clavicula_right.value: cls.clavicula_right.value, # clavicula_right Full_Body_Instance.femur_left.value: cls.femur_left.value, # femur_left Full_Body_Instance.femur_right.value: cls.femur_right.value, # femur_right - Full_Body_Instance.pelvis_left.value: cls.hip_left.value, # hip_left - Full_Body_Instance.pelvis_right.value: cls.hip_right.value, # hip_right - Full_Body_Instance.channel.value: cls.spinal_cord.value, # spinal_cord + # `hip_left`/`hip_right` do not exist on this enum; the labels are called + # pelvis_left/pelvis_right (50/51), matching `50: pelvis_left` in the inverse map. + Full_Body_Instance.pelvis_left.value: cls.pelvis_left.value, # hip_left + Full_Body_Instance.pelvis_right.value: cls.pelvis_right.value, # hip_right Full_Body_Instance.gluteus_maximus_left.value: cls.gluteus_maximus_left.value, # gluteus_maximus_left Full_Body_Instance.gluteus_maximus_right.value: cls.gluteus_maximus_right.value, # gluteus_maximus_right Full_Body_Instance.gluteus_medius_left.value: cls.gluteus_medius_left.value, # gluteus_medius_left diff --git a/TPTBox/mesh3D/mesh.py b/TPTBox/mesh3D/mesh.py index 9bd3e2d2..c0bbee4e 100644 --- a/TPTBox/mesh3D/mesh.py +++ b/TPTBox/mesh3D/mesh.py @@ -123,7 +123,7 @@ def __init__(self, int_arr: np.ndarray | Image_Reference) -> None: # Force dtype to uint if np.issubdtype(int_arr.dtype, np.floating): print("input is of type float, converting to int") - int_arr.astype(np.uint16) + int_arr = int_arr.astype(np.uint16) # astype returns a new array; the result was discarded # calculate bounding box cutout bbox_crop = np_bbox_binary(int_arr, px_dist=2) x1, y1, z1 = bbox_crop[0].start, bbox_crop[1].start, bbox_crop[2].start diff --git a/TPTBox/registration/_deepali/deepali_model.py b/TPTBox/registration/_deepali/deepali_model.py index 4646179f..523256e6 100644 --- a/TPTBox/registration/_deepali/deepali_model.py +++ b/TPTBox/registration/_deepali/deepali_model.py @@ -276,8 +276,9 @@ def __init__( super().__init__( source=source.to_deepali(), target=fix.to_deepali(), - source_seg=to_nii(moving_seg, True).to_deepali() if fixed_seg is not None else None, - target_seg=to_nii(fixed_seg, True).to_deepali() if moving_seg is not None else None, + # each branch used to test the *other* variable + source_seg=to_nii(moving_seg, True).to_deepali() if moving_seg is not None else None, + target_seg=to_nii(fixed_seg, True).to_deepali() if fixed_seg is not None else None, source_pset=source_pset, target_pset=target_pset, source_landmarks=source_landmarks, diff --git a/TPTBox/registration/_ridged_points/point_registration.py b/TPTBox/registration/_ridged_points/point_registration.py index b25f535c..f158481d 100644 --- a/TPTBox/registration/_ridged_points/point_registration.py +++ b/TPTBox/registration/_ridged_points/point_registration.py @@ -377,8 +377,10 @@ def load_(cls, w: tuple) -> Point_Registration: ) = w a: Has_Grid b: Has_Grid - self._img_fixed = nii_to_sitk(a.make_nii()) - self._img_moving = nii_to_sitk(b.make_nii()) + # get_dump() writes (version, moving, fixed, ...); these were assigned the wrong way + # round, so every reloaded registration resampled into the wrong space. + self._img_moving = nii_to_sitk(a.make_nii()) + self._img_fixed = nii_to_sitk(b.make_nii()) assert version == 1, f"Version mismatch {version=}" return self diff --git a/TPTBox/segmentation/VibeSeg/inference_nnunet.py b/TPTBox/segmentation/VibeSeg/inference_nnunet.py index 5232c728..d1402e51 100644 --- a/TPTBox/segmentation/VibeSeg/inference_nnunet.py +++ b/TPTBox/segmentation/VibeSeg/inference_nnunet.py @@ -370,8 +370,8 @@ def to_int(a: str, k: int | None = None): for k, v in mapping_.items(): key = to_int(k) value = to_int(v, key) - if k != value: - mapping[k] = value + if key != value: # `k` is the raw string: map_labels_ needs the int key + mapping[key] = value unknown_strings[v] = value logger.print(f"{unknown_strings}") logger.print(f"{mapping=}") diff --git a/TPTBox/segmentation/nnUnet_utils/predictor.py b/TPTBox/segmentation/nnUnet_utils/predictor.py index 4e95e710..f61a317a 100755 --- a/TPTBox/segmentation/nnUnet_utils/predictor.py +++ b/TPTBox/segmentation/nnUnet_utils/predictor.py @@ -866,7 +866,7 @@ def add_slicer(self, s: tuple[slice, ...]) -> None: if self.max_s is None: self.max_s = [s.stop for s in s[1:]] else: - self.max_s = [max(s.stop, m) for s, m in zip(s[1:], self.min_s)] + self.max_s = [max(s.stop, m) for s, m in zip(s[1:], self.max_s)] # was accumulating from min_s assert len(s) - 1 == len(self.meta_slice) self.slicers.append(s) diff --git a/TPTBox/spine/snapshot2D/snapshot_modular.py b/TPTBox/spine/snapshot2D/snapshot_modular.py index b9e11dfb..3cb04e6f 100755 --- a/TPTBox/spine/snapshot2D/snapshot_modular.py +++ b/TPTBox/spine/snapshot2D/snapshot_modular.py @@ -700,7 +700,7 @@ def plot_sag_centroids( v[0] * zms[0], c, d, - color=cmap(color - 1 % LABEL_MAX % cmap.N), + color=cmap((color - 1) % LABEL_MAX % cmap.N), ) ) if "text_sag" in ctd.info: @@ -811,7 +811,7 @@ def plot_cor_centroids( v[0] * zms[0], c, d, - color=cmap(color - 1 % LABEL_MAX % cmap.N), + color=cmap((color - 1) % LABEL_MAX % cmap.N), ) ) if "text_cor" in ctd.info: diff --git a/TPTBox/spine/spinestats/angles.py b/TPTBox/spine/spinestats/angles.py index 7c7f832f..e10fdca6 100644 --- a/TPTBox/spine/spinestats/angles.py +++ b/TPTBox/spine/spinestats/angles.py @@ -732,7 +732,7 @@ def plot_compute_lordosis_and_kyphosis( vert = round((id1.value + id2.value) / 2) while (vert, 50) not in poi and vert != 0: vert -= 1 - if (vert, 50) not in poi: + if (vert, 50) in poi: # the body needs the key to be PRESENT; `not in` made it dead cord = poi[vert, 50] text_out.append((vert, (f"{str(name).split('_')[-1]}: {v:.1f}°", 15, cord[1]))) diff --git a/TPTBox/spine/spinestats/body_quadrants.py b/TPTBox/spine/spinestats/body_quadrants.py index 8cf287b5..7ebbbe5d 100644 --- a/TPTBox/spine/spinestats/body_quadrants.py +++ b/TPTBox/spine/spinestats/body_quadrants.py @@ -98,7 +98,10 @@ def make_quadrants( subreg_id=[ Location.Vertebra_Corpus, Location.Vertebra_Direction_Inferior, - Location.Vertebra_Direction_Inferior, + # Vertebra_Direction_Right is read below; it used to be a second copy of + # Inferior, so the lookup always raised KeyError and every vertebra was + # skipped -> make_quadrants silently returned an all-zero image. + Location.Vertebra_Direction_Right, ], buffer_file=poi_buffer, ) From 3f5162a728ad96d08476e7b814e7a713fac8b9d5 Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 5 Aug 2026 06:54:40 +0000 Subject: [PATCH 7/8] Cut import time 35%, stop leaking figures, and drop redundant volume copies MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Import cost: nii_wrapper_math imported peak_signal_noise_ratio and structural_similarity from skimage.metrics at module level, but they are used only inside NII.ssim and NII.psnr. That import pulls in scipy.stats, which accounted for ~35% of `import TPTBox`. Moved both into their methods, following the existing skimage.exposure precedent in nii_wrapper.py. import TPTBox 948 ms -> 613 ms (scipy.stats no longer loaded at all) Memory: - vertebra_direction allocated `subreg_iso.get_array() * 0` once per vertebra: get_array() copies the whole volume and `* 0` allocates a second one. Now a single np.zeros of the same shape and dtype. Two sites. - set_dtype called get_array() (a full copy) twice on the smallest_int/uint path; it is now fetched at most once. Resource leaks: - _help.py created a figure that was never closed - one leaked figure per call. - snapshot_modular used a bare plt.close(), which closes the *current* figure rather than the one just saved, and was skipped entirely if savefig raised. Both now close their own figure in a finally block. Encoding: the Logger opened its .log with the platform default encoding, so print_statistic's "±" raised UnicodeEncodeError under a C/POSIX locale (Docker, cron, CI). Its logs directory is also created with parents=True, exist_ok=True, which parallel jobs were racing on. The same missing encoding= is fixed for the POI, BIDS-sidecar and DICOM JSON readers. Co-Authored-By: Claude Opus 5 --- TPTBox/core/bids_files.py | 2 +- TPTBox/core/dicom/dicom2nii_utils.py | 4 ++-- TPTBox/core/nii_wrapper.py | 8 ++++++-- TPTBox/core/nii_wrapper_math.py | 8 ++++++-- TPTBox/core/poi_fun/_help.py | 19 +++++++++++-------- TPTBox/core/poi_fun/poi_abstract.py | 2 +- TPTBox/core/poi_fun/save_load.py | 4 ++-- TPTBox/core/poi_fun/vertebra_direction.py | 6 ++++-- TPTBox/logger/log_file.py | 8 ++++---- TPTBox/spine/snapshot2D/snapshot_modular.py | 12 ++++++++---- 10 files changed, 45 insertions(+), 28 deletions(-) diff --git a/TPTBox/core/bids_files.py b/TPTBox/core/bids_files.py index 882445fa..801748e0 100755 --- a/TPTBox/core/bids_files.py +++ b/TPTBox/core/bids_files.py @@ -1336,7 +1336,7 @@ def open_json(self) -> dict: Raises: KeyError: If no JSON file is registered in :attr:`file`. """ - with open(self.file["json"]) as f: + with open(self.file["json"], encoding="utf-8") as f: return json.load(f) def open_poi(self, nii: TPTBox.Image_Reference | None = None) -> TPTBox.POI: diff --git a/TPTBox/core/dicom/dicom2nii_utils.py b/TPTBox/core/dicom/dicom2nii_utils.py index cc835b4d..41ce673e 100755 --- a/TPTBox/core/dicom/dicom2nii_utils.py +++ b/TPTBox/core/dicom/dicom2nii_utils.py @@ -284,7 +284,7 @@ def test_name_conflict(json_ob: dict, file: str | Path) -> bool: ``False`` otherwise (file does not exist or content matches). """ if Path(file).exists(): - with open(file) as f: + with open(file, encoding="utf-8") as f: js = json.load(f) if "grid" in js: del js["grid"] @@ -328,7 +328,7 @@ def load_json(file: str | Path) -> dict: Returns: Parsed contents of the JSON file. """ - with open(file) as file_handel: + with open(file, encoding="utf-8") as file_handel: return json.load(file_handel) diff --git a/TPTBox/core/nii_wrapper.py b/TPTBox/core/nii_wrapper.py index 6fd809fa..a1e84ce3 100755 --- a/TPTBox/core/nii_wrapper.py +++ b/TPTBox/core/nii_wrapper.py @@ -739,8 +739,10 @@ def set_dtype(self, dtype: type | Literal['smallest_int', 'smallest_uint'] = np. The NII with the new dtype (``self`` when ``inplace=True``, a new NII otherwise). """ sel = self if inplace else self.copy() + arr = None # get_array() copies the whole volume; fetch it at most once if dtype in ("smallest_uint", "smallest_int"): - dtype = _smallest_int_dtype(self.get_array(), unsigned=dtype == "smallest_uint") + arr = self.get_array() + dtype = _smallest_int_dtype(arr, unsigned=dtype == "smallest_uint") if self.__unpacked: self._unpack() sel._arr = sel._arr.astype(dtype) @@ -748,7 +750,9 @@ def set_dtype(self, dtype: type | Literal['smallest_int', 'smallest_uint'] = np. else: sel.nii.set_data_dtype(dtype) if sel.nii.get_data_dtype() != self.dtype: #type: ignore - sel.nii = Nifti1Image(self.get_array().astype(dtype,casting=casting,order=order),self.affine,self.header) + if arr is None: + arr = self.get_array() + sel.nii = Nifti1Image(arr.astype(dtype,casting=casting,order=order),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: diff --git a/TPTBox/core/nii_wrapper_math.py b/TPTBox/core/nii_wrapper_math.py index 7d94b39e..ebabc60a 100755 --- a/TPTBox/core/nii_wrapper_math.py +++ b/TPTBox/core/nii_wrapper_math.py @@ -5,8 +5,6 @@ from typing import TYPE_CHECKING, Union import numpy as np -from skimage.metrics import peak_signal_noise_ratio as psnr -from skimage.metrics import structural_similarity as ssim from typing_extensions import Self from TPTBox.core.np_utils import np_dice @@ -454,6 +452,9 @@ def ssim(self, nii:NII_Proxy, min_v = 0)->float: Returns: float: SSIM score in the range [-1, 1] (1 = identical). """ + # imported here: skimage.metrics pulls in scipy.stats, ~35% of `import TPTBox` + from skimage.metrics import structural_similarity as ssim + img_1 = nii.get_array() - min_v img_2 = self.get_array() - min_v img_1 = img_1/ img_1.max() # out-of-place: /= fails on integer arrays @@ -478,6 +479,9 @@ def psnr(self,nii: NII_Proxy,min_v=0)->float: Returns: float: PSNR score in dB (higher is better; inf when images are identical). """ + # imported here: skimage.metrics pulls in scipy.stats, ~35% of `import TPTBox` + from skimage.metrics import peak_signal_noise_ratio as psnr + img_1 = nii.get_array() - min_v img_2 = self.get_array() - min_v img_1 = img_1/ img_1.max() # out-of-place: /= fails on integer arrays diff --git a/TPTBox/core/poi_fun/_help.py b/TPTBox/core/poi_fun/_help.py index 334d0f51..ed6a6b6b 100644 --- a/TPTBox/core/poi_fun/_help.py +++ b/TPTBox/core/poi_fun/_help.py @@ -147,11 +147,14 @@ def make_spine_plot(pois: POI, body_spline: np.ndarray, vert_nii: NII, filenames vert_nii = vert_nii.reorient().rescale(pois.zoom) body_center_list = list(np.array(pois.values())) # fitting a curve to the poi and getting it's first derivative - plt.figure(figsize=(10, 10)) - plt.imshow( - np.swapaxes(np.max(vert_nii.get_array(), axis=vert_nii.get_axis(direction="R")), 0, 1), - cmap=plt.cm.gray, # type: ignore - ) - plt.plot(np.asarray(body_center_list)[:, 0], np.asarray(body_center_list)[:, 1]) - plt.plot(np.asarray(body_spline[:, 0]), np.asarray(body_spline[:, 1]), "-") - plt.savefig(filenames) + fig = plt.figure(figsize=(10, 10)) + try: + plt.imshow( + np.swapaxes(np.max(vert_nii.get_array(), axis=vert_nii.get_axis(direction="R")), 0, 1), + cmap=plt.cm.gray, # type: ignore + ) + plt.plot(np.asarray(body_center_list)[:, 0], np.asarray(body_center_list)[:, 1]) + plt.plot(np.asarray(body_spline[:, 0]), np.asarray(body_spline[:, 1]), "-") + plt.savefig(filenames) + finally: + plt.close(fig) # otherwise every call leaks a figure diff --git a/TPTBox/core/poi_fun/poi_abstract.py b/TPTBox/core/poi_fun/poi_abstract.py index 194f03db..7b1bd418 100755 --- a/TPTBox/core/poi_fun/poi_abstract.py +++ b/TPTBox/core/poi_fun/poi_abstract.py @@ -66,7 +66,7 @@ def __init__( ) -> None: """Placeholder class to move string names to integers with multiple definitions.""" if path is not None: - with open(path) as f: + with open(path, encoding="utf-8") as f: info = json.load(f) region = info["region"] subregion = info["subregion"] diff --git a/TPTBox/core/poi_fun/save_load.py b/TPTBox/core/poi_fun/save_load.py index 01f6c066..5af3a985 100644 --- a/TPTBox/core/poi_fun/save_load.py +++ b/TPTBox/core/poi_fun/save_load.py @@ -279,7 +279,7 @@ def _open_file(ctd_path: Union[Path, str, bids_files.BIDS_FILE]) -> dict | list: # --- 1) try JSON --- try: - with path.open("r") as f: + with path.open("r", encoding="utf-8") as f: return json.load(f) except json.JSONDecodeError: pass # not JSON → continue @@ -778,7 +778,7 @@ def _load_landmark_txt(path: Path) -> list: label_name = {} label_group_id = 1 current_group: str | None = None - with path.open("r") as f: + with path.open("r", encoding="utf-8") as f: for raw_line in f: line = raw_line.strip() if not line: diff --git a/TPTBox/core/poi_fun/vertebra_direction.py b/TPTBox/core/poi_fun/vertebra_direction.py index 112cf2cf..88af2670 100644 --- a/TPTBox/core/poi_fun/vertebra_direction.py +++ b/TPTBox/core/poi_fun/vertebra_direction.py @@ -111,7 +111,8 @@ def calc_orientation_of_vertebra_PIR( plane_coords = plane_coords.astype(int) # create_subregion # 1 where the selected subreg is, else 0 - select = subreg_iso.get_array() * 0 + # get_array() copies the whole volume and `* 0` allocates a second one, per iteration + select = np.zeros(subreg_iso.shape, dtype=subreg_iso.dtype) select[plane_coords[:, :, 0], plane_coords[:, :, 1], plane_coords[:, :, 2]] = 1 out[out == 0] += (target_labels * select * reg_label)[out == 0] @@ -403,7 +404,8 @@ def calc_center_spinal_cord( plane_coords = plane_coords.astype(int) # create_subregion # 1 where the selected subreg is, else 0 - select = subreg_iso.get_array() * 0 + # get_array() copies the whole volume and `* 0` allocates a second one, per iteration + select = np.zeros(subreg_iso.shape, dtype=subreg_iso.dtype) select[plane_coords[:, :, 0], plane_coords[:, :, 1], plane_coords[:, :, 2]] = 1 out += target_labels * select * reg_label diff --git a/TPTBox/logger/log_file.py b/TPTBox/logger/log_file.py index b6db6d6b..e4abb7b6 100755 --- a/TPTBox/logger/log_file.py +++ b/TPTBox/logger/log_file.py @@ -305,10 +305,10 @@ def __init__( # Creates logs folder if not existent log_path = Path(path).joinpath("logs") - if not Path.exists(log_path): - Path.mkdir(log_path) - # Open log file - self.f = open(log_path.joinpath(log_filename_full), "w") # noqa: SIM115 + log_path.mkdir(parents=True, exist_ok=True) # exist_ok: parallel jobs race here + # Open log file. encoding is explicit: print_statistic emits "±", which raises + # UnicodeEncodeError under a C/POSIX locale (Docker, cron, CI). + self.f = open(log_path.joinpath(log_filename_full), "w", encoding="utf-8") # noqa: SIM115 # calls close() if program terminates self._finalizer = weakref.finalize(self.f, self.close) self.default_verbose = default_verbose diff --git a/TPTBox/spine/snapshot2D/snapshot_modular.py b/TPTBox/spine/snapshot2D/snapshot_modular.py index 3cb04e6f..c5618e50 100755 --- a/TPTBox/spine/snapshot2D/snapshot_modular.py +++ b/TPTBox/spine/snapshot2D/snapshot_modular.py @@ -1363,10 +1363,14 @@ def create_snapshot( # noqa: C901 if not isinstance(snp_path, list): snp_path = [str(snp_path)] - for path in snp_path: - fig.savefig(str(path)) - print("[*] Snapshot saved:", path) if verbose else None - plt.close() + try: + for path in snp_path: + fig.savefig(str(path)) + print("[*] Snapshot saved:", path) if verbose else None + finally: + # close THIS figure (bare plt.close() closes the current one, which may be another + # figure entirely) and do it even if savefig raises, so failures do not leak + plt.close(fig) return snp_path From 0b0d14cea9287298b3ad86d3860a811ba65e6337 Mon Sep 17 00:00:00 2001 From: iback Date: Wed, 5 Aug 2026 06:57:25 +0000 Subject: [PATCH 8/8] Remove a stale duplicate package and narrow over-broad import guards TPTBox/registration/ridged_intensity/ was meant to be renamed to _ridged_intensity/ in e919218 ("rename folder to _ to show that this is not the recommended way to import things"), but the old directory stayed tracked. Both copies then received the docstring pass in 81db55a, while only the underscore copy received 1b8e0fb ("fix bug for very elongated segmentations") - so the duplicate still carries the pre-fix `w = max(target.shape[2:])` and the un-scaled delta comparison. Nothing imports it and its __init__.py is empty, so it is removed, along with the stale __pycache__-only leftovers at registration/{deepali,deformable,ridged_points}/. The optional-dependency guards in registration/__init__.py and _deepali/__init__.py caught bare `Exception`, which makes a real NameError or AttributeError inside those modules indistinguishable from "torch is not installed" - the symbol just silently disappears from the public API. That is how the stale duplicate above, and the NameError in deepali_trainer, went unnoticed. Narrowed to ImportError; all five exported names still resolve. Also drops the duplicate MODES definition in nii_wrapper.py, which shadowed the identical import 13 lines above it, and two leftover debug prints - one in POI buffer loading, one emitted once per bisection step per point from ray_casting. Co-Authored-By: Claude Opus 5 --- TPTBox/core/nii_wrapper.py | 1 - TPTBox/core/poi.py | 1 - TPTBox/core/poi_fun/ray_casting.py | 1 - TPTBox/registration/__init__.py | 6 +- TPTBox/registration/_deepali/__init__.py | 4 +- .../registration/ridged_intensity/__init__.py | 0 .../ridged_intensity/affine_deepali.py | 531 ------------------ 7 files changed, 5 insertions(+), 539 deletions(-) delete mode 100644 TPTBox/registration/ridged_intensity/__init__.py delete mode 100644 TPTBox/registration/ridged_intensity/affine_deepali.py diff --git a/TPTBox/core/nii_wrapper.py b/TPTBox/core/nii_wrapper.py index a1e84ce3..d656c990 100755 --- a/TPTBox/core/nii_wrapper.py +++ b/TPTBox/core/nii_wrapper.py @@ -72,7 +72,6 @@ if TYPE_CHECKING: from stl.mesh import Mesh from torch import device -MODES = Literal["constant", "nearest", "reflect", "wrap"] _unpacked_nii = tuple[np.ndarray, AFFINE, nib.nifti1.Nifti1Header] _formatwarning = warnings.formatwarning diff --git a/TPTBox/core/poi.py b/TPTBox/core/poi.py index 45e06a35..057e541b 100755 --- a/TPTBox/core/poi.py +++ b/TPTBox/core/poi.py @@ -1023,7 +1023,6 @@ def wrap(*args, **kwargs): len_pref = 0 if buffer_file is not None and Path(buffer_file).exists(): assert extend_to is None - print("load") extend_to = POI.load(buffer_file) len_pref = len(extend_to) kwargs["extend_to"] = extend_to diff --git a/TPTBox/core/poi_fun/ray_casting.py b/TPTBox/core/poi_fun/ray_casting.py index 1155a98f..e645ac62 100644 --- a/TPTBox/core/poi_fun/ray_casting.py +++ b/TPTBox/core/poi_fun/ray_casting.py @@ -102,7 +102,6 @@ def max_distance_ray_cast_convex_npfast( y = start_coord[1] + norm_vec[1] * mid z = start_coord[2] + norm_vec[2] * mid val = trilinear_interpolate(region_array, x, y, z) - print(f"Raycast check at distance {mid:.2f}: value={val:.4f}") if val > 0.5: min_v = mid else: diff --git a/TPTBox/registration/__init__.py b/TPTBox/registration/__init__.py index 8da19776..e638c694 100755 --- a/TPTBox/registration/__init__.py +++ b/TPTBox/registration/__init__.py @@ -3,20 +3,20 @@ try: from ._ridged_points.point_registration import Point_Registration, ridged_points_from_poi, ridged_points_from_subreg_vert -except Exception: +except ImportError: pass try: from TPTBox.registration._deformable.deformable_reg import Deformable_Registration from TPTBox.registration._deformable.multilabel_segmentation import Template_Registration from ._deepali.spine_rigid_elements_reg import Rigid_Elements_Registration -except Exception: +except ImportError: pass try: from ._deepali.deepali_model import General_Registration -except Exception: +except ImportError: pass __all__ = [ "Deformable_Registration", diff --git a/TPTBox/registration/_deepali/__init__.py b/TPTBox/registration/_deepali/__init__.py index 8ccedd39..1910a508 100644 --- a/TPTBox/registration/_deepali/__init__.py +++ b/TPTBox/registration/_deepali/__init__.py @@ -3,11 +3,11 @@ try: from .spine_rigid_elements_reg import Rigid_Elements_Registration -except Exception: +except ImportError: pass try: from .deepali_model import General_Registration -except Exception: +except ImportError: pass diff --git a/TPTBox/registration/ridged_intensity/__init__.py b/TPTBox/registration/ridged_intensity/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/TPTBox/registration/ridged_intensity/affine_deepali.py b/TPTBox/registration/ridged_intensity/affine_deepali.py deleted file mode 100644 index c3db154d..00000000 --- a/TPTBox/registration/ridged_intensity/affine_deepali.py +++ /dev/null @@ -1,531 +0,0 @@ -from __future__ import annotations - -from abc import ABCMeta, abstractmethod - -# pip install hf-deepali -from collections.abc import Sequence -from copy import deepcopy -from typing import Literal, Union - -import torch -import torch.optim -from deepali import spatial -from deepali.core import PathStr, Sampling -from deepali.data import Image as deepaliImage -from deepali.losses import ( - PairwiseImageLoss, -) -from deepali.losses.functional import ncc_loss -from torch import Tensor -from torch.nn import Module -from tqdm import tqdm - -from TPTBox import Image_Reference -from TPTBox.core.internal.deep_learning_utils import DEVICES, get_device -from TPTBox.registration._deepali.deepali_model import General_Registration -from TPTBox.registration._deepali.deepali_trainer import PairwiseImageLoss - - -class PairwiseSegImageLoss(Module, metaclass=ABCMeta): - r"""Base class of pairwise image dissimilarity criteria for segmentation masks.""" - - @abstractmethod - def forward(self, source: Tensor, target: Tensor, mask: [Tensor] | None = None) -> Tensor: - r"""Evaluate image dissimilarity loss. - - Args: - source: Warped source image/segmentation tensor. - target: Fixed target image/segmentation tensor. - mask: Optional spatial mask restricting the loss computation. - - Returns: - Scalar loss tensor. - """ - raise NotImplementedError(f"{type(self).__name__}.forward()") - - -def center_of_mass(tensor: torch.Tensor) -> torch.Tensor: - """Compute the centre of mass of a spatial tensor. - - Args: - tensor: Arbitrary-shape float tensor treated as a spatial density map. - - Returns: - 1-D tensor of length ``tensor.ndim`` with the weighted-average - coordinate along each axis. - """ - grid = torch.meshgrid([torch.arange(s, device=tensor.device) for s in tensor.shape], indexing="ij") - t = tensor / tensor.sum() - com = torch.stack([(t * g).sum() for g in grid]) - return com - - -class Tether_single(PairwiseImageLoss): - """Centre-of-mass tethering loss for a single binary channel. - - Penalises the Euclidean distance between the centre of mass of the warped - source and the fixed target. The penalty is zeroed out when the distance - is smaller than 10 voxels. - """ - - def forward( - self, - source: torch.Tensor, - target: torch.Tensor, - mask: torch.Tensor | None = None, # noqa: ARG002 - ) -> torch.Tensor: # noqa: ARG002 - """Compute the centre-of-mass distance loss between source and target.""" - com_fixed = center_of_mass(target) - com_warped = center_of_mass(source) - l_com = torch.norm(com_fixed - com_warped) - if l_com < 10: - l_com = source.sum() * 0 - l_com = torch.nan_to_num(l_com, nan=0) - return l_com # type: ignore - - -def center_of_mass_cc(tensor: torch.Tensor) -> torch.Tensor: - """Compute the centre of mass for each channel in a ``(B, C, X, Y, Z)`` tensor. - - Returns a tensor of shape ``(B, C, 3)`` with ``(x, y, z)`` coordinates per channel. - """ - dtype = tensor.dtype - B, C, *spatial_shape = tensor.shape - tensor = tensor.float() - grid = torch.meshgrid([torch.arange(s, device=tensor.device) for s in spatial_shape], indexing="ij") # each g is (X, Y, Z) - grid = torch.stack(grid, dim=0) # (3, X, Y, Z) - - # Flatten spatial dims - tensor_flat = tensor.view(B, C, -1) # (B, C, X*Y*Z) - grid_flat = grid.view(3, -1) # (3, X*Y*Z) - - # Normalize tensor - norm = tensor_flat.sum(dim=-1, keepdim=True) # (B, C, 1) - norm[norm == 0] = 1 # avoid division by zero - - com = torch.einsum("bcn,nm->bcm", tensor_flat, grid_flat.T.to(tensor_flat.dtype)) / norm # (B, C, 3) - return com.to(dtype) - - -class Tether_Seg(PairwiseSegImageLoss): - """Per-channel centre-of-mass tethering loss for multi-channel segmentation tensors. - - Computes the mean normalised Euclidean distance between the centre of mass - of each channel in the warped source and the fixed target. Channels whose - displacement is smaller than ``delta`` (relative to the largest spatial - dimension) are ignored. - - Args: - delta: Minimum relative displacement threshold below which a channel - contributes zero loss. - """ - - def __init__(self, delta: float = 1, *args, **kwargs): - self.delta = delta - super().__init__(*args, **kwargs) - - def forward( - self, - source: torch.Tensor, # shape: (B, C, X, Y, Z) - target: torch.Tensor, # shape: (B, C, X, Y, Z) - mask: torch.Tensor | None = None, # noqa: ARG002 - ) -> torch.Tensor: - """Compute the mean per-channel normalised centre-of-mass distance.""" - w = max(target.shape[2:]) - com_fixed = center_of_mass_cc(target) # (B, C, 3) - com_warped = center_of_mass_cc(source) # (B, C, 3) - - l_com = torch.norm(com_fixed - com_warped, dim=-1) / w # (B, C) - - # Zero out channels with small displacement (<10) or NaNs - l_com = torch.where(l_com < self.delta, torch.zeros_like(l_com), l_com) - l_com = torch.nan_to_num(l_com, nan=0.0) - - return l_com.mean() # type: ignore - - -class Tether(PairwiseImageLoss): - """Centre-of-mass tethering loss with optional per-class and memory features. - - Args: - delta: Minimum voxel distance below which the loss is zeroed. Prevents - penalising already-well-aligned structures. - uniq: If ``True``, compute separate centre-of-mass distances per unique - label value instead of the whole foreground mass. - remember: If ``True``, skip computing the loss for ``remember_c`` - iterations after the loss is zeroed (early stopping memory). - remember_c: Number of iterations to skip after the loss drops to zero. - max_v: Scale factor mapping continuous output values to integer labels. - """ - - def __init__( - self, - delta: float = 10, - uniq: bool = False, - remember: bool = False, - remember_c: int = 10, - max_v: float = 1, - *args, - **kwargs, - ) -> None: - self.delta = delta - self.uniq = uniq - self.remember = remember - self.remember_c = remember_c - self.count = 0 - self.max_v = max_v - super().__init__(*args, **kwargs) - - def forward( - self, - source: torch.Tensor, - target: torch.Tensor, - mask: torch.Tensor | None = None, # noqa: ARG002 - ) -> torch.Tensor: # noqa: ARG002 - """Compute the centre-of-mass tethering loss.""" - if self.count != 0: - self.count -= 1 - return torch.zeros(1, device=source.device) - if self.uniq: - loss = torch.zeros(1, device=source.device) - k = 0 - target = (target * self.max_v).round(decimals=0) - source = (source * self.max_v).round(decimals=0) - u = torch.unique(target) - for i in u: - if i == 0: - continue - com_fixed = center_of_mass(target == i) - com_warped = center_of_mass(source == i) - l_com = torch.norm(com_fixed - com_warped) - l_com = torch.nan_to_num(l_com, nan=0) - # print(l_com) - if l_com > self.delta: - loss += l_com - k += 1 - # print(loss / k, k, len(u)) - if k == 0: - if self.remember: - self.count = 10 - return loss - return loss / k - else: - com_fixed = center_of_mass(target != 0) - com_warped = center_of_mass(source != 0) - l_com = torch.norm(com_fixed - com_warped) - if l_com < self.delta: - l_com = torch.zeros(1, device=source.device) - if self.remember: - self.count = 10 - l_com = torch.nan_to_num(l_com, nan=0) - return l_com # type: ignore - - -def subsample_coords(coords: torch.Tensor, k: int) -> torch.Tensor: - """Return a random subset of ``k`` rows from ``coords``, or the full tensor if it has ≤ ``k`` rows. - - Samples without replacement via ``torch.randperm``; works entirely on-device. - """ - n = coords.size(0) - if n <= k: - return coords - idx = torch.randperm(n, device=coords.device)[:k] - return coords[idx] - - -class DISTANCE_to_TARGET(PairwiseImageLoss): - """Chamfer-style distance loss penalising mislabelled voxels. - - For each foreground class, computes the mean minimum distance from - incorrectly predicted voxels to the nearest correct ground-truth voxel of - the same class. - - Args: - max_v: Scale factor to convert continuous predictions to integer labels. - res_gt: Spatial downsampling factor applied to the ground-truth when - computing nearest-neighbour distances (reduces memory usage). - """ - - def __init__( - self, - max_v: float = 1, - res_gt: int = 4, - *args, - **kwargs, - ) -> None: - self.max_v = max_v - self.res_gt = res_gt - super().__init__(*args, **kwargs) - - def forward( - self, - source: torch.Tensor, - target: torch.Tensor, - mask: torch.Tensor | None = None, # noqa: ARG002 - ) -> torch.Tensor: - """Chamfer-style distance loss for mis-labelled voxels. - - Parameters - ---------- - source : (D, H, W)[, …] torch.Tensor - Model prediction in label form (one channel per voxel). - target : (D, H, W)[, …] torch.Tensor - Ground-truth labels. - max_v : float, default 1 - Same scale factor you use elsewhere to map the continuous range [0, 1] - back to integer labels. Set to 1 if `source` and `target` are already - integer encoded. - - Returns: - ------- - torch.Tensor scalar - The mean distance (in voxel units) from every wrongly predicted voxel - to the nearest correct voxel of the same class in the target. - """ - max_v = self.max_v - device = source.device - # Discretise - src = (source * max_v).round().short() # .long() - tgt = (target * max_v).round().short() # .long() - - classes = torch.unique(tgt) - classes = classes[classes != 0] # skip background label 0 - - if classes.numel() == 0: - return torch.zeros(1, device=device) - - per_class_losses = [] - - for c in classes: - wrong_mask = (src == c) & (tgt != c) # voxels we predicted as c but shouldn't - if not wrong_mask.any(): - continue # no penalty if we never made that error - res_gt = self.res_gt - gt_mask = tgt[..., ::res_gt, ::res_gt, ::res_gt] == c - if not gt_mask.any(): - # Optional: if the class is missing in GT you could add - # a constant penalty or skip it. Here we skip. - continue - - # Coordinates of voxels - wrong_coords = torch.nonzero(wrong_mask, as_tuple=False).float() - # print(gt_mask.shape) - # - gt_coords = torch.nonzero(gt_mask, as_tuple=False).float() - - # Pairwise distances (N_wrong, N_gt); differentiable - d = torch.cdist(subsample_coords(wrong_coords, 5000), gt_coords) - min_dists = d.min(dim=1).to_numpy() # (N_wrong,) - - per_class_losses.append(min_dists.mean()) - - if not per_class_losses: - # Nothing to penalise - perfect overlap - return torch.zeros(1, device=device) - - # Average over foreground classes - return torch.stack(per_class_losses).mean() - - -class Rigid_Registration_with_Tether(General_Registration): - """Rigid (or affine) image registration with an optional centre-of-mass tether. - - Extends :class:`~TPTBox.registration._deepali.deepali_model.General_Registration` - with a patience-based early-stopping strategy and an additional - centre-of-mass regularisation term that keeps the warped image anchored - near the fixed image. - - Args: - fixed_image: Reference (target) image. - moving_image: Image to be registered. - reference_image: Optional reference image defining the output grid. - device: PyTorch device for computation. - gpu: GPU index used when ``ddevice`` is ``"cuda"``. - ddevice: Device type string (``"cpu"``, ``"cuda"``, or ``"mps"``). - fixed_mask: Optional foreground mask for the fixed image. - moving_mask: Optional foreground mask for the moving image. - normalize_strategy: Intensity normalisation strategy - (``"auto"``, ``"CT"``, ``"MRI"``, or ``None``). - pyramid_levels: Number of multi-resolution pyramid levels. - finest_level: Index of the finest pyramid level (0 = full resolution). - coarsest_level: Index of the coarsest pyramid level. - pyramid_finest_spacing: Voxel spacing at the finest pyramid level. - pyramid_min_size: Minimum spatial size per axis at the coarsest level. - dims: Spatial dimensions to include in the transform. - align: If ``True``, initialise the transform to align image centres. - transform_name: Name of the spatial transform class from DeepALI. - transform_args: Extra keyword arguments for the transform constructor. - transform_init: Optional path to a pre-computed flow field for - warm-starting. - optim_name: PyTorch optimiser class name. - lr: Learning rate. - optim_args: Extra keyword arguments for the optimiser (excluding ``lr``). - smooth_grad: Gradient smoothing factor. - verbose: Verbosity level (0 = silent). - max_steps: Maximum optimisation steps per pyramid level. - max_history: History length for convergence checking. - min_value: Minimum loss value for early stopping. - min_delta: Minimum loss improvement for early stopping. - loss_terms: Tuple of loss functions (image similarity + regulariser). - weights: Corresponding loss weights. - patience: Number of steps without improvement before stopping. - patience_delta: Minimum loss change to reset patience counter. - desc: Description string shown in the progress bar. - """ - - def __init__( - self, - fixed_image: Image_Reference, - moving_image: Image_Reference, - reference_image: Image_Reference | None = None, - device: Union[torch.device, str, int] | None = None, - gpu=0, - ddevice: DEVICES = "cuda", - # foreground_mask - fixed_mask=None, - moving_mask=None, - # normalize - normalize_strategy: Literal["auto", "CT", "MRI"] | None = None, - # Pyramid - pyramid_levels: int | None = None, # 1/None = no pyramid; int: number of stacks, tuple from to (0 is finest) - finest_level: int = 0, - coarsest_level: int | None = None, - pyramid_finest_spacing: Sequence[int] | torch.Tensor | None = None, - pyramid_min_size=16, - dims=("x", "y", "z"), - align=False, - transform_name: str = "RigidTransform", # Names that are defined in deepali.spatial.LINEAR_TRANSFORMS and deepali.spatialNONRIGID_TRANSFORMS. Override on_make_transform for finer control - transform_args: dict | None = None, - transform_init: PathStr | None = None, # reload initial flowfield from file - optim_name="Adam", # Optimizer name defined in torch.optim. or override on_optimizer finer control - lr=0.01, # Learning rate - optim_args=None, # args of Optimizer with out lr - smooth_grad=0.0, - verbose=0, - max_steps: int | Sequence[int] = 250, # Early stopping. override on_converged finer control - max_history: int | None = None, - min_value=0.0, # Early stopping. override on_converged finer control - min_delta=0.0, # Early stopping. override on_converged finer control - loss_terms=(ncc_loss, None), - weights=(1, 0.001), - patience=100, - patience_delta=0.0, - desc="RRwT", - ) -> None: - self.patience = patience - self.patience_delta = patience_delta - if device is None: - device = get_device(ddevice, gpu) - self.best = 1000000000 - self.best2 = 1000000000 - self.early_stopping = 0 - self.desc = desc - super().__init__( - fixed_image, - moving_image, - reference_image, - device=device, - fixed_mask=fixed_mask, - moving_mask=moving_mask, - normalize_strategy=normalize_strategy, - pyramid_levels=pyramid_levels, - finest_level=finest_level, - coarsest_level=coarsest_level, - pyramid_finest_spacing=pyramid_finest_spacing, - pyramid_min_size=pyramid_min_size, - dims=dims, - align=align, - transform_name=transform_name, - transform_args=transform_args, - transform_init=transform_init, - optim_name=optim_name, - lr=lr, - optim_args=optim_args, - smooth_grad=smooth_grad, - verbose=verbose, - max_steps=max_steps, - max_history=max_history, - min_value=min_value, - min_delta=min_delta, - loss_terms=loss_terms, - weights=weights, - ) - - def run_level( - self, - grid_transform: spatial.SequentialTransform, - target_image: deepaliImage, - source_image: deepaliImage, - opt: torch.optim.Optimizer, - lr_sq, # noqa: ARG002 - level, # noqa: ARG002 - max_steps: int, - sampling: Union[Sampling, str] = Sampling.LINEAR, # noqa: ARG002 - ) -> None: - """Optimise the transform at a single pyramid level. - - Args: - grid_transform: The spatial transform being optimised. - target_image: Fixed target image at the current resolution. - source_image: Moving source image at the current resolution. - opt: Initialised PyTorch optimiser. - lr_sq: Learning-rate scheduler (unused, kept for API compatibility). - level: Current pyramid level index (unused). - max_steps: Maximum number of gradient steps for this level. - sampling: Interpolation mode for image warping (unused at this - level; linear sampling is hard-coded). - """ - loss_list = [] - self.loss_values = loss_list - self.transform = grid_transform - transformer = spatial.ImageTransformer(grid_transform).to(self.device) - loss = next(iter(self.loss_terms.values())) - lambda_mse, lambda_com = self.weights.values() - pbar = tqdm(range(max_steps)) - for _ in pbar: - if self.on_converged(): - break - warped_batch = transformer(source_image) - l_mse = loss(warped_batch, target_image) - if lambda_com != 0: - # Compute center-of-mass loss - com_fixed = center_of_mass(target_image) - com_warped = center_of_mass(warped_batch) - l_com = torch.norm(com_fixed - com_warped) - if l_com < 10: - l_com = 0 - else: - l_com = 0 - l = lambda_mse * l_mse + lambda_com * l_com # Weighted sum of losses - - loss_list.append(l.item()) - opt.zero_grad() - l.backward() - opt.step() - pbar.desc = f"{self.desc} loss={l_mse.item() * lambda_mse:.5f}, center_of_mass={l_com * lambda_com:.5f}, {self.early_stopping=}, {self.best=}" - - def on_converged(self) -> bool: - """Check patience-based early-stopping convergence criterion. - - Returns: - ``True`` when the patience counter exceeds the configured patience - limit and the best transform has been restored; ``False`` otherwise. - """ - values = self.loss_values - if not values: - return False - value = values[-1] - if value <= self.best - self.patience_delta: - self.early_stopping = 0 - self.best = value - else: - self.early_stopping += 1 - - if value <= self.best2: - self.best_transform = deepcopy(self.transform) - self.best2 = value - - if self.early_stopping <= self.patience: - return False - self.transform = self.best_transform - return True