nnUNet API utilities¶

class LabelManagerAdapter:
    """
    Adapter around nnU-Net LabelManager private state mutations.

    Centralizes private attribute access (`_all_labels`, `_regions`) so calling
    code does not directly mutate them.
    """

    def __init__(self, label_manager: LabelManager):
        """Initializes the adapter with a given LabelManager.

        Args:
            label_manager (LabelManager): The nnU-Net LabelManager to adapt.
        """
        self.label_manager = label_manager

    def build_subset_state(
        self, class_idx: int | list[int]
    ) -> LabelSubsetState:
        """Build the temporary label mapping state for a class subset.

        Args:
            class_idx (int | list[int]): The index or indices of the
                classes to include.

        Returns:
            LabelSubsetState: The constructed subset state for inference.
        """
        if isinstance(class_idx, int):
            class_idx = [class_idx]
        used_labels = [0, *class_idx]
        correspondence_dict = {i: label for i, label in enumerate(used_labels)}

        if self.label_manager._has_regions:
            new_regions = [
                region
                for region in self.label_manager._regions
                if any(region_label in used_labels for region_label in region)
            ]
            if new_regions:
                used_labels = sorted(
                    np.unique(np.concatenate(new_regions)).tolist()
                )

        return LabelSubsetState(
            prediction_indices=[0, *class_idx],
            used_labels=used_labels,
            correspondence_dict=correspondence_dict,
        )

    @contextmanager
    def apply_subset_state(self, subset_state: LabelSubsetState):
        """
        Temporarily apply subset labels and always restore original state.

        Args:
            subset_state (LabelSubsetState): The state to apply during the context.

        Yields:
            None: Allows execution within the modified label manager state.
        """
        old_labels = self.label_manager._all_labels
        old_regions = self.label_manager._regions
        self.label_manager._all_labels = subset_state.used_labels
        try:
            yield
        finally:
            self.label_manager._all_labels = old_labels
            self.label_manager._regions = old_regions

@dataclass
class LabelSubsetState:
    """Holds temporary label-manager state used for class-subset inference.

    Attributes:
        prediction_indices (list[int]): Indices of the classes to be kept in the
            prediction.
        used_labels (list[int]): List of unique labels that are active in the
            current subset.
        correspondence_dict (dict[int, int]): Mapping from subset label indices
            to original labels.
    """

    prediction_indices: list[int]
    used_labels: list[int]
    correspondence_dict: dict[int, int]

class SeriesLoader:
    """
    Load and cache medical image series as SimpleITK volumes.

    This helper provides:
    - Caching: each unique series path is read at most once.
    - Optional DICOM loading via `read_dicom_as_sitk`.
    - Optional on-access post-processing controlled by a suffix in the requested path.

    The path string passed to `__getitem__` (and used inside `series_paths`) can include
    simple modifiers:
    - `"/path/to/seg.nii.gz=3"`: returns a binary mask `(volume == 3)` cast to `sitkInt32`.
    - `"/path/to/4d.nii.gz[0]"`: returns the slice/volume at the given index.

    Args:
        series_paths (list[list[str]]): Nested list of series identifiers grouped by
            "stage". Each inner list is the set of series to be used at that stage.
        is_dicom (bool, optional): If `True`, each series path is treated as a DICOM
            directory and read with `read_dicom_as_sitk`. If `False`, each series
            path is read using `sitk.ReadImage`.
        bvalue_for_filtering (int | None, optional): If provided when `is_dicom=True`,
            filters DICOM series by the specified b-value using `filter_by_bvalue`.
            Defaults to None.
    """

    def __init__(
        self,
        series_paths: list[list[str]],
        is_dicom: bool = False,
        bvalue_for_filtering: int | None = None,
    ):
        """
        Create a loader and pre-compute the unique series paths.

        Note:
            Only the base path (without modifiers) is cached. For example,
            requesting `"image.nii.gz=1"` and `"image.nii.gz=2"` will load
            `"image.nii.gz"` once and apply post-processing per request.
        """
        self.series_paths = series_paths
        self.is_dicom = is_dicom
        self.bvalue_for_filtering = bvalue_for_filtering

        self.n_stages = len(self.series_paths)
        self.unique_series_paths = []

        for series_path in series_paths:
            for s in series_path:
                if s not in self.unique_series_paths:
                    self.unique_series_paths.append(s)

        self.loaded_volumes = {}
        self.good_file_paths = {}
        self.loaded_resampled_volumes = {}
        logger.info(f"Identified {len(self.unique_series_paths)} unique series")

    def __getitem__(self, path: str) -> tuple[sitk.Image, list[str] | None]:
        """
        Load a series (if needed) and return the (optionally processed) volume.

        Args:
            path (str): Series path, optionally suffixed with `=<label>` or `[<index>]`.

        Returns:
            A tuple `(volume, good_file_paths)` where:
                - `volume` is the loaded SimpleITK image, potentially post-processed.
                - `good_file_paths` is only populated for DICOM inputs and corresponds to
                the list of files considered valid by `read_dicom_as_sitk`.
        """

        path, equal, index = self.get_info(path)
        if path not in self.loaded_volumes:
            if self.is_dicom:
                if os.path.isdir(path):
                    dcm_img = read_dicom_as_sitk(
                        path, bvalue_for_filtering=self.bvalue_for_filtering
                    )
                    self.loaded_volumes[path] = dcm_img[0]
                    self.good_file_paths[path] = dcm_img[1]
                else:
                    self.loaded_volumes[path] = sitk.ReadImage(path)
                    self.good_file_paths[path] = None
            else:
                self.loaded_volumes[path] = sitk.ReadImage(path)
                self.good_file_paths[path] = None
        volume, good_file_paths = (
            self.loaded_volumes[path],
            self.good_file_paths[path],
        )
        volume = self.post_process(volume, equal=equal, index=index)
        return volume, good_file_paths

    def __setitem__(self, key: str, value: sitk.Image):
        """Manually set/cache a volume for a given path.

        Args:
            key (str): The series identifier or path.
            value (sitk.Image): The SimpleITK image to cache.
        """
        self.loaded_volumes[key] = value
        self.good_file_paths[key] = None

    def register(
        self,
        image: sitk.Image,
        stage: int,
        image_file_name: str = "prediction.nii.gz",
    ):
        """Registers an image for a specific stage.

        Args:
            image (sitk.Image): The image to register.
            stage (int): The stage index to associate the image with.
            image_file_name (str, optional): image file name presumed to be inside
                the stage directory. Defaults to "prediction.nii.gz".
        """
        key = os.path.join(f"stage_{stage}", image_file_name)
        self[key] = image

    def get_info(self, path: str) -> tuple[str, int | None, int | None]:
        """
        Parse a series path and extract any post-processing modifier.

        Supported syntaxes:
        - `"<path>=<int>"`: equality comparison against the given integer label.
        - `"<path>[<int>]"`: SimpleITK index selection.

        Args:
            path (str): Raw path string.

        Returns:
            A tuple `(base_path, equal, index)` where:
            - `base_path` is the path without the modifier.
            - `equal` is the parsed integer after `=` (or `None`).
            - `index` is the parsed integer inside `[...]` (or `None`).
        """
        equal, index = None, None
        if "=" in path:
            path, equal = path.split("=")
            equal = int(equal)
        elif "[" in path:
            path, index = path.split("[")
            index = int(index.split("]")[0])
        return path, equal, index

    def post_process(
        self,
        volume: sitk.Image,
        equal: int | None = None,
        index: int | None = None,
    ) -> sitk.Image:
        """
        Apply an optional post-processing operation to a loaded volume.

        Args:
            volume (sitk.Image): Input SimpleITK image.
            equal (int | None): If provided, returns `(volume == equal)` cast
                to `sitkInt32`.
            index (int | None): If provided (and `equal` is `None`), returns
                `volume[index]`.

        Returns:
            The processed (or original) image.
        """
        if equal is not None:
            return volume == equal
        elif index is not None:
            select_f = sitk.VectorIndexSelectionCastImageFilter()
            select_f.SetIndex(index)
            return select_f.Execute(volume)
        return sitk.Cast(volume, sitk.sitkFloat32)

    def get_volumes(self, stage: int) -> list[sitk.Image]:
        """
        Return the loaded (and post-processed) volumes for a given stage.

        Args:
            stage (int): Index into `series_paths`.

        Returns:
            List of `sitk.Image` volumes for the stage.
        """
        return [self.get_resampled(s, stage) for s in self.series_paths[stage]]

    def get_file_paths(self, stage: int) -> list[list[str]]:
        """Return the DICOM file-path lists for a given stage.

        Args:
            stage (int): Index into `series_paths`.

        Returns:
            List of file-path lists returned by `read_dicom_as_sitk` for the stage.
            For non-DICOM inputs, entries will be `None`.
        """
        return [self[s][1] for s in self.series_paths[stage]]

    def might_be_mask(self, image: sitk.Image) -> bool:
        """
        Makes an educated guess on whether ``image`` is a mask.

        Args:
            image (sitk.Image): image.

        Returns:
            Boolean which is True if image is an integer _and_ has fewer than
                50 unique values.
        """
        is_mask = False
        if "int" in image.GetPixelIDTypeAsString().lower():
            stats_filter = sitk.LabelShapeStatisticsImageFilter()
            stats_filter.Execute(image)
            if stats_filter.GetNumberOfLabels() < 50:
                is_mask = True
        return is_mask

    def get_resampled(self, path: str, stage: int) -> sitk.Image:
        """
        Similar to __getitem__ but resamples to first series in stage.

        Args:
            path (str): Series path, optionally suffixed with `=<label>` or `[<index>]`.
            stage (int): Index into `series_paths`.

        Returns:
            sitk.Image: The loaded image, resampled to match the first series in the stage.
        """
        first_path = self.series_paths[stage][0]
        identifier = (path, first_path)
        if identifier not in self.loaded_resampled_volumes:
            if path == first_path:
                resampled_image = self[path][0]
            else:
                resampled_image = resample_image_to_target(
                    self[path][0],
                    self[first_path][0],
                    is_mask=self.might_be_mask(self[path][0]),
                )

            self.loaded_resampled_volumes[identifier] = resampled_image

        return self.loaded_resampled_volumes[identifier]