Skip to content

Commit a4df5a4

Browse files
authored
Merge branch 'main' into robert-pc-pr-stitching-refactor
2 parents 2350a5a + e012157 commit a4df5a4

5 files changed

Lines changed: 601 additions & 3 deletions

File tree

‎TPTBox/core/bids_files.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -243,10 +243,10 @@ def save_buffer(f: Path, buffer_name: str) -> list[Path]:
243243
print("[ ] Dose not exist:", (folder), f"{' ':20}") if verbose else None
244244
continue
245245
if (folder / buffer_name).exists():
246-
import datetime
246+
import datetime as dt
247247

248-
file_mod_time = datetime.datetime.fromtimestamp(os.path.getmtime(folder / buffer_name))
249-
today = datetime.datetime.today()
248+
file_mod_time = dt.datetime.fromtimestamp(os.path.getmtime(folder / buffer_name))
249+
today = dt.datetime.today()
250250

251251
age = today - file_mod_time
252252
if age.days >= int(max_age_days):

‎TPTBox/core/nii_wrapper.py‎

Lines changed: 174 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@
5858
np_unique,
5959
np_unique_withoutzero,
6060
np_volume,
61+
np_voronoi_labels,
6162
)
6263
from TPTBox.core.vert_constants import (
6364
AFFINE,
@@ -1808,6 +1809,59 @@ def dilate_msk_(self, n_pixel: int = 5, labels: LABEL_REFERENCE = None, connecti
18081809
"""In-place variant of `dilate_msk`."""
18091810
return self.dilate_msk(n_pixel=n_pixel, labels=labels, connectivity=connectivity, mask=mask, inplace=True, verbose=verbose,use_crop=use_crop,ignore_direction=ignore_direction)
18101811

1812+
def voronoi_labels(self, labels: LABEL_REFERENCE = None, max_distance: float | None = None, use_mm: bool = True, signed_background: bool = True, tie_to_zero: bool = True, inplace=False, verbose: logging = True) -> Self:
1813+
"""Assigns every voxel the label of the nearest labelled region (a Voronoi partition).
1814+
1815+
Args:
1816+
labels (LABEL_REFERENCE, optional): Labels to partition by. Everything else is treated as
1817+
background. Defaults to None (all labels present in the segmentation).
1818+
max_distance (float | None, optional): Leave background voxels further than this from any
1819+
region at 0. In mm when `use_mm` is True, else in voxels. Defaults to None (no limit).
1820+
use_mm (bool, optional): If True, distances are measured in mm using this image's voxel
1821+
spacing, so the partition is correct on anisotropic images. If False, voxels are
1822+
treated as isotropic. Defaults to True.
1823+
signed_background (bool, optional): If True, voxels filled in from the background get the
1824+
negated label of their region, so the input mask stays recoverable as `out > 0`. The
1825+
returned NII then has a signed dtype. Defaults to True.
1826+
tie_to_zero (bool, optional): If True, voxels equidistant to two different regions are set
1827+
to 0 instead of being assigned to one of them. Defaults to True.
1828+
inplace (bool, optional): Whether to modify this NII in place or return a new object.
1829+
Defaults to False.
1830+
verbose (bool, optional): Whether to print a message. Defaults to True.
1831+
1832+
Returns:
1833+
NII: The partition, signed when `signed_background` is set.
1834+
1835+
Notes:
1836+
Runtime is independent of the number of labels (one Euclidean feature transform for the
1837+
whole volume), unlike `dilate_msk`.
1838+
Set `signed_background=False, tie_to_zero=False` and a `max_distance` for a plain,
1839+
unsigned expansion of the labels (`np_expand_labels`).
1840+
1841+
Examples:
1842+
>>> nii.voronoi_labels(max_distance=10) # doctest: +SKIP
1843+
# every voxel within 1 cm of a labelled region takes its label
1844+
"""
1845+
assert self.seg
1846+
log.print("voronoi labels", end="\r", verbose=verbose)
1847+
out = np_voronoi_labels(
1848+
self.get_seg_array(),
1849+
label_ref=labels,
1850+
max_distance=max_distance,
1851+
zoom=self.zoom if use_mm else None,
1852+
signed_background=signed_background,
1853+
tie_to_zero=tie_to_zero,
1854+
)
1855+
if not signed_background:
1856+
# keep the image dtype; the negated labels of a signed partition would not fit an unsigned one
1857+
out = out.astype(self.dtype)
1858+
log.print("Voronoi partition of", len(np_unique_withoutzero(np.abs(out))), "labels", verbose=verbose)
1859+
return self.set_array(out, inplace=inplace)
1860+
1861+
def voronoi_labels_(self, labels: LABEL_REFERENCE = None, max_distance: float | None = None, use_mm: bool = True, signed_background: bool = True, tie_to_zero: bool = True, verbose: logging = True) -> Self:
1862+
"""In-place variant of `voronoi_labels`."""
1863+
return self.voronoi_labels(labels=labels, max_distance=max_distance, use_mm=use_mm, signed_background=signed_background, tie_to_zero=tie_to_zero, inplace=True, verbose=verbose)
1864+
18111865

18121866
def fill_holes(self, labels: LABEL_REFERENCE = None, slice_wise_dim: int|str | None = None, verbose:logging=False, inplace=False,use_crop=True) -> Self: # noqa: ARG002
18131867
"""Fills holes in segmentation.
@@ -2996,6 +3050,126 @@ def relabel_by_position_(self, axis: DIRECTIONS = "I", offset: int = 0, verbose:
29963050
"""In-place version of :meth:`relabel_by_position`."""
29973051
return self.relabel_by_position(axis=axis, offset=offset, inplace=True, verbose=verbose)
29983052

3053+
def _split_axis(self, axis: DIRECTIONS | int) -> tuple[int, bool]:
3054+
"""Resolves ``axis`` to an array axis and whether that direction runs with increasing index."""
3055+
if isinstance(axis, int):
3056+
return axis, True
3057+
# get_axis falls back to the opposite letter, so recover which way the axis actually runs.
3058+
return self.get_axis(axis), axis in self.orientation
3059+
3060+
def _split_label_profile(self, labels: LABEL_REFERENCE, ax: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
3061+
"""Returns the seg array, the mask of ``labels`` and its per-slice voxel count along ``ax``."""
3062+
arr = self.get_seg_array()
3063+
selected = arr > 0 if labels is None else np_isin(arr, labels)
3064+
counts = selected.sum(axis=tuple(i for i in range(arr.ndim) if i != ax))
3065+
return arr, selected, counts
3066+
3067+
def _split_label_at_index(self, arr: np.ndarray, selected: np.ndarray, ax: int, region: slice, offset: int, inplace: bool) -> Self:
3068+
"""Adds ``offset`` to the selected voxels that lie inside ``region`` along axis ``ax``."""
3069+
part = np.zeros_like(selected)
3070+
part[(slice(None),) * ax + (region,)] = True
3071+
part &= selected
3072+
if offset != 0 and part.any():
3073+
arr = arr.astype(np.promote_types(arr.dtype, np.min_scalar_type(int(arr.max()) + offset)), copy=False)
3074+
arr[part] += offset
3075+
return self.set_array(arr, inplace=inplace)
3076+
3077+
def partition_segmentation_by_volume(
3078+
self,
3079+
labels: LABEL_REFERENCE = None,
3080+
axis: DIRECTIONS | int = "P",
3081+
fraction: float = 0.5,
3082+
offset: int = 100,
3083+
inplace: bool = False,
3084+
) -> Self:
3085+
"""Splits labels along an axis at the plane that cuts off a given fraction of their volume.
3086+
3087+
Slices are walked inwards from the ``axis``-most end of the mask, summing voxels until
3088+
``fraction`` of the total volume is covered; that leading slab gets ``offset`` added to its
3089+
label values, the rest is left untouched. With the default ``axis="P"`` and
3090+
``fraction=1/3`` the posterior third (by volume, not by extent) is split off.
3091+
3092+
The direction is resolved against :attr:`orientation`, so no reorientation round-trip is
3093+
needed. Passing an ``int`` instead selects that array axis, pointing at high indices.
3094+
3095+
Args:
3096+
labels (LABEL_REFERENCE, optional): Label(s) to split. Defaults to None (every label).
3097+
axis (DIRECTIONS | int, optional): Anatomical direction the split-off part lies in,
3098+
or a raw array axis. Defaults to ``"P"``.
3099+
fraction (float, optional): Volume fraction that ends up in the split-off part, in
3100+
``[0, 1]``. Defaults to 0.5.
3101+
offset (int, optional): Added to the labels of the split-off part. Defaults to 100.
3102+
inplace (bool, optional): If True, modifies this NII in place. Defaults to False.
3103+
3104+
Returns:
3105+
NII: The segmentation with the split-off part relabeled to ``label + offset``.
3106+
3107+
Examples:
3108+
>>> ethmoid.partition_segmentation_by_volume([7, 8], "P", 1 / 3) # doctest: +SKIP
3109+
# posterior third of labels 7 and 8 becomes 107 and 108
3110+
"""
3111+
assert 0 <= fraction <= 1, f"fraction has to be in [0,1], got {fraction}"
3112+
ax, forward = self._split_axis(axis)
3113+
arr, selected, counts = self._split_label_profile(labels, ax)
3114+
total = int(counts.sum())
3115+
if total == 0:
3116+
return self if inplace else self.copy()
3117+
n = len(counts)
3118+
# walk inwards from the axis-most end until the slab holds fraction of the volume
3119+
order = counts[::-1] if forward else counts
3120+
target = fraction * total
3121+
taken = 0 if target <= 0 else min(int(np.searchsorted(np.cumsum(order), target)) + 1, n)
3122+
region = slice(n - taken, None) if forward else slice(None, taken)
3123+
return self._split_label_at_index(arr, selected, ax, region, offset, inplace)
3124+
3125+
def partition_segmentation_by_center_of_mass(
3126+
self,
3127+
labels: LABEL_REFERENCE = None,
3128+
axis: DIRECTIONS | int = "L",
3129+
offset: int = 100,
3130+
inplace: bool = False,
3131+
) -> Self:
3132+
"""Splits labels along an axis at their center of mass.
3133+
3134+
The center of mass of ``labels`` is projected onto ``axis`` and rounded to the nearest
3135+
slice; everything from there towards ``axis`` gets ``offset`` added to its label values.
3136+
With the default ``axis="L"`` this separates the left half of a structure from its right
3137+
half, and flipping the letter to ``"R"`` splits off the other side. The slice the plane
3138+
lands on always goes to the ``axis`` side, so the two halves overlap in that one slice
3139+
when the center of mass sits exactly on a voxel center.
3140+
3141+
The direction is resolved against :attr:`orientation` and the rounding is done in the
3142+
anatomical frame, so the split is the same no matter which orientation the image is
3143+
stored in and no reorientation round-trip is needed. Passing an ``int`` instead selects
3144+
that array axis, pointing at high indices.
3145+
3146+
Args:
3147+
labels (LABEL_REFERENCE, optional): Label(s) to split. Defaults to None (every label).
3148+
axis (DIRECTIONS | int, optional): Anatomical direction the split-off part lies in,
3149+
or a raw array axis. Defaults to ``"L"``.
3150+
offset (int, optional): Added to the labels of the split-off part. Defaults to 100.
3151+
inplace (bool, optional): If True, modifies this NII in place. Defaults to False.
3152+
3153+
Returns:
3154+
NII: The segmentation with the split-off part relabeled to ``label + offset``.
3155+
3156+
Examples:
3157+
>>> ethmoid.partition_segmentation_by_center_of_mass([7, 15], "L") # doctest: +SKIP
3158+
# left half of labels 7 and 15 becomes 107 and 115
3159+
"""
3160+
ax, forward = self._split_axis(axis)
3161+
arr, selected, counts = self._split_label_profile(labels, ax)
3162+
total = int(counts.sum())
3163+
if total == 0:
3164+
return self if inplace else self.copy()
3165+
n = len(counts)
3166+
com = float((counts * np.arange(n)).sum()) / total
3167+
# round in the anatomical frame, not the array frame, so the result does not depend on
3168+
# whether the axis happens to be stored forwards or backwards
3169+
c = int(np.clip(np.rint(com if forward else n - 1 - com), 0, n))
3170+
region = slice(c, None) if forward else slice(None, n - c)
3171+
return self._split_label_at_index(arr, selected, ax, region, offset, inplace)
3172+
29993173
def filter_connected_components_by_bbox_chain(
30003174
self,
30013175
margin_mm: float = 0.0,

‎TPTBox/core/np_utils.py‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@
2424
distance_transform_edt,
2525
gaussian_filter,
2626
generate_binary_structure,
27+
maximum_filter,
28+
minimum_filter,
2729
sobel,
2830
)
2931
from scipy.spatial.distance import cdist
@@ -598,6 +600,145 @@ def np_dilate_msk_euclid(arr: np.ndarray, n_pixel: int = 3, use_crop=True, label
598600
return arr
599601

600602

603+
def _np_voronoi_ties(arr: UINTARRAY, indices: np.ndarray, nearest: np.ndarray, zoom: Sequence[float] | None) -> np.ndarray:
604+
"""Finds the voxels that are equidistant to two differently labelled regions.
605+
606+
``distance_transform_edt`` picks one nearest source per voxel and never reports that the choice
607+
was arbitrary, so the ties have to be recovered afterwards. Only voxels next to a change in the
608+
nearest-label assignment can sit on a boundary between two regions, so the search is restricted
609+
to those and then each of their neighbours' nearest sources is tested for being at exactly the
610+
same distance. Any neighbour's source is a real foreground voxel, so a hit is always a genuine
611+
tie -- the test can miss one, but never invent one.
612+
613+
Args:
614+
arr (UINTARRAY): The labelled input array.
615+
indices (np.ndarray): Feature transform of ``arr``, shape ``(ndim, *arr.shape)``.
616+
nearest (np.ndarray): ``arr`` gathered at ``indices``.
617+
zoom (Sequence[float] | None): Voxel spacing used for the distances.
618+
619+
Returns:
620+
np.ndarray: Boolean mask of the tied voxels.
621+
"""
622+
ties = np.zeros(arr.shape, dtype=bool)
623+
boundary = maximum_filter(nearest, size=3) != minimum_filter(nearest, size=3)
624+
pts = np.argwhere(boundary)
625+
if len(pts) == 0:
626+
return ties
627+
sampling = np.ones(arr.ndim) if zoom is None else np.asarray(zoom, dtype=float)
628+
limit = np.asarray(arr.shape) - 1
629+
key = (slice(None), *pts.T)
630+
dist2 = (((indices[key].T - pts) * sampling) ** 2).sum(1)
631+
label = nearest[tuple(pts.T)]
632+
tied = np.zeros(len(pts), dtype=bool)
633+
for offset in itertools.product((-1, 0, 1), repeat=arr.ndim):
634+
if not any(offset):
635+
continue
636+
neighbour = np.clip(pts + offset, 0, limit)
637+
source = indices[(slice(None), *neighbour.T)]
638+
cand2 = (((source.T - pts) * sampling) ** 2).sum(1)
639+
tied |= (arr[tuple(source)] != label) & (np.abs(cand2 - dist2) <= 1e-9 * np.maximum(dist2, 1.0))
640+
ties[tuple(pts[tied].T)] = True
641+
return ties
642+
643+
644+
def np_voronoi_labels(
645+
arr: UINTARRAY,
646+
label_ref: LABEL_REFERENCE = None,
647+
max_distance: float | None = None,
648+
zoom: Sequence[float] | None = None,
649+
signed_background: bool = True,
650+
tie_to_zero: bool = True,
651+
) -> INTARRAY:
652+
"""Assigns every voxel the label of the nearest labelled region (a Voronoi partition).
653+
654+
Computed with a single Euclidean feature transform of the background: the nearest background
655+
voxel of ``arr == 0`` is exactly the nearest labelled voxel of ``arr``, so one pass yields the
656+
assignment for the whole volume, independent of how many labels there are.
657+
658+
By default the partition is signed: voxels that were already labelled keep their positive
659+
label, voxels filled in from the background carry the *negative* label of the region they were
660+
assigned to, so the original mask can still be recovered from the result. Voxels that are
661+
exactly equidistant to two different regions are set to 0 rather than being broken arbitrarily.
662+
663+
For a bounded expansion that overwrites the input in place and crops to the mask first, see
664+
:func:`np_dilate_msk_euclid`; this function is the unbounded, spacing-aware variant.
665+
666+
Args:
667+
arr (UINTARRAY): Labelled input array with background 0.
668+
label_ref (int | list[int] | None, optional): Labels to partition by. Everything else is
669+
treated as background. Defaults to None (all labels found in ``arr``).
670+
max_distance (float | None, optional): Leave background voxels further than this from any
671+
region at 0. In the units of ``zoom``, i.e. voxels when ``zoom`` is None.
672+
Defaults to None (no limit).
673+
zoom (Sequence[float] | None, optional): Voxel spacing, so distances are measured in mm on
674+
anisotropic images. Defaults to None (isotropic voxels).
675+
signed_background (bool, optional): If True, background voxels get the negated label of
676+
their region. Defaults to True.
677+
tie_to_zero (bool, optional): If True, voxels equidistant to two different regions are set
678+
to 0 instead of being assigned to one of them. Defaults to True.
679+
680+
Returns:
681+
INTARRAY: The partition, signed when ``signed_background`` is set.
682+
683+
Examples:
684+
>>> arr = np.zeros((5, 1, 1), dtype=np.uint8)
685+
>>> arr[0] = 1
686+
>>> arr[4] = 2
687+
>>> np_voronoi_labels(arr).ravel().tolist()
688+
[1, -1, 0, -2, 2]
689+
"""
690+
assert 2 <= arr.ndim <= 3, f"expected 2D or 3D, but got {arr.ndim}"
691+
assert np.min(arr) >= 0, f"expected non-negative labels, got {np.min(arr)}"
692+
assert zoom is None or len(zoom) == arr.ndim, f"zoom has to have {arr.ndim} entries, got {zoom}"
693+
if label_ref is not None:
694+
arr = np.where(np_isin(arr, label_ref), arr, 0)
695+
background = arr == 0
696+
if background.all():
697+
return np.zeros(arr.shape, dtype=np.int8 if signed_background else arr.dtype)
698+
699+
want_distance = max_distance is not None
700+
result = distance_transform_edt(background, sampling=zoom, return_distances=want_distance, return_indices=True)
701+
distances, indices = result if want_distance else (None, result)
702+
nearest = arr[tuple(indices)]
703+
704+
out = nearest.astype(np.min_scalar_type(-int(arr.max()) - 1) if signed_background else arr.dtype, copy=True)
705+
if signed_background:
706+
out[background] = -out[background]
707+
if want_distance:
708+
out[background & (distances > max_distance)] = 0
709+
if tie_to_zero:
710+
out[_np_voronoi_ties(arr, indices, nearest, zoom)] = 0
711+
return out
712+
713+
714+
def np_expand_labels(
715+
arr: UINTARRAY,
716+
distance: float | None = None,
717+
label_ref: LABEL_REFERENCE = None,
718+
zoom: Sequence[float] | None = None,
719+
) -> UINTARRAY:
720+
"""Grows every label into the surrounding background, up to ``distance``.
721+
722+
Each background voxel within reach takes the label of the nearest region; labelled voxels are
723+
left alone, so labels never eat into each other. This is :func:`np_voronoi_labels` with the
724+
plain, unsigned semantics -- see there for the signed partition and the handling of voxels that
725+
are equidistant to two regions.
726+
727+
Args:
728+
arr (UINTARRAY): Labelled input array with background 0.
729+
distance (float | None, optional): How far to grow, in the units of ``zoom``, i.e. voxels
730+
when ``zoom`` is None. Defaults to None (fill the whole volume).
731+
label_ref (int | list[int] | None, optional): Labels to grow. Everything else is treated as
732+
background. Defaults to None (all labels found in ``arr``).
733+
zoom (Sequence[float] | None, optional): Voxel spacing, so ``distance`` is in mm on
734+
anisotropic images. Defaults to None (isotropic voxels).
735+
736+
Returns:
737+
UINTARRAY: A copy of ``arr`` with the labels grown.
738+
"""
739+
return np_voronoi_labels(arr, label_ref=label_ref, max_distance=distance, zoom=zoom, signed_background=False, tie_to_zero=False)
740+
741+
601742
def np_dilate_msk(
602743
arr: np.ndarray,
603744
label_ref: LABEL_REFERENCE = None,

0 commit comments

Comments
 (0)