Nnunet output

nnunet_output

MaskSavingStrategy

Bases: str, enum.Enum

Enum for mask saving strategies.

Attributes:

Name	Type	Description
LABEL_IMAGE	str	No overlaps allowed.
SPARSE_MASK	str	Allows overlaps, but is lossy if overlaps exist.
REGION_MASK	str	Work around that creates a new region for each overlap.

MaskSavingStrategyError

Bases: imgtools.io.nnunet_output.nnUNetOutputError

Raised when an invalid mask saving strategy is provided.

MissingROIsError

MissingROIsError(
    sample_number: str,
    expected_rois: list[str],
    found_rois: list[list[str]],
)

Bases: imgtools.io.nnunet_output.nnUNetOutputError

Raised when a VectorMask does not contain all required ROI keys.

Source code in src/imgtools/io/nnunet_output.py

def __init__(
    self,
    sample_number: str,
    expected_rois: list[str],
    found_rois: list[list[str]],
) -> None:
    msg = (
        f"Not all required ROI names found in sample {sample_number}. "
        f"Expected: {expected_rois}. Found: {found_rois}"
    )
    super().__init__(msg)
    self.sample_number = sample_number
    self.expected_rois = expected_rois
    self.found_rois = found_rois

NoSegmentationImagesError

NoSegmentationImagesError(sample_number: str)

Bases: imgtools.io.nnunet_output.nnUNetOutputError

Raised when no segmentation images are found in a sample.

Source code in src/imgtools/io/nnunet_output.py

def __init__(self, sample_number: str) -> None:
    msg = f"No segmentation images found in sample {sample_number}"
    super().__init__(msg)
    self.sample_number = sample_number

nnUNetOutput

Bases: pydantic.BaseModel

Configuration model for saving medical imaging outputs in nnUNet format.

This class provides a standardized configuration for saving medical images, supporting various file formats and output organization strategies.

Attributes:

Name	Type	Description
directory	pathlib.Path	Directory where output files will be saved. Must exist and be writable.
filename_format	str	Format string for output filenames with placeholders for metadata values.
existing_file_mode	imgtools.io.writers.ExistingFileMode	How to handle existing files (FAIL, SKIP, OVERWRITE).
extra_context	typing.Dict[str, typing.Any]	Additional metadata to include when saving files.

Examples:

>>> from imgtools.io import nnUNetOutput
>>> from imgtools.io.writers import ExistingFileMode
>>> output = nnUNetOutput(
...     directory="results/patient_scans",
...     existing_file_mode=ExistingFileMode.SKIP,
... )
>>> output(scan_list)  # Save all scans in the list

Methods:

Name	Description
default	Create a default instance of SampleOutput.
finalize_dataset	Finalize dataset by generating preprocessing scripts and dataset JSON configuration.
model_post_init	Initialize the writer after model initialization.
validate_directory	Validate that the output directory exists or can be created, and is writable.

writer property

writer: imgtools.io.writers.AbstractBaseWriter

Get the writer instance.

default classmethod

default() -> imgtools.io.nnunet_output.nnUNetOutput

Create a default instance of SampleOutput.

Source code in src/imgtools/io/nnunet_output.py

@classmethod
def default(cls) -> nnUNetOutput:
    """Create a default instance of SampleOutput."""
    return cls(
        directory=Path("output"),
        dataset_name="Dataset",
        roi_keys=["ROI_1", "ROI_2"],
        mask_saving_strategy=MaskSavingStrategy.LABEL_IMAGE,
        existing_file_mode=ExistingFileMode.FAIL,
        extra_context={},
    )

finalize_dataset

finalize_dataset() -> None

Finalize dataset by generating preprocessing scripts and dataset JSON configuration.

Source code in src/imgtools/io/nnunet_output.py

def finalize_dataset(self) -> None:
    """Finalize dataset by generating preprocessing scripts and dataset JSON configuration."""

    generate_nnunet_scripts(self.directory, self.dataset_id)

    index_df = pd.read_csv(self.writer.index_file)
    _image_modalities = index_df["Modality"].unique()

    # Construct channel names mapping
    channel_names = {
        channel_num.lstrip("0") or "0": modality
        for modality, channel_num in MODALITY_MAP.items()
        if modality in _image_modalities
    }

    # Count the number of training cases
    num_training_cases = sum(
        1
        for file in (self.directory / "imagesTr").iterdir()
        if file.is_file()
    )

    # Construct labels
    labels: dict[str, int | list[int]] = {"background": 0}
    if self.mask_saving_strategy is MaskSavingStrategy.REGION_MASK:
        n_components = len(self.roi_keys)
        max_val = 2**n_components

        for component_index in range(n_components):
            indices = [
                value
                for value in range(1, max_val)
                if (value >> component_index) & 1
            ]
            labels[self.roi_keys[component_index]] = indices
        regions_class_order = tuple(idx + 1 for idx in range(n_components))
    else:
        labels = {
            **{label: i + 1 for i, label in enumerate(self.roi_keys)},
        }
        regions_class_order = None

    generate_dataset_json(
        self.directory,
        channel_names=channel_names,
        labels=labels,
        num_training_cases=num_training_cases,
        file_ending=".nii.gz",
        regions_class_order=regions_class_order,
    )

model_post_init

model_post_init(__context) -> None

Initialize the writer after model initialization.

Source code in src/imgtools/io/nnunet_output.py

def model_post_init(self, __context) -> None:  # type: ignore # noqa: ANN001
    """Initialize the writer after model initialization."""
    # Create required directories
    for subdir in ["nnUNet_results", "nnUNet_preprocessed", "nnUNet_raw"]:
        (self.directory / subdir).mkdir(parents=True, exist_ok=True)

    # Determine the next available dataset ID
    existing_ids = {
        int(folder.name[7:10])
        for folder in (self.directory / "nnUNet_raw").glob("Dataset*")
        if folder.name[7:10].isdigit()
    }
    self.dataset_id = min(set(range(1, 1000)) - existing_ids)

    # Update root directory to the specific dataset folder
    self.directory = (
        self.directory
        / "nnUNet_raw"
        / f"Dataset{self.dataset_id:03d}_{self.dataset_name}"
    )

    self._file_name_format = (
        "{DirType}{SplitType}/{Dataset}_{SampleID}.nii.gz"
    )

    self._writer = NIFTIWriter(
        root_directory=self.directory,
        existing_file_mode=self.existing_file_mode,
        filename_format=self._file_name_format,
        context=self.extra_context,
    )

validate_directory classmethod

validate_directory(v: str | pathlib.Path) -> pathlib.Path

Validate that the output directory exists or can be created, and is writable.

Source code in src/imgtools/io/nnunet_output.py

@field_validator("directory")
@classmethod
def validate_directory(cls, v: str | Path) -> Path:
    """Validate that the output directory exists or can be created, and is writable."""
    return validate_directory(v)

nnUNetOutputError

Bases: Exception

Base class for errors related to sample data.