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.
split_dataset	Split successfully processed cases into nnUNet train and test folders.
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.update(
            {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}/{PatientID}_{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,
    )

split_dataset

split_dataset(
    successful_results: typing.List[
        imgtools.autopipeline.ProcessSampleResult
    ],
) -> None

Split successfully processed cases into nnUNet train and test folders.

Files must already be under imagesTr and labelsTr. Selected cases are moved to imagesTs and labelsTs. The number of test cases is ceil(test_set_ratio * n_cases). A ratio of 1.0 moves every case to the test set.

Parameters:

Name	Type	Description	Default
successful_results	typing.List[imgtools.autopipeline.ProcessSampleResult]	Results for successfully processed cases.	required

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

def split_dataset(
    self,
    successful_results: List[ProcessSampleResult],
) -> None:
    """
    Split successfully processed cases into nnUNet train and test folders.

    Files must already be under imagesTr and labelsTr. Selected cases are moved
    to imagesTs and labelsTs. The number of test cases is
    ``ceil(test_set_ratio * n_cases)``. A ratio of 1.0 moves every case to the
    test set.

    Parameters
    ----------
    successful_results: List[ProcessSampleResult]
        Results for successfully processed cases.

    """
    if self.test_set_ratio == 0.0 or not successful_results:
        return

    n_cases = len(successful_results)
    n_test = ceil(self.test_set_ratio * n_cases)
    dir_map = {"labelsTr": "labelsTs", "imagesTr": "imagesTs"}

    pool = sorted(successful_results, key=lambda r: r.sample_id)
    rng = random.Random(self.random_seed)
    test_set = rng.sample(pool, n_test)

    for sample in test_set:
        sample.output_files = [
            self._move_file_to_test_split(path, dir_map)
            for path in sample.output_files
        ]

    self._patch_index_for_test_split({s.sample_id for s in test_set})

    logger.info(
        "Test set split: moved %d of %d cases to imagesTs/labelsTs "
        "(ratio=%.2f, seed=%d)",
        n_test,
        n_cases,
        self.test_set_ratio,
        self.random_seed,
    )
    if n_test == n_cases:
        logger.warning(
            "All %d successful cases were moved to the test set; "
            "imagesTr will be empty and numTraining will be 0.",
            n_cases,
        )

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.