diff --git a/MIGRATION.md b/MIGRATION.md index dc1e327..0a58834 100644 --- a/MIGRATION.md +++ b/MIGRATION.md @@ -81,6 +81,18 @@ spineps sample -i scan.nii.gz --model-semantic t2w --model-instance instance --n | `Segmentation_Model_NNunet` | `SegmentationModelNNunet` | | `Segmentation_Model_Unet3D` | `SegmentationModelUnet3D` | +### Removed + +| Removed | Why / what to do instead | +| --- | --- | +| `--model-semantic auto` (CLI) and `process_dataset(model_semantic=None)` | The auto-selection it depended on (`spineps.seg_utils.find_best_matching_model`) was never implemented and always raised `NotImplementedError`. Name a model explicitly. | +| `spineps.seg_utils.find_best_matching_model` | See above. | +| `spineps.utils.image` (vendored spinalcordtoolbox `Image`) and `spineps.utils.generate_disc_labels` | Standalone disc-label export, wired to no entry point. Derive disc labels from the vertebra mask with `TPTBox` instead. | +| `spineps.architectures_new.unet2D` | The instance model is 3D only; `PLNet(do2D=True)` now raises. | +| `spineps.example` scripts | Never shipped in the wheel; see the README for usage examples. | +| `spineps.seg_pipeline.pipeline_revision` | Centroid metadata no longer records a git revision (see below). | +| `SPINEPS_TURN_OF_CITATION_REMINDER` | Renamed to `SPINEPS_NO_CITATION_REMINDER` (and it now actually works). | + `process_dataset`, `get_semantic_model`, `get_instance_model`, `get_labeling_model`, `predict_semantic_mask`, `predict_instance_mask` and the other phase functions keep their names. @@ -95,3 +107,45 @@ spineps sample -i scan.nii.gz --model-semantic t2w --model-instance instance --n (semantic sliding-window step), and `--tta` / `--no-tta` (toggle test-time mirroring; `SegmentationModel.set_test_time_augmentation(...)` in Python). - Clearer errors: invalid paths / missing models now raise `FileNotFoundError` / `ValueError` instead of bare `AssertionError`, and `spineps sample -h` / `dataset -h` no longer crash. + +## Fixed in 2.0 — output may change + +These are bug fixes, so results can differ from 1.x. Each was wrong before. + +- **Endplate labels reach the semantic mask.** The superior/inferior split was computed and then thrown away + by a binarising `extract_label`, so the `msk` output carried endplate voxels labelled `1` instead of + `Vertebral_Body_Endplate_Superior` (52) / `_Inferior` (53). If you worked around this, remove the workaround. +- **The semantic bounding-box clean keeps what it incorporates.** It grew a region to take in nearby connected + components and then cropped to the *largest* component's box anyway, deleting the rest. Spines split across + components (gaps, implants) keep more of the mask now. +- **No phantom disc.** `detect_and_solve_merged_vertebra` offset every voxel including the background, adding a + volume-sized fake IVD to the height-sorted list the split-C2 heuristic reads. +- **Small connected-component cleaning actually runs.** With `only_delete=False`, the neighbourhood mask was + destroyed before it was used, so nothing was ever deleted or relabelled despite the log saying otherwise. +- **Merged-corpus splitting.** `get_separating_components` returned its two parts *after* dilating them in + place, so they overlapped and the separating plane was derived from smeared centers of mass. +- **Incompatible models stop the run.** `process_dataset` logged "stop program" and then carried on; it now + raises `ValueError` unless `ignore_model_compatibility=True` / `--ignore-model-compatibility`. +- **Non-overlapping instance partners.** The couple search detected two partners that agree with the anchor but + not with each other, logged that it was skipping them, and used both anyway. +- **Version metadata.** `ctd.info["version"]` is the installed package version. It used to shell out to `git` + with no working directory, so it recorded whatever repository you happened to be standing in (or + `"Version not found"`), and `ctd.info["revision"]` is gone. +- **Labeling no longer crashes** on an empty instance mask, or with `disable_c1=False` and no subregion mask. +- **`import spineps` no longer creates a directory** inside the installed package; the fallback models folder is + created on demand. + +## Faster and smaller + +No output change -- pinned by voxel-identical regression tests. + +- The instance phase keeps each cutout prediction where it lives instead of in a dense + `(n_vertebrae, 3, *volume)` array, and compares candidates on their overlapping region only. +- The endplate splitter runs per vertebra on numpy arrays and grows each dilation by one voxel per round + instead of re-dilating from scratch. +- `clean_cc_artifacts` builds connected components one label at a time and works inside each component's + bounding box. +- The input volume is read from disk once per image instead of up to three times. + +On a 24-vertebra, 8.1M-voxel whole-spine volume: instance phase 2.30s -> 1.00s (peak RSS +313MB -> +157MB), +combined post-processing 3.34s -> 0.59s. diff --git a/README.md b/README.md index 1930861..1489bdc 100644 --- a/README.md +++ b/README.md @@ -177,6 +177,9 @@ spineps sample -i --model-semantic --model-instance (replacing `` with the model you want to use). You can also call SPINEPS from Python — see [Using the Code](#using-the-code). +SPINEPS prints a short citation reminder on first use and at exit. Set `SPINEPS_NO_CITATION_REMINDER=1` +(or `true`/`yes`/`on`) to silence it. + ### Issues - import issues: try installing via the requirements again, somethings it doesn't install everything @@ -276,7 +279,7 @@ To that end, we are using TPTBox (see https://github.com/Hendrik-code/TPTBox) | argument | explanation | | :--- | --------- | | --directory, -i, -d | Absolute path to the dataset directory, preferably a BIDS dataset (required) | -| --model-semantic, -ms | The model used for the semantic segmentation, or `auto` to select automatically by modality (default: t2w) | +| --model-semantic, -ms | The model used for the semantic segmentation (default: t2w) | | --model-instance, -mv, -mi | The model used for the vertebra instance segmentation (default: instance) | | --model-labeling, -ml | The (optional) VERIDAH model used for vertebra labeling (default: t2w_labeling) | | --rawdata-name, -rn | Sets the name of the rawdata folder of the dataset (default: "rawdata") @@ -320,9 +323,11 @@ In the subregion segmentation: | 47 | Inferior_Articular_Left | | 48 | Inferior_Articular_Right | | 49 | Vertebra_Corpus_border | +| 52 | Vertebral_Body_Endplate_Superior | +| 53 | Vertebral_Body_Endplate_Inferior | | 60 | Spinal_Cord | | 61 | Spinal_Canal | -| 62 | Endplate | +| 62 | Endplate (only where the plate could not be assigned to a vertebra) | | 100 | Vertebra_Disc | | 26 | Sacrum | diff --git a/docs/api/architectures.md b/docs/api/architectures.md index 6b9923d..3fdbf66 100644 --- a/docs/api/architectures.md +++ b/docs/api/architectures.md @@ -20,7 +20,7 @@ Network architectures and the vertebra label definitions used by the models. ## spineps.architectures_new.pl_unet -The newer PyTorch Lightning U-Net wrapper (2D or 3D), used by [`spineps.seg_model`](models.md). +The newer PyTorch Lightning U-Net wrapper, used by [`spineps.seg_model`](models.md). ::: spineps.architectures_new.pl_unet @@ -28,10 +28,6 @@ The newer PyTorch Lightning U-Net wrapper (2D or 3D), used by [`spineps.seg_mode ::: spineps.architectures_new.unet3D -## spineps.architectures_new.unet2D - -::: spineps.architectures_new.unet2D - ## spineps.architectures_new.dice ::: spineps.architectures_new.dice diff --git a/docs/api/utils.md b/docs/api/utils.md index 6df1055..ccbda8f 100644 --- a/docs/api/utils.md +++ b/docs/api/utils.md @@ -1,6 +1,6 @@ # Utilities -Image processing, the vertebra-labeling path solver, disc labeling and other helpers. +Image processing, the vertebra-labeling path solver and other helpers. ## spineps.utils.resolution @@ -17,10 +17,6 @@ behaviour is consistent across MRI and CT resolutions. ::: spineps.utils.find_min_cost_path -## spineps.utils.generate_disc_labels - -::: spineps.utils.generate_disc_labels - ## spineps.utils.filepaths ::: spineps.utils.filepaths diff --git a/pyproject.toml b/pyproject.toml index 8cb271e..f146b1b 100755 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,7 +94,6 @@ exclude = [ "dist", "node_modules", "venv", - "spineps/utils/image.py", # vendored from spinalcordtoolbox ".toml", ] line-length = 140 diff --git a/spineps/architectures_new/pl_unet.py b/spineps/architectures_new/pl_unet.py index 06fab73..1b07e3a 100644 --- a/spineps/architectures_new/pl_unet.py +++ b/spineps/architectures_new/pl_unet.py @@ -13,7 +13,6 @@ from torch.optim import Adam, lr_scheduler from .dice import MemoryEfficientSoftDiceLoss -from .unet2D import Unet2D from .unet3D import Unet3D @@ -42,36 +41,41 @@ def _tb_logger(module: pl.LightningModule) -> TensorBoardLogger: class PLNet(pl.LightningModule): - """LightningModule training a 2D or 3D U-Net with a combined cross-entropy, Dice and L2 loss. + """LightningModule training a 3D U-Net with a combined cross-entropy, Dice and L2 loss. - Wraps :class:`Unet2D` or :class:`Unet3D` and handles the training/validation loops, loss computation, - Dice metric logging and optimizer configuration. + Wraps :class:`Unet3D` and handles the training/validation loops, loss computation, Dice metric + logging and optimizer configuration. """ - def __init__(self, opt: Namespace | None = None, do2D: bool = False, num_channels=11, dim=8, *args: Any, **kwargs: Any) -> None: # ruff: ignore[unused-method-argument] + # The previous suppression comment here used a directive form ruff does not recognise, so it never + # silenced anything; ARG002 needs a real noqa. + def __init__(self, opt: Namespace | None = None, do2D: bool = False, num_channels=11, dim=8, *args: Any, **kwargs: Any) -> None: # noqa: ARG002 """Build the network and configure losses, metrics and training hyperparameters. Args: opt (Namespace): Configuration namespace providing ``channelwise``, ``n_epoch``, ``lr``, ``lr_end_factor``, ``l2_reg_w`` and ``dsc_loss_w``. - do2D (bool): If ``True``, use the 2D U-Net; otherwise the 3D U-Net. + do2D (bool): Kept so 2.x checkpoints still restore their saved hyperparameters. Only 3D is + supported; passing ``True`` raises. *args (Any): Unused positional arguments. **kwargs (Any): Unused keyword arguments. """ if opt is None: opt = Namespace(**kwargs) + if do2D: + raise NotImplementedError("PLNet only supports the 3D U-Net; the 2D variant was removed in SPINEPS 2.0") + super().__init__() self.save_hyperparameters() - arch = Unet2D if do2D else Unet3D - self.network = arch( + self.network = Unet3D( dim=dim, dim_mults=(1, 2, 4, 8), out_dim=4, channels=1 if not opt.channelwise else num_channels, ) - self.do2D = do2D + self.do2D = False self.n_epoch = opt.n_epoch self.start_lr = opt.lr self.linear_end_factor = opt.lr_end_factor @@ -314,4 +318,4 @@ def __str__(self): Returns: str: ``"Unet_2D"`` or ``"Unet_3D"``. """ - return f"Unet_{'2D' if self.do2D else '3D'}" + return "Unet_3D" diff --git a/spineps/architectures_new/unet2D.py b/spineps/architectures_new/unet2D.py deleted file mode 100644 index 4d84da6..0000000 --- a/spineps/architectures_new/unet2D.py +++ /dev/null @@ -1,570 +0,0 @@ -"""2D U-Net architecture with time, label and embedding conditioning for diffusion-style models.""" - -from __future__ import annotations - -import itertools -import math -from functools import partial - -import torch -from einops import rearrange -from torch import nn - - -def default(val, d): - """Return ``val`` if it is not ``None``, otherwise the default ``d``. - - Args: - val: The value to use when it is not ``None``. - d: The fallback value, or a zero-argument callable producing the fallback. - - Returns: - ``val`` when it is not ``None``; otherwise ``d()`` if ``d`` is a function, else ``d``. - """ - from inspect import isfunction - - if val is not None: - return val - return d() if isfunction(d) else d - - -# sinusoidal positional embeds - - -class SinusoidalPosEmb(nn.Module): - """Fixed sinusoidal positional embedding used to encode the diffusion time step.""" - - def __init__(self, dim): - """Initialize the embedding. - - Args: - dim (int): Output embedding dimension. Half is used for the sine and half for the cosine components. - """ - super().__init__() - self.dim = dim - - def forward(self, x): - """Compute the sinusoidal embedding for the given scalar values. - - Args: - x (torch.Tensor): Tensor of shape ``(b,)`` with the values (e.g. time steps) to embed. - - Returns: - torch.Tensor: Embedding of shape ``(b, dim)``. - """ - device = x.device - half_dim = self.dim // 2 - emb = math.log(10000) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, device=device) * -emb) - emb = x[:, None] * emb[None, :] - emb = torch.cat((emb.sin(), emb.cos()), dim=-1) - return emb - - -class LearnedSinusoidalPosEmb(nn.Module): - """following @crowsonkb 's lead with learned sinusoidal pos emb""" - - """ https://github.com/crowsonkb/v-diffusion-jax/blob/master/diffusion/models/danbooru_128.py#L8 """ - - def __init__(self, dim): - """Initialize the learned embedding. - - Args: - dim (int): Output embedding dimension; must be even. ``dim // 2`` learnable frequencies are used. - - Raises: - AssertionError: If ``dim`` is not even. - """ - super().__init__() - assert (dim % 2) == 0 - half_dim = dim // 2 - self.weights = nn.Parameter(torch.randn(half_dim)) # type: ignore - - def forward(self, x): - """Compute the learned Fourier embedding for the given scalar values. - - Args: - x (torch.Tensor): Tensor of shape ``(b,)`` with the values (e.g. time steps) to embed. - - Returns: - torch.Tensor: Embedding of shape ``(b, dim + 1)`` concatenating the input with its sine and cosine features. - """ - x = rearrange(x, "b -> b 1") - freqs = x * rearrange(self.weights, "d -> 1 d") * 2 * math.pi - fouriered = torch.cat((freqs.sin(), freqs.cos()), dim=-1) - fouriered = torch.cat((x, fouriered), dim=-1) - return fouriered - - -class PreNorm(nn.Module): - """Apply layer normalization to the input before passing it through a wrapped module.""" - - def __init__(self, dim, fn): - """Initialize the pre-normalization wrapper. - - Args: - dim (int): Number of channels to normalize over. - fn (nn.Module): Module applied to the normalized input. - """ - super().__init__() - self.fn = fn - self.norm = LayerNorm(dim) - - def forward(self, x, *args, **kwargs): - """Normalize ``x`` and forward it (with any extra arguments) through the wrapped module. - - Args: - x (torch.Tensor): Input tensor of shape ``(b, c, h, w)``. - *args: Positional arguments forwarded to the wrapped module. - **kwargs: Keyword arguments forwarded to the wrapped module. - - Returns: - torch.Tensor: Output of the wrapped module applied to the normalized input. - """ - x = self.norm(x) - return self.fn(x, *args, **kwargs) - - -class Residual(nn.Module): - """Add a skip connection around a wrapped module.""" - - def __init__(self, fn): - """Initialize the residual wrapper. - - Args: - fn (nn.Module): Module whose output is added to its input. - """ - super().__init__() - self.fn = fn - - def forward(self, x, *args, **kwargs): - """Forward ``x`` through the wrapped module and add the input as a residual. - - Args: - x (torch.Tensor): Input tensor. - *args: Positional arguments forwarded to the wrapped module. - **kwargs: Keyword arguments forwarded to the wrapped module. - - Returns: - torch.Tensor: ``fn(x) + x``. - """ - return self.fn(x, *args, **kwargs) + x - - -class LayerNorm(nn.Module): - """Channel-wise layer normalization for 4D ``(b, c, h, w)`` tensors with learnable scale and bias.""" - - def __init__(self, dim, eps=1e-5): - """Initialize the layer normalization. - - Args: - dim (int): Number of channels to normalize over. - eps (float): Small constant added to the variance for numerical stability. - """ - super().__init__() - self.eps = eps - self.g = nn.Parameter(torch.ones(1, dim, 1, 1)) # type: ignore - self.b = nn.Parameter(torch.zeros(1, dim, 1, 1)) # type: ignore - - def forward(self, x): - """Normalize ``x`` over the channel dimension and apply the learnable scale and bias. - - Args: - x (torch.Tensor): Input tensor of shape ``(b, c, h, w)``. - - Returns: - torch.Tensor: Normalized tensor of the same shape as ``x``. - """ - var = torch.var(x, dim=1, unbiased=False, keepdim=True) - mean = torch.mean(x, dim=1, keepdim=True) - return (x - mean) / (var + self.eps).sqrt() * self.g + self.b - - -class Block(nn.Module): - """Convolution, group normalization and SiLU activation, with optional FiLM-style scale/shift modulation.""" - - def __init__(self, dim, dim_out, groups=8): - """Initialize the block. - - Args: - dim (int): Number of input channels. - dim_out (int): Number of output channels. - groups (int): Number of groups for group normalization. - """ - super().__init__() - self.proj = nn.Conv2d(dim, dim_out, 3, padding=1) - self.norm = nn.GroupNorm(groups, dim_out) - self.act = nn.SiLU() - - def forward(self, x, scale_shift=None): - """Apply convolution, normalization, optional modulation and activation. - - Args: - x (torch.Tensor): Input tensor of shape ``(b, dim, h, w)``. - scale_shift (tuple[torch.Tensor, torch.Tensor] | None): Optional ``(scale, shift)`` tensors used to - modulate the normalized features as ``x * (scale + 1) + shift``. - - Returns: - torch.Tensor: Output tensor of shape ``(b, dim_out, h, w)``. - """ - x = self.proj(x) - x = self.norm(x) - - if scale_shift is not None: - scale, shift = scale_shift - x = x * (scale + 1) + shift - - x = self.act(x) - return x - - -class ResnetBlock(nn.Module): - """Residual block of two convolutional blocks with optional time-embedding conditioning.""" - - def __init__(self, dim, dim_out, *, time_emb_dim=None, groups=8): - """Initialize the residual block. - - Args: - dim (int): Number of input channels. - dim_out (int): Number of output channels. - time_emb_dim (int | None): Dimension of the time embedding. If given, an MLP produces per-channel - scale and shift values; if ``None``, no time conditioning is applied. - groups (int): Number of groups for the group normalization in each block. - """ - super().__init__() - self.mlp = nn.Sequential(nn.SiLU(), nn.Linear(time_emb_dim, dim_out * 2)) if time_emb_dim is not None else None - - self.block1 = Block(dim, dim_out, groups=groups) - self.block2 = Block(dim_out, dim_out, groups=groups) - self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity() - - def forward(self, x, time_emb=None): - """Apply the two blocks with optional time conditioning and a residual connection. - - Args: - x (torch.Tensor): Input tensor of shape ``(b, dim, h, w)``. - time_emb (torch.Tensor | None): Optional time embedding of shape ``(b, time_emb_dim)`` used to derive - the scale/shift modulation of the first block. - - Returns: - torch.Tensor: Output tensor of shape ``(b, dim_out, h, w)``. - """ - scale_shift = None - if self.mlp is not None and time_emb is not None: - time_emb = self.mlp(time_emb) - time_emb = rearrange(time_emb, "b c -> b c 1 1") - scale_shift = time_emb.chunk(2, dim=1) - - h = self.block1(x, scale_shift=scale_shift) - - h = self.block2(h) - - return h + self.res_conv(x) - - -class LinearAttention(nn.Module): - """Multi-head linear attention over spatial positions with linear complexity in the number of pixels.""" - - def __init__(self, dim, heads=4, dim_head=32): - """Initialize the linear attention module. - - Args: - dim (int): Number of input and output channels. - heads (int): Number of attention heads. - dim_head (int): Channel dimension per attention head. - """ - super().__init__() - self.scale = dim_head**-0.5 - self.heads = heads - hidden_dim = dim_head * heads - self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False) - - self.to_out = nn.Sequential(nn.Conv2d(hidden_dim, dim, 1), LayerNorm(dim)) - - def forward(self, x): - """Apply linear attention to the spatial feature map. - - Args: - x (torch.Tensor): Input tensor of shape ``(b, dim, h, w)``. - - Returns: - torch.Tensor: Output tensor of shape ``(b, dim, h, w)``. - """ - _b, _c, h, w = x.shape - qkv = self.to_qkv(x).chunk(3, dim=1) - q, k, v = (rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads) for t in qkv) - - q = q.softmax(dim=-2) - k = k.softmax(dim=-1) - - q = q * self.scale - context = torch.einsum("b h d n, b h e n -> b h d e", k, v) - - out = torch.einsum("b h d e, b h d n -> b h e n", context, q) - out = rearrange(out, "b h c (x y) -> b (h c) x y", h=self.heads, x=h, y=w) - return self.to_out(out) - - -class Attention(nn.Module): - """Standard multi-head softmax self-attention over spatial positions.""" - - def __init__(self, dim, heads=4, dim_head=32): - """Initialize the attention module. - - Args: - dim (int): Number of input and output channels. - heads (int): Number of attention heads. - dim_head (int): Channel dimension per attention head. - """ - super().__init__() - self.scale = dim_head**-0.5 - self.heads = heads - hidden_dim = dim_head * heads - self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False) - self.to_out = nn.Conv2d(hidden_dim, dim, 1) - - def forward(self, x): - """Apply full self-attention to the spatial feature map. - - Args: - x (torch.Tensor): Input tensor of shape ``(b, dim, h, w)``. - - Returns: - torch.Tensor: Output tensor of shape ``(b, dim, h, w)``. - """ - _b, _c, h, w = x.shape - qkv = self.to_qkv(x).chunk(3, dim=1) - q, k, v = (rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads) for t in qkv) - q = q * self.scale - - sim = torch.einsum("b h d i, b h d j -> b h i j", q, k) - sim = sim - sim.amax(dim=-1, keepdim=True).detach() - attn = sim.softmax(dim=-1) - - out = torch.einsum("b h i j, b h d j -> b h i d", attn, v) - out = rearrange(out, "b h (x y) d -> b (h d) x y", x=h, y=w) - return self.to_out(out) - - -class Unet2D(nn.Module): - """2D U-Net with residual blocks, (linear) attention and time/label/embedding conditioning. - - The network down-samples through a configurable number of resolution stages, applies a bottleneck with - full attention and up-samples again using skip connections. The diffusion time step is encoded via a - sinusoidal (or learned-sinusoidal) embedding, and the model can additionally be conditioned on a class - label and/or an external embedding vector. Optional patching folds spatial patches into the channel - dimension to reduce the spatial resolution processed by the network. - """ - - def __init__( - self, - dim, - init_dim=None, - out_dim=None, - dim_mults=(1, 2, 4, 8), - channels=1, - conditional_dimensions=0, - resnet_block_groups=8, - learned_variance=False, - learned_sinusoidal_cond=False, - learned_sinusoidal_dim=0, - conditional_label_size=0, - conditional_embedding_size=0, - patch_size=1, # Improving Diffusion Model Efficiency Through Patching https://arxiv.org/abs/2207.04316; 1 means deactivated (Note: Increases Training difficulty by a lot!) - ): - """Build the 2D U-Net layers. - - Args: - dim (int): Base feature dimension used to derive the channel widths and the time embedding size. - init_dim (int | None): Channels produced by the initial convolution. Defaults to ``dim``. - out_dim (int | None): Number of output channels before patch unfolding. Defaults to ``channels`` - (doubled when ``learned_variance`` is set). - dim_mults (tuple): Channel multipliers applied to ``dim`` for the successive resolution stages. - channels (int): Number of image channels. - conditional_dimensions (int): Number of additional conditioning channels concatenated to the input. - resnet_block_groups (int): Number of groups for the group normalization inside the residual blocks. - learned_variance (bool): If ``True``, double the default output channels to also predict a variance. - learned_sinusoidal_cond (bool): If ``True``, use a learned sinusoidal time embedding instead of a fixed one. - learned_sinusoidal_dim (int): Dimension of the learned sinusoidal embedding when enabled. - conditional_label_size (int): Number of classes for label conditioning; 0 disables it. - conditional_embedding_size (int): Size of an external embedding concatenated to the time embedding; 0 disables it. - patch_size (int): Spatial patch size folded into channels; 1 disables patching. - """ - super().__init__() - self.patch_size = patch_size - self.learned_variance = learned_variance - self.conditional_label_size = conditional_label_size - self.conditional_dimensions = conditional_dimensions - # determine dimensions - - self.channels = channels - - init_dim = default(init_dim, dim) - self.init_conv = nn.Conv2d((channels + conditional_dimensions) * patch_size * patch_size, init_dim, 7, padding=3) - - dims = [init_dim, *(int(dim * m) for m in dim_mults)] - in_out = list(itertools.pairwise(dims)) - - res_block = partial(ResnetBlock, groups=resnet_block_groups) - - # time embeddings - - time_dim = dim * 4 - - self.learned_sinusoidal_cond = learned_sinusoidal_cond - - if learned_sinusoidal_cond: - sinus_pos_emb = LearnedSinusoidalPosEmb(learned_sinusoidal_dim) - fourier_dim = learned_sinusoidal_dim + 1 - else: - sinus_pos_emb = SinusoidalPosEmb(dim) - fourier_dim = dim - - self.time_mlp = nn.Sequential(sinus_pos_emb, nn.Linear(fourier_dim, time_dim), nn.GELU(), nn.Linear(time_dim, time_dim)) - - if conditional_label_size != 0: - self.label_emb = nn.Embedding(conditional_label_size, time_dim) - - self.conditional_embedding_size = conditional_embedding_size - if conditional_embedding_size: - time_dim += conditional_embedding_size - # layers - - self.downs = nn.ModuleList([]) - self.ups = nn.ModuleList([]) - num_resolutions = len(in_out) - - for ind, (dim_in, dim_out) in enumerate(in_out): - is_last = ind >= (num_resolutions - 1) - - self.downs.append( - nn.ModuleList( - [ - res_block(dim_in, dim_out, time_emb_dim=time_dim), - res_block(dim_out, dim_out, time_emb_dim=time_dim), - Residual(PreNorm(dim_out, LinearAttention(dim_out))), - nn.Conv2d(dim_out, dim_out, 4, 2, 1) if not is_last else nn.Identity(), - ] - ) - ) - - mid_dim = dims[-1] - self.mid_block1 = res_block(mid_dim, mid_dim, time_emb_dim=time_dim) - self.mid_attn = Residual(PreNorm(mid_dim, Attention(mid_dim))) - self.mid_block2 = res_block(mid_dim, mid_dim, time_emb_dim=time_dim) - - for ind, (dim_in, dim_out) in enumerate(reversed(in_out)): - is_last = ind == (len(in_out) - 1) - - self.ups.append( - nn.ModuleList( - [ - res_block(dim_out * 2, dim_in, time_emb_dim=time_dim), - res_block(dim_in, dim_in, time_emb_dim=time_dim), - Residual(PreNorm(dim_in, LinearAttention(dim_in))), - nn.ConvTranspose2d(dim_in, dim_in, 4, 2, 1) if not is_last else nn.Identity(), - ] - ) - ) - - default_out_dim = channels * (1 if not learned_variance else 2) - self.out_dim = default(out_dim, default_out_dim) * patch_size * patch_size - - self.final_res_block = res_block(dim * 2, dim, time_emb_dim=time_dim) - self.final_conv = nn.Conv2d(dim, self.out_dim, 1) - - # Improving Diffusion Model Efficiency Through Patching https://arxiv.org/abs/2207.04316 (Note: Increases Training difficulty by a lot!) - def to_patches(self, x): - """Fold ``patch_size x patch_size`` spatial patches into the channel dimension. - - Args: - x (torch.Tensor): Input tensor of shape ``(B, C, H, W)`` with ``H`` and ``W`` divisible by ``patch_size``. - - Returns: - torch.Tensor: Tensor of shape ``(B, C * patch_size**2, H // patch_size, W // patch_size)``. - """ - p = self.patch_size - B, C, H, W = x.shape - x = x.permute(0, 2, 3, 1).reshape(B, H, W // p, C * p) - x = x.permute(0, 2, 1, 3).reshape(B, W // p, H // p, C * p * p) - return x.permute(0, 3, 2, 1) - - def from_patches(self, x): - """Invert :meth:`to_patches`, unfolding the channel dimension back into spatial patches. - - Args: - x (torch.Tensor): Patched tensor of shape ``(B, C, H, W)``. - - Returns: - torch.Tensor: Tensor of shape ``(B, C // patch_size**2, H * patch_size, W * patch_size)``. - """ - p = self.patch_size - B, C, H, W = x.shape - - x = x.permute(0, 3, 2, 1).reshape(B, W, H * p, C // p) - x = x.permute(0, 2, 1, 3).reshape(B, H * p, W * p, C // (p * p)) - return x.permute(0, 3, 1, 2) - - def forward(self, x, time=None, label=None, embedding=None) -> torch.Tensor: - """Run the U-Net forward pass. - - Args: - x (torch.Tensor): Input image tensor of shape ``(b, channels (+ conditional_dimensions), h, w)``. - time (torch.Tensor | None): Diffusion time steps of shape ``(b,)``. Defaults to a tensor of ones. - label (torch.Tensor | None): Class labels of shape ``(b,)``; required if the model was built with - ``conditional_label_size != 0``. - embedding (torch.Tensor | None): External conditioning embedding; required if the model was built - with ``conditional_embedding_size != 0``. - - Returns: - torch.Tensor: Output tensor of shape ``(b, out_dim, h, w)``. - - Raises: - AssertionError: If a required ``label`` or ``embedding`` is not provided. - """ - if self.patch_size != 1: - x = self.to_patches(x) - - x = self.init_conv(x) - r = x.clone() - - if time is None: - time = torch.ones((1,), device=x.device) - - t = self.time_mlp(time) - if hasattr(self, "label_emb"): - assert label is not None, "This UNet requires a class label" - t = t + self.label_emb(label) - - if self.conditional_embedding_size != 0: - assert embedding is not None, "This UNet requires a embedding" - # This is a general implementation, you my specialize this to your needs. The + operator instead of cat is possible. - t = torch.cat([embedding, t], dim=-1) - - h = [] - - for block1, block2, attn, downsample in self.downs: # type: ignore - x = block1(x, t) - x = block2(x, t) - x = attn(x) - h.append(x) - x = downsample(x) - - x = self.mid_block1(x, t) - x = self.mid_attn(x) - x = self.mid_block2(x, t) - - for block1, block2, attn, upsample in self.ups: # type: ignore - x = torch.cat((x, h.pop()), dim=1) - x = block1(x, t) - x = block2(x, t) - x = attn(x) - x = upsample(x) - - x = torch.cat((x, r), dim=1) - - x = self.final_res_block(x, t) - x = self.final_conv(x) - if self.patch_size != 1: - x = self.from_patches(x) - return x diff --git a/spineps/entrypoint.py b/spineps/entrypoint.py index cc1b30a..5b0fef2 100755 --- a/spineps/entrypoint.py +++ b/spineps/entrypoint.py @@ -182,7 +182,7 @@ def entry_point(): "--model-semantic", "-ms", default="t2w", - help="The model used for the subregion segmentation. Pass 'auto' to auto-select a model by modality, or an absolute path to the model folder", + help="The model used for the subregion segmentation. You can also pass an absolute path to the model folder", ) parser_dataset.add_argument( "--model-instance", @@ -333,8 +333,8 @@ def run_sample(opt: Namespace): def run_dataset(opt: Namespace): """Run the segmentation pipeline over a whole (preferably BIDS) dataset directory. - Resolves the semantic, instance and (optional) labeling models (``"auto"`` defers model selection to the - pipeline), then calls :func:`process_dataset`, optionally under a cProfiler. + Resolves the semantic, instance and (optional) labeling models, then calls :func:`process_dataset`, + optionally under a cProfiler. Args: opt (Namespace): Parsed CLI arguments from the ``dataset`` subcommand (dataset directory, rawdata and @@ -355,9 +355,7 @@ def run_dataset(opt: Namespace): raise NotADirectoryError(f"-directory is not a directory, got {input_dir}") # Model semantic - if opt.model_semantic == "auto": - model_semantic = None - elif "/" in str(opt.model_semantic): + if "/" in str(opt.model_semantic): model_semantic = get_actual_model(opt.model_semantic, use_cpu=opt.cpu).load() else: model_semantic = get_semantic_model(opt.model_semantic, use_cpu=opt.cpu).load() diff --git a/spineps/example/get_gpu.py b/spineps/example/get_gpu.py deleted file mode 100644 index 5c18504..0000000 --- a/spineps/example/get_gpu.py +++ /dev/null @@ -1,82 +0,0 @@ -"""Helpers for selecting idle GPUs and thread-aware logging when running SPINEPS in parallel.""" # noqa: INP001 - -from __future__ import annotations - -import time - -import GPUtil -from TPTBox import No_Logger - -logger = No_Logger() - - -def get_gpu(verbose: bool = False, max_load: float = 0.3, max_memory: float = 0.4): - """Return the IDs of currently available GPUs below the given load and memory thresholds. - - Args: - verbose (bool): If ``True``, print the current GPU utilization before querying. - max_load (float): Maximum allowed compute load (0-1) for a GPU to count as available. - max_memory (float): Maximum allowed memory usage (0-1) for a GPU to count as available. - - Returns: - list[int]: Up to four available GPU IDs ordered by load. - """ - GPUtil.showUtilization() if verbose else None - device_ids = GPUtil.getAvailable( - order="load", - limit=4, - maxLoad=max_load, - maxMemory=max_memory, - includeNan=False, - excludeID=[], - excludeUUID=[], - ) - return device_ids - - -def intersection(lst1, lst2): - """Return the set intersection of two iterables. - - Args: - lst1: First iterable. - lst2: Second iterable. - - Returns: - set: Elements present in both ``lst1`` and ``lst2``. - """ - return set(lst1).intersection(lst2) - - -def get_free_gpus(blocked_gpus=None, max_load: float = 0.3, max_memory: float = 0.4): - """Poll the GPUs repeatedly and return those consistently free and not explicitly blocked. - - Availability is sampled 15 times (a short sleep between samples) and intersected so that only GPUs that stay - idle across all samples are returned. - - Args: - blocked_gpus (dict[int, bool] | None): Mapping of GPU ID to a blocked flag; a GPU is excluded when its flag - is not ``False``. Defaults to ``{0: False, 1: False, 2: False, 3: False}``. - max_load (float): Maximum allowed compute load (0-1) for the initial availability query. - max_memory (float): Maximum allowed memory usage (0-1) for the initial availability query. - - Returns: - list[int]: IDs of GPUs that are consistently available and not blocked. - """ - if blocked_gpus is None: - blocked_gpus = {0: False, 1: False, 2: False, 3: False} - cached_list = get_gpu(max_load=max_load, max_memory=max_memory) - for _ in range(15): - time.sleep(0.25) - cached_list = intersection(cached_list, get_gpu()) - gpulist = [i for i in list(cached_list) if i not in blocked_gpus or blocked_gpus[i] is False] - return gpulist - - -def thread_print(fold, *text): - """Print a message prefixed with the fold identifier of the calling thread. - - Args: - fold: Identifier of the fold/thread used as the message prefix. - *text: Values to print after the prefix. - """ - logger.print(f"Fold [{fold}]: ", *text) diff --git a/spineps/example/helper_parallel.py b/spineps/example/helper_parallel.py deleted file mode 100755 index 820fad1..0000000 --- a/spineps/example/helper_parallel.py +++ /dev/null @@ -1,52 +0,0 @@ -"""CLI entry point running the SPINEPS pipeline on a single image, intended to be launched in parallel.""" # noqa: INP001 - -from __future__ import annotations - -import sys -from pathlib import Path - -file = Path(__file__).resolve() -sys.path.append(str(file.parents[1])) -sys.path.append(str(file.parents[2])) - -import argparse # noqa: E402 - -from TPTBox import BIDS_FILE # noqa: E402 - -from spineps.get_models import get_instance_model, get_semantic_model # noqa: E402 -from spineps.seg_run import segment_image # noqa: E402 - -# Example -# python /spineps/example/helper_parallel.py -i PATH/TO/IMG.nii.gz -ds DATASET-PATH -der derivatives -ms [t1w,t2w,vibe] -mv instance - -if __name__ == "__main__": - main_parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) - main_parser.add_argument("-i", type=str) - main_parser.add_argument("-ds", type=str) - main_parser.add_argument("-der", default="derivatives", type=str) - main_parser.add_argument("-ms", default="t2w", type=str) - main_parser.add_argument("-mv", default="instance", type=str) - main_parser.add_argument("-snap", default=None, type=str) - - opt = main_parser.parse_args() - - input_bids_file = BIDS_FILE(file=opt.i, dataset=opt.ds) - - ms = get_semantic_model(opt.ms) - mv = get_instance_model(opt.mv) - if opt.snap is not None: - Path(opt.snap).mkdir(exist_ok=True, parents=True) - segment_image( - img_ref=input_bids_file, - derivative_name=opt.der, - model_semantic=ms, - model_instance=mv, - override_semantic=False, - override_instance=False, - save_debug_data=False, - verbose=False, - ignore_compatibility_issues=False, # If true, we do not check if the file ending match like _T2w.nii.gz for T2w images - ignore_bids_filter=False, # If true, we do not check if BIDS compliant - save_raw=False, # Save output as they are produced by the model - snapshot_copy_folder=opt.snap, - ) diff --git a/spineps/example/template_roll_out.py b/spineps/example/template_roll_out.py deleted file mode 100755 index 4e79069..0000000 --- a/spineps/example/template_roll_out.py +++ /dev/null @@ -1,116 +0,0 @@ -"""Template script for batch-running the SPINEPS pipeline over a BIDS dataset (edit the TODO markers).""" # noqa: INP001 - -from __future__ import annotations - -import sys -from pathlib import Path - -file = Path(__file__).resolve() -sys.path.append(str(file.parents[1])) -sys.path.append(str(file.parents[2])) - -import time # noqa: E402 - -import numpy as np # noqa: E402 -from TPTBox import BIDS_FILE, NII, POI, BIDS_Global_info, No_Logger # noqa: E402 - -from spineps.get_models import get_instance_model, get_semantic_model # noqa: E402 -from spineps.seg_run import ErrCode, segment_image # noqa: E402 - -# INPUT -in_ds = Path("DATASET_PATH") # TODO change this to the path to your dataset folder -raw = "rawdata" # TODO change this to your rawdata directory name -der = "derivatives" # TODO change this to your derivatives directory name - -head_logger = ( - No_Logger() -) # (in_ds, log_filename="source-run-spineps", default_verbose=True) # TODO uncomment this to save a logger across the files - - -block = "" # TODO put i.e. 101 in here if your dataset is split into blocks and you want to specify one -parent_raw = str(Path(raw).joinpath(str(block))) -parent_der = str(Path(der).joinpath(str(block))) - -model_semantic = get_semantic_model("Model_name") # TODO put the modelname here -model_instance = get_instance_model("Model_name") # TODO put the modelname here - -bids_ds = BIDS_Global_info(datasets=[in_ds], parents=[parent_raw, parent_der], verbose=False) - -execution_times = [] - - -def injection_function(seg_nii: NII): - """Post-process the semantic segmentation mask before instance segmentation (placeholder hook). - - Passed as ``lambda_semantic`` to the pipeline; customize it to modify the semantic mask. By default it is - a no-op returning the mask unchanged. - - Args: - seg_nii (NII): Semantic segmentation mask produced by the pipeline. - - Returns: - NII: The (optionally modified) semantic segmentation mask. - """ - # TODO do something with semantic mask - return seg_nii - - -for name, subject in bids_ds.enumerate_subjects(sort=True): - logger = head_logger.add_sub_logger(name=name) - q = subject.new_query() - q.flatten() - q.filter("part", "inphase", required=False) - q.filter("chunk", "LWS") - q.unflatten() - q.filter_format("T2w") - q.filter_filetype("nii.gz") - families = q.loop_dict(sort=True, key_addendum=["part"]) - for f in families: - fid = f.family_id - - # TODO if you have a list of ids, check here if this family fits - - # TODO adapt this to your needs - if "T2w_part-inphase" not in f: - logger.print(f"{fid}: T2w_part-inphase not found, skip") - if "T2w" not in f: - logger.print(f"{fid}: T2w without part- not found, skip") - continue - - start_time = time.perf_counter() - ref: BIDS_FILE = f["T2w_part-inphase"][0] if "T2w_part-inphase" in f else f["T2w"][0] - # Call to the pipeline - output_paths, errcode = segment_image( - img_ref=ref, - derivative_name=der, - model_semantic=model_semantic, - model_instance=model_instance, - override_semantic=False, - override_instance=False, - lambda_semantic=injection_function, - save_debug_data=False, - verbose=False, - ) - # TODO measures time for each sample - end_time = time.perf_counter() - execution_time = end_time - start_time - logger.print(f"Inference time is: {execution_time}") - execution_times.append(execution_time) - - if errcode not in [ErrCode.OK, ErrCode.ALL_DONE]: - logger.print(f"{fid}: Pipeline threw errorcode {errcode}") - # TODO continue? assert? - - # TODO if you want to do something directly with the outputs again, you can load them like this - # Load Outputs - img_nii = ref.open_nii() - seg_nii = NII.load(output_paths["out_spine"], seg=True) # semantic mask - vert_nii = NII.load(output_paths["out_vert"], seg=True) # instance mask - ctd = POI.load(output_paths["out_ctd"]) # centroid file - - # TODO do something with outputs, potentially saving them again to the output paths - -if len(execution_times) > 0: - head_logger.print( - f"\nExecution times:\n{execution_times}\nRange:{min(execution_times)}, {max(execution_times)}\nAvg {np.average(execution_times)}" - ) diff --git a/spineps/get_models.py b/spineps/get_models.py index cbb11bd..7272851 100755 --- a/spineps/get_models.py +++ b/spineps/get_models.py @@ -182,9 +182,12 @@ def check_available_models( config_paths = search_path(models_folder, query="**/inference_config.json", suppress=True) global _modelid2folder_semantic, _modelid2folder_instance, _modelid2folder_labeling # noqa: PLW0603 - _modelid2folder_semantic = semantic # id to model_folder - _modelid2folder_instance = instances # id to model_folder - _modelid2folder_labeling = labeling + # Copies, not aliases: these dicts are seeded from the download registry and then filled with the + # locally found model folders. Mutating the registry itself would permanently replace its release + # URLs with local paths for the rest of the process. + _modelid2folder_semantic = dict(semantic) # id to model_folder + _modelid2folder_instance = dict(instances) # id to model_folder + _modelid2folder_labeling = dict(labeling) for cp in tqdm(config_paths, desc="Checking models"): model_folder = cp.parent model_folder_name = model_folder.parent.name.lower() if "nnUNetPlans" in model_folder.name else model_folder.name.lower() diff --git a/spineps/lab_model.py b/spineps/lab_model.py index 598638c..906f7da 100755 --- a/spineps/lab_model.py +++ b/spineps/lab_model.py @@ -112,7 +112,7 @@ def __init__( """ super().__init__(model_folder, inference_config, use_cpu, default_verbose, default_allow_tqdm) assert len(self.inference_config.expected_inputs) == 1, "Unet3D cannot expect more than one input" - self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + self.device = torch.device("cuda:0" if torch.cuda.is_available() and not use_cpu else "cpu") self.final_size: tuple[int, int, int] = DEFAULT_CLASSIFIER_INPUT_SIZE self.totensor = ToTensor() self.transform = Compose( @@ -334,9 +334,13 @@ def run_given_center_pos( img_v = img.set_array(arr_cut).reorient_(("I", "P", "L")) seg_v = seg.set_array(sem_cut).reorient_(("I", "P", "L")) + # Read the patches back in the model orientation: the cutouts above are still in the input's axis + # order, and the crop below (and set_array_) assume (I, P, L). + arr_cut = img_v.get_array() + sem_cut = seg_v.get_seg_array() if angle is not None and angle != 0: - arr_cut = rotate_patch_sagitally(img_v.get_array(), -angle, msk=False) - sem_cut = rotate_patch_sagitally(seg_v.get_seg_array(), -angle, msk=True) + arr_cut = rotate_patch_sagitally(arr_cut, -angle, msk=False) + sem_cut = rotate_patch_sagitally(sem_cut, -angle, msk=True) # crop down to final cutout size (200, 160, 32) arr_cut = arr_cut[ @@ -422,9 +426,10 @@ def _run_array(self, img_arr: np.ndarray, seg_arr: np.ndarray | torch.Tensor | N model_input = model_input.to(torch.float32) model_input = model_input.to(self.device) - self.predictor.eval() - self.predictor.to(self.device) - logits_dict = self.predictor.forward(model_input) - logits_soft = {k: self.predictor.softmax(v)[0].detach().cpu().numpy() for k, v in logits_dict.items()} + # eval()/to(device) are done once in load(); the autograd graph built without inference_mode was + # allocated and thrown away for every single vertebra. + with torch.inference_mode(): + logits_dict = self.predictor.forward(model_input) + logits_soft = {k: self.predictor.softmax(v)[0].detach().cpu().numpy() for k, v in logits_dict.items()} pred_cls = {k: np.argmax(v, 0) for k, v in logits_soft.items()} return logits_soft, pred_cls diff --git a/spineps/phase_instance.py b/spineps/phase_instance.py index 652f39f..e24b1de 100755 --- a/spineps/phase_instance.py +++ b/spineps/phase_instance.py @@ -3,6 +3,7 @@ from __future__ import annotations from collections.abc import Sequence +from typing import NamedTuple import numpy as np from TPTBox import NII, Location, Log_Type @@ -11,7 +12,6 @@ np_center_of_mass, np_connected_components, np_count_nonzero, - np_dice, np_dilate_msk, np_erode_msk, np_filter_connected_components, @@ -23,7 +23,7 @@ from spineps.seg_enums import ErrCode, OutputType from spineps.seg_model import SegmentationModel -from spineps.seg_pipeline import IVD_LABEL_OFFSET, logger +from spineps.seg_pipeline import IVD_LABEL_OFFSET, debug_put, logger from spineps.utils.proc_functions import clean_cc_artifacts from spineps.utils.resolution import ( INFERIOR_AXIS_PIR, @@ -46,6 +46,71 @@ MIN_NEIGHBORS_FOR_VOLUME_CHECK = 2 # A corpus is considered merged when its volume exceeds the neighbor average by this ratio. MERGED_CORPUS_VOLUME_RATIO = 1.5 +# Minimum Dice for two cutout predictions to be considered the same vertebra. +PREDICTION_COUPLE_DICE_THRESHOLD = 0.3 +# A couple overlapping already-established vertebrae by more than this fraction is discarded. +MAX_ESTABLISHED_OVERLAP = 0.6 + + +class SparsePrediction(NamedTuple): + """One cutout prediction, stored where it actually lives instead of in a whole-volume array. + + Each cutout only ever covers ``cutout_size`` voxels, so the dense ``(n_coms, 3, *volume)`` array this + replaces was almost entirely zeros -- hundreds of megabytes for a whole-spine scan, and every Dice + comparison below had to read all of it. + + Attributes: + bbox: Location of ``mask`` inside the cropped volume, one slice per axis. + mask: Boolean cutout mask, shaped like ``bbox``. + """ + + bbox: tuple[slice, slice, slice] + mask: np.ndarray + + @property + def size(self) -> int: + """Number of foreground voxels.""" + return int(np.count_nonzero(self.mask)) + + +def _bbox_overlap(a: tuple[slice, ...], b: tuple[slice, ...]) -> tuple[slice, ...] | None: + """Intersection of two bounding boxes, or None when they do not overlap.""" + out = [] + for sa, sb in zip(a, b): + start = max(sa.start, sb.start) + stop = min(sa.stop, sb.stop) + if start >= stop: + return None + out.append(slice(start, stop)) + return tuple(out) + + +def _local(bbox: tuple[slice, ...], window: tuple[slice, ...]) -> tuple[slice, ...]: + """Express ``window`` (absolute coordinates) relative to the start of ``bbox``.""" + return tuple(slice(w.start - b.start, w.stop - b.start) for b, w in zip(bbox, window)) + + +def sparse_dice(a: SparsePrediction, b: SparsePrediction) -> float: + """Dice score between two sparsely stored binary predictions. + + Identical to ``np_dice`` on the two full-volume masks: outside the intersection of the bounding + boxes at least one operand is zero everywhere, so it cannot contribute to the intersection. + + Args: + a (SparsePrediction): First prediction. + b (SparsePrediction): Second prediction. + + Returns: + float: The Dice score; 1.0 when both masks are empty, matching ``np_dice``'s NaN handling. + """ + denom = a.size + b.size + if denom == 0: + return 1.0 + window = _bbox_overlap(a.bbox, b.bbox) + if window is None: + return 0.0 + intersect = int(np.count_nonzero(a.mask[_local(a.bbox, window)] & b.mask[_local(b.bbox, window)])) + return (2.0 * intersect) / denom def predict_instance_mask( @@ -103,7 +168,7 @@ def predict_instance_mask( # shp = seg_nii.shape seg_nii_rdy = seg_nii.reorient(verbose=logger) - debug_data["inst_uncropped_Subreg_nii_a_PIR"] = seg_nii_rdy.copy() + debug_put(debug_data, "inst_uncropped_Subreg_nii_a_PIR", seg_nii_rdy.copy) # Padding? if pad_size > 0: @@ -118,14 +183,14 @@ def predict_instance_mask( logger.print( "Vertebra seg_nii_uncropped", seg_nii_uncropped.zoom, seg_nii_uncropped.orientation, seg_nii_uncropped.shape, verbose=verbose ) - debug_data["inst_uncropped_Subreg_nii_b_zms"] = seg_nii_uncropped.copy() + debug_put(debug_data, "inst_uncropped_Subreg_nii_b_zms", seg_nii_uncropped.copy) uncropped_vert_mask = np.zeros(seg_nii_uncropped.shape, dtype=seg_nii_uncropped.dtype) logger.print("Vertebra uncropped_vert_mask empty", uncropped_vert_mask.shape, verbose=verbose) crop = seg_nii_rdy.compute_crop(dist=INSTANCE_CROP_MARGIN_MM / min(seg_nii_rdy.zoom)) seg_nii_rdy.apply_crop_(crop) logger.print(f"Crop down from {uncropped_vert_mask.shape} to {seg_nii_rdy.shape}", verbose=verbose) logger.print("Vertebra seg_nii_rdy", seg_nii_rdy.zoom, seg_nii_rdy.orientation, seg_nii_rdy.shape, verbose=verbose) - debug_data["inst_cropped_Subreg_nii_b"] = seg_nii_rdy.copy() + debug_put(debug_data, "inst_cropped_Subreg_nii_b", seg_nii_rdy.copy) # # make threshold in actual mm corpus_border_threshold = int(corpus_border_threshold / expected_zms[1]) @@ -138,7 +203,7 @@ def predict_instance_mask( logger.print(f"no corpus ({Location.Vertebra_Corpus_border.value}) labels in this segmentation, cannot proceed", Log_Type.FAIL) return None, ErrCode.EMPTY # get all the 3vert predictions - vert_predictions, hierarchical_existing_predictions, n_corpus_coms = collect_vertebra_predictions( + vert_predictions, n_corpus_coms = collect_vertebra_predictions( seg_nii=seg_nii_rdy, model=model, corpus_size_cleaning=corpus_size_cleaning if proc_corpus_clean else 0, @@ -159,12 +224,12 @@ def predict_instance_mask( whole_vert_nii, debug_data, errcode = from_vert3_predictions_make_vert_mask( seg_nii_rdy, vert_predictions, - hierarchical_existing_predictions, + n_corpus_coms, vert_size_threshold, debug_data=debug_data, proc_inst_clean_small_cc_artifacts=proc_inst_clean_small_cc_artifacts, ) - del vert_predictions, hierarchical_existing_predictions + del vert_predictions if errcode != ErrCode.OK: return None, errcode logger.print("Merged predictions into vert mask") @@ -177,7 +242,7 @@ def predict_instance_mask( if proc_inst_fill_3d_holes: whole_vert_nii.fill_holes_(verbose=logger) - debug_data["inst_cropped_vert_arr_c_proc"] = whole_vert_nii.copy() + debug_put(debug_data, "inst_cropped_vert_arr_c_proc", whole_vert_nii.copy) n_vert_bodies = len(uniq_labels) logger.print(f"Predicted {n_vert_bodies} vertebrae") if n_vert_bodies < n_corpus_coms: @@ -195,7 +260,7 @@ def predict_instance_mask( uncropped_vert_mask[crop] = vert_nii_cleaned.get_seg_array() logger.print(f"Uncrop back from {vert_nii_cleaned.shape} to {uncropped_vert_mask.shape}", verbose=verbose) whole_vert_nii_uncropped = seg_nii_uncropped.set_array(uncropped_vert_mask) - debug_data["inst_uncropped_vert_arr_a"] = whole_vert_nii_uncropped.copy() + debug_put(debug_data, "inst_uncropped_vert_arr_a", whole_vert_nii_uncropped.copy) # Uncrop again if pad_size > 0: @@ -297,9 +362,13 @@ def get_corpus_coms( stats_by_height = dict(sorted(stats.items(), key=lambda x: x[1][0])) stats_by_height_keys = list(stats_by_height.keys()) + key_position = {k: n for n, k in enumerate(stats_by_height_keys)} - for vl in stats_by_height_keys: - idx = stats_by_height_keys.index(vl) + # Iterate a snapshot: the working list is rebuilt below whenever a same-height neighbour is merged away. + for vl in list(stats_by_height_keys): + if vl not in stats_by_height: + continue + idx = key_position[vl] statsvl = stats_by_height[vl] is_ivd = statsvl[1] @@ -324,13 +393,14 @@ def get_corpus_coms( stats_by_height.pop(vl) stats_by_height = dict(sorted(stats_by_height.items(), key=lambda x: x[1][0])) stats_by_height_keys = list(stats_by_height.keys()) + key_position = {k: n for n, k in enumerate(stats_by_height_keys)} continue logger.print("Merged corpi, try to fix it", verbose=verbose) neighbor_verts = { stats_by_height_keys[idx + i]: stats_by_height[stats_by_height_keys[idx + i]] for i in NEIGHBOR_OFFSETS - if (idx + i) < len(stats_by_height_keys) and (idx + i) >= 0 and stats_by_height_keys[idx + i] < 99 + if (idx + i) < len(stats_by_height_keys) and (idx + i) >= 0 and stats_by_height_keys[idx + i] < IVD_LABEL_OFFSET } logger.print("neighbor_vert_labels", neighbor_verts, verbose=verbose) @@ -360,6 +430,9 @@ def get_corpus_coms( logger.print("Splitting by plane") plane_split_nii = get_plane_split(segvert, corpus_nii, spart, tpart, spart_dil, tpart_dil) split_vert = split_by_plane(segvert, plane_split_nii) + # NOTE: `stats` / `stats_by_height` deliberately keep their pre-split values; the loop + # only ever splits off one extra corpus per detected alternation error, and the final + # centers of mass below are recomputed from `corpus_cc` anyway. corpus_cc[split_vert == 2] = corpus_cc.max() + 1 except Exception as e: logger.print(f"Separating Corpi failed with exception {e}", Log_Type.FAIL) @@ -404,22 +477,26 @@ def get_separating_components( vol_old = vol.copy() iterations = 0 while True: - vol_erode = np_erode_msk(vol, n_pixel=1, connectivity=connectivity) + # np_erode_msk mutates its input and returns the same object, so erode a copy. Without it `vol`, + # `vol_old` and `vol_erode` all end up aliasing one array after the first iteration, "the iteration + # before" is lost, and the subreg_cc_n == 0 branch below could never find its two parts. + vol_erode = np_erode_msk(vol.copy(), n_pixel=1, connectivity=connectivity) subreg_cc, subreg_cc_n = np_connected_components(vol_erode, connectivity=check_connectivity) if subreg_cc_n > 1: vol = subreg_cc break elif subreg_cc_n == 0: # np.max(subreg_cc)# 1 not in np_unique(subreg_cc) - vol_dilated = np_dilate_msk(vol, n_pixel=1, connectivity=connectivity, mask=vol.copy()) + # Dilate a copy: `vol` is read again two lines down, and np_dilate_msk works in place. + vol_dilated = np_dilate_msk(vol.copy(), n_pixel=1, connectivity=connectivity, mask=vol.copy()) # use iteration before to get other CC vol[vol_old != 0] = 2 # all possible voxels are 2 vol[vol_dilated == 1] = 1 - if 2 not in np_volume(vol): + volume = np_volume(vol) + if 1 not in volume or 2 not in volume: raise Exception( # noqa: TRY002 - f"cannot split volume into two parts after {iterations} iterations, all values are 0 after erosion." + f"cannot split volume into two parts after {iterations} iterations, got regions {volume}." ) - volume = np_volume(vol) dil_iter = 0 while volume[1] / (volume[1] + volume[2]) < 0.5: vol_dilated = np_dilate_msk(vol_dilated, n_pixel=1, connectivity=connectivity, mask=vol.copy()) @@ -456,8 +533,11 @@ def get_separating_components( if spart.sum() == 0 or tpart.sum() == 0: raise Exception("S or T are empty") # noqa: TRY002 - spart_dil = np_dilate_msk(spart, n_pixel=1, connectivity=connectivity) - tpart_dil = np_dilate_msk(tpart, n_pixel=1, connectivity=connectivity) + # Dilate copies: np_dilate_msk works in place and returns its input, so dilating `spart`/`tpart` + # directly grew them too -- the function then returned two overlapping blobs as "the two separated + # components", and get_plane_split took its normal vector between their smeared centers of mass. + spart_dil = np_dilate_msk(spart.copy(), n_pixel=1, connectivity=connectivity) + tpart_dil = np_dilate_msk(tpart.copy(), n_pixel=1, connectivity=connectivity) stpart = (spart_dil + (tpart_dil * 2)).astype(np.uint8) while 3 not in np_volume(stpart): spart_dil = np_dilate_msk(spart_dil, n_pixel=1, connectivity=connectivity) @@ -636,7 +716,7 @@ def collect_vertebra_predictions( instance_batch_size: int = 4, amp: bool = False, verbose: bool = False, -) -> tuple[np.ndarray | None, list[str], int]: +) -> tuple[dict[tuple[int, int], SparsePrediction] | None, int]: """Run the instance model on a cutout around each corpus center of mass and collect per-label predictions. Computes corpus centers of mass, and for each one extracts a ``cutout_size`` window (nudged inferiorly until @@ -660,9 +740,9 @@ def collect_vertebra_predictions( verbose (bool, optional): Emit additional progress logging. Defaults to False. Returns: - tuple[np.ndarray | None, list[str], int]: A hierarchical prediction array of shape - ``(n_corpus_coms, 3, *seg_shape)``, a list of ``"comidx_label"`` identifiers for the predictions actually - produced, and the number of corpus centers of mass. Returns ``(None, [], 0)`` if no corpus is found. + tuple[dict[tuple[int, int], SparsePrediction] | None, int]: The produced predictions keyed by + ``(corpus-com index, label index)``, and the number of corpus centers of mass. Returns ``(None, 0)`` + if no corpus is found. """ corpus_coms = get_corpus_coms( seg_nii, @@ -671,24 +751,16 @@ def collect_vertebra_predictions( verbose=verbose, ) if corpus_coms is None: - return None, [], 0 + return None, 0 n_corpus_coms = len(corpus_coms) if n_corpus_coms < 3: logger.print(f"Too few vertebra semantically segmented ({n_corpus_coms}), might have bad result", Log_Type.WARNING) - # return None, [], 0 - - shp = ( - # n_corpus_coms, - # 3 - seg_nii.shape[0], - seg_nii.shape[1], - seg_nii.shape[2], - ) - hierarchical_existing_predictions = [] - # Holds only binary {0, 1} per-label masks, so uint8 is sufficient (the source dtype can be wider, - # which would needlessly inflate this n_coms x 3 x volume array and slow the Dice comparisons below). - hierarchical_predictions = np.zeros((n_corpus_coms, 3, *shp), dtype=np.uint8) + + shp = seg_nii.shape + # Each prediction is kept where it lives (see SparsePrediction); the dense (n_coms, 3, *volume) array + # this replaces was almost all zeros and dominated both peak memory and the Dice comparisons below. + predictions: dict[tuple[int, int], SparsePrediction] = {} # relabel to the labels expected by the model # {41: 1, 42: 2, 43: 3, 44: 4, 45: 5, 46: 6, 47: 7, 48: 8, 49: 9, 50: 9, Location.Dens_axis.value: 9} @@ -718,7 +790,7 @@ def collect_vertebra_predictions( # Calc cutout cut_nii, cutout_coords, paddings = nii_calc_crop_around_centerpoint(com, seg_arr_c, cutout_size) # cut_nii = seg_nii_for_cut.set_array(arr_cut, verbose=False).reorient_() - debug_data[f"inst_cutout_vert_nii_{com_idx}_cut"] = cut_nii + debug_put(debug_data, f"inst_cutout_vert_nii_{com_idx}_cut", lambda c=cut_nii: c) cut_niis.append(cut_nii) cut_meta.append((com_idx, com, cutout_coords, paddings)) @@ -735,7 +807,7 @@ def collect_vertebra_predictions( for (com_idx, com, cutout_coords, paddings), results in zip(cut_meta, batched_results): vert_cut_nii = results[OutputType.seg].reorient_() - debug_data[f"inst_cutout_vert_nii_{com_idx}_pred"] = vert_cut_nii.copy() + debug_put(debug_data, f"inst_cutout_vert_nii_{com_idx}_pred", vert_cut_nii.copy) vert_cut_nii = post_process_single_3vert_prediction( vert_cut_nii, None, @@ -743,14 +815,14 @@ def collect_vertebra_predictions( largest_cc=proc_inst_largest_k_cc, # type:ignore ) vert_labels = vert_cut_nii.unique() # 1,2,3 - debug_data[f"inst_cutout_vert_nii_{com_idx}_proc"] = vert_cut_nii.copy() + debug_put(debug_data, f"inst_cutout_vert_nii_{com_idx}_proc", vert_cut_nii.copy) cutout_sizes = tuple(cutout_coords[i].stop - cutout_coords[i].start for i in range(len(cutout_coords))) pad_cutout = tuple(slice(paddings[i][0], paddings[i][0] + cutout_sizes[i]) for i in range(len(paddings))) arr = vert_cut_nii.get_seg_array() cutout_vals = arr[pad_cutout] - # Write straight into the (already fully-allocated) hierarchical_predictions slice instead of building - # full-volume-sized temporaries per vertebra/label: everything outside cutout_coords is 0 either way. + # Store the cutout-sized mask together with where it belongs, rather than scattering it into a + # full-volume array. local_com = tuple(int(com[i]) - cutout_coords[i].start for i in range(3)) seg_at_com = cutout_vals[local_com] if seg_at_com == 0: @@ -759,9 +831,8 @@ def collect_vertebra_predictions( mask = cutout_vals == l labelindex = l - 1 if mask.any(): - hierarchical_predictions[com_idx, labelindex][cutout_coords] = mask.astype(np.uint8) - hierarchical_existing_predictions.append(str_id_com_label(com_idx, labelindex)) - return hierarchical_predictions, hierarchical_existing_predictions, n_corpus_coms + predictions[(com_idx, labelindex)] = SparsePrediction(cutout_coords, mask) + return predictions, n_corpus_coms def post_process_single_3vert_prediction( @@ -789,29 +860,16 @@ def post_process_single_3vert_prediction( return vert_nii -def str_id_com_label(com_idx: int, label: int) -> str: - """Build the string identifier for a single (corpus-com, label) prediction. - - Args: - com_idx (int): Index of the corpus center of mass. - label (int): Label index within that center's three-vertebra prediction. - - Returns: - str: The identifier ``"{com_idx}_{label}"``. - """ - return str(com_idx) + "_" + str(label) - - def from_vert3_predictions_make_vert_mask( seg_nii: NII, - vert_predictions: np.ndarray, # already hierarchical [com_idx, l, map] - hierarchical_existing_predictions: list[str], # list of actually used vert predictions + vert_predictions: dict[tuple[int, int], SparsePrediction], + n_corpus_coms: int, vert_size_threshold: int, debug_data: dict, proc_inst_clean_small_cc_artifacts: bool = True, verbose: bool = False, ) -> tuple[NII, dict, ErrCode]: - """Merge the hierarchical three-vertebra predictions into a single vertebra instance mask. + """Merge the per-cutout three-vertebra predictions into a single vertebra instance mask. Each per-label prediction looks among neighboring predictions (center index -2 to +2, all three labels) for its most-agreeing partners (by Dice), forming prediction couples. The couples are then merged into one @@ -819,8 +877,8 @@ def from_vert3_predictions_make_vert_mask( Args: seg_nii (NII): Reference segmentation providing shape and spatial metadata. - vert_predictions (np.ndarray): Hierarchical predictions of shape ``(com_idx, label, *shape)``. - hierarchical_existing_predictions (list[str]): Identifiers of the predictions that were actually produced. + vert_predictions (dict[tuple[int, int], SparsePrediction]): Predictions keyed by ``(com index, label index)``. + n_corpus_coms (int): Number of corpus centers of mass the predictions were made around. vert_size_threshold (int): Voxel threshold for removing small instance artifacts. debug_data (dict): Dictionary for collecting intermediate results. proc_inst_clean_small_cc_artifacts (bool, optional): Whether to delete small instance artifacts. Defaults to True. @@ -835,16 +893,15 @@ def from_vert3_predictions_make_vert_mask( # idx is always in the order in predictions (so bottom2up corpus CC) # arcus_coms sorted bottom to top - hierarchical_predictions = vert_predictions # search space: all neighboring predictions # all search for up to two other predictions with best agreement - coupled_predictions = create_prediction_couples(hierarchical_predictions, hierarchical_existing_predictions) + coupled_predictions = create_prediction_couples(vert_predictions, n_corpus_coms) logger.print("Coupled predictions", coupled_predictions, verbose=verbose) return merge_coupled_predictions( seg_nii, coupled_predictions=coupled_predictions, - hierarchical_predictions=hierarchical_predictions, + predictions=vert_predictions, debug_data=debug_data, proc_clean_small_cc_artifacts=proc_inst_clean_small_cc_artifacts, vert_size_threshold=vert_size_threshold, @@ -853,35 +910,31 @@ def from_vert3_predictions_make_vert_mask( def create_prediction_couples( - hierarchical_predictions: np.ndarray, - hierarchical_existing_predictions, + predictions: dict[tuple[int, int], SparsePrediction], + n_predictions: int, verbose: bool = False, ) -> dict: - """Form and rank prediction couples across all hierarchical predictions. + """Form and rank prediction couples across all cutout predictions. For every (center index, label) prediction, finds its best-agreeing partners and groups them into a couple, averaging the agreement scores of duplicate couples. The result is sorted so that larger, higher-agreement couples come first (key = ``(len(couple) + 1) * mean_agreement``). Args: - hierarchical_predictions (np.ndarray): Hierarchical predictions of shape ``(com_idx, label, *shape)``. - hierarchical_existing_predictions (list[str]): Identifiers of the predictions that were actually produced. + predictions (dict[tuple[int, int], SparsePrediction]): Predictions keyed by ``(com index, label index)``. + n_predictions (int): Number of corpus centers of mass. verbose (bool, optional): Emit additional progress logging. Defaults to False. Returns: dict: Mapping from each couple (a tuple of ``(com_idx, label)`` members) to its mean agreement score, ordered by descending size-weighted agreement. """ - n_predictions = hierarchical_predictions.shape[0] - # Set for O(1) membership in the inner candidate search (called 3 * n_predictions times). - existing_predictions = set(hierarchical_existing_predictions) - coupled_predictions = {} # TODO try to calculate list of candidates here, take the predictions and then parallelize the find_prediction_couple for idx in range(n_predictions): for pred in range(3): - couple, agreement = find_prediction_couple(idx, pred, hierarchical_predictions, existing_predictions, n_predictions, verbose) + couple, agreement = find_prediction_couple(idx, pred, predictions, n_predictions, verbose) if couple is None: continue if couple not in coupled_predictions: @@ -899,39 +952,24 @@ def create_prediction_couples( return coupled_predictions -def parallel_dice(anchor, pred, cand_loc: tuple) -> tuple[float, tuple]: - """Compute the Dice score between two masks, tagged with a candidate location. - - Args: - anchor (np.ndarray): Anchor prediction mask. - pred (np.ndarray): Candidate prediction mask to compare against. - cand_loc: Candidate location identifier carried through unchanged. - - Returns: - tuple[float, Any]: The Dice score between ``anchor`` and ``pred`` and the passed-through ``cand_loc``. - """ - return float(np_dice(anchor, pred)), cand_loc - - def find_prediction_couple( idx, pred, - hierarchical_predictions: np.ndarray, - hierarchical_existing_predictions, + predictions: dict[tuple[int, int], SparsePrediction], n_predictions, verbose: bool = False, ) -> tuple[tuple | None, float]: """Find the best-agreeing partner predictions for one anchor prediction. Considers candidate predictions within +/-2 of the anchor's center index (all three labels, excluding the - anchor itself), ranks them by Dice with the anchor, and keeps up to the two best whose Dice exceeds 0.3. - The anchor itself is appended, and the members are returned sorted by center index. + anchor itself), ranks them by Dice with the anchor, and keeps up to the two best whose Dice exceeds + ``PREDICTION_COUPLE_DICE_THRESHOLD``. If two partners are kept but do not overlap each other, only the better + one survives. The anchor itself is appended, and the members are returned sorted by center index. Args: idx (int): Center-of-mass index of the anchor prediction. pred (int): Label index of the anchor prediction. - hierarchical_predictions (np.ndarray): Hierarchical predictions of shape ``(com_idx, label, *shape)``. - hierarchical_existing_predictions (list[str]): Identifiers of the predictions that were actually produced. + predictions (dict[tuple[int, int], SparsePrediction]): Predictions keyed by ``(com index, label index)``. n_predictions (int): Total number of corpus centers of mass. verbose (bool, optional): Emit additional progress logging. Defaults to False. @@ -939,24 +977,18 @@ def find_prediction_couple( tuple[tuple | None, float]: The couple (a sorted tuple of ``(com_idx, label)`` members including the anchor) and its mean partner agreement. Returns ``(None, 0)`` if the anchor prediction does not exist. """ - if str_id_com_label(idx, pred) not in hierarchical_existing_predictions: - logger.print(f"{str_id_com_label(idx, pred)} not in predictions {hierarchical_existing_predictions}", verbose=verbose) + if (idx, pred) not in predictions: + logger.print(f"({idx}, {pred}) not in predictions {sorted(predictions)}", verbose=verbose) return None, 0 - anchor = hierarchical_predictions[idx][pred] + anchor = predictions[(idx, pred)] dices = {} min_idx = max(0, idx - 2) max_idx = min(idx + 2, n_predictions) - list_of_candidates = [ - (i, l) - for i in range(min_idx, max_idx + 1) - for l in [0, 1, 2] - if i != idx and str_id_com_label(i, l) in hierarchical_existing_predictions - ] - - # list_of_candidates = np.array(np.meshgrid(idx_candidates, [0, 1, 2])).T.reshape(-1, 2) + list_of_candidates = [(i, l) for i in range(min_idx, max_idx + 1) for l in [0, 1, 2] if i != idx and (i, l) in predictions] + for cand_loc in list_of_candidates: - dices[tuple(cand_loc)] = float(np_dice(anchor, hierarchical_predictions[cand_loc[0]][cand_loc[1]])) + dices[cand_loc] = sparse_dice(anchor, predictions[cand_loc]) # find k best partners dices = dict(sorted(dices.items(), key=lambda item: item[1], reverse=True)) @@ -964,23 +996,17 @@ def find_prediction_couple( best_k = best_k[:2] couple = [] - dice_threshold = 0.3 - if len(best_k) > 0 and dices[best_k[0]] > dice_threshold: + if len(best_k) > 0 and dices[best_k[0]] > PREDICTION_COUPLE_DICE_THRESHOLD: couple.append(best_k[0]) - if len(best_k) > 1 and dices[best_k[1]] > dice_threshold: + if len(best_k) > 1 and dices[best_k[1]] > PREDICTION_COUPLE_DICE_THRESHOLD: couple.append(best_k[1]) - # if dices[best_k[2]] > dice_threshold: - # couple.append(best_k[2]) if len(couple) == 2: - # sort out if the other two do not overlap over threshold - dice_partners = float( - np_dice( - hierarchical_predictions[best_k[0][0]][best_k[0][1]], - hierarchical_predictions[best_k[1][0]][best_k[1][1]], - ) - ) - if dice_partners < dice_threshold: - logger.print(couple, " was skipped because the partners do not overlap", verbose=verbose) + # The two partners agree with the anchor but not with each other, so they cannot both be the + # same vertebra -- keep only the better one (best_k is sorted by descending Dice). + dice_partners = sparse_dice(predictions[best_k[0]], predictions[best_k[1]]) + if dice_partners < PREDICTION_COUPLE_DICE_THRESHOLD: + logger.print(couple[1], "was dropped because the partners do not overlap", verbose=verbose) + couple = couple[:1] agreement = 0 if len(couple) > 0: @@ -996,7 +1022,7 @@ def find_prediction_couple( def merge_coupled_predictions( seg_nii: NII, coupled_predictions, - hierarchical_predictions: np.ndarray, + predictions: dict[tuple[int, int], SparsePrediction], debug_data: dict, proc_clean_small_cc_artifacts: bool = True, vert_size_threshold: int = 0, @@ -1007,12 +1033,13 @@ def merge_coupled_predictions( Iterates over the couples in priority order, summing their member maps and thresholding by voxel agreement (requiring overlap from at least two members unless the couple is small or low-agreement). Each accepted couple is written as a new instance label into voxels not yet claimed; couples overlapping established - vertebrae by more than 60% are skipped. Small connected-component artifacts are optionally cleaned afterwards. + vertebrae by more than ``MAX_ESTABLISHED_OVERLAP`` are skipped. Small connected-component artifacts are + optionally cleaned afterwards. Args: seg_nii (NII): Reference segmentation providing shape and spatial metadata. coupled_predictions (dict): Mapping from couple to mean agreement, ordered by priority. - hierarchical_predictions (np.ndarray): Hierarchical predictions of shape ``(com_idx, label, *shape)``. + predictions (dict[tuple[int, int], SparsePrediction]): Predictions keyed by ``(com index, label index)``. debug_data (dict): Dictionary for collecting intermediate results. proc_clean_small_cc_artifacts (bool, optional): Whether to delete small instance artifacts. Defaults to True. vert_size_threshold (int, optional): Voxel threshold for removing small instance artifacts. Defaults to 0. @@ -1024,7 +1051,6 @@ def merge_coupled_predictions( """ whole_vert_nii = seg_nii.copy() whole_vert_arr = np.zeros(whole_vert_nii.shape, dtype=np.uint16) # this is fixed segmentations from vert - combine = np.zeros(whole_vert_nii.shape, dtype=whole_vert_nii.dtype) # reused scratch buffer, reset per couple below idx = 1 for k, overall_agreement in coupled_predictions.items(): @@ -1032,27 +1058,35 @@ def merge_coupled_predictions( take_no_overlap = len(k) <= 2 if overall_agreement < 0.3 + 0.15 * (4 - len(k)): take_no_overlap = True - combine.fill(0) + # A couple only ever covers the union of its members' cutouts, so accumulate and write there + # instead of allocating and scanning whole-volume buffers per couple. + member_bboxes = [predictions[cid].bbox for cid in k] + window = tuple( + slice(min(b[axis].start for b in member_bboxes), max(b[axis].stop for b in member_bboxes)) + for axis in range(len(member_bboxes[0])) + ) + combine = np.zeros(tuple(s.stop - s.start for s in window), dtype=np.uint8) for cid in k: - combine += hierarchical_predictions[cid[0]][cid[1]] + member = predictions[cid] + combine[_local(window, member.bbox)] += member.mask m = 1 if take_no_overlap else 2 - # m = min(max(1, np.max(combine)), 2) # type:ignore - combine[combine < m] = 0 - combine[combine != 0] = idx + selected = combine >= m - count_new = np_count_nonzero(combine) + count_new = int(np.count_nonzero(selected)) if count_new == 0: logger.print("ZERO instance mask failure on vertebra instance creation", Log_Type.FAIL) return seg_nii, debug_data, ErrCode.EMPTY - count_cut = np_count_nonzero((combine != 0) & (whole_vert_arr == 0)) + target = whole_vert_arr[window] + free = selected & (target == 0) + count_cut = int(np.count_nonzero(free)) relative_overlap = (count_new - count_cut) / count_new - if relative_overlap > 0.6: + if relative_overlap > MAX_ESTABLISHED_OVERLAP: logger.print(k, f" was skipped because it overlaps {round(relative_overlap, 4)} with established verts", verbose=verbose) continue - whole_vert_arr[whole_vert_arr == 0] = combine[whole_vert_arr == 0] + target[free] = idx idx += 1 - debug_data["inst_crop_vert_arr_a_raw"] = seg_nii.set_array(whole_vert_arr) + debug_put(debug_data, "inst_crop_vert_arr_a_raw", lambda: seg_nii.set_array(whole_vert_arr)) if np_is_empty(whole_vert_arr): logger.print("Vert mask empty, will skip", Log_Type.FAIL) diff --git a/spineps/phase_labeling.py b/spineps/phase_labeling.py index 1af46d0..b7a8312 100644 --- a/spineps/phase_labeling.py +++ b/spineps/phase_labeling.py @@ -54,6 +54,8 @@ T13_LABEL = 28 # Crop margin in millimeters kept around the vertebrae before labeling. LABELING_CROP_MARGIN_MM = 128 +# A dens (odontoid process) overlapping an instance by at least this physical volume identifies that instance as C2. +MIN_DENS_VOLUME_MM3 = 250 def perform_labeling_step( @@ -84,21 +86,25 @@ def perform_labeling_step( if model.predictor is None: model.load() + vert_nii_u = vert_nii.unique() + if len(vert_nii_u) == 0: + logger.on_fail("perform_labeling_step: instance mask is empty, nothing to label") + return vert_nii + if subreg_nii is not None: # crop for corpus instead of whole vertebra corpus_nii = subreg_nii.extract_label((Location.Vertebra_Corpus, Location.Vertebra_Corpus_border, Location.Dens_axis)) vert_nii_c = vert_nii * corpus_nii else: vert_nii_c = vert_nii - vert_nii_u = vert_nii.unique() force_c2 = None force_c1 = None - if not disable_c1: + if not disable_c1 and subreg_nii is not None: dense = vert_nii_c * subreg_nii.extract_label(Location.Dens_axis.value) volumes = dense.volumes(in_mm3=True) # dict[label, volume_mm3] if volumes: max_label, max_volume = max(volumes.items(), key=lambda x: x[1]) - if max_volume > 250: + if max_volume > MIN_DENS_VOLUME_MM3: force_c2 = max_label # Force C1 if the preceding vertebra exists if force_c2 != 1: diff --git a/spineps/phase_post.py b/spineps/phase_post.py index 1a3b6b9..634d3f7 100644 --- a/spineps/phase_post.py +++ b/spineps/phase_post.py @@ -5,6 +5,7 @@ import heapq import numpy as np +from scipy.ndimage import binary_dilation, generate_binary_structure from TPTBox import NII, Location, Log_Type, v_idx2name, v_name2idx from TPTBox.core.np_utils import ( np_bbox_binary, @@ -22,7 +23,7 @@ ) from spineps.phase_labeling import VertLabelingClassifier, perform_labeling_step -from spineps.seg_pipeline import ENDPLATE_LABEL_OFFSET, IVD_LABEL_OFFSET, logger, vertebra_subreg_labels +from spineps.seg_pipeline import ENDPLATE_LABEL_OFFSET, IVD_LABEL_OFFSET, debug_put, logger, vertebra_subreg_labels from spineps.utils.compat import zip_strict from spineps.utils.proc_functions import fix_wrong_posterior_instance_label from spineps.utils.resolution import REFERENCE_VOXEL_VOLUME_MM3, REFERENCE_ZOOM, isotropic_area_to_voxels @@ -183,7 +184,7 @@ def phase_postprocess_combined( logger.print("vert_uncropped volumes", vert_uncropped.volumes()) logger.print("seg_uncropped", seg_uncropped.unique()) - debug_data["vert_arr_return_final"] = vert_uncropped.copy() + debug_put(debug_data, "vert_arr_return_final", vert_uncropped.copy) return seg_uncropped, vert_uncropped @@ -236,7 +237,10 @@ def mask_cleaning_other( subreg_arr[deletion_map == 1] = 0 n_vert_pixels = np_count_nonzero(vert_arr_cleaned) - n_subreg_vert_pixels = subreg_vert_nii.volumes()[1] + n_subreg_vert_pixels = subreg_vert_nii.volumes().get(1, 0) + if n_subreg_vert_pixels == 0 or n_vert_bodies == 0: + logger.print("No vertebra voxels to reconcile between the instance and semantic masks", Log_Type.WARNING) + return whole_vert_nii.set_array(vert_arr_cleaned), seg_nii.set_array(subreg_arr) n_vert_pixel_per_vertebra = n_subreg_vert_pixels / n_vert_bodies n_difference_pixels = n_subreg_vert_pixels - n_vert_pixels if n_difference_pixels > 0: @@ -309,7 +313,7 @@ def assign_missing_cc( ) subreg_arr_vert_rest = reference_arr.copy() subreg_arr_vert_rest[target_arr_ != 0] = 0 - deletion_map = np.zeros(reference_arr.shape) + deletion_map = np.zeros_like(reference_arr, dtype=np.uint8) label_rest = np_unique(subreg_arr_vert_rest) if len(label_rest) == 1 and label_rest[0] == 0: @@ -365,6 +369,85 @@ def assign_missing_cc( return target_arr, reference_arr, deletion_map +def _split_endplates(seg_t: NII, vert_t: NII, ep_labels: list[int], verbose: bool = True) -> NII: + """Divide the endplate band into superior and inferior plates by growing each vertebra into it. + + Each vertebra instance is dilated one voxel per round; endplate voxels it reaches, and that no other + vertebra has claimed yet, are labelled from the endplate instance id already stored in ``vert_t``: + the current vertebra's own plate becomes ``Vertebral_Body_Endplate_Inferior`` and the previous + vertebra's becomes ``Vertebral_Body_Endplate_Superior``. Whatever is still unclaimed when the rounds + end keeps the generic ``Endplate`` label. + + This is the hot loop of the whole post-processing phase, so it runs on numpy arrays inside each + vertebra's own window rather than on whole-volume NII operators (every NII operator copies the array + twice), and each round dilates the previous round's mask by one instead of re-dilating the original + by ``dil``. + + Args: + seg_t (NII): Subregion semantic mask in PIR, providing the endplate voxels and the output grid. + vert_t (NII): Vertebra instance mask in PIR, already carrying the endplate instance ids. + ep_labels (list[int]): Semantic labels that count as endplate. + verbose (bool): If True, report the detected fraction while iterating. + + Returns: + NII: A mask holding only the superior/inferior/unassigned endplate labels. + """ + inferior = Location.Vertebral_Body_Endplate_Inferior.value + superior = Location.Vertebral_Body_Endplate_Superior.value + + ep_arr = seg_t.extract_label(ep_labels).get_seg_array().astype(bool) + vert_vals = vert_t.get_seg_array() + # int32 so `out + plates` cannot wrap: the intermediate values are vertebra ids offset by + # ENDPLATE_LABEL_OFFSET and can be summed where two rounds touch the same voxel. + out_arr = np.zeros(seg_t.shape, dtype=np.int32) + total = int(ep_arr.sum()) + if total == 0: + return seg_t.set_array(out_arr) + + # vert_t.unique() is sorted, so this is the same set the old `if i >= LIMIT: break` selected. + vert_labels_to_split = [int(i) for i in vert_t.unique() if i < INSTANCE_LABEL_LIMIT] + # Per-vertebra window: its bounding box grown by the largest dilation we can apply, so a window-local + # dilation is identical to the global one everywhere it can matter. + windows: dict[int, tuple[slice, ...]] = {} + grown: dict[int, np.ndarray] = {} + for i in vert_labels_to_split: + label_mask = vert_vals == i + bbox = np_bbox_binary(label_mask, px_dist=MAX_ENDPLATE_DILATION) + windows[i] = bbox + grown[i] = label_mask[bbox] + # connectivity=3 (26-neighbourhood), matching NII.dilate_msk's default. scipy's binary dilation is the + # same operation as TPTBox's np_dilate_msk on a single-label binary mask, but not a per-voxel Python loop. + struct = generate_binary_structure(3, 3) + + pref = 1 + old_vol = -1 + for _dil in range(1, MAX_ENDPLATE_DILATION): + new_vol = int(np.count_nonzero((out_arr == inferior) | (out_arr == superior))) + logger.print(rf"{new_vol / total * 100:.2f}% endplates detected", end="\r") if verbose else None + if old_vol == new_vol and old_vol != 0: + break + old_vol = new_vol + if total == new_vol: + logger.print("Found all Endplates ") + break + for i in vert_labels_to_split: + bbox = windows[i] + grown[i] = binary_dilation(grown[i], structure=struct) + out_w = out_arr[bbox] + unclaimed = (out_w != inferior) & (out_w != superior) + reached = ep_arr[bbox] & grown[i] & unclaimed + plates = np.where(reached, vert_vals[bbox], 0) + plates = np_map_labels(plates, {i + ENDPLATE_LABEL_OFFSET: inferior, pref + ENDPLATE_LABEL_OFFSET: superior}) + out_arr[bbox] = out_w + plates + pref = i + + # whatever no vertebra reached keeps the generic endplate label + leftover = ep_arr & (out_arr != inferior) & (out_arr != superior) + out_arr[leftover] += Location.Endplate.value + keep = (out_arr == inferior) | (out_arr == superior) | (out_arr == Location.Endplate.value) + return seg_t.set_array(np.where(keep, out_arr, 0)) + + def add_ivd_ep_vert_label(whole_vert_nii: NII, seg_nii: NII, include_sacrum=False, verbose=True) -> tuple[np.ndarray, np.ndarray]: """Attach intervertebral-disc and endplate instance labels and split endplates into superior/inferior. @@ -393,16 +476,10 @@ def add_ivd_ep_vert_label(whole_vert_nii: NII, seg_nii: NII, include_sacrum=Fals vert_arr = vert_t.get_seg_array() subreg_arr = seg_t.get_seg_array() - coms_vert_dict = {} - for l in vert_labels: - vert_l = vert_arr.copy() - vert_l[vert_l != l] = 0 - vert_l[subreg_arr != 49] = 0 # com of corpus region - vert_l[vert_l != 0] = 1 - try: - coms_vert_dict[l] = np_center_of_mass(vert_l)[1][1] # center_of_mass(vert_l)[1] - except Exception: - coms_vert_dict[l] = 0 + # One pass over the corpus voxels for every label at once, instead of a whole-volume copy per label. + corpus_instances = np.where(subreg_arr == Location.Vertebra_Corpus_border.value, vert_arr, 0) + corpus_coms = np_center_of_mass(corpus_instances) + coms_vert_dict = {l: (corpus_coms[l][1] if l in corpus_coms else 0) for l in vert_labels} coms_vert_y = list(coms_vert_dict.values()) coms_vert_labels = list(coms_vert_dict.keys()) @@ -411,17 +488,15 @@ def add_ivd_ep_vert_label(whole_vert_nii: NII, seg_nii: NII, include_sacrum=Fals n_ivd_unique = 0 if Location.Vertebra_Disc.value in seg_t.unique(): # Map IVDS - subreg_cc = seg_t.get_connected_components(labels=Location.Vertebra_Disc.value) - subreg_cc_n = len(subreg_cc.unique()) - subreg_cc = subreg_cc.get_seg_array() - cc_labelset = list(range(1, subreg_cc_n + 1)) + subreg_cc = seg_t.get_connected_components(labels=Location.Vertebra_Disc.value).get_seg_array() mapping_cc_to_vert_label = {} + # All component centroids in one pass; the per-component `subreg_cc == c` built a full-volume + # boolean for every disc. + cc_coms = np_center_of_mass(subreg_cc) coms_ivd_dict = {} - for c in cc_labelset: - if c == 0: - continue - com_y = np_center_of_mass(subreg_cc == c)[1][1] # center_of_mass(c_l)[1] + for c, com in cc_coms.items(): + com_y = com[1] if com_y < min(coms_vert_y): label = min(coms_vert_labels) - 1 @@ -459,16 +534,11 @@ def add_ivd_ep_vert_label(whole_vert_nii: NII, seg_nii: NII, include_sacrum=Fals # FIXME Problem: For some reason Endplate are mapped to the IVD in MRI aka the superior endplate hat the IVD of vertebra above instead of below. if Location.Endplate.value in u or has_split_endplates: # MAP Endplate - ep_cc = seg_t.get_connected_components(labels=ep_labels) - ep_cc_n = len(ep_cc.unique()) - ep_cc = ep_cc.get_seg_array() - cc_ep_labelset = list(range(1, ep_cc_n + 1)) + ep_cc = seg_t.get_connected_components(labels=ep_labels).get_seg_array() mapping_ep_cc_to_vert_label = {} - coms_ivd_dict = {} - for c in cc_ep_labelset: - if c == 0: - continue - com_y = np_center_of_mass(ep_cc == c)[1][1] + ep_cc_coms = np_center_of_mass(ep_cc) + for c, com in ep_cc_coms.items(): + com_y = com[1] nearest_lower = ( find_nearest_lower(coms_vert_y, com_y) if not has_split_endplates @@ -489,52 +559,7 @@ def add_ivd_ep_vert_label(whole_vert_nii: NII, seg_nii: NII, include_sacrum=Fals # This code sets the IDs to the respective IVD instead of vertebra disc! has_split_endplates is True for CT vert_arr[subreg_arr == Location.Endplate.value] = subreg_ep[subreg_arr == Location.Endplate.value] vert_t.set_array_(vert_arr) - # divide into upper and lower endplate - out = seg_t * 0 - pref = 1 - old_vol = -1 - # seg_t and vert_t are not modified in this loop, so compute these invariants once. - endplate_nii = seg_t.extract_label(ep_labels) - total = endplate_nii.sum() - vert_labels_to_split = vert_t.unique() - for dil in range(1, MAX_ENDPLATE_DILATION): - curr = out.extract_label([Location.Vertebral_Body_Endplate_Inferior.value, Location.Vertebral_Body_Endplate_Superior.value]) - new_vol = curr.sum() - logger.print(rf"{new_vol / total * 100:.2f}% endplates detected", end="\r") if verbose else None - if old_vol == new_vol and old_vol != 0: - break - old_vol = new_vol - if total == new_vol: - logger.print("Found all Endplates ") - break - for i in vert_labels_to_split: - if i >= INSTANCE_LABEL_LIMIT: - break - curr = out.extract_label( - [Location.Vertebral_Body_Endplate_Inferior.value, Location.Vertebral_Body_Endplate_Superior.value] - ) - v = vert_t.extract_label(i).dilate_msk(dil, verbose=False) - end = endplate_nii * v - end *= -curr + 1 # type: ignore - plates = vert_t * end - plates.map_labels_( - { - i + ENDPLATE_LABEL_OFFSET: Location.Vertebral_Body_Endplate_Inferior.value, - pref + ENDPLATE_LABEL_OFFSET: Location.Vertebral_Body_Endplate_Superior.value, - }, - verbose=False, - ) - out += plates - pref = i - curr = out.extract_label([Location.Vertebral_Body_Endplate_Inferior.value, Location.Vertebral_Body_Endplate_Superior.value]) - - end = seg_t.extract_label(ep_labels) - end *= -curr + 1 - # end += end.dilate_msk(3) - out += end * Location.Endplate.value - seg_t = out.extract_label( - [Location.Vertebral_Body_Endplate_Inferior.value, Location.Vertebral_Body_Endplate_Superior.value, Location.Endplate.value] - ) + seg_t = _split_endplates(seg_t, vert_t, ep_labels, verbose=verbose) else: # Endplates are already split semantically. # Assign endplate instance IDs while preserving the semantic labels. @@ -567,14 +592,14 @@ def find_nearest_lower(seq, x) -> float: def find_nearest_higher(seq, x) -> float: - """Return the largest element of ``seq`` strictly smaller than ``x``, or the minimum if none exists. + """Return the smallest element of ``seq`` strictly larger than ``x``, or the maximum if none exists. Args: seq (Sequence[float]): Values to search. x (float): Reference value. Returns: - float: The greatest element below ``x``, or ``min(seq)`` when no element is below ``x``. + float: The smallest element above ``x``, or ``max(seq)`` when no element is above ``x``. """ values_higher = [item for item in seq if item > x] if len(values_higher) == 0: @@ -653,9 +678,12 @@ def assign_vertebra_inconsistency( if ccl == 0: continue cc_map = np_extract_label(subreg_cc, ccl, inplace=False) - vert_arr_cc = vert_arr.copy() - vert_arr_cc += 1 - vert_arr_cc[cc_map == 0] = 0 + # An articular process is tiny next to the volume; work inside its bounding box instead of + # copying and incrementing the whole instance array once per connected component. + cc_bbox = np_bbox_binary(cc_map) + cc_map_c = cc_map[cc_bbox] + vert_arr_cc = vert_arr[cc_bbox] + 1 + vert_arr_cc[cc_map_c == 0] = 0 gt_volume = np_volume(vert_arr_cc) k_keys_sorted = heapq.nlargest(2, gt_volume, key=gt_volume.__getitem__) @@ -667,7 +695,7 @@ def assign_vertebra_inconsistency( if biggest_volume[1] * ARTICULAR_DOMINANCE_RATIO > second_volume[1]: to_label = biggest_volume[0] - 1 # int(list(gt_volume.keys())[argmax] - 1) - vert_arr[cc_map == 1] = to_label + vert_arr[cc_bbox][cc_map_c == 1] = to_label logger.print( f"set cc to {to_label}, with volume decision {gt_volume}, based on {biggest_volume}, {second_volume}", ) @@ -695,7 +723,9 @@ def detect_and_solve_merged_vertebra(seg_nii: NII, vert_nii: NII) -> tuple[NII, stats = {} # Map IVDS subreg_cc: NII = seg_sem.get_connected_components(labels=Location.Vertebra_Disc.value) - subreg_cc += 100 + # Offset only the foreground: a plain `+= OFFSET` would also lift the background out of 0 and + # add a phantom component spanning the whole volume to the stats below. + subreg_cc[subreg_cc > 0] += IVD_LABEL_OFFSET coms = subreg_cc.center_of_masses() volumes = subreg_cc.volumes() @@ -712,6 +742,8 @@ def detect_and_solve_merged_vertebra(seg_nii: NII, vert_nii: NII) -> tuple[NII, stats_by_height_keys = list(stats_by_height.keys()) # detect C2 split into two components + if len(stats_by_height_keys) < 2: + return seg_nii, vert_nii first_key, second_key = stats_by_height_keys[0], stats_by_height_keys[1] first_stats, second_stats = stats_by_height[first_key], stats_by_height[second_key] if first_stats[1] is False and second_stats[1] is False: # noqa: SIM102 diff --git a/spineps/phase_semantic.py b/spineps/phase_semantic.py index 39d1638..61d0000 100755 --- a/spineps/phase_semantic.py +++ b/spineps/phase_semantic.py @@ -7,7 +7,7 @@ from spineps.seg_enums import ErrCode, OutputType from spineps.seg_model import SegmentationModel -from spineps.seg_pipeline import fill_holes_labels, logger +from spineps.seg_pipeline import debug_put, fill_holes_labels, logger from spineps.utils.proc_functions import clean_cc_artifacts from spineps.utils.resolution import REFERENCE_VOXEL_VOLUME_MM3, REFERENCE_ZOOM, mm3_to_voxels, mm_to_voxels @@ -73,14 +73,16 @@ def predict_semantic_mask( logger.print("Post-process semantic mask...") - debug_data["sem_raw"] = seg_nii.copy() + debug_put(debug_data, "sem_raw", seg_nii.copy) if seg_nii.is_empty: logger.print("Subregion mask is empty, skip this", Log_Type.FAIL) return seg_nii, softmax_logits, ErrCode.EMPTY + # Both helpers write through set_array_ and return the same object, and the result is rebound to + # seg_nii either way, so the defensive copies these calls used to make were pure whole-volume waste. if proc_remove_inferior_beyond_canal: - seg_nii = remove_nonsacrum_beyond_canal_height(seg_nii=seg_nii.copy()) + seg_nii = remove_nonsacrum_beyond_canal_height(seg_nii) if proc_clean_small_cc_artifacts: seg_nii.set_array_( @@ -110,14 +112,14 @@ def predict_semantic_mask( # Do two iterations of both processing if enabled to make sure if proc_remove_inferior_beyond_canal: - seg_nii = remove_nonsacrum_beyond_canal_height(seg_nii=seg_nii.copy()) + seg_nii = remove_nonsacrum_beyond_canal_height(seg_nii) if proc_clean_beyond_largest_bounding_box: - seg_nii = semantic_bounding_box_clean(seg_nii=seg_nii.copy()) + seg_nii = semantic_bounding_box_clean(seg_nii) if proc_remove_inferior_beyond_canal and proc_clean_beyond_largest_bounding_box: - seg_nii = remove_nonsacrum_beyond_canal_height(seg_nii=seg_nii.copy()) - seg_nii = semantic_bounding_box_clean(seg_nii=seg_nii.copy()) + seg_nii = remove_nonsacrum_beyond_canal_height(seg_nii) + seg_nii = semantic_bounding_box_clean(seg_nii) if proc_fill_3d_holes: seg_nii = seg_nii.fill_holes_(fill_holes_labels, verbose=logger) @@ -210,16 +212,18 @@ def semantic_bounding_box_clean(seg_nii: NII) -> NII: break seg_bin_arr = seg_binary.get_seg_array() - crop = (p_slice, i_slice, r_slice) - seg_bin_clean_arr = np.zeros(seg_bin_arr.shape) - seg_bin_clean_arr[crop] = 1 + # The region to keep is the union of every incorporated component's bounding box, not just the + # largest component's -- otherwise growing the region has no effect on the result at all. + crop = tuple(slice(min(b[axis].start for b in bboxes), max(b[axis].stop for b in bboxes)) for axis in range(len(bboxes[0]))) + seg_bin_clean_arr = np.zeros(seg_bin_arr.shape, dtype=bool) + seg_bin_clean_arr[crop] = True - # everything below biggest k get always removed + # every component that was never incorporated is dropped, even where it falls inside the union bbox largest_k_arr = seg_bin_largest_k_cc_nii.get_seg_array() - seg_bin_clean_arr[largest_k_arr == 0] = 0 + seg_bin_clean_arr &= np.isin(largest_k_arr, incorporated) seg_arr = seg_nii.get_seg_array() - seg_arr[seg_bin_clean_arr != 1] = 0 + seg_arr[~seg_bin_clean_arr] = 0 seg_nii.set_array_(seg_arr) seg_nii.reorient_(ori) cleaned_ks = [l for l in range(2, max_k + 1) if l not in incorporated] diff --git a/spineps/seg_model.py b/spineps/seg_model.py index f0c93ec..e25f84b 100755 --- a/spineps/seg_model.py +++ b/spineps/seg_model.py @@ -529,6 +529,9 @@ def load(self, folds: tuple[str, ...] | None = None) -> Self: # noqa: ARG002 raise FileNotFoundError( f"expected exactly one '*weights*.ckpt' checkpoint in {self.model_folder}, found {len(chktpath)}: {chktpath}" ) + # Two U-Net wrappers are kept on purpose: released checkpoints exist for both the legacy + # `spineps.architectures` PLNet and the current `spineps.architectures_new` one, and only the + # weights themselves say which. Try the legacy layout first and fall back on a shape mismatch. try: model = PLNet.load_from_checkpoint(checkpoint_path=chktpath[0], weights_only=False, map_location=self.device) except RuntimeError: diff --git a/spineps/seg_pipeline.py b/spineps/seg_pipeline.py index df49115..c943717 100755 --- a/spineps/seg_pipeline.py +++ b/spineps/seg_pipeline.py @@ -2,7 +2,10 @@ from __future__ import annotations -import subprocess +from collections.abc import Callable +from functools import lru_cache +from importlib.metadata import PackageNotFoundError +from importlib.metadata import version as _package_version from typing import Any from scipy.ndimage import center_of_mass @@ -24,6 +27,31 @@ IVD_LABEL_RANGE = range(IVD_LABEL_OFFSET, IVD_LABEL_OFFSET + _MAX_DERIVED_LABELS_PER_TYPE) ENDPLATE_LABEL_RANGE = range(ENDPLATE_LABEL_OFFSET, ENDPLATE_LABEL_OFFSET + _MAX_DERIVED_LABELS_PER_TYPE) + +class NoOpDebugSink(dict): + """Dict-like sink that discards writes; used to skip retaining debug data when it won't be saved.""" + + def __setitem__(self, key, value): + pass + + +def debug_put(debug_data: dict, key: str, factory: Callable[[], Any]) -> None: + """Store a debug value, skipping its construction entirely when the sink discards writes. + + ``NoOpDebugSink`` drops the value, but Python still evaluates the argument -- so a plain + ``debug_data[key] = nii.copy()`` paid for a whole-volume copy even with ``save_debug_data=False``. + Pass a zero-argument factory instead and the copy never happens. + + Args: + debug_data (dict): The debug sink; a plain dict stores, a ``NoOpDebugSink`` discards. + key (str): Name to store the value under. + factory (Callable[[], Any]): Builds the value, called only when it will actually be kept. + """ + if isinstance(debug_data, NoOpDebugSink): + return + debug_data[key] = factory() + + fill_holes_labels = [ Location.Vertebra_Corpus_border.value, Location.Spinal_Canal.value, @@ -87,48 +115,24 @@ def predict_centroids_from_both( ctd.info["source"] = "MRI Segmentation Pipeline" ctd.info["version"] = pipeline_version() ctd.info["models"] = models_repr - ctd.info["revision"] = pipeline_revision() ctd.info["timestamp"] = format_time_short(get_time()) for pname, pvalue in parameter.items(): ctd.info[pname] = str(pvalue) return ctd +@lru_cache(maxsize=1) def pipeline_version() -> str: - """Return the pipeline version string derived from the git commit count on ``main``. + """Return the installed SPINEPS version. - Returns: - str: A version like ``"v1."``, or ``"Version not found"`` if git is unavailable. - """ - try: - label = subprocess.check_output(["git", "rev-list", "--count", "main"]).strip() - label = str(label).replace("'", "") - while not label[0].isdigit(): - label = label[1:] - except Exception: - return "Version not found" - return "v1." + str(label) - - -def pipeline_revision() -> str: - """Return the current git revision string for the pipeline. + Read from the package metadata, which poetry-dynamic-versioning derives from the git tag at build time. + (This used to shell out to ``git`` without a ``cwd``, so it reported whatever repository the caller + happened to be standing in.) Returns: - str: ``"::"``; either part is empty if the corresponding git call fails. + str: The installed version, or ``"unknown"`` if SPINEPS is not installed as a distribution. """ - label = "" - rev = "" try: - label = subprocess.check_output(["git", "describe", "--always"]).strip() - except Exception: - pass - try: - rev = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("ascii").strip() - except Exception: - pass - return str(label) + "::" + str(rev) - - -if __name__ == "__main__": - print(pipeline_version()) - print(pipeline_revision()) + return _package_version("SPINEPS") + except PackageNotFoundError: + return "unknown" diff --git a/spineps/seg_run.py b/spineps/seg_run.py index e0fe52b..7e5e55d 100755 --- a/spineps/seg_run.py +++ b/spineps/seg_run.py @@ -20,23 +20,16 @@ from spineps.phase_semantic import predict_semantic_mask from spineps.seg_enums import Acquisition, ErrCode, Modality from spineps.seg_model import SegmentationModel -from spineps.seg_pipeline import logger, predict_centroids_from_both -from spineps.seg_utils import Modality_Pair, check_input_model_compatibility, check_model_modality_acquisition, find_best_matching_model +from spineps.seg_pipeline import NoOpDebugSink, logger, predict_centroids_from_both +from spineps.seg_utils import Modality_Pair, check_input_model_compatibility, check_model_modality_acquisition from spineps.utils.citation_reminder import citation_reminder -class _NoOpDebugDict(dict): - """Dict-like sink that discards writes; used to skip retaining debug data when it won't be saved.""" - - def __setitem__(self, key, value): - pass - - @citation_reminder def process_dataset( # noqa: C901 dataset_path: Path, model_instance: SegmentationModel, - model_semantic: list[SegmentationModel] | SegmentationModel | None = None, + model_semantic: list[SegmentationModel] | SegmentationModel, model_labeling: VertLabelingClassifier | None = None, # rawdata_name: str = "rawdata", @@ -90,8 +83,8 @@ def process_dataset( # noqa: C901 Args: dataset_path (Path): Path to the BIDS dataset. model_instance (SegmentationModel): Model for the vertebra (instance) segmentation. - model_semantic (list[SegmentationModel] | SegmentationModel | None, optional): Models for the subregion (semantic) - segmentation, one per modality pair. If None, attempts to find a matching model for each modality. Defaults to None. + model_semantic (list[SegmentationModel] | SegmentationModel): Model(s) for the subregion (semantic) segmentation. + Pass a list with one model per modality pair, or a single model used for all of them. model_labeling (VertLabelingClassifier | None, optional): Classifier used to label the vertebra instances. Defaults to None. rawdata_name (str, optional): Name of the rawdata folder. Defaults to "rawdata". derivative_name (str, optional): Name of the derivatives output folder. Defaults to "derivatives_seg". @@ -142,8 +135,7 @@ def process_dataset( # noqa: C901 Defaults to False. ignore_bids_filter (bool, optional): If true, disables the BIDS query filters and processes all niftys found. Defaults to False. tta (bool | None, optional): If not None, forces test-time augmentation (mirroring) on/off for the semantic - model(s), covering both explicitly-passed and auto-resolved models. If None, uses each model's configured - setting. Defaults to None. + model(s). If None, uses each model's configured setting. Defaults to None. log_inference_time (bool, optional): If true, logs the inference time of each step. Defaults to True. verbose (bool, optional): If true, prints verbose information. Defaults to False. """ @@ -160,17 +152,13 @@ def process_dataset( # noqa: C901 elif snapshot_copy_folder is False: snapshot_copy_folder = None - if model_semantic is None: - model_semantic = [find_best_matching_model(m, expected_resolution=None) for m in modalities] - logger.print("Found matching models:") - for idx, m in enumerate(model_semantic): - logger.print("-", str(modalities[idx]), ":", str(m.modelid())) - del idx, m if not isinstance(model_semantic, list): - model_semantic = [model_semantic] + model_semantic = [model_semantic] * len(modalities) + if len(model_semantic) != len(modalities): + raise ValueError(f"need one semantic model per modality pair, got {len(model_semantic)} models for {len(modalities)} modalities") - # Optionally force test-time augmentation (mirroring) on/off for the semantic model(s); covers both - # explicitly-passed and auto-resolved models. Load eagerly so the toggle reaches the predictor. + # Optionally force test-time augmentation (mirroring) on/off for the semantic model(s). + # Load eagerly so the toggle reaches the predictor. if tta is not None: for m in model_semantic: if m is not None: @@ -187,9 +175,18 @@ def process_dataset( # noqa: C901 if not compatible and not ignore_model_compatibility: logger.print("Compatibility issues (see above), stop program", Log_Type.FAIL) + raise ValueError( + "the given model(s) do not support the requested modality/acquisition pairs (see the warnings above); " + "pass ignore_model_compatibility=True (CLI: --ignore-model-compatibility) to run anyway" + ) - # Activate logger - args = locals() + # Activate logger. Log the plain options plus the model ids -- the model objects stringify to their + # entire inference config, which floods the log file. + _not_logged = ("model_instance", "model_semantic", "model_labeling", "compatible") + args = {k: v for k, v in locals().items() if k not in _not_logged} + args["model_instance"] = model_instance.modelid() + args["model_semantic"] = [m.modelid() for m in model_semantic] + args["model_labeling"] = model_labeling.modelid() if model_labeling is not None else None if save_log_data: logger = Logger(dataset_path, log_filename="segmentation_pipeline", default_verbose=True, log_arguments=args, prefix="SegPipeline") logger.print(f"Processing dataset in {dataset_path}", Log_Type.BOLD) @@ -202,7 +199,7 @@ def process_dataset( # noqa: C901 processed_seen_counter = 0 processed_alldone_counter = 0 processed_counter = 0 - not_properly_processed: list[str] = [] + not_properly_processed: list[tuple[ErrCode, str]] = [] for s_idx, (name, subject) in enumerate(bids_ds.enumerate_subjects(sort=True)): logger.print() @@ -447,7 +444,6 @@ def segment_image( # noqa: C901 out_snap = output_paths["out_snap"] out_ctd = output_paths["out_ctd"] out_snap2 = output_paths["out_snap2"] - out_raw = output_paths["out_raw"] out_debug = output_paths["out_debug"] if isinstance(snapshot_copy_folder, Path): snapshot_copy_folder.mkdir(parents=True, exist_ok=True) @@ -467,8 +463,8 @@ def segment_image( # noqa: C901 return output_paths, ErrCode.ALL_DONE done_something = False - # Avoid retaining full-volume debug copies for the whole run when they'll never be saved (see seg_run.py:699). - debug_data_run: dict[str, NII] = {} if save_debug_data else _NoOpDebugDict() + # Avoid retaining -- and, via debug_put, even building -- full-volume debug copies when they'll never be saved. + debug_data_run: dict[str, NII] = {} if save_debug_data else NoOpDebugSink() if Modality.CT in model_semantic.modalities(): proc_normalize_input = False # Never normalize input if it is an CT @@ -476,8 +472,13 @@ def segment_image( # noqa: C901 if model_semantic.inference_config.has_c1: vertebra_instance_labeling_offset = 1 - compatible = check_input_model_compatibility(img_ref, model=model_semantic) - compatible_labeling = check_input_model_compatibility(img_ref, model=model_labeling) if model_labeling is not None else True + # Load the volume once and hand it to the compatibility checks: BIDS_FILE.open_nii() does not cache, + # so the input used to be read from disk up to three times per image. + input_nii = _nii if _nii is not None else img_ref.open_nii() + compatible = check_input_model_compatibility(img_ref, model=model_semantic, img_nii=input_nii) + compatible_labeling = ( + check_input_model_compatibility(img_ref, model=model_labeling, img_nii=input_nii) if model_labeling is not None else True + ) if not (compatible and compatible_labeling): if not ignore_compatibility_issues: return output_paths, ErrCode.COMPATIBILITY @@ -491,7 +492,6 @@ def segment_image( # noqa: C901 with logger: if verbose: model_semantic.logger.default_verbose = True - input_nii = _nii if _nii is not None else img_ref.open_nii() input_nii.seg = False input_nii_ = input_nii.copy() if timing: @@ -592,6 +592,8 @@ def segment_image( # noqa: C901 seg_nii_modelres.save(out_spine_raw, verbose=logger) if save_softmax_logits and isinstance(softmax_logits, np.ndarray): save_nparray(softmax_logits, out_logits) + # Both are whole-volume and finished with; drop them before the instance stage allocates. + del input_preprocessed, softmax_logits done_something = True if timing: logger.print(f"Predict semantic took: {perf_counter() - start_time2:.2f} seconds", Log_Type.OK, verbose=log_inference_time) @@ -701,20 +703,15 @@ def segment_image( # noqa: C901 # save debug if save_debug_data: - if debug_data_run is None: - logger.print("Save_debug_data: no debug data found", Log_Type.WARNING) - else: - out_debug.parent.mkdir(parents=True, exist_ok=True) - for k, v in debug_data_run.items(): - v.reorient_(input_nii_.orientation).save( - out_debug.joinpath(k + f"_{input_format}.nii.gz"), make_parents=True, verbose=False - ) - logger.print(f"Saved debug data into {out_debug}/*", Log_Type.OK) - if timing: - logger.print( - f"Save debug data took: {perf_counter() - start_time2:.2f} seconds", Log_Type.OK, verbose=log_inference_time - ) - start_time2 = perf_counter() + out_debug.parent.mkdir(parents=True, exist_ok=True) + for k, v in debug_data_run.items(): + v.reorient_(input_nii_.orientation).save( + out_debug.joinpath(k + f"_{input_format}.nii.gz"), make_parents=True, verbose=False + ) + logger.print(f"Saved debug data into {out_debug}/*", Log_Type.OK) + if timing: + logger.print(f"Save debug data took: {perf_counter() - start_time2:.2f} seconds", Log_Type.OK, verbose=log_inference_time) + start_time2 = perf_counter() # Snapshot if not out_snap.exists() or done_something: @@ -736,7 +733,7 @@ def segment_image( # noqa: C901 start_time2 = perf_counter() elif not out_snap2.exists(): logger.print(f"Copying snapshot into {snapshot_copy_folder!s}") - out_snap2.parent.mkdir(exist_ok=True) + out_snap2.parent.mkdir(parents=True, exist_ok=True) shutil.copy(out_snap, out_snap2) logger.print(f"Pipeline took: {perf_counter() - start_time:.2f} seconds", Log_Type.OK, verbose=log_inference_time) @@ -816,14 +813,6 @@ def output_paths_from_input( make_parent=False, ) out_vert_raw = out_raw.joinpath(out_vert_raw.name) - out_unc = img_ref.get_changed_path( - bids_format="uncertainty", - parent=derivative_name, - info={"seg": "spine", "mod": img_ref.format}, - non_strict_mode=non_strict_mode, - make_parent=False, - ) - out_unc = out_raw.joinpath(out_unc.name) out_logits = img_ref.get_changed_path( file_type="npz", bids_format="logit", @@ -845,7 +834,6 @@ def output_paths_from_input( "out_spine_raw": out_spine_raw, "out_vert": out_vert, "out_vert_raw": out_vert_raw, - "out_unc": out_unc, "out_logits": out_logits, "out_snap": out_snap, "out_ctd": out_ctd, diff --git a/spineps/seg_utils.py b/spineps/seg_utils.py index 6b60769..01f4b78 100755 --- a/spineps/seg_utils.py +++ b/spineps/seg_utils.py @@ -4,7 +4,7 @@ from typing import Union -from TPTBox import BIDS_FILE, ZOOMS, Log_Type +from TPTBox import BIDS_FILE, NII, Log_Type from spineps.seg_enums import Acquisition, Modality from spineps.seg_model import SegmentationModel @@ -15,27 +15,6 @@ Modality_Pair = tuple[Union[list[Modality], Modality], Acquisition] -def find_best_matching_model( - modality_pair: Modality_Pair, - expected_resolution: ZOOMS | None, # actual resolution here? -) -> SegmentationModel: - """Select the segmentation model best matching a modality/acquisition pair and resolution. - - Not yet implemented: intended to iterate over model configs and pick the one best matching the requested resolution. - - Args: - modality_pair (Modality_Pair): The desired ``(modality(ies), acquisition)`` pair. - expected_resolution (ZOOMS | None): The desired voxel resolution, or None. - - Returns: - SegmentationModel: The best-matching model (once implemented). - - Raises: - NotImplementedError: Always, as this function is not yet implemented; also for an unmapped modality pair. - """ - raise NotImplementedError("find_best_matching_model()") - - def check_model_modality_acquisition( model: SegmentationModel, mod_pair: Modality_Pair, @@ -106,6 +85,7 @@ def check_input_model_compatibility( ignore_acquisition: bool = False, ignore_labelkey: bool = False, verbose: bool = True, + img_nii: NII | None = None, ) -> bool: """Check whether an input image file is compatible with a model's expected modality, acquisition, and naming. @@ -120,6 +100,7 @@ def check_input_model_compatibility( ignore_acquisition (bool): If True, tolerate an acquisition mismatch. ignore_labelkey (bool): If True, tolerate an unexpected ``label`` key in the filename. verbose (bool): If True, log warnings describing incompatibilities. + img_nii (NII | None): The already-loaded image, to avoid re-reading it from disk just for the plane check. Returns: bool: True if the input is compatible with the model (after applying the ignore flags), otherwise False. @@ -148,8 +129,9 @@ def check_input_model_compatibility( compatible = False else: add_ignore_text(logger_texts) - if has_seg_key and allowed_format not in Modality.format_keys(Modality.SEG): - logger_texts.append("Input acquisition not segmentation, but found a 'seg'-key.") + # `allowed_format` is a list, so the old `allowed_format not in Modality.format_keys(...)` was always True. + if has_seg_key and Modality.SEG not in model_modalities: + logger_texts.append("Found a 'seg'-key in the filename, but the model does not take a segmentation as input.") if not ignore_modality: compatible = False else: @@ -173,7 +155,8 @@ def check_input_model_compatibility( logger_texts.append("Probably a debug file (debug in name or parent).") compatible = False - img_nii = img_ref.open_nii() + if img_nii is None: + img_nii = img_ref.open_nii() if img_nii.get_plane() not in ["iso", *allowed_acq]: logger_texts.append(f"input {img_nii.get_plane()=} is not 'iso' or one of the expected {allowed_acq}.") compatible = False diff --git a/spineps/utils/citation_reminder.py b/spineps/utils/citation_reminder.py index 42cf08c..eab6b2f 100644 --- a/spineps/utils/citation_reminder.py +++ b/spineps/utils/citation_reminder.py @@ -12,16 +12,25 @@ ARXIV_LINK = "https://arxiv.org/abs/2402.16368" +# Set this environment variable to any of the values below to silence the reminder entirely. +OPT_OUT_ENV_VAR = "SPINEPS_NO_CITATION_REMINDER" +_OPT_OUT_VALUES = frozenset({"1", "true", "yes", "on"}) + has_reminded_citation = False +def reminder_disabled() -> bool: + """Return whether the user opted out of the citation reminder via the environment.""" + return os.environ.get(OPT_OUT_ENV_VAR, "").strip().lower() in _OPT_OUT_VALUES + + def citation_reminder(func): """Decorator to remind users to cite SPINEPS.""" @functools.wraps(func) def wrapper(*args, **kwargs): global has_reminded_citation # noqa: PLW0603 - if not has_reminded_citation and os.environ.get("SPINEPS_TURN_OF_CITATION_REMINDER", "FALSE") != "TRUE": + if not has_reminded_citation and not reminder_disabled(): print_citation_reminder() has_reminded_citation = True return func(*args, **kwargs) @@ -45,4 +54,14 @@ def print_citation_reminder(): console.line() -atexit.register(print_citation_reminder) +def _print_citation_reminder_at_exit() -> None: + """Repeat the reminder on interpreter exit, but only if SPINEPS actually ran and the user did not opt out. + + Registering ``print_citation_reminder`` directly made merely importing ``spineps`` print the banner, and + ignored the opt-out environment variable entirely. + """ + if has_reminded_citation and not reminder_disabled(): + print_citation_reminder() + + +atexit.register(_print_citation_reminder_at_exit) diff --git a/spineps/utils/filepaths.py b/spineps/utils/filepaths.py index 9b582d3..b256a57 100755 --- a/spineps/utils/filepaths.py +++ b/spineps/utils/filepaths.py @@ -14,8 +14,9 @@ spineps_environment_path_override = None # Path( # "/DATA/NAS/ongoing_projects/hendrik/mri_usage/models/" # ) # None # You can put an absolute path to the model weights here instead of using environment variable -spineps_environment_path_backup = Path(__file__).parent.parent.joinpath("models") # EDIT this to use this instead of environment variable -spineps_environment_path_backup.mkdir(exist_ok=True) +# EDIT this to use this instead of the environment variable. Created on demand by +# get_mri_segmentor_models_dir(), never at import time (that would write into site-packages). +spineps_environment_path_backup = Path(__file__).parent.parent.joinpath("models") def get_mri_segmentor_models_dir() -> Path: @@ -25,26 +26,33 @@ def get_mri_segmentor_models_dir() -> Path: Path: Path to the overall models folder Raises: - RuntimeError: If no models directory could be determined from the environment variable, override or backup. - FileNotFoundError: If the resolved models directory does not exist. + RuntimeError: If no models directory could be determined, or the fallback directory cannot be created. + FileNotFoundError: If the directory named by 'SPINEPS_SEGMENTOR_MODELS' does not exist. """ - folder_path = ( - os.environ.get("SPINEPS_SEGMENTOR_MODELS") - if spineps_environment_path_override is None or not spineps_environment_path_override.exists() - else spineps_environment_path_override - ) - if folder_path is None and spineps_environment_path_backup is not None: - folder_path = spineps_environment_path_backup - - if folder_path is None: + if spineps_environment_path_override is not None and spineps_environment_path_override.exists(): + return spineps_environment_path_override + + from_env = os.environ.get("SPINEPS_SEGMENTOR_MODELS") + if from_env is not None: + folder_path = Path(from_env) + if not folder_path.exists(): + raise FileNotFoundError(f"Environment variable 'SPINEPS_SEGMENTOR_MODELS' = {folder_path} does not exist") + return folder_path + + if spineps_environment_path_backup is None: raise RuntimeError( "Environment variable 'SPINEPS_SEGMENTOR_MODELS' is not defined. Setup the environment variable as stated " "in the readme or set the override in utils.filepaths.py" ) - folder_path = Path(folder_path) - if not folder_path.exists(): - raise FileNotFoundError(f"Environment variable 'SPINEPS_SEGMENTOR_MODELS' = {folder_path} does not exist") - return folder_path + try: + spineps_environment_path_backup.mkdir(parents=True, exist_ok=True) + except OSError as e: + raise RuntimeError( + f"Environment variable 'SPINEPS_SEGMENTOR_MODELS' is not defined and the fallback models directory " + f"{spineps_environment_path_backup} could not be created ({e}). Set the environment variable as stated " + "in the readme." + ) from e + return spineps_environment_path_backup def filepath_model(model_folder_name: str, model_dir: str | Path | None = None) -> Path: diff --git a/spineps/utils/find_min_cost_path.py b/spineps/utils/find_min_cost_path.py index c8e0e79..b5db9d1 100644 --- a/spineps/utils/find_min_cost_path.py +++ b/spineps/utils/find_min_cost_path.py @@ -184,15 +184,18 @@ def find_most_probably_sequence( # noqa: C901 # define regions n_classes = shape[1] assert min_start_class < n_classes - regions_ranges = None + if n_classes < regions[-1]: + warn(f"n_classes < defined regions, got {n_classes} and {regions}", stacklevel=3) + # Local copy closed by the class-axis end: appending to `regions` would mutate the caller's list + # (and DEFAULT_REGION_STARTS itself for any caller that forgets to copy it). + region_bounds = [*regions, n_classes] + # Built unconditionally: minCostAlgo indexes it whenever allow_skip_at_region is non-empty, which is + # independent of whether region_rel_cost was given. + regions_ranges = [(region_bounds[i], region_bounds[i + 1] - 1) for i in range(len(region_bounds) - 1)] if region_rel_cost is not None: - if n_classes < regions[-1]: - warn(f"n_classes < defined regions, got {n_classes} and {regions}", stacklevel=3) - regions.append(n_classes) - regions_ranges = [(regions[i], regions[i + 1] - 1) for i in range(len(regions) - 1)] region_rel_shape = region_rel_cost.shape - assert region_rel_shape[1] == ((len(regions) - 1) * 2), ( - f"expected region_rel_cost with shape {((len(regions) - 1) * 2)}, but got {region_rel_shape[1]}" + assert region_rel_shape[1] == ((len(region_bounds) - 1) * 2), ( + f"expected region_rel_cost with shape {((len(region_bounds) - 1) * 2)}, but got {region_rel_shape[1]}" ) # softmax (deprecated, handled elsewhere) @@ -220,7 +223,7 @@ def add_option_path(options, r, c, extracost): def minCostAlgo(r, c): logger.print(f"Called vert {r}, label {c}") # get current region - region_cur = c_to_region_idx(c, regions) + region_cur = c_to_region_idx(c, region_bounds) # start point if c == -1 and r == -1: # go over each possible start column @@ -311,7 +314,7 @@ def t13_cost_single(r, c): cost_add = 0 if vertt13_cost is not None: vt13_cost = vertt13_cost[r][1] - if c == 18: + if c == T12_CLASS_IDX: cost_add += vt13_cost return cost_add @@ -334,7 +337,7 @@ def rel_cost(r, c, pnext, region_cur): logger.print(f"Added F {rel_cost} to vert {r}, label {c}, {internal_to_real_path(pnext)}") cost_add += rel_cost # break - elif last == 1 and (c_to_region_idx(pnext[-1][1], regions) >= region_cur + 1): # or pnext[-1][1] == c): + elif last == 1 and (c_to_region_idx(pnext[-1][1], region_bounds) >= region_cur + 1): # or pnext[-1][1] == c): logger.print(f"Added L {rel_cost} to vert {r}, label {c}, {internal_to_real_path(pnext)}") cost_add += rel_cost return cost_add diff --git a/spineps/utils/generate_disc_labels.py b/spineps/utils/generate_disc_labels.py deleted file mode 100644 index 7e91ed2..0000000 --- a/spineps/utils/generate_disc_labels.py +++ /dev/null @@ -1,266 +0,0 @@ -""" -This script generates discs labels using SPINEPS' vertebrae segmentation - -Author: Nathan Molinier -""" - -from __future__ import annotations - -import argparse -from pathlib import Path - -import cc3d -import numpy as np - -from spineps.utils.compat import zip_strict -from spineps.utils.image import Image - -DISCS_MAP = { - 2: 1, - 102: 3, - 103: 4, - 104: 5, - 105: 6, - 106: 7, - 107: 8, - 108: 9, - 109: 10, - 110: 11, - 111: 12, - 112: 13, - 113: 14, - 114: 15, - 115: 16, - 116: 17, - 117: 18, - 118: 19, - 119: 20, - 120: 21, - 121: 22, - 122: 23, - 123: 24, - 124: 25, -} - - -def get_parser() -> argparse.ArgumentParser: - """Build the command-line argument parser for disc-label generation. - - Returns: - argparse.ArgumentParser: Parser accepting the input vertebrae label path and the optional output path. - """ - # parse command line arguments - parser = argparse.ArgumentParser(description="Generate discs labels from spineps' vertebrae segmentation.") - parser.add_argument( - "--path-vert", - type=str, - required=True, - help='Path to the SPINEPS vertebrae labels. Example: "//sub-amuALT_T2w_label-vert_dseg.nii.gz" (Required)', - ) - parser.add_argument( - "--path-out", - type=str, - default="", - help="Output path of the discs label. " - 'Example: "//sub-amuALT_T2w_label-discs_dlabel.nii.gz". ' - 'By default, the structure "_label-discs_dlabel" will be used.', - ) - return parser - - -def main(): - """Run the disc-label generation CLI. - - Parses arguments, loads the SPINEPS vertebrae segmentation, derives single-voxel disc labels from it and - writes the result to the chosen (or default) output path. - """ - # Load parser - parser = get_parser() - args = parser.parse_args() - - # Fetch paths - path_in = Path(args.path_vert).resolve() - path_out = Path(args.path_out).resolve() if args.path_out else default_name_discs(path_in) - - # Check if output folder exists - if not path_out.parent.exists(): - path_out.parent.mkdir(parents=True) - - # Extract discs labels - vert_image = Image(str(path_in)) - print("-" * 80) - print(f"Creating discs label using SPINEPS prediction: {path_in}") - print("-" * 80) - discs_nii_clean = extract_discs_label(vert_image, mapping=DISCS_MAP) - - # Save discs labels - discs_nii_clean.save(str(path_out)) - print("-" * 80) - print(f"Discs label: {path_out} was created.") - print("-" * 80) - - -def default_name_discs(path_in: Path | str, suffix="_label-discs_dlabel") -> Path: - """Derive the default output path for disc labels by swapping in a disc suffix. - - Args: - path_in: Path to the input vertebrae label file (may include compound extensions like ``.nii.gz``). - suffix (str, optional): Suffix inserted before the extension. Defaults to ``"_label-discs_dlabel"``. - - Returns: - Path: The default output path with the disc suffix applied. - """ - # Fetch suffixes - path_obj = Path(path_in) - ext = "".join(path_obj.suffixes) - - # Add suffix - path_out = Path(str(path_in).replace(ext, suffix + ext)) - return path_out - - -def extract_discs_label(label: Image, mapping: dict) -> Image: - """Derive single-voxel disc labels from a vertebrae segmentation. - - Remaps vertebra label values to disc values, locates each disc's posterior tip by shifting a centerline - (interpolated through the disc centroids) posteriorly and picking the closest segmented voxel, inserts disc 2 - between discs 1 and 3 when both are present, and writes one labeled voxel per disc into the image. - - Args: - label (Image): Vertebrae segmentation image; its data is replaced in place with the disc labels. - mapping (dict): Mapping from vertebra label values to disc label values. - - Returns: - Image: The image holding the disc labels, restored to its original orientation. - """ - # Store input orientation - orig_orientation = label.orientation - - # Use RSP orientation - label.change_orientation("RSP") - - # Extract only discs segmentations based on mapping - data = label.data - data_discs_seg = np.zeros_like(data) - for seg_value, discs_value in mapping.items(): - data_discs_seg[np.where(data == seg_value)] = discs_value - - # Deal with disc 1 obtained with the first vertebrae (Highest vertical coordinate) - if 1 in data_discs_seg: - # If the first vertebrae is present identify label disc 1 at the top - vert1_seg = np.array(np.where(data_discs_seg == 1)) - disc1_idx = np.argmin(vert1_seg[1]) # find min on the S-I axis - disc1_coord = vert1_seg[:, disc1_idx] - data_discs_seg[np.where(data_discs_seg == 1)] = 0 - data_discs_seg[disc1_coord[0], disc1_coord[1], disc1_coord[2]] = 1 - - ## Identify the posterior tip of the disc - # Extract the center of mass of every discs segmentation to create discs labels - # Centroids are sorted based on the vertical axis - discs_centroids, discs_bb = extract_centroids_3d(data_discs_seg) - - # Generate a centerline between the discs by doing linear interpolation - yvals = np.linspace(discs_centroids[0, 1], discs_centroids[-1, 1], round(8 * len(discs_centroids))) - xvals = np.interp(yvals, discs_centroids[:, 1], discs_centroids[:, 0]) - zvals = np.interp(yvals, discs_centroids[:, 1], discs_centroids[:, 2]) - centerline = np.concatenate((np.expand_dims(xvals, axis=1), np.expand_dims(yvals, axis=1), np.expand_dims(zvals, axis=1)), axis=1) - - # Shift the centerline to the posterior direction until there is no intersection with the - # discs segmentations - # Find the min coordinate of the discs segmentation on the A-P axis - min_seg_ap = np.min(np.where(data_discs_seg > 0)[2]) - max_centroid_ap = np.max(discs_centroids[:, 2]) - offset = 5 - shift = (max_centroid_ap - min_seg_ap + offset) if min_seg_ap >= offset else (max_centroid_ap - min_seg_ap) - - centerline_shifted = np.copy(centerline) - centerline_shifted[:, 2] = centerline_shifted[:, 2] - shift - - # For each segmented disc, identify the closest voxel to this shifted centerline - discs_list = closest_point_seg_to_line(data_discs_seg, centerline_shifted, discs_bb) - - # Add disc 2 between disc 1 and 3 - if 1 and 3 in discs_list[:, -1]: - disc1_coord = discs_list[discs_list[:, -1] == 1] - disc2_coord = discs_list[discs_list[:, -1] == 3] - disc2_coord[0, 1] = (disc2_coord[0, 1] + disc1_coord[0, 1]) // 2 - disc2_coord[0, -1] = 2 - discs_list = np.insert(discs_list, 1, disc2_coord, axis=0) - - # Create output Image - data_discs = np.zeros_like(data) - for x, y, z, v in discs_list: - data_discs[x, y, z] = v - label.data = data_discs - return label.change_orientation(orig_orientation) - - -def extract_centroids_3d(arr: np.ndarray) -> tuple[np.ndarray, np.ndarray]: - """Extract connected-component centroids and bounding boxes from a 3D array, sorted along the vertical axis. - - Args: - arr (np.ndarray): 3D label array (assumed RSP orientation, so axis 1 is the superior-inferior axis). - - Returns: - tuple[np.ndarray, np.ndarray]: Integer centroid coordinates and the matching bounding boxes, both sorted - by the vertical (axis-1) coordinate, with the background component removed. - """ - stats = cc3d.statistics(cc3d.connected_components(arr)) - centroids = stats["centroids"][1:] # Remove backgroud <0> - bounding_boxes = stats["bounding_boxes"][1:] - - # Sort according to the vertical axis because RSP orientation - sort_args = np.argsort(centroids[:, 1]) - - centroids_sorted = centroids[sort_args] - bb_sorted = np.array(bounding_boxes)[sort_args] - return centroids_sorted.astype(int), bb_sorted - - -def project_point_on_line(point: np.ndarray, line: np.ndarray) -> tuple[np.ndarray, float]: - """Project a point onto a polyline by finding the closest line point. - - Copied from https://github.com/spinalcordtoolbox/spinalcordtoolbox. - - Args: - point (np.ndarray): Coordinates of the point, ``numpy.array([x, y, z])``. - line (np.ndarray): Coordinates of the points composing the line. - - Returns: - tuple[np.ndarray, float]: The closest point on the line and the squared distance to it. - """ - # Calculate distances between the referenced point and the line then keep the closest point - dist = np.sum((line - point) ** 2, axis=1) - - return line[np.argmin(dist)], np.min(dist) - - -def closest_point_seg_to_line(discs_seg: np.ndarray, centerline: np.ndarray, bounding_boxes: np.ndarray) -> np.ndarray: - """Find, per disc, the segmented voxel closest to a reference centerline. - - Args: - discs_seg (np.ndarray): Disc-labeled segmentation array. - centerline (np.ndarray): Coordinates of the points composing the reference line. - bounding_boxes (np.ndarray): Bounding box (slice tuple) for each disc, used to isolate it. - - Returns: - np.ndarray: Array of ``[x, y, z, disc_value]`` rows, one per disc, giving the closest voxel and its label. - """ - discs_list = [] - for x, y, z in bounding_boxes: - zer = np.zeros_like(discs_seg) - zer[x, y, z] = discs_seg[x, y, z] # isolate disc - # Loop on all the pixels of the segmentation - min_dist = np.inf - nonzero = np.where(zer > 0) - for u, v, w in zip_strict(nonzero[0], nonzero[1], nonzero[2]): - _, dist = project_point_on_line(np.array([u, v, w]), centerline) - if dist < min_dist: - min_dist = dist - min_point = np.array([u, v, w, discs_seg[u, v, w]]) - discs_list.append(min_point) - return np.array(discs_list) - - -if __name__ == "__main__": - main() diff --git a/spineps/utils/image.py b/spineps/utils/image.py deleted file mode 100644 index 4b34d1f..0000000 --- a/spineps/utils/image.py +++ /dev/null @@ -1,697 +0,0 @@ -from __future__ import annotations - -import logging -import os -from copy import deepcopy - -import nibabel as nib -import numpy as np - -logger = logging.getLogger(__name__) - - -class Image: - """ - Compact version of SCT's Image Class (https://github.com/spinalcordtoolbox/spinalcordtoolbox/blob/master/spinalcordtoolbox/image.py#L245) - Create an object that behaves similarly to nibabel's image object. Useful additions include: dims, change_orientation and getNonZeroCoordinates. - """ - - def __init__(self, param=None, hdr=None, orientation=None, absolutepath=None, dim=None): # noqa: ARG002 - """ - :param param: string indicating a path to a image file or an `Image` object. - """ - - # initialization of all parameters - self.affine = None - self.data = None - self._path = None - self.ext = "" - - if absolutepath is not None: - self._path = os.path.abspath(absolutepath) # noqa: PTH100 - - # Case 1: load an image from file - if isinstance(param, str): - self.loadFromPath(param) - # Case 2: create a copy of an existing `Image` object - elif isinstance(param, type(self)): - self.copy(param) - # Case 3: create a blank image from a list of dimensions - elif isinstance(param, list): - self.data = np.zeros(param) - self.hdr = hdr.copy() if hdr is not None else nib.Nifti1Header() - self.hdr.set_data_shape(self.data.shape) - # Case 4: create an image from an existing data array - elif isinstance(param, (np.ndarray, np.generic)): - self.data = param - self.hdr = hdr.copy() if hdr is not None else nib.Nifti1Header() - self.hdr.set_data_shape(self.data.shape) - else: - raise TypeError("Image constructor takes at least one argument.") - - # Fix any mismatch between the array's datatype and the header datatype - self.fix_header_dtype() - - @property - def dim(self): - return get_dimension(self) - - @property - def orientation(self): - return get_orientation(self) - - @property - def absolutepath(self): - """ - Storage path (either actual or potential) - - Notes: - - - As several tools perform chdir() it's very important to have absolute paths - - When set, if relative: - - - If it already existed, it becomes a new basename in the old dirname - - Else, it becomes absolute (shortcut) - - Usually not directly touched (use `Image.save`), but in some cases it's - the best way to set it. - """ - return self._path - - @absolutepath.setter - def absolutepath(self, value): - if value is None: - self._path = None - return - elif not os.path.isabs(value) and self._path is not None: # noqa: PTH117 - value = os.path.join(os.path.dirname(self._path), value) # noqa: PTH118, PTH120 - elif not os.path.isabs(value): # noqa: PTH117 - value = os.path.abspath(value) # noqa: PTH100 - self._path = value - - @property - def header(self): - return self.hdr - - @header.setter - def header(self, value): - self.hdr = value - - def __deepcopy__(self, memo): - return type(self)( - deepcopy(self.data, memo), - deepcopy(self.hdr, memo), - deepcopy(self.orientation, memo), - deepcopy(self.absolutepath, memo), - deepcopy(self.dim, memo), - ) - - def copy(self, image=None): - if image is not None: - self.affine = deepcopy(image.affine) - self.data = deepcopy(image.data) - self.hdr = deepcopy(image.hdr) - self._path = deepcopy(image._path) - else: - return deepcopy(self) - - def loadFromPath(self, path): - """ - This function load an image from an absolute path using nibabel library - - :param path: path of the file from which the image will be loaded - :return: - """ - - self.absolutepath = os.path.abspath(path) # noqa: PTH100 - im_file = nib.load(self.absolutepath, mmap=True) - self.affine = im_file.affine.copy() - self.data = np.asanyarray(im_file.dataobj) - self.hdr = im_file.header.copy() - if path != self.absolutepath: - logger.debug("Loaded %s (%s) orientation %s shape %s", path, self.absolutepath, self.orientation, self.data.shape) - else: - logger.debug("Loaded %s orientation %s shape %s", path, self.orientation, self.data.shape) - - def change_orientation(self, orientation, inverse=False): - """ - Change orientation on image (in-place). - - :param orientation: orientation string (SCT "from" convention) - - :param inverse: if you think backwards, use this to specify that you actually\ - want to transform *from* the specified orientation, not *to*\ - it. - - """ - change_orientation(self, orientation, self, inverse=inverse) - return self - - def getNonZeroCoordinates(self, sorting=None, reverse_coord=False): - """ - This function return all the non-zero coordinates that the image contains. - Coordinate list can also be sorted by x, y, z, or the value with the parameter sorting='x', sorting='y', sorting='z' or sorting='value' - If reverse_coord is True, coordinate are sorted from larger to smaller. - - Removed Coordinate object - """ - n_dim = 1 - n_dim = 3 if self.dim[3] == 1 else 4 - if self.dim[2] == 1: - n_dim = 2 - - if n_dim == 3: - x, y, z = (self.data > 0).nonzero() - list_coordinates = [[x[i], y[i], z[i], self.data[x[i], y[i], z[i]]] for i in range(len(x))] - elif n_dim == 2: - try: - x, y = (self.data > 0).nonzero() - list_coordinates = [[x[i], y[i], 0, self.data[x[i], y[i]]] for i in range(len(x))] - except ValueError: - x, y, z = (self.data > 0).nonzero() - list_coordinates = [[x[i], y[i], 0, self.data[x[i], y[i], 0]] for i in range(len(x))] - - if sorting is not None: - if reverse_coord not in [True, False]: - raise ValueError("reverse_coord parameter must be a boolean") - - if sorting == "x": - list_coordinates = sorted(list_coordinates, key=lambda el: el[0], reverse=reverse_coord) - elif sorting == "y": - list_coordinates = sorted(list_coordinates, key=lambda el: el[1], reverse=reverse_coord) - elif sorting == "z": - list_coordinates = sorted(list_coordinates, key=lambda el: el[2], reverse=reverse_coord) - elif sorting == "value": - list_coordinates = sorted(list_coordinates, key=lambda el: el[3], reverse=reverse_coord) - else: - raise ValueError("sorting parameter must be either 'x', 'y', 'z' or 'value'") - - return list_coordinates - - def change_type(self, dtype): - """ - Change data type on image. - - Note: the image path is voided. - """ - change_type(self, dtype, self) - return self - - def fix_header_dtype(self): - """ - Change the header dtype to the match the datatype of the array. - """ - # Using bool for nibabel headers is unsupported, so use uint8 instead: - # `nibabel.spatialimages.HeaderDataError: data dtype "bool" not supported` - dtype_data = self.data.dtype - if dtype_data is bool: - dtype_data = np.uint8 - - dtype_header = self.hdr.get_data_dtype() - if dtype_header != dtype_data: - logger.warning( - f"Image header specifies datatype '{dtype_header}', but array is of type " # noqa: G004 - f"'{dtype_data}'. Header metadata will be overwritten to use '{dtype_data}'." - ) - self.hdr.set_data_dtype(dtype_data) - - def save(self, path=None, dtype=None, verbose=1, mutable=False): - """ - Write an image in a nifti file - - :param path: Where to save the data, if None it will be taken from the\ - absolutepath member.\ - If path is a directory, will save to a file under this directory\ - with the basename from the absolutepath member. - - :param dtype: if not set, the image is saved in the same type as input data\ - if 'minimize', image storage space is minimized\ - (2, 'uint8', np.uint8, "NIFTI_TYPE_UINT8"),\ - (4, 'int16', np.int16, "NIFTI_TYPE_INT16"),\ - (8, 'int32', np.int32, "NIFTI_TYPE_INT32"),\ - (16, 'float32', np.float32, "NIFTI_TYPE_FLOAT32"),\ - (32, 'complex64', np.complex64, "NIFTI_TYPE_COMPLEX64"),\ - (64, 'float64', np.float64, "NIFTI_TYPE_FLOAT64"),\ - (256, 'int8', np.int8, "NIFTI_TYPE_INT8"),\ - (512, 'uint16', np.uint16, "NIFTI_TYPE_UINT16"),\ - (768, 'uint32', np.uint32, "NIFTI_TYPE_UINT32"),\ - (1024,'int64', np.int64, "NIFTI_TYPE_INT64"),\ - (1280, 'uint64', np.uint64, "NIFTI_TYPE_UINT64"),\ - (1536, 'float128', _float128t, "NIFTI_TYPE_FLOAT128"),\ - (1792, 'complex128', np.complex128, "NIFTI_TYPE_COMPLEX128"),\ - (2048, 'complex256', _complex256t, "NIFTI_TYPE_COMPLEX256"), - - :param mutable: whether to update members with newly created path or dtype - """ - if mutable: # do all modifications in-place - # Case 1: `path` not specified - if path is None: - if self.absolutepath: # Fallback to the original filepath - path = self.absolutepath - else: - raise ValueError("Don't know where to save the image (no absolutepath or path parameter)") - # Case 2: `path` points to an existing directory - elif os.path.isdir(path): # noqa: PTH112 - if self.absolutepath: # Use the original filename, but save to the directory specified by `path` - path = os.path.join(os.path.abspath(path), os.path.basename(self.absolutepath)) # noqa: PTH100, PTH118, PTH119 - else: - raise ValueError("Don't know where to save the image (path parameter is dir, but absolutepath is missing)") - # Case 3: `path` points to a file (or a *nonexistent* directory) so use its value as-is - # (We're okay with letting nonexistent directories slip through, because it's difficult to distinguish - # between nonexistent directories and nonexistent files. Plus, `nibabel` will catch any further errors.) - else: - pass - - if os.path.isfile(path) and verbose: # noqa: PTH113 - logger.warning("File %s already exists. Will overwrite it.", path) - if os.path.isabs(path): # noqa: PTH117 - logger.debug("Saving image to %s orientation %s shape %s", path, self.orientation, self.data.shape) - else: - logger.debug( - "Saving image to %s (%s) orientation %s shape %s", - path, - os.path.abspath(path), # noqa: PTH100 - self.orientation, - self.data.shape, - ) - - # Now that `path` has been set and log messages have been written, we can assign it to the image itself - self.absolutepath = os.path.abspath(path) # noqa: PTH100 - - if dtype is not None: - self.change_type(dtype) - - if self.hdr is not None: - self.hdr.set_data_shape(self.data.shape) - self.fix_header_dtype() - - # nb. that copy() is important because if it were a memory map, save() would corrupt it - dataobj = self.data.copy() - affine = None - header = self.hdr.copy() if self.hdr is not None else None - nib.save(nib.nifti1.Nifti1Image(dataobj, affine, header), self.absolutepath) - if not os.path.isfile(self.absolutepath): # noqa: PTH113 - raise RuntimeError(f"Couldn't save image to {self.absolutepath}") - else: - # if we're not operating in-place, then make any required modifications on a throw-away copy - self.copy().save(path, dtype, verbose, mutable=True) - return self - - -class SlicerOneAxis: - """ - Image slicer to use when you don't care about the 2D slice orientation, - and don't want to specify them. - The slicer will just iterate through the right axis that corresponds to - its specification. - - Can help getting ranges and slice indices. - - Copied from https://github.com/spinalcordtoolbox/spinalcordtoolbox/image.py - """ - - def __init__(self, im, axis="IS"): - opposite_character = {"L": "R", "R": "L", "A": "P", "P": "A", "I": "S", "S": "I"} - axis_labels = "LRPAIS" - if len(axis) != 2: - raise ValueError() - if axis[0] not in axis_labels: - raise ValueError() - if axis[1] not in axis_labels: - raise ValueError() - if axis[0] != opposite_character[axis[1]]: - raise ValueError() - - for idx_axis in range(2): - dim_nr = im.orientation.find(axis[idx_axis]) - if dim_nr != -1: - break - if dim_nr == -1: - raise ValueError() - - # SCT convention - from_dir = im.orientation[dim_nr] - self.direction = +1 if axis[0] == from_dir else -1 - self.nb_slices = im.dim[dim_nr] - self.im = im - self.axis = axis - self._slice = lambda idx: tuple([(idx if x in axis else slice(None)) for x in im.orientation]) - - def __len__(self): - return self.nb_slices - - def __getitem__(self, idx): - """ - - :return: an image slice, at slicing index idx - :param idx: slicing index (according to the slicing direction) - """ - if isinstance(idx, slice): - raise NotImplementedError() - - if idx >= self.nb_slices: - raise IndexError(f"I just have {self.nb_slices} slices!") - - if self.direction == -1: - idx = self.nb_slices - 1 - idx - - return self.im.data[self._slice(idx)] - - -def get_dimension(im_file, verbose=1): # noqa: ARG001 - """ - Copied from https://github.com/spinalcordtoolbox/spinalcordtoolbox/ - - Get dimension from Image or nibabel object. Manages 2D, 3D or 4D images. - - :param: im_file: Image or nibabel object - :return: nx, ny, nz, nt, px, py, pz, pt - """ - if not isinstance(im_file, (nib.nifti1.Nifti1Image, Image)): - raise TypeError("The provided image file is neither a nibabel.nifti1.Nifti1Image instance nor an Image instance") - # initializating ndims [nx, ny, nz, nt] and pdims [px, py, pz, pt] - ndims = [1, 1, 1, 1] - pdims = [1, 1, 1, 1] - data_shape = im_file.header.get_data_shape() - zooms = im_file.header.get_zooms() - for i in range(min(len(data_shape), 4)): - ndims[i] = data_shape[i] - pdims[i] = zooms[i] - return *ndims, *pdims - - -def change_orientation(im_src, orientation, im_dst=None, inverse=False): - """ - Copied from https://github.com/spinalcordtoolbox/spinalcordtoolbox/ - - :param im_src: source image - :param orientation: orientation string (SCT "from" convention) - :param im_dst: destination image (can be the source image for in-place - operation, can be unset to generate one) - :param inverse: if you think backwards, use this to specify that you actually - want to transform *from* the specified orientation, not *to* it. - :return: an image with changed orientation - - .. note:: - - the resulting image has no path member set - - if the source image is < 3D, it is reshaped to 3D and the destination is 3D - """ - - if len(im_src.data.shape) < 3: - pass # Will reshape to 3D - elif len(im_src.data.shape) == 3: - pass # OK, standard 3D volume - elif len(im_src.data.shape) == 4: - pass # OK, standard 4D volume - elif len(im_src.data.shape) == 5 and im_src.header.get_intent()[0] == "vector": - pass # OK, physical displacement field - else: - raise NotImplementedError("Don't know how to change orientation for this image") - - im_src_orientation = im_src.orientation - im_dst_orientation = orientation - if inverse: - im_src_orientation, im_dst_orientation = im_dst_orientation, im_src_orientation - - perm, inversion = _get_permutations(im_src_orientation, im_dst_orientation) - - if im_dst is None: - im_dst = im_src.copy() - im_dst._path = None - - im_src_data = im_src.data - if len(im_src_data.shape) < 3: - im_src_data = im_src_data.reshape(tuple(list(im_src_data.shape) + ([1] * (3 - len(im_src_data.shape))))) - - # Update data by performing inversions and swaps - - # axes inversion (flip) - data = im_src_data[:: inversion[0], :: inversion[1], :: inversion[2]] - - # axes manipulations (transpose) - if perm == [1, 0, 2]: - data = np.swapaxes(data, 0, 1) - elif perm == [2, 1, 0]: - data = np.swapaxes(data, 0, 2) - elif perm == [0, 2, 1]: - data = np.swapaxes(data, 1, 2) - elif perm == [2, 0, 1]: - data = np.swapaxes(data, 0, 2) # transform [2, 0, 1] to [1, 0, 2] - data = np.swapaxes(data, 0, 1) # transform [1, 0, 2] to [0, 1, 2] - elif perm == [1, 2, 0]: - data = np.swapaxes(data, 0, 2) # transform [1, 2, 0] to [0, 2, 1] - data = np.swapaxes(data, 1, 2) # transform [0, 2, 1] to [0, 1, 2] - elif perm == [0, 1, 2]: - # do nothing - pass - else: - raise NotImplementedError() - - # Update header - - im_src_aff = im_src.hdr.get_best_affine() - aff = nib.orientations.inv_ornt_aff(np.array((perm, inversion)).T, im_src_data.shape) - im_dst_aff = np.matmul(im_src_aff, aff) - - im_dst.header.set_qform(im_dst_aff) - im_dst.header.set_sform(im_dst_aff) - im_dst.header.set_data_shape(data.shape) - im_dst.data = data - - return im_dst - - -def _get_permutations(im_src_orientation, im_dst_orientation): - """ - Copied from https://github.com/spinalcordtoolbox/spinalcordtoolbox/ - - :param im_src_orientation str: Orientation of source image. Example: 'RPI' - :param im_dest_orientation str: Orientation of destination image. Example: 'SAL' - :return: list of axes permutations and list of inversions to achieve an orientation change - """ - - opposite_character = {"L": "R", "R": "L", "A": "P", "P": "A", "I": "S", "S": "I"} - - perm = [0, 1, 2] - inversion = [1, 1, 1] - for i, character in enumerate(im_src_orientation): - try: - perm[i] = im_dst_orientation.index(character) - except ValueError: - perm[i] = im_dst_orientation.index(opposite_character[character]) - inversion[i] = -1 - - return perm, inversion - - -def get_orientation(im): - """ - Copied from https://github.com/spinalcordtoolbox/spinalcordtoolbox/ - - :param im: an Image - :return: reference space string (ie. what's in Image.orientation) - """ - res = "".join(nib.orientations.aff2axcodes(im.hdr.get_best_affine())) - return orientation_string_nib2sct(res) - - -def orientation_string_nib2sct(s): - """ - Copied from https://github.com/spinalcordtoolbox/spinalcordtoolbox/ - - :return: SCT reference space code from nibabel one - """ - opposite_character = {"L": "R", "R": "L", "A": "P", "P": "A", "I": "S", "S": "I"} - return "".join([opposite_character[x] for x in s]) - - -def change_type(im_src, dtype, im_dst=None): # noqa: C901 - """ - Change the voxel type of the image - - :param dtype: if not set, the image is saved in standard type\ - if 'minimize', image space is minimize\ - if 'minimize_int', image space is minimize and values are approximated to integers\ - (2, 'uint8', np.uint8, "NIFTI_TYPE_UINT8"),\ - (4, 'int16', np.int16, "NIFTI_TYPE_INT16"),\ - (8, 'int32', np.int32, "NIFTI_TYPE_INT32"),\ - (16, 'float32', np.float32, "NIFTI_TYPE_FLOAT32"),\ - (32, 'complex64', np.complex64, "NIFTI_TYPE_COMPLEX64"),\ - (64, 'float64', np.float64, "NIFTI_TYPE_FLOAT64"),\ - (256, 'int8', np.int8, "NIFTI_TYPE_INT8"),\ - (512, 'uint16', np.uint16, "NIFTI_TYPE_UINT16"),\ - (768, 'uint32', np.uint32, "NIFTI_TYPE_UINT32"),\ - (1024,'int64', np.int64, "NIFTI_TYPE_INT64"),\ - (1280, 'uint64', np.uint64, "NIFTI_TYPE_UINT64"),\ - (1536, 'float128', _float128t, "NIFTI_TYPE_FLOAT128"),\ - (1792, 'complex128', np.complex128, "NIFTI_TYPE_COMPLEX128"),\ - (2048, 'complex256', _complex256t, "NIFTI_TYPE_COMPLEX256"), - :return: - - Copied from https://github.com/spinalcordtoolbox/spinalcordtoolbox/ - """ - - if im_dst is None: - im_dst = im_src.copy() - im_dst._path = None - - if dtype is None: - return im_dst - - # get min/max from input image - min_in = np.nanmin(im_src.data) - max_in = np.nanmax(im_src.data) - - # find optimum type for the input image - if dtype in ("minimize", "minimize_int"): - # warning: does not take intensity resolution into account, neither complex voxels - - # check if voxel values are real or integer - isinteger = True - if dtype == "minimize": - for vox in im_src.data.flatten(): - if int(vox) != vox: - isinteger = False - break - - if isinteger: - if min_in >= 0: # unsigned - if max_in <= np.iinfo(np.uint8).max: - dtype = np.uint8 - elif max_in <= np.iinfo(np.uint16): - dtype = np.uint16 - elif max_in <= np.iinfo(np.uint32).max: - dtype = np.uint32 - elif max_in <= np.iinfo(np.uint64).max: - dtype = np.uint64 - else: - raise ValueError("Maximum value of the image is to big to be represented.") - else: # noqa: PLR5501 - if max_in <= np.iinfo(np.int8).max and min_in >= np.iinfo(np.int8).min: - dtype = np.int8 - elif max_in <= np.iinfo(np.int16).max and min_in >= np.iinfo(np.int16).min: - dtype = np.int16 - elif max_in <= np.iinfo(np.int32).max and min_in >= np.iinfo(np.int32).min: - dtype = np.int32 - elif max_in <= np.iinfo(np.int64).max and min_in >= np.iinfo(np.int64).min: - dtype = np.int64 - else: - raise ValueError("Maximum value of the image is to big to be represented.") - else: # noqa: PLR5501 - # if max_in <= np.finfo(np.float16).max and min_in >= np.finfo(np.float16).min: - # type = 'np.float16' # not supported by nibabel - if max_in <= np.finfo(np.float32).max and min_in >= np.finfo(np.float32).min: - dtype = np.float32 - elif max_in <= np.finfo(np.float64).max and min_in >= np.finfo(np.float64).min: - dtype = np.float64 - - dtype = to_dtype(dtype) - else: - dtype = to_dtype(dtype) - - # if output type is int, check if it needs intensity rescaling - if "int" in dtype.name: - # get min/max from output type - min_out = np.iinfo(dtype).min - max_out = np.iinfo(dtype).max - # before rescaling, check if there would be an intensity overflow - - if (min_in < min_out) or (max_in > max_out): - # This condition is important for binary images since we do not want to scale them - logger.warning( - f"To avoid intensity overflow due to convertion to +{dtype.name}+, intensity will be rescaled to the maximum quantization scale" # noqa: G004 - ) - # rescale intensity - data_rescaled = im_src.data * (max_out - min_out) / (max_in - min_in) - im_dst.data = data_rescaled - (data_rescaled.min() - min_out) - - # change type of data in both numpy array and nifti header - im_dst.data = getattr(np, dtype.name)(im_dst.data) - im_dst.hdr.set_data_dtype(dtype) - return im_dst - - -def to_dtype(dtype): - """ - Take a dtypeification and return an np.dtype - - :param dtype: dtypeification (string or np.dtype or None are supported for now) - :return: dtype or None - - Copied from https://github.com/spinalcordtoolbox/spinalcordtoolbox/ - """ - # TODO add more or filter on things supported by nibabel - - if dtype is None: - return None - if isinstance(dtype, type) and isinstance(dtype(0).dtype, np.dtype): - return dtype(0).dtype - if isinstance(dtype, np.dtype): - return dtype - if isinstance(dtype, str): - return np.dtype(dtype) - - raise TypeError(f"data type {dtype.__class__}: {dtype} not understood") - - -def zeros_like(img, dtype=None): - """ - - :param img: reference image - :param dtype: desired data type (optional) - :return: an Image with the same shape and header, filled with zeros - - Similar to numpy.zeros_like(), the goal of the function is to show the developer's - intent and avoid doing a copy, which is slower than initialization with a constant. - - Copied from https://github.com/spinalcordtoolbox/spinalcordtoolbox/image.py - """ - zimg = Image(np.zeros_like(img.data), hdr=img.hdr.copy()) - if dtype is not None: - zimg.change_type(dtype) - return zimg - - -def empty_like(img, dtype=None): - """ - :param img: reference image - :param dtype: desired data type (optional) - :return: an Image with the same shape and header, whose data is uninitialized - - Similar to numpy.empty_like(), the goal of the function is to show the developer's - intent and avoid touching the allocated memory, because it will be written to - afterwards. - - Copied from https://github.com/spinalcordtoolbox/spinalcordtoolbox/image.py - """ - dst = change_type(img, dtype) - return dst - - -def find_zmin_zmax(im, threshold=0.1): - """ - Find the min (and max) z-slice index below which (and above which) slices only have voxels below a given threshold. - - :param im: Image object - :param threshold: threshold to apply before looking for zmin/zmax, typically corresponding to noise level. - :return: [zmin, zmax] - - Copied from https://github.com/spinalcordtoolbox/spinalcordtoolbox/image.py - """ - slicer = SlicerOneAxis(im, axis="IS") - - # Make sure image is not empty - if not np.any(slicer): - logger.error("Input image is empty") - - # Iterate from bottom to top until we find data - for zmin in range(len(slicer)): - if np.any(slicer[zmin] > threshold): - break - - # Conversely from top to bottom - for zmax in range(len(slicer) - 1, zmin, -1): - if np.any(slicer[zmax] > threshold): - break - - return zmin, zmax diff --git a/spineps/utils/proc_functions.py b/spineps/utils/proc_functions.py index 101b867..e4ec311 100755 --- a/spineps/utils/proc_functions.py +++ b/spineps/utils/proc_functions.py @@ -125,49 +125,65 @@ def clean_cc_artifacts( f"cc_size_threshold size does not match number of given labels to clean, got {len(labels)} and {len(cc_size_threshold)}. Specifiy only an int for cc_size_threshold to use it for all labels" ) - subreg_cc, subreg_cc_stats = connected_components_3d(result_arr, connectivity=1) - cc_to_clean = {} for lidx, label in enumerate(tqdm(labels, desc=f"{logger._get_logger_prefix()} cleaning...", disable=not verbose)): + # One label at a time: each costs a full-volume component array, and asking for all of them up + # front made this the memory peak of the instance merge (~25 labels). Components come from the + # untouched input, exactly as when they were all computed before the loop started. + subreg_cc, subreg_cc_stats = connected_components_3d(mask_arr, connectivity=1, label_ref=[label]) + if label not in subreg_cc: + continue idx = [i for i, v in enumerate(subreg_cc_stats[label]["voxel_counts"]) if v < cc_size_threshold[lidx] and v > 0] if len(idx) > 0: cc_to_clean[label] = idx + bounding_boxes = subreg_cc_stats[label]["bounding_boxes"] + mask_cc = subreg_cc[label] for cc_idx in idx: - # extract cc label - mask_cc = subreg_cc[label] - mask_cc_l = mask_cc.copy() + # The components handled here are by definition small, so everything below runs inside the + # component's own bounding box (padded so the 1-voxel dilation still fits) instead of over + # the whole volume. + bbox = _padded_bbox(bounding_boxes[cc_idx], mask_cc.shape, pad=2) + mask_cc_l = mask_cc[bbox].copy() mask_cc_l[mask_cc_l != cc_idx] = 0 + cc_voxels = mask_cc_l != 0 log_string = "" if verbose: cc_volume = np_count_nonzero(mask_cc_l) - cc_centroid = center_of_mass(mask_cc_l) - cc_centroid = [int(c) + 1 for c in cc_centroid] # type: ignore + cc_centroid = [int(c + bbox[d].start) + 1 for d, c in enumerate(center_of_mass(mask_cc_l))] log_string = f"Label {label}, cc{cc_idx}, at {cc_centroid}, volume {cc_volume}: " if only_delete: logger.print(log_string + "deleted") if verbose else None # dilated mask nothing in original mask, just delete it - result_arr[mask_cc_l != 0] = 0 + result_arr[bbox][cc_voxels] = 0 continue - dilated_m = np_dilate_msk(mask_cc_l, n_pixel=1) - dilated_m[mask_cc_l != 0] = 0 + # np_dilate_msk mutates its input and returns the same object. The old code relied on that + # accidentally: `dilated_m[mask_cc_l != 0] = 0` re-read the *already dilated* mask, so it zeroed + # the shell as well as the component -- leaving an empty neighbourhood, an always-taken "delete" + # branch, and a `mask_cc_l` that no longer selected anything. The whole relabel/delete path was a + # no-op. Dilate a copy so the component mask stays intact. + dilated_m = np_dilate_msk(mask_cc_l.copy(), n_pixel=1) + dilated_m[cc_voxels] = 0 neighbor_voxel_count = np_count_nonzero(dilated_m) - mult = mask_arr * dilated_m + mask_arr_c = mask_arr[bbox] + mult = mask_arr_c * dilated_m if np_count_nonzero(mult) <= int(neighbor_voxel_count * neighbor_factor_2_delete): logger.print(log_string + "deleted") if verbose else None # dilated mask nothing in original mask, just delete it - result_arr[mask_cc_l != 0] = 0 + result_arr[bbox][cc_voxels] = 0 else: # majority voting dilated_m[dilated_m != 0] = 1 - mult = mask_arr * dilated_m + mult = mask_arr_c * dilated_m volumes = np_volume(mult) nlabels = list(volumes.keys()) volumes_values = list(volumes.values()) newlabel = nlabels[np.argmax(volumes_values)] # type: ignore - result_arr[mask_cc_l != 0] = newlabel + result_arr[bbox][cc_voxels] = newlabel logger.print(log_string + f"labeled as {newlabel}") if verbose else None + # release this label's full-volume component map before building the next one + del mask_cc, subreg_cc, subreg_cc_stats n_to_clean = {k: len(v) for k, v in cc_to_clean.items()} # By clearning: look at surrounding neighbor pixels. If too few, remove cc. otherwise, do majority voting if len(n_to_clean) != 0: @@ -175,7 +191,17 @@ def clean_cc_artifacts( return result_arr -def connected_components_3d(mask_image: np.ndarray, connectivity: int = 3, verbose: bool = False) -> tuple[dict, dict]: # noqa: ARG001 +def _padded_bbox(bbox: tuple[slice, ...], shape: tuple[int, ...], pad: int) -> tuple[slice, ...]: + """Grow a component bounding box by ``pad`` voxels per side, clamped to the array bounds.""" + return tuple(slice(max(s.start - pad, 0), min(s.stop + pad, shape[d])) for d, s in enumerate(bbox)) + + +def connected_components_3d( + mask_image: np.ndarray, + connectivity: int = 3, + verbose: bool = False, # noqa: ARG001 + label_ref: int | list[int] | None = None, +) -> tuple[dict, dict]: """Compute 3D connected components per label together with their statistics. Args: @@ -183,14 +209,17 @@ def connected_components_3d(mask_image: np.ndarray, connectivity: int = 3, verbo connectivity (int, optional): Voxel connectivity in range [1, 3]. For 2D images 2 and 3 are equivalent. Defaults to 3. verbose (bool, optional): Currently unused. Defaults to False. + label_ref (int | list[int] | None, optional): Restrict the computation to these labels. Each label costs + one full-volume component array, so passing only the labels you need matters. Defaults to None (all). Returns: tuple[dict, dict]: A dict mapping each label to its connected-component array, and a dict mapping each - label to its ``cc3d`` component statistics. + label to its ``cc3d`` component statistics (including per-component ``bounding_boxes``). """ subreg_cc = np_connected_components_per_label( mask_image, connectivity=connectivity, + label_ref=label_ref, ) subreg_cc_stats = {k: cc3d.statistics(v) for k, v in subreg_cc.items()} return subreg_cc, subreg_cc_stats @@ -226,8 +255,17 @@ def fix_wrong_posterior_instance_label(seg_sem: NII, seg_inst: NII, logger: Logg instance_labels = [i for i in seg_inst.unique() if 1 <= i <= MAX_VERTEBRA_INSTANCE_LABEL] for vert in instance_labels: - inst_vert = seg_inst.extract_label(vert) - # sem_vert = seg_sem.apply_mask(inst_vert) + # Everything below concerns one vertebra, so crop to its bounding box (+1, the margin the inner + # per-component crops use) once instead of running connected components and several crops over the + # whole volume for each of ~25 instances. + inst_vert_full = seg_inst.extract_label(vert) + try: + vert_crop = inst_vert_full.compute_crop(dist=1) + except ValueError: # label vanished, nothing to reassign + continue + inst_vert = inst_vert_full.apply_crop(vert_crop) + seg_inst_c = seg_inst.apply_crop(vert_crop) + seg_sem_c = seg_sem.apply_crop(vert_crop) # Check if multiple CC exist inst_vert_cc: NII = inst_vert.filter_connected_components(max_count_component=3, keep_label=False) @@ -242,7 +280,7 @@ def fix_wrong_posterior_instance_label(seg_sem: NII, seg_inst: NII, logger: Logg crop = inst_vert_cc_i.compute_crop(dist=1) inst_vert_cc_i_c = inst_vert_cc_i.apply_crop(crop) - cc_sem_vert = seg_sem.apply_crop(crop).apply_mask(inst_vert_cc_i_c) + cc_sem_vert = seg_sem_c.apply_crop(crop).apply_mask(inst_vert_cc_i_c) # cc_vert is semantic mask of only that cc of instance cc_sem_vert_labels = cc_sem_vert.unique() @@ -251,7 +289,7 @@ def fix_wrong_posterior_instance_label(seg_sem: NII, seg_inst: NII, logger: Logg [i in [Location.Arcus_Vertebrae.value, Location.Spinosus_Process.value] for i in cc_sem_vert_labels] ): # neighbor that have non arcus/spinosus label? - neighbor_instance_labels = seg_inst.apply_crop(crop).get_seg_array() + neighbor_instance_labels = seg_inst_c.apply_crop(crop).get_seg_array() neighbor_instance_labels[inst_vert_cc_i_c.get_seg_array() == 1] = 0 neighbor_instance_labels = np_unique_withoutzero(neighbor_instance_labels) # which instance labels does it touch @@ -260,7 +298,7 @@ def fix_wrong_posterior_instance_label(seg_sem: NII, seg_inst: NII, logger: Logg if len(neighbor_instance_labels) == 1 and neighbor_instance_labels[0] != vert: to_label = neighbor_instance_labels[0] logger.print(f"vert {vert}, cc_k {i} relabel to instance {to_label}") - seg_inst_arr_proc[inst_vert_cc_i.get_seg_array() == 1] = to_label + seg_inst_arr_proc[vert_crop][inst_vert_cc_i.get_seg_array() == 1] = to_label seg_inst_proc = seg_inst.set_array(seg_inst_arr_proc).reorient_(orientation) return seg_inst_proc diff --git a/unit_tests/test_bugfixes.py b/unit_tests/test_bugfixes.py new file mode 100644 index 0000000..eb6c4fc --- /dev/null +++ b/unit_tests/test_bugfixes.py @@ -0,0 +1,338 @@ +# Call 'python -m unittest' on this folder +"""Targeted tests for the bugs found in the 2.0 audit pass. + +Each test names the defect it pins down, so a regression is self-explaining. +""" + +from __future__ import annotations + +import os +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +import nibabel as nib +import numpy as np +from TPTBox import NII, No_Logger +from TPTBox.tests.test_utils import get_test_mri + +from spineps.phase_instance import find_prediction_couple +from spineps.phase_labeling import perform_labeling_step +from spineps.phase_post import detect_and_solve_merged_vertebra +from spineps.phase_semantic import semantic_bounding_box_clean +from spineps.seg_pipeline import pipeline_version +from spineps.utils.citation_reminder import OPT_OUT_ENV_VAR, reminder_disabled +from spineps.utils.find_min_cost_path import DEFAULT_REGION_STARTS, find_most_probably_sequence + +logger = No_Logger() + + +def _nii(arr: np.ndarray, zoom: float = 1.0) -> NII: + """Wrap an array as a ``(P, I, R)`` oriented segmentation without permuting it. + + Axis 0 grows posteriorly, axis 1 inferiorly, axis 2 to the right -- so the index arithmetic in the + tests below reads the same way the pipeline code does. + """ + affine = np.array([[0, 0, zoom, 0], [-zoom, 0, 0, 0], [0, -zoom, 0, 0], [0, 0, 0, 1.0]]) + nii = NII(nib.Nifti1Image(arr.astype(np.uint8), affine=affine), seg=True) + assert nii.orientation == ("P", "I", "R"), nii.orientation + return nii + + +class Test_Merged_Vertebra_Background(unittest.TestCase): + """The IVD components must be offset without dragging the background out of 0. + + A plain ``subreg_cc += OFFSET`` turned every background voxel into one giant phantom "IVD" + whose center of mass sits in the middle of the volume. When the anatomy lives in the inferior + half, that phantom sorts *above* every real structure and takes the first slot in the + height-sorted list -- which is exactly the slot the split-C2 heuristic inspects. + """ + + @staticmethod + def _fixture() -> tuple[NII, NII]: + # PIR: axis 1 is the inferior axis, so "high index" == inferior == low in the body. + shape = (12, 40, 12) + seg = np.zeros(shape, dtype=np.uint8) + vert = np.zeros(shape, dtype=np.uint8) + # Everything sits in the inferior half so the background centroid lands above it all. + # A small upper vertebra (1) stacked directly onto a large one (2) -> should be merged. + seg[3:9, 24:27, 3:9] = 49 # Vertebra_Corpus_border, instance 1 + vert[3:9, 24:27, 3:9] = 1 + seg[3:9, 27:37, 3:9] = 49 # instance 2, clearly larger + vert[3:9, 27:37, 3:9] = 2 + seg[3:9, 37:39, 3:9] = 100 # an IVD below them, so the disc branch has something to find + return _nii(seg), _nii(vert) + + def test_top_two_instances_are_merged(self): + seg_nii, vert_nii = self._fixture() + detect_and_solve_merged_vertebra(seg_nii, vert_nii) + self.assertNotIn(1, vert_nii.unique(), "the small top instance should have been merged into its neighbour") + self.assertIn(2, vert_nii.unique()) + + +class Test_Semantic_Bounding_Box_Clean(unittest.TestCase): + """The region kept must be the union of the incorporated components' boxes, not just the largest.""" + + def test_incorporated_component_survives(self): + shape = (16, 60, 16) + arr = np.zeros(shape, dtype=np.uint8) + arr[6:10, 10:40, 6:10] = 49 # largest component + arr[6:10, 42:52, 6:10] = 49 # second component, below it, within the inferior margin + seg = _nii(arr, zoom=1.0) + n_second = int((arr[6:10, 42:52, 6:10] != 0).sum()) + kept = semantic_bounding_box_clean(seg.copy()).get_seg_array() + # The old code cropped to the largest component's box only, so the tail of an incorporated + # component beyond that box was silently deleted. + self.assertEqual(int((kept[6:10, 42:52, 6:10] != 0).sum()), n_second, "an incorporated component must survive whole") + self.assertGreater(kept[6:10, 10:40, 6:10].sum(), 0) + + def test_far_component_is_dropped(self): + shape = (16, 60, 16) + arr = np.zeros(shape, dtype=np.uint8) + arr[6:10, 4:34, 6:10] = 49 # largest component, superior + arr[1:3, 56:59, 1:3] = 49 # far away in every axis -> never incorporated + seg = _nii(arr, zoom=1.0) + cleaned = semantic_bounding_box_clean(seg.copy()) + self.assertEqual(cleaned.get_seg_array()[1:3, 56:59, 1:3].sum(), 0) + + +class Test_Prediction_Couple_Partner_Drop(unittest.TestCase): + """Two partners that agree with the anchor but not with each other cannot both be kept.""" + + @staticmethod + def _sparse(spans: dict) -> dict: + """Build sparse predictions from ``{(com, label): (start, stop)}`` spans along the first axis.""" + from spineps.phase_instance import SparsePrediction + + return { + key: SparsePrediction((slice(a, b), slice(0, 1), slice(0, 1)), np.ones((b - a, 1, 1), dtype=bool)) + for key, (a, b) in spans.items() + } + + def test_overlapping_partners_are_both_kept(self): + preds = self._sparse({(1, 1): (0, 6), (0, 1): (0, 5), (2, 1): (1, 7)}) + couple, agreement = find_prediction_couple(1, 1, preds, 3) + self.assertEqual(len(couple), 3, couple) + self.assertGreater(agreement, 0) + + def test_non_overlapping_partner_is_dropped(self): + # Both partners clear the Dice threshold against the anchor, but they are disjoint from each + # other -- so they cannot both be the same vertebra as the anchor. + preds = self._sparse({(1, 1): (0, 10), (0, 1): (0, 4), (2, 1): (6, 10)}) + couple, _agreement = find_prediction_couple(1, 1, preds, 3) + self.assertEqual(len(couple), 2, f"the weaker of two disagreeing partners must be dropped, got {couple}") + self.assertIn((1, 1), couple, "the anchor always stays in its own couple") + + +class Test_Sparse_Dice_Matches_Dense(unittest.TestCase): + """`sparse_dice` must agree with `np_dice` on the equivalent full-volume masks, exactly.""" + + def test_random_boxes(self): + from TPTBox.core.np_utils import np_dice + + from spineps.phase_instance import SparsePrediction, sparse_dice + + rng = np.random.default_rng(0) + shape = (24, 20, 18) + for _ in range(60): + preds = [] + for _side in range(2): + starts = [int(rng.integers(0, s - 4)) for s in shape] + sizes = [int(rng.integers(2, min(9, s - st))) for st, s in zip(starts, shape)] + bbox = tuple(slice(st, st + sz) for st, sz in zip(starts, sizes)) + mask = rng.random(tuple(sizes)) < 0.6 + preds.append(SparsePrediction(bbox, mask)) + dense = [] + for pr in preds: + full = np.zeros(shape, dtype=np.uint8) + full[pr.bbox] = pr.mask + dense.append(full) + self.assertAlmostEqual(sparse_dice(preds[0], preds[1]), float(np_dice(dense[0], dense[1])), places=12) + + def test_disjoint_and_empty(self): + from spineps.phase_instance import SparsePrediction, sparse_dice + + a = SparsePrediction((slice(0, 2), slice(0, 2), slice(0, 2)), np.ones((2, 2, 2), dtype=bool)) + far = SparsePrediction((slice(9, 11), slice(0, 2), slice(0, 2)), np.ones((2, 2, 2), dtype=bool)) + empty = SparsePrediction((slice(0, 2), slice(0, 2), slice(0, 2)), np.zeros((2, 2, 2), dtype=bool)) + self.assertEqual(sparse_dice(a, far), 0.0) + self.assertEqual(sparse_dice(empty, empty), 1.0, "np_dice returns 1.0 when both masks are empty") + self.assertEqual(sparse_dice(a, a), 1.0) + + +class Test_Min_Cost_Path_Arguments(unittest.TestCase): + def test_region_skip_without_rel_cost(self): + """`allow_skip_at_region` used to dereference a `regions_ranges` that was never built.""" + rng = np.random.default_rng(0) + cost = rng.random((4, 24)) + fcost, fpath, _mcp = find_most_probably_sequence(cost, allow_skip_at_region=[0], region_rel_cost=None) + self.assertEqual(len(fpath), 4) + self.assertIsInstance(fcost, float) + + def test_regions_argument_is_not_mutated(self): + before = list(DEFAULT_REGION_STARTS) + regions = list(DEFAULT_REGION_STARTS) + rng = np.random.default_rng(1) + find_most_probably_sequence(rng.random((3, 24)), regions=regions, region_rel_cost=None) + self.assertEqual(regions, before, "the caller's region list must not be appended to") + self.assertEqual(list(DEFAULT_REGION_STARTS), before) + + +class Test_Labeling_Guards(unittest.TestCase): + def test_empty_instance_mask_returns_unchanged(self): + from unit_tests.test_inference_mocked import Labeling_Model_Dummy + + mri, _subreg, vert, _label = get_test_mri() + empty = vert.set_array(np.zeros_like(vert.get_seg_array())) + model = Labeling_Model_Dummy().load() + out = perform_labeling_step(model, mri, empty, subreg_nii=None) + self.assertEqual(len(out.unique()), 0) + + def test_no_subreg_with_c1_enabled_does_not_crash(self): + from unittest.mock import MagicMock + + from unit_tests.test_inference_mocked import Labeling_Model_Dummy, _fake_run_all_seg_instances + + mri, _subreg, vert, _label = get_test_mri() + model = Labeling_Model_Dummy().load() + model.run_all_seg_instances = MagicMock(side_effect=_fake_run_all_seg_instances) + # disable_c1=False previously dereferenced the (None) subregion mask. + out = perform_labeling_step(model, mri, vert.copy(), subreg_nii=None, disable_c1=False) + self.assertIsInstance(out, NII) + + +class Test_Process_Dataset_Compatibility(unittest.TestCase): + def test_incompatible_model_raises(self): + """`process_dataset` used to log "stop program" and then carry on regardless.""" + from spineps.seg_enums import Acquisition, Modality + from spineps.seg_run import process_dataset + from unit_tests.test_inference_mocked import SegmentationModelDummy + + model = SegmentationModelDummy().load() + with tempfile.TemporaryDirectory() as d, self.assertRaises(ValueError): + # The dummy model is sagittal T2w/SEG/T1w; asking for an axial CT cannot work. + process_dataset( + dataset_path=Path(d), + model_instance=model, + model_semantic=model, + modalities=(Modality.CT, Acquisition.ax), + save_log_data=False, + ) + + +class Test_Pipeline_Version(unittest.TestCase): + def test_version_is_not_read_from_the_callers_git_repo(self): + pipeline_version.cache_clear() + with patch("spineps.seg_pipeline._package_version", return_value="2.0.0") as m: + self.assertEqual(pipeline_version(), "2.0.0") + m.assert_called_once() + pipeline_version.cache_clear() + + +class Test_Citation_Opt_Out(unittest.TestCase): + def test_env_var_disables_reminder(self): + old = os.environ.get(OPT_OUT_ENV_VAR) + try: + os.environ[OPT_OUT_ENV_VAR] = "1" + self.assertTrue(reminder_disabled()) + os.environ[OPT_OUT_ENV_VAR] = "TRUE" + self.assertTrue(reminder_disabled()) + os.environ.pop(OPT_OUT_ENV_VAR) + self.assertFalse(reminder_disabled()) + finally: + if old is None: + os.environ.pop(OPT_OUT_ENV_VAR, None) + else: + os.environ[OPT_OUT_ENV_VAR] = old + + +class Test_Separating_Components(unittest.TestCase): + """`get_separating_components` splits a merged corpus into two parts. + + ``np_erode_msk`` / ``np_dilate_msk`` mutate and return their input. Dilating ``spart``/``tpart`` in + place therefore grew the very arrays the function returns as "the two separated components", so it + handed back two overlapping blobs -- and ``get_plane_split`` derived its separating plane from their + smeared centers of mass. + """ + + def test_dumbbell_splits_into_disjoint_parts(self): + from spineps.phase_instance import get_separating_components + + arr = np.zeros((24, 12, 12), dtype=np.uint8) + arr[2:10, 2:10, 2:10] = 1 # first body + arr[14:22, 2:10, 2:10] = 1 # second body + arr[10:14, 5:7, 5:7] = 1 # thin bridge that erosion breaks + spart, tpart, spart_dil, tpart_dil, stpart = get_separating_components(arr, connectivity=3) + self.assertGreater(spart.sum(), 0) + self.assertGreater(tpart.sum(), 0) + self.assertEqual((spart & tpart).sum(), 0, "the two parts must be disjoint") + self.assertIn(3, np.unique(stpart), "the dilated parts must end up touching") + # the dilations are strictly larger than the parts they came from + self.assertGreater(spart_dil.sum(), spart.sum()) + self.assertGreater(tpart_dil.sum(), tpart.sum()) + + def test_unsplittable_shape_fails_with_a_readable_error(self): + """A uniform bar has no waist to split at; the fallback branch used to die with a bare KeyError.""" + from spineps.phase_instance import get_separating_components + + arr = np.zeros((20, 16, 16), dtype=np.uint8) + arr[4:16, 6:10, 6:10] = 1 + with self.assertRaises(Exception) as ctx: + get_separating_components(arr, connectivity=3) + self.assertNotIsInstance(ctx.exception, KeyError, "the failure must name the problem, not blow up on a missing key") + self.assertTrue(str(ctx.exception), "the exception must carry a message") + + +class Test_Endplate_Labels_Reach_The_Semantic_Mask(unittest.TestCase): + """The split superior/inferior endplate labels must survive into the returned semantic mask. + + ``NII.extract_label`` binarises unless ``keep_label=True``, so the final extract in + ``add_ivd_ep_vert_label`` collapsed the whole split to 1 and the semantic mask came back with + endplates labelled ``1`` -- a label that means nothing in the subregion space. + """ + + def test_superior_and_inferior_plates_are_present(self): + from TPTBox import Location + + from unit_tests.test_regression_golden import run_post_phase_synthetic + + seg_cleaned, _vert = run_post_phase_synthetic() + labels = seg_cleaned.unique() + self.assertIn(Location.Vertebral_Body_Endplate_Superior.value, labels) + self.assertIn(Location.Vertebral_Body_Endplate_Inferior.value, labels) + self.assertNotIn(1, labels, "1 is not a subregion label; it was the binarised endplate mask") + + +class Test_Fix_Wrong_Posterior_Instance_Label(unittest.TestCase): + """A detached arcus fragment must be relabelled to the single instance it touches. + + The function now crops to each vertebra's bounding box before running connected components; the + per-component windows and the write-back have to stay aligned with the full volume. + """ + + def test_detached_arcus_is_reassigned(self): + from TPTBox import Location + + from spineps.utils.proc_functions import fix_wrong_posterior_instance_label + + shape = (40, 40, 20) + sem = np.zeros(shape, dtype=np.uint8) + inst = np.zeros(shape, dtype=np.uint8) + + # instance 1: a corpus high up (small I index == superior) + sem[6:16, 4:12, 6:14] = Location.Vertebra_Corpus_border.value + inst[6:16, 4:12, 6:14] = 1 + # instance 2: a corpus below it, with its own arcus + sem[6:16, 20:30, 6:14] = Location.Vertebra_Corpus_border.value + sem[16:22, 20:30, 8:12] = Location.Arcus_Vertebrae.value + inst[6:22, 20:30, 6:14] = 2 + # a stray arcus-only fragment carrying instance 1's label but sitting on instance 2 + sem[22:25, 22:26, 9:11] = Location.Arcus_Vertebrae.value + inst[22:25, 22:26, 9:11] = 1 + + sem_nii = _nii(sem) + inst_nii = _nii(inst) + out = fix_wrong_posterior_instance_label(sem_nii, inst_nii, logger=logger).get_seg_array() + self.assertTrue(np.all(out[22:25, 22:26, 9:11] == 2), "the stray arcus should follow the instance it touches") + self.assertTrue(np.all(out[6:16, 4:12, 6:14] == 1), "the real instance-1 corpus must be untouched") diff --git a/unit_tests/test_disc_labels.py b/unit_tests/test_disc_labels.py deleted file mode 100644 index c6371dd..0000000 --- a/unit_tests/test_disc_labels.py +++ /dev/null @@ -1,41 +0,0 @@ -# Call 'python -m unittest' on this folder -# coverage run -m unittest -# coverage report -# coverage html -from __future__ import annotations - -import contextlib -import io -import sys -import unittest -from pathlib import Path -from unittest import mock - -import numpy as np -from TPTBox import No_Logger -from typing_extensions import Self - -from spineps.utils.generate_disc_labels import Image, main - -logger = No_Logger() - - -class Test_DiscLabels(unittest.TestCase): - def test_main_without_args(self): - # --path-vert is required; omitting it must exit via argparse, not proceed. - with ( - self.assertRaises(SystemExit) as cm, - mock.patch.object(sys, "argv", ["generate_disc_labels"]), - contextlib.redirect_stderr(io.StringIO()), - ): - main() - self.assertEqual(cm.exception.code, 2) - - def test_image(self): - img = Image(param=np.array([0, 0, 0, 0])) - - img.dim # noqa: B018 - img.orientation # noqa: B018 - img.absolutepath # noqa: B018 - img.copy() - img.change_type(np.uint8) diff --git a/unit_tests/test_generate_disc_labels_extra.py b/unit_tests/test_generate_disc_labels_extra.py deleted file mode 100644 index acf066c..0000000 --- a/unit_tests/test_generate_disc_labels_extra.py +++ /dev/null @@ -1,175 +0,0 @@ -# Call 'python -m unittest' on this folder -# coverage run -m unittest -# coverage report -# coverage html -from __future__ import annotations - -import unittest -from pathlib import Path - -import numpy as np -import numpy.testing as npt - -from spineps.utils.generate_disc_labels import ( - DISCS_MAP, - closest_point_seg_to_line, - default_name_discs, - extract_centroids_3d, - project_point_on_line, -) - - -class Test_ProjectPointOnLine(unittest.TestCase): - def test_point_on_axis_aligned_line(self): - # Line along the x-axis; an off-line point projects to the nearest vertex. - line = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0], [3.0, 0.0, 0.0]]) - point = np.array([1.4, 5.0, 0.0]) - closest, dist = project_point_on_line(point, line) - # Nearest vertex is x=1 (1.4 rounds toward 1), distance^2 = 0.4^2 + 5.0^2 = 25.16 - npt.assert_allclose(closest, np.array([1.0, 0.0, 0.0])) - self.assertAlmostEqual(dist, 25.16) - - def test_point_exactly_on_vertex(self): - line = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]]) - point = np.array([2.0, 0.0, 0.0]) - closest, dist = project_point_on_line(point, line) - npt.assert_allclose(closest, np.array([2.0, 0.0, 0.0])) - self.assertAlmostEqual(dist, 0.0) - - def test_snaps_to_nearest_vertex_not_interpolated(self): - # The function returns the closest *vertex*, it does not interpolate along the segment. - line = np.array([[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 2.0, 0.0], [0.0, 3.0, 0.0]]) - point = np.array([0.0, 1.6, 0.0]) - closest, dist = project_point_on_line(point, line) - # 1.6 is closer to vertex 2 than vertex 1 -> [0, 2, 0], distance^2 = 0.4^2 = 0.16 - npt.assert_allclose(closest, np.array([0.0, 2.0, 0.0])) - self.assertAlmostEqual(dist, 0.16) - - def test_returns_squared_distance(self): - # Verify the returned distance is the squared euclidean distance (not the root). - line = np.array([[0.0, 0.0, 0.0]]) - point = np.array([3.0, 4.0, 0.0]) - closest, dist = project_point_on_line(point, line) - npt.assert_allclose(closest, np.array([0.0, 0.0, 0.0])) - self.assertAlmostEqual(dist, 25.0) # 3^2 + 4^2, not 5 - - -class Test_ExtractCentroids3d(unittest.TestCase): - def test_two_components_sorted_by_vertical_axis(self): - arr = np.zeros((10, 10, 10), dtype=int) - arr[1:3, 1:3, 1:3] = 5 # lower along axis 1 (S-I axis in RSP) - arr[1:3, 6:8, 1:3] = 7 # higher along axis 1 - centroids, _bounding_boxes = extract_centroids_3d(arr) - - self.assertEqual(len(centroids), 2) - # Centroid of a 2x2x2 block starting at (1,1,1) is (1,1,1) after int truncation; same for (1,6,1). - npt.assert_array_equal(centroids, np.array([[1, 1, 1], [1, 6, 1]])) - # Sorted ascending along axis 1. - self.assertTrue(np.all(np.diff(centroids[:, 1]) >= 0)) - # Integer dtype is guaranteed by the implementation. - self.assertTrue(np.issubdtype(centroids.dtype, np.integer)) - - def test_background_component_removed(self): - arr = np.zeros((6, 6, 6), dtype=int) - arr[2:4, 2:4, 2:4] = 3 - centroids, bounding_boxes = extract_centroids_3d(arr) - # Only one foreground component, background (label 0) must be dropped. - self.assertEqual(len(centroids), 1) - self.assertEqual(len(bounding_boxes), 1) - npt.assert_array_equal(centroids[0], np.array([2, 2, 2])) - - def test_sorting_independent_of_insertion_order(self): - arr = np.zeros((12, 12, 4), dtype=int) - arr[0:2, 8:10, 0:2] = 1 # high axis 1 - arr[0:2, 0:2, 0:2] = 2 # low axis 1 - arr[0:2, 4:6, 0:2] = 3 # mid axis 1 - centroids, _ = extract_centroids_3d(arr) - # Regardless of which label was written first, output is sorted by axis-1 coord. - npt.assert_array_equal(centroids[:, 1], np.array([0, 4, 8])) - - def test_bounding_boxes_match_components(self): - arr = np.zeros((8, 8, 8), dtype=int) - arr[1:3, 1:3, 1:3] = 4 - _, bounding_boxes = extract_centroids_3d(arr) - self.assertEqual(len(bounding_boxes), 1) - # cc3d returns slice tuples; the bounding box must contain exactly the block we placed. - bb = bounding_boxes[0] - sub = arr[bb[0], bb[1], bb[2]] - self.assertEqual(sub.shape, (2, 2, 2)) - self.assertTrue(np.all(sub == 4)) - - -class Test_ClosestPointSegToLine(unittest.TestCase): - def test_picks_closest_voxel_and_preserves_label(self): - arr = np.zeros((10, 10, 10), dtype=int) - arr[1:3, 1:3, 1:3] = 5 - arr[1:3, 6:8, 1:3] = 7 - _, bounding_boxes = extract_centroids_3d(arr) - - # Centerline far in the +z direction -> nearest voxel of each disc is the one with max z. - centerline = np.array([[1.0, 1.0, 100.0], [1.0, 6.0, 100.0]]) - result = closest_point_seg_to_line(arr, centerline, bounding_boxes) - - self.assertEqual(result.shape, (2, 4)) - # Each row is [x, y, z, disc_value]; z must be 2 (top of the [1,3) z-range), labels preserved. - npt.assert_array_equal(result, np.array([[1, 1, 2, 5], [1, 6, 2, 7]])) - - def test_label_value_in_last_column(self): - arr = np.zeros((6, 6, 6), dtype=int) - arr[2:4, 2:4, 2:4] = 9 - _, bounding_boxes = extract_centroids_3d(arr) - centerline = np.array([[0.0, 0.0, 0.0]]) - result = closest_point_seg_to_line(arr, centerline, bounding_boxes) - self.assertEqual(result.shape, (1, 4)) - # Closest voxel to the origin within the [2,4) block is (2,2,2), value 9. - npt.assert_array_equal(result[0], np.array([2, 2, 2, 9])) - - def test_single_voxel_disc(self): - arr = np.zeros((5, 5, 5), dtype=int) - arr[3, 1, 4] = 11 # a single labelled voxel - _, bounding_boxes = extract_centroids_3d(arr) - centerline = np.array([[0.0, 0.0, 0.0], [10.0, 10.0, 10.0]]) - result = closest_point_seg_to_line(arr, centerline, bounding_boxes) - npt.assert_array_equal(result, np.array([[3, 1, 4, 11]])) - - -class Test_DefaultNameDiscs(unittest.TestCase): - def test_default_suffix_with_compound_extension(self): - out = default_name_discs("/data/sub-amu_T2w_dseg.nii.gz") - self.assertEqual(out, Path("/data/sub-amu_T2w_dseg_label-discs_dlabel.nii.gz")) - - def test_custom_suffix(self): - out = default_name_discs(Path("/data/foo.nii.gz"), suffix="_disc") - self.assertEqual(out, Path("/data/foo_disc.nii.gz")) - - def test_single_extension(self): - out = default_name_discs("/data/foo.mha") - self.assertEqual(out, Path("/data/foo_label-discs_dlabel.mha")) - - def test_accepts_path_object_input(self): - out = default_name_discs(Path("/data/scan.nii")) - self.assertIsInstance(out, Path) - self.assertEqual(out.name, "scan_label-discs_dlabel.nii") - - -class Test_DiscsMap(unittest.TestCase): - def test_mapping_known_values(self): - # Spot-check the static vertebra->disc remapping table. - self.assertEqual(DISCS_MAP[2], 1) - self.assertEqual(DISCS_MAP[102], 3) - self.assertEqual(DISCS_MAP[124], 25) - - def test_mapping_is_consecutive_for_thoracolumbar_block(self): - # Keys 102..124 map to consecutive disc values 3..25. - block_keys = list(range(102, 125)) - values = [DISCS_MAP[k] for k in block_keys] - self.assertEqual(values, list(range(3, 26))) - - def test_disc_value_2_is_not_directly_mapped(self): - # Disc 2 is inserted between 1 and 3 by extract_discs_label, so it is absent from the map values. - self.assertNotIn(2, DISCS_MAP.values()) - self.assertEqual(len(DISCS_MAP), 24) - - -if __name__ == "__main__": - unittest.main() diff --git a/unit_tests/test_proc_functions.py b/unit_tests/test_proc_functions.py index e8741d4..1cdaaf1 100644 --- a/unit_tests/test_proc_functions.py +++ b/unit_tests/test_proc_functions.py @@ -8,6 +8,7 @@ import unittest from pathlib import Path +import numpy as np from TPTBox import Log_Type, No_Logger from TPTBox.tests.test_utils import get_test_mri @@ -75,3 +76,40 @@ def test_clean_artifacts_zeros(self): l3_cleaned = clean_cc_artifacts( l3, logger=logger, labels=[41, 42, 43, 44, 45, 46, 47, 48, 49], ignore_missing_labels=ignore_missing_labels ) + + +class Test_Clean_CC_Artifacts_Branches(unittest.TestCase): + """Pin both cleaning branches: a small component next to a big one is relabeled, an isolated one deleted. + + ``clean_cc_artifacts`` now works inside each component's padded bounding box instead of over the whole + volume, so the neighbourhood dilation and the majority vote need to keep giving the same answers. + """ + + @staticmethod + def _mask() -> np.ndarray: + arr = np.zeros((20, 20, 20), dtype=np.uint8) + arr[2:12, 2:12, 2:12] = 1 # big label-1 body + arr[12:14, 5:7, 5:7] = 2 # small label-2 speck glued to it -> majority vote says 1 + arr[17:19, 17:19, 17:19] = 2 # isolated label-2 speck -> deleted + return arr + + def test_relabel_and_delete(self): + arr = self._mask() + out = clean_cc_artifacts(arr, logger=logger, labels=[2], cc_size_threshold=100, only_delete=False, verbose=False) + self.assertTrue(np.all(out[12:14, 5:7, 5:7] == 1), "the attached speck should inherit its neighbour's label") + self.assertTrue(np.all(out[17:19, 17:19, 17:19] == 0), "the isolated speck should be deleted") + self.assertTrue(np.all(out[2:12, 2:12, 2:12] == 1), "the big component must be untouched") + + def test_only_delete_removes_both(self): + arr = self._mask() + out = clean_cc_artifacts(arr, logger=logger, labels=[2], cc_size_threshold=100, only_delete=True, verbose=False) + self.assertEqual(int((out == 2).sum()), 0) + self.assertTrue(np.all(out[2:12, 2:12, 2:12] == 1)) + + def test_component_touching_the_volume_edge(self): + """The bounding-box crop must clamp at the array bounds exactly like the full-volume code did.""" + arr = np.zeros((20, 20, 20), dtype=np.uint8) + arr[2:12, 2:12, 2:12] = 1 + arr[0:2, 0:2, 0:2] = 2 # in the corner, so the padded bbox is clipped + out = clean_cc_artifacts(arr, logger=logger, labels=[2], cc_size_threshold=100, only_delete=False, verbose=False) + self.assertEqual(int((out == 2).sum()), 0) diff --git a/unit_tests/test_regression_golden.py b/unit_tests/test_regression_golden.py new file mode 100644 index 0000000..d8a7af4 --- /dev/null +++ b/unit_tests/test_regression_golden.py @@ -0,0 +1,180 @@ +# Call 'python -m unittest' on this folder +"""Golden-output regression tests for the two whole-volume post-model phases. + +The instance merge (``phase_instance``) and the combined post-processing (``phase_post``) are the two +places where memory/speed optimisations are most likely to silently change a label. These tests pin +their exact output on a deterministic fixture, so any refactor that is supposed to be behaviour +preserving has to prove it. + +If a change here is *intended*, regenerate the digests with:: + + python -m unit_tests.test_regression_golden +""" + +from __future__ import annotations + +import hashlib +import unittest + +import nibabel as nib +import numpy as np +from TPTBox import NII, Location, No_Logger +from TPTBox.tests.test_utils import get_test_mri +from typing_extensions import Self + +from spineps.phase_instance import predict_instance_mask +from spineps.phase_post import phase_postprocess_combined +from spineps.seg_enums import ErrCode, OutputType +from spineps.seg_model import SegmentationModel +from spineps.utils.seg_modelconfig import Segmentation_Inference_Config + +logger = No_Logger() + + +def digest(arr: np.ndarray) -> str: + """Stable content digest of an array, including its shape and dtype kind.""" + a = np.ascontiguousarray(arr) + h = hashlib.sha256() + h.update(repr((a.shape, a.dtype.kind)).encode()) + h.update(a.astype(np.int64).tobytes()) + return h.hexdigest()[:32] + + +class ThirdsInstanceModel(SegmentationModel): + """Deterministic stand-in for the instance model. + + Splits each cutout into three bands along the inferior axis and labels them 1/2/3, which is the + shape of a real three-vertebra prediction -- enough to drive the couple search and the merge. + """ + + def __init__(self, cutout_size: tuple[int, int, int] = (24, 24, 24)) -> None: + self.logger = No_Logger() + config = Segmentation_Inference_Config( + logger=self.logger, + modality=["SEG"], + acquisition="sag", + log_name="ThirdsInstanceModel", + modeltype="unet", + model_expected_orientation=("P", "I", "R"), + available_folds=1, + inference_augmentation=False, + resolution_range=[1.5, 1.5, 1.5], + default_step_size=0.5, + labels={1: 1, 2: 2, 3: 3}, + expected_inputs=["seg"], + cutout_size=cutout_size, + ) + super().__init__(__file__, config, default_verbose=False, default_allow_tqdm=False) + + def load(self, folds: tuple[str, ...] | None = None) -> Self: # noqa: ARG002 + self.predictor = object() + return self + + def run(self, input_nii: list[NII], verbose: bool = False) -> dict[OutputType, NII | None]: # noqa: ARG002 + nii = input_nii[0] + arr = nii.get_seg_array() + out = np.zeros_like(arr) + height = arr.shape[1] + for band, lo in enumerate((0, height // 3, 2 * height // 3)): + hi = height if band == 2 else (band + 1) * height // 3 + band_slice = out[:, lo:hi, :] + band_slice[arr[:, lo:hi, :] != 0] = band + 1 + return {OutputType.seg: nii.set_array(out), OutputType.softmax_logits: None} + + +def run_instance_phase() -> NII: + """Run the instance phase on the shared fixture and return the vertebra mask.""" + _mri, subreg, _vert, _label = get_test_mri() + model = ThirdsInstanceModel().load() + vert_nii, errcode = predict_instance_mask(subreg, model, debug_data={}, verbose=False) + assert errcode == ErrCode.OK, errcode + assert vert_nii is not None + return vert_nii + + +def run_post_phase() -> tuple[NII, NII]: + """Run the combined post-processing on the shared fixture.""" + mri, subreg, vert, _label = get_test_mri() + return phase_postprocess_combined(mri, subreg, vert, model_labeling=None, debug_data={}) + + +def synthetic_spine() -> tuple[NII, NII, NII]: + """A deterministic multi-vertebra spine with discs and endplates, in (P, I, R). + + The shared TPTBox fixture only has three vertebrae, which barely exercises the endplate splitter -- + the loop that dominates post-processing. This one has six, each with its own corpus, arch, disc and + endplate band, so superior/inferior plate assignment actually has neighbours to disagree about. + """ + shape = (40, 120, 40) + affine = np.array([[0, 0, 1.0, 0], [-1.0, 0, 0, 0], [0, -1.0, 0, 0], [0, 0, 0, 1.0]]) + seg = np.zeros(shape, dtype=np.uint8) + vert = np.zeros(shape, dtype=np.uint8) + img = np.zeros(shape, dtype=np.float32) + + pitch = 18 # vertebra + disc period along the inferior axis + for n in range(6): + top = 6 + n * pitch + body = slice(top, top + 12) + seg[10:26, body, 12:28] = Location.Vertebra_Corpus_border.value + seg[26:32, body, 16:24] = Location.Arcus_Vertebrae.value # posterior elements + vert[10:32, body, 12:28] = n + 1 + # endplate band directly below the corpus, then the disc below that + seg[10:26, top + 12 : top + 14, 12:28] = Location.Endplate.value + seg[10:26, top + 14 : top + 18, 12:28] = Location.Vertebra_Disc.value + seg[30:34, 4:112, 18:22] = Location.Spinal_Canal.value + img[seg != 0] = 800.0 + + def wrap(arr, is_seg): + return NII(nib.Nifti1Image(arr, affine=affine), seg=is_seg) + + return wrap(img, False), wrap(seg, True), wrap(vert, True) + + +def run_post_phase_synthetic() -> tuple[NII, NII]: + """Run the combined post-processing on the six-vertebra synthetic spine.""" + mri, subreg, vert = synthetic_spine() + return phase_postprocess_combined(mri, subreg, vert, model_labeling=None, debug_data={}) + + +# Regenerate with `python -m unit_tests.test_regression_golden` after an intentional change. +GOLDEN_INSTANCE = "d03a9f8081523184108c6f2698737ed5" +GOLDEN_POST_SEG = "98d67f4a9a4ad6c565e0d647b0466441" +GOLDEN_POST_VERT = "8928f04c5539c1eff12672743df80f61" +GOLDEN_SYNTH_SEG = "ab6a357788a70bbe4c53b9dc2a253e3e" +GOLDEN_SYNTH_VERT = "f0850727dd21c001bd661814103ca770" + + +class Test_Golden_Instance_Phase(unittest.TestCase): + def test_instance_mask_unchanged(self): + vert_nii = run_instance_phase() + self.assertEqual(digest(vert_nii.get_seg_array()), GOLDEN_INSTANCE) + + +class Test_Golden_Post_Phase(unittest.TestCase): + def test_postprocess_output_unchanged(self): + seg_cleaned, vert_cleaned = run_post_phase() + self.assertEqual(digest(seg_cleaned.get_seg_array()), GOLDEN_POST_SEG) + self.assertEqual(digest(vert_cleaned.get_seg_array()), GOLDEN_POST_VERT) + + def test_postprocess_synthetic_spine_unchanged(self): + seg_cleaned, vert_cleaned = run_post_phase_synthetic() + self.assertEqual(digest(seg_cleaned.get_seg_array()), GOLDEN_SYNTH_SEG) + self.assertEqual(digest(vert_cleaned.get_seg_array()), GOLDEN_SYNTH_VERT) + + def test_synthetic_spine_splits_endplates(self): + """Sanity check on the fixture itself: the splitter must produce both plate labels.""" + seg_cleaned, vert_cleaned = run_post_phase_synthetic() + labels = seg_cleaned.unique() + self.assertIn(Location.Vertebral_Body_Endplate_Inferior.value, labels) + self.assertIn(Location.Vertebral_Body_Endplate_Superior.value, labels) + self.assertGreaterEqual(len([v for v in vert_cleaned.unique() if v < 40]), 5) + + +if __name__ == "__main__": + seg_c, vert_c = run_post_phase() + seg_s, vert_s = run_post_phase_synthetic() + print(f'GOLDEN_INSTANCE = "{digest(run_instance_phase().get_seg_array())}"') + print(f'GOLDEN_POST_SEG = "{digest(seg_c.get_seg_array())}"') + print(f'GOLDEN_POST_VERT = "{digest(vert_c.get_seg_array())}"') + print(f'GOLDEN_SYNTH_SEG = "{digest(seg_s.get_seg_array())}"') + print(f'GOLDEN_SYNTH_VERT = "{digest(vert_s.get_seg_array())}"')