nnInteractive 2.0.0__py3-none-any.whl

nnInteractive/utils/crop.py

@@ -0,0 +1,268 @@
+from typing import Sequence, Optional
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+
+def crop_and_pad_into_buffer(
+    target_tensor: torch.Tensor, bbox: Sequence[Sequence[int]], source_tensor, source_leading_slice=None
+) -> None:
+    """
+    Copies a sub-region from source_tensor into target_tensor based on a bounding box.
+
+    Args:
+        target_tensor (torch.Tensor): A preallocated tensor that will be updated.
+        bbox (sequence of [int, int]): A bounding box for each dimension of the source tensor
+            that is covered by the bbox. The bbox is defined as [start, end) (half-open interval)
+            and may extend outside the source tensor. If source_tensor has more dimensions than
+            len(bbox), the leading dimensions will be fully included.
+        source_tensor: The tensor (or blosc2 NDArray) to copy data from.
+        source_leading_slice: Optional slice to apply to the first leading dimension of the source
+            instead of slice(None). Useful for reading a subset of channels from a blosc2 NDArray
+            without decompressing channel 0.
+
+    Behavior:
+        For each dimension that the bbox covers (i.e. the last len(bbox) dims of source_tensor):
+            - Compute the overlapping region between the bbox and the source tensor.
+            - Determine the corresponding indices in the target tensor where the data will be copied.
+        For any extra leading dimensions (i.e. source_tensor.ndim > len(bbox)):
+            - Use slice(None) to include the entire dimension (or source_leading_slice for the
+              first leading dim when provided).
+        If source_tensor and target_tensor are on different devices, only the overlapping subregion
+        is transferred to the device of target_tensor.
+    """
+    total_dims = source_tensor.ndim
+    bbox_dims = len(bbox)
+    # Compute the number of leading dims that are not covered by bbox.
+    leading_dims = total_dims - bbox_dims
+
+    source_slices = []
+    target_slices = []
+
+    # For the leading dimensions, include the entire dimension.
+    for d in range(leading_dims):
+        if d == 0 and source_leading_slice is not None:
+            source_slices.append(source_leading_slice)
+        else:
+            source_slices.append(slice(None))
+        target_slices.append(slice(None))
+
+    # Process the dimensions covered by the bbox.
+    for d in range(bbox_dims):
+        box_start, box_end = bbox[d]
+        d_source = leading_dims + d
+        source_size = source_tensor.shape[d_source]
+
+        # Compute the overlapping region in source coordinates.
+        copy_start_source = max(box_start, 0)
+        copy_end_source = min(box_end, source_size)
+        copy_size = copy_end_source - copy_start_source
+
+        # Compute the corresponding indices in the target tensor.
+        copy_start_target = max(0, -box_start)
+        copy_end_target = copy_start_target + copy_size
+
+        source_slices.append(slice(copy_start_source, copy_end_source))
+        target_slices.append(slice(copy_start_target, copy_end_target))
+
+    # Extract the overlapping region from the source.
+    sub_source = source_tensor[tuple(source_slices)]
+    # Convert non-torch (e.g. blosc2 NDArray or numpy) to torch tensor.
+    if not isinstance(sub_source, torch.Tensor):
+        sub_source = torch.from_numpy(np.asarray(sub_source))
+    # Transfer only this subregion to the target tensor's device.
+    sub_source = sub_source.to(target_tensor.device) if isinstance(target_tensor, torch.Tensor) else sub_source.cpu()
+    # Write the data into the preallocated target_tensor.
+    target_tensor[tuple(target_slices)] = sub_source
+
+
+def paste_tensor(target, source, bbox, channel_idx=None):
+    """
+    Paste a source tensor into a target tensor using a given bounding box.
+
+    Both tensors are assumed to be 3D (or 4D when channel_idx is provided for the target).
+    The bounding box is specified in the coordinate system of the target as:
+      [[x1, x2], [y1, y2], [z1, z2]]
+    and its size is assumed to be equal to the shape of the source tensor.
+    The bbox may exceed the boundaries of the target tensor.
+
+    Args:
+        target: The target tensor (torch.Tensor of shape (T0, T1, T2)) or blosc2 NDArray of
+                shape (C, T0, T1, T2) when channel_idx is provided.
+        source: The source tensor of shape (S0, S1, S2). It must be the same size as the bbox.
+        bbox (list or tuple): List of intervals for each dimension: [[x1, x2], [y1, y2], [z1, z2]].
+        channel_idx (int, optional): If provided, paste into this channel of a 4D target.
+            For torch.Tensor targets, delegates to target[channel_idx]. For blosc2 NDArray
+            targets, performs a numpy read-modify-write on the specified channel.
+
+    Returns:
+        The target after pasting in the source (for torch.Tensor targets).
+    """
+    # When channel_idx is given and target is a torch Tensor, delegate to the channel view.
+    if channel_idx is not None and isinstance(target, torch.Tensor):
+        return paste_tensor(target[channel_idx], source, bbox)
+
+    if channel_idx is not None:
+        # target is a 4D blosc2 NDArray; write to target[channel_idx] at bbox.
+        target_shape = target.shape[1:]  # spatial dims
+
+        target_indices = []
+        source_indices = []
+
+        for i, (b0, b1) in enumerate(bbox):
+            t_start = max(b0, 0)
+            t_end = min(b1, target_shape[i])
+            if t_start >= t_end:
+                return  # no overlap in this dimension
+            s_start = t_start - b0
+            s_end = s_start + (t_end - t_start)
+            target_indices.append((t_start, t_end))
+            source_indices.append((s_start, s_end))
+
+        src = source[
+            source_indices[0][0] : source_indices[0][1],
+            source_indices[1][0] : source_indices[1][1],
+            source_indices[2][0] : source_indices[2][1],
+        ]
+        if isinstance(src, torch.Tensor):
+            src = src.cpu().numpy().astype(np.float16)
+        else:
+            src = np.asarray(src).astype(np.float16)
+
+        target[
+            (
+                channel_idx,
+                slice(target_indices[0][0], target_indices[0][1]),
+                slice(target_indices[1][0], target_indices[1][1]),
+                slice(target_indices[2][0], target_indices[2][1]),
+            )
+        ] = src
+        return
+
+    target_shape = target.shape  # (T0, T1, T2)
+
+    # For each dimension compute:
+    #   - The valid region in the target: [t_start, t_end)
+    #   - The corresponding region in the source: [s_start, s_end)
+    target_indices = []
+    source_indices = []
+
+    for i, (b0, b1) in enumerate(bbox):
+        # Determine valid region in target tensor:
+        t_start = max(b0, 0)
+        t_end = min(b1, target_shape[i])
+        # If there's no overlap in any dimension, nothing gets pasted.
+        if t_start >= t_end:
+            return target
+
+        # Determine corresponding indices in the source tensor.
+        # The source's coordinate 0 corresponds to b0 in the target.
+        s_start = t_start - b0
+        s_end = s_start + (t_end - t_start)
+
+        target_indices.append((t_start, t_end))
+        source_indices.append((s_start, s_end))
+
+    # Paste the corresponding region from source into target.
+    if isinstance(target, torch.Tensor):
+        target[
+            target_indices[0][0] : target_indices[0][1],
+            target_indices[1][0] : target_indices[1][1],
+            target_indices[2][0] : target_indices[2][1],
+        ] = source[
+            source_indices[0][0] : source_indices[0][1],
+            source_indices[1][0] : source_indices[1][1],
+            source_indices[2][0] : source_indices[2][1],
+        ].to(
+            target.device
+        )
+    else:
+        target[
+            target_indices[0][0] : target_indices[0][1],
+            target_indices[1][0] : target_indices[1][1],
+            target_indices[2][0] : target_indices[2][1],
+        ] = source[
+            source_indices[0][0] : source_indices[0][1],
+            source_indices[1][0] : source_indices[1][1],
+            source_indices[2][0] : source_indices[2][1],
+        ].cpu()
+
+    return target
+
+
+def crop_to_valid(img, bbox):
+    """
+    Crops the image to the part of the bounding box that lies within the image.
+    Supports a 4D tensor of shape (C, X, Y, Z). The bounding box is specified as
+    [[x1, x2], [y1, y2], [z1, z2]] with half-open intervals.
+
+    Args:
+        img: Input tensor (or blosc2 NDArray) of shape (C, X, Y, Z).
+        bbox (list or tuple): Bounding box as a list of three intervals for spatial dims:
+                              [[x1, x2], [y1, y2], [z1, z2]].
+
+    Returns:
+        cropped: Cropped data of shape (C, cropped_x, cropped_y, cropped_z).
+        pad (list of tuples): A list [(pad_x_left, pad_x_right),
+                                     (pad_y_left, pad_y_right),
+                                     (pad_z_left, pad_z_right)]
+                              indicating how much padding needs to be applied on each side.
+    """
+    # Only spatial dimensions (X, Y, Z) are cropped; channels are preserved.
+    spatial_dims = img.shape[1:]  # (X, Y, Z)
+    crop_indices = []
+    pad = []  # for each spatial dimension
+
+    for i, (start, end) in enumerate(bbox):
+        dim_size = spatial_dims[i]
+        # Clamp the indices to the valid range for cropping.
+        crop_start = max(start, 0)
+        crop_end = min(end, dim_size)
+        crop_indices.append((crop_start, crop_end))
+        # Calculate padding if the bbox goes out-of-bound.
+        pad_left = -start if start < 0 else 0
+        pad_right = end - dim_size if end > dim_size else 0
+        pad.append((pad_left, pad_right))
+
+    # Crop the image on spatial dimensions, leaving the channel dimension intact.
+    cropped = img[
+        :,
+        crop_indices[0][0] : crop_indices[0][1],
+        crop_indices[1][0] : crop_indices[1][1],
+        crop_indices[2][0] : crop_indices[2][1],
+    ]
+    return cropped, pad
+
+
+def pad_cropped(cropped: torch.Tensor, pad):
+    """
+    Pads the cropped image using the given pad amounts.
+    Supports a 4D tensor of shape (C, X, Y, Z) and applies padding only on the spatial dimensions.
+    For 3D (volumetric) padding, F.pad expects a 5D tensor with shape (N, C, X, Y, Z).
+    Hence, we temporarily add a dummy batch dimension.
+
+    Args:
+        cropped (torch.Tensor): Cropped tensor of shape (C, X, Y, Z).
+        pad (list of tuples): List of padding for each spatial dimension, in order (x, y, z):
+                              [(pad_x_left, pad_x_right),
+                               (pad_y_left, pad_y_right),
+                               (pad_z_left, pad_z_right)].
+
+    Returns:
+        padded (torch.Tensor): Padded tensor of shape (C, desired_x, desired_y, desired_z),
+                               where the spatial dimensions match the bbox size.
+    """
+    # F.pad for 3D data expects a 5D input (N, C, X, Y, Z) and a pad tuple of length 6:
+    # (pad_z_left, pad_z_right, pad_y_left, pad_y_right, pad_x_left, pad_x_right)
+    need_unsqueeze = cropped.dim() == 4
+    if need_unsqueeze:
+        cropped = cropped.unsqueeze(0)  # Now shape is (1, C, X, Y, Z)
+
+    # Reverse the pad list (currently in order x, y, z) to match F.pad's expected order: z, y, x.
+    pad_rev = pad[::-1]
+    pad_flat = [p for pair in pad_rev for p in pair]
+    padded = F.pad(cropped, pad_flat)
+
+    if need_unsqueeze:
+        padded = padded.squeeze(0)
+    return padded

nnInteractive/utils/erosion_dilation.py

@@ -0,0 +1,48 @@
+import torch
+import torch.nn.functional as F
+
+from nnunetv2.utilities.helpers import empty_cache
+from torch.backends import cudnn
+
+
+@torch.inference_mode()
+def iterative_3x3_same_padding_pool3d(x, kernel_size: int, use_min_pool: bool = False):
+    """
+    Applies 3D max pooling with manual asymmetric padding such that
+    the output shape is the same as the input shape.
+
+    Args:
+        x (Tensor): Input tensor of shape (N, C, D, H, W)
+        kernel_size (int or tuple): Kernel size for the pooling.
+            If int, the same kernel size is used for all three dimensions.
+
+    Returns:
+        Tensor: Output tensor with the same (D, H, W) dimensions as the input.
+    """
+    benchmark = cudnn.benchmark
+    cudnn.benchmark = False
+
+    assert kernel_size % 2 == 1, "Only works with odd kernels"
+
+    # Compute asymmetric padding for each dimension:
+    pad_front = (kernel_size - 1) // 2
+    pad_back = (kernel_size - 1) - pad_front
+
+    # For 3D (input shape: [N, C, D, H, W]), F.pad expects the padding in the order:
+    # (pad_left, pad_right, pad_top, pad_bottom, pad_front, pad_back)
+    x = F.pad(x, (pad_front, pad_back, pad_front, pad_back, pad_front, pad_back), mode="replicate")
+
+    iters = (kernel_size - 1) // 2
+    # Apply max pooling with no additional padding.
+    if not use_min_pool:
+        for _ in range(iters):
+            x = F.max_pool3d(x, kernel_size=3, stride=1, padding=0)
+        empty_cache(x.device)
+        cudnn.benchmark = benchmark
+        return x
+    else:
+        for _ in range(iters):
+            x = -F.max_pool3d(-x, kernel_size=3, stride=1, padding=0)
+        empty_cache(x.device)
+        cudnn.benchmark = benchmark
+        return x

nnInteractive/utils/inference_helpers.py

@@ -0,0 +1,45 @@
+import re
+from typing import List, Tuple, Union
+
+
+def version_to_tuple(version: str) -> Tuple[int, ...]:
+    return tuple(int(i) for i in re.findall(r"\d+", version))
+
+
+def parse_channel_pair(channel_name: str, raw_channels) -> Tuple[int, int]:
+    if not isinstance(raw_channels, (tuple, list)) or len(raw_channels) != 2:
+        raise ValueError(
+            f"Invalid channel mapping for '{channel_name}': expected a pair [pos, neg], got {raw_channels}."
+        )
+    return int(raw_channels[0]), int(raw_channels[1])
+
+
+def infer_num_interaction_channels_from_mapping(channel_mapping: dict) -> int:
+    max_positive_index = -1
+    max_negative_magnitude = 0
+
+    for k, v in channel_mapping.items():
+        if k == "prev_seg":
+            indices = [int(v)]
+        else:
+            pos_ch, neg_ch = parse_channel_pair(k, v)
+            indices = [pos_ch, neg_ch]
+
+        for idx in indices:
+            if idx >= 0:
+                max_positive_index = max(max_positive_index, idx)
+            else:
+                max_negative_magnitude = max(max_negative_magnitude, abs(idx))
+
+    # Positive indexing is 0-based, while negative indexing is 1-based-from-end.
+    return max(max_positive_index + 1, max_negative_magnitude, 1)
+
+
+def transform_coordinates_noresampling(
+    coords_orig: Union[List[int], Tuple[int, ...]],
+    nnunet_preprocessing_crop_bbox: List[Tuple[int, int]],
+) -> Tuple[int, ...]:
+    """
+    Converts coordinates in the original uncropped image to the internal cropped representation.
+    """
+    return tuple(coords_orig[d] - nnunet_preprocessing_crop_bbox[d][0] for d in range(len(coords_orig)))

nnInteractive/utils/os_shennanigans.py

@@ -0,0 +1,16 @@
+import platform
+import sys
+
+
+def is_linux_kernel_6_11():
+    if sys.platform != "linux":
+        return False
+
+    kernel_version = platform.release()  # e.g., '6.11.0-24-generic'
+    version_parts = kernel_version.split(".")
+    try:
+        major = int(version_parts[0])
+        minor = int(version_parts[1])
+        return major == 6 and minor == 11
+    except (IndexError, ValueError):
+        return False

nnInteractive/utils/rounding.py

@@ -0,0 +1,13 @@
+from math import ceil, floor
+
+
+def round_to_nearest_odd(number: float):
+    assert number > 0
+    cl = ceil(number)
+    fl = floor(number)
+    if cl % 2 == 1:
+        return cl
+    elif fl % 2 == 1:
+        return fl
+    else:
+        return round(number) + 1