Skip to content

refactor: extract the nnU-Net transform tail shared by five trainers #64

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Hendrik-code merged 5 commits into from
Aug 27, 2026
+134 −255
Merged

refactor: extract the nnU-Net transform tail shared by five trainers #64

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
c79e102
refactor: extract the nnU-Net transform tail shared by five trainers
Hendrik-code Aug 20, 2026
6dc67ad
Merge branch 'hm/dedup-tails' into hm/dedup-trainer-tail
Hendrik-code Aug 27, 2026
5bdaecf
Merge branch 'hm/dedup-tails' into hm/dedup-trainer-tail
Hendrik-code Aug 27, 2026
60d8b46
Merge branch 'main' into hm/dedup-trainer-tail
Hendrik-code Aug 27, 2026
2945572
Merge remote-tracking branch 'origin/hm/dedup-trainer-tail' into hm/d…
Hendrik-code Aug 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
178 changes: 28 additions & 150 deletions smauglab/trainers/nnUNetTrainerDAExt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,23 +8,19 @@
import torch
from batchgeneratorsv2.helpers.scalar_type import RandomScalar
from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform
from batchgeneratorsv2.transforms.nnunet.random_binary_operator import ApplyRandomBinaryOperatorTransform
from batchgeneratorsv2.transforms.nnunet.remove_connected_components import RemoveRandomConnectedComponentFromOneHotEncodingTransform
from batchgeneratorsv2.transforms.nnunet.seg_to_onehot import MoveSegAsOneHotToDataTransform
from batchgeneratorsv2.transforms.spatial.spatial import SpatialTransform
from batchgeneratorsv2.transforms.utils.compose import ComposeTransforms
from batchgeneratorsv2.transforms.utils.deep_supervision_downsampling import DownsampleSegForDSTransform
from batchgeneratorsv2.transforms.utils.nnunet_masking import MaskImageTransform
from batchgeneratorsv2.transforms.utils.pseudo2d import Convert2DTo3DTransform, Convert3DTo2DTransform
from batchgeneratorsv2.transforms.utils.random import RandomTransform
from batchgeneratorsv2.transforms.utils.remove_label import RemoveLabelTansform
from batchgeneratorsv2.transforms.utils.seg_to_regions import ConvertSegmentationToRegionsTransform
from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer
from nnunetv2.utilities.helpers import dummy_context
from torch import autocast

from smauglab import configs
from smauglab.trainers.utils import DownsampleSegForDSTransformCustom
from smauglab.trainers.utils import DownsampleSegForDSTransformCustom, nnunet_tail_transforms
from smauglab.transforms.cpu.transforms import AugTransforms
from smauglab.transforms.gpu.transforms import AugTransformsGPU

Expand Down Expand Up @@ -96,54 +92,16 @@ def get_training_transforms(
if do_dummy_2d_data_aug:
transforms.append(Convert2DTo3DTransform())

if use_mask_for_norm is not None and any(use_mask_for_norm):
transforms.append(
MaskImageTransform(
apply_to_channels=[i for i in range(len(use_mask_for_norm)) if use_mask_for_norm[i]],
channel_idx_in_seg=0,
set_outside_to=0,
)
transforms.extend(
nnunet_tail_transforms(
use_mask_for_norm=use_mask_for_norm,
deep_supervision_scales=deep_supervision_scales,
is_cascaded=is_cascaded,
foreground_labels=foreground_labels,
regions=regions,
ignore_label=ignore_label,
)

transforms.append(RemoveLabelTansform(-1, 0))

# The following augmentations are related to special nnunet executions
if is_cascaded:
assert foreground_labels is not None, "We need foreground_labels for cascade augmentations"
transforms.append(
MoveSegAsOneHotToDataTransform(source_channel_idx=1, all_labels=foreground_labels, remove_channel_from_source=True)
)
transforms.append(
RandomTransform(
ApplyRandomBinaryOperatorTransform(
channel_idx=list(range(-len(foreground_labels), 0)), strel_size=(1, 8), p_per_label=1
),
apply_probability=0.4,
)
)
transforms.append(
RandomTransform(
RemoveRandomConnectedComponentFromOneHotEncodingTransform(
channel_idx=list(range(-len(foreground_labels), 0)),
fill_with_other_class_p=0,
dont_do_if_covers_more_than_x_percent=0.15,
p_per_label=1,
),
apply_probability=0.2,
)
)

if regions is not None:
# the ignore label must also be converted
transforms.append(
ConvertSegmentationToRegionsTransform(
regions=[*list(regions), ignore_label] if ignore_label is not None else regions, channel_in_seg=0
)
)

if deep_supervision_scales is not None:
transforms.append(DownsampleSegForDSTransform(ds_scales=deep_supervision_scales))

)
return ComposeTransforms(transforms)


Expand Down Expand Up @@ -220,57 +178,16 @@ def get_training_transforms(
if do_dummy_2d_data_aug:
transforms.append(Convert2DTo3DTransform())

if use_mask_for_norm is not None and any(use_mask_for_norm):
transforms.append(
MaskImageTransform(
apply_to_channels=[i for i in range(len(use_mask_for_norm)) if use_mask_for_norm[i]],
channel_idx_in_seg=0,
set_outside_to=0,
)
)

transforms.append(RemoveLabelTansform(-1, 0))

# The following augmentations are related to special nnunet executions
if is_cascaded:
assert foreground_labels is not None, "We need foreground_labels for cascade augmentations"
transforms.append(
MoveSegAsOneHotToDataTransform(source_channel_idx=1, all_labels=foreground_labels, remove_channel_from_source=True)
)
transforms.append(
RandomTransform(
ApplyRandomBinaryOperatorTransform(
channel_idx=list(range(-len(foreground_labels), 0)), strel_size=(1, 8), p_per_label=1
),
apply_probability=0.4,
)
)
transforms.append(
RandomTransform(
RemoveRandomConnectedComponentFromOneHotEncodingTransform(
channel_idx=list(range(-len(foreground_labels), 0)),
fill_with_other_class_p=0,
dont_do_if_covers_more_than_x_percent=0.15,
p_per_label=1,
),
apply_probability=0.2,
)
)

if regions is not None:
# the ignore label must also be converted
transforms.append(
ConvertSegmentationToRegionsTransform(
regions=[*list(regions), ignore_label] if ignore_label is not None else regions, channel_in_seg=0
)
transforms.extend(
nnunet_tail_transforms(
use_mask_for_norm=use_mask_for_norm,
deep_supervision_scales=None,
is_cascaded=is_cascaded,
foreground_labels=foreground_labels,
regions=regions,
ignore_label=ignore_label,
)

# transforms.append(ZscoreNormalization())

# NOTE: DownsampleSegForDSTransform is now handled in train_step for GPU augmentations
# if deep_supervision_scales is not None:
# transforms.append(DownsampleSegForDSTransform(ds_scales=deep_supervision_scales))

)
return ComposeTransforms(transforms)

@staticmethod
Expand Down Expand Up @@ -399,55 +316,16 @@ def get_training_transforms(
)
)

if use_mask_for_norm is not None and any(use_mask_for_norm):
transforms.append(
MaskImageTransform(
apply_to_channels=[i for i in range(len(use_mask_for_norm)) if use_mask_for_norm[i]],
channel_idx_in_seg=0,
set_outside_to=0,
)
transforms.extend(
nnunet_tail_transforms(
use_mask_for_norm=use_mask_for_norm,
deep_supervision_scales=None,
is_cascaded=is_cascaded,
foreground_labels=foreground_labels,
regions=regions,
ignore_label=ignore_label,
)

transforms.append(RemoveLabelTansform(-1, 0))

# The following augmentations are related to special nnunet executions
if is_cascaded:
assert foreground_labels is not None, "We need foreground_labels for cascade augmentations"
transforms.append(
MoveSegAsOneHotToDataTransform(source_channel_idx=1, all_labels=foreground_labels, remove_channel_from_source=True)
)
transforms.append(
RandomTransform(
ApplyRandomBinaryOperatorTransform(
channel_idx=list(range(-len(foreground_labels), 0)), strel_size=(1, 8), p_per_label=1
),
apply_probability=0.4,
)
)
transforms.append(
RandomTransform(
RemoveRandomConnectedComponentFromOneHotEncodingTransform(
channel_idx=list(range(-len(foreground_labels), 0)),
fill_with_other_class_p=0,
dont_do_if_covers_more_than_x_percent=0.15,
p_per_label=1,
),
apply_probability=0.2,
)
)

if regions is not None:
# the ignore label must also be converted
transforms.append(
ConvertSegmentationToRegionsTransform(
regions=[*list(regions), ignore_label] if ignore_label is not None else regions, channel_in_seg=0
)
)

# NOTE: DownsampleSegForDSTransform is now handled in train_step for GPU augmentations
# if deep_supervision_scales is not None:
# transforms.append(DownsampleSegForDSTransform(ds_scales=deep_supervision_scales))

)
return ComposeTransforms(transforms)

def train_step(self, batch: dict) -> dict:
Expand Down
123 changes: 19 additions & 104 deletions smauglab/trainers/nnUNetTrainerTest.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,23 +5,15 @@
import torch
from batchgeneratorsv2.helpers.scalar_type import RandomScalar
from batchgeneratorsv2.transforms.base.basic_transform import BasicTransform
from batchgeneratorsv2.transforms.nnunet.random_binary_operator import ApplyRandomBinaryOperatorTransform
from batchgeneratorsv2.transforms.nnunet.remove_connected_components import RemoveRandomConnectedComponentFromOneHotEncodingTransform
from batchgeneratorsv2.transforms.nnunet.seg_to_onehot import MoveSegAsOneHotToDataTransform
from batchgeneratorsv2.transforms.spatial.spatial import SpatialTransform
from batchgeneratorsv2.transforms.utils.compose import ComposeTransforms
from batchgeneratorsv2.transforms.utils.deep_supervision_downsampling import DownsampleSegForDSTransform
from batchgeneratorsv2.transforms.utils.nnunet_masking import MaskImageTransform
from batchgeneratorsv2.transforms.utils.pseudo2d import Convert2DTo3DTransform, Convert3DTo2DTransform
from batchgeneratorsv2.transforms.utils.random import RandomTransform
from batchgeneratorsv2.transforms.utils.remove_label import RemoveLabelTansform
from batchgeneratorsv2.transforms.utils.seg_to_regions import ConvertSegmentationToRegionsTransform
from nnunetv2.training.nnUNetTrainer.nnUNetTrainer import nnUNetTrainer
from nnunetv2.utilities.helpers import dummy_context
from torch import autocast

from smauglab import configs
from smauglab.trainers.utils import DownsampleSegForDSTransformCustom
from smauglab.trainers.utils import DownsampleSegForDSTransformCustom, nnunet_tail_transforms
from smauglab.transforms.cpu.transforms import AugTransformsTest
from smauglab.transforms.gpu.transforms import AugTransformsGPU

Expand Down Expand Up @@ -74,54 +66,16 @@ def get_training_transforms(
if do_dummy_2d_data_aug:
transforms.append(Convert2DTo3DTransform())

if use_mask_for_norm is not None and any(use_mask_for_norm):
transforms.append(
MaskImageTransform(
apply_to_channels=[i for i in range(len(use_mask_for_norm)) if use_mask_for_norm[i]],
channel_idx_in_seg=0,
set_outside_to=0,
)
transforms.extend(
nnunet_tail_transforms(
use_mask_for_norm=use_mask_for_norm,
deep_supervision_scales=deep_supervision_scales,
is_cascaded=is_cascaded,
foreground_labels=foreground_labels,
regions=regions,
ignore_label=ignore_label,
)

transforms.append(RemoveLabelTansform(-1, 0))

# The following augmentations are related to special nnunet executions
if is_cascaded:
assert foreground_labels is not None, "We need foreground_labels for cascade augmentations"
transforms.append(
MoveSegAsOneHotToDataTransform(source_channel_idx=1, all_labels=foreground_labels, remove_channel_from_source=True)
)
transforms.append(
RandomTransform(
ApplyRandomBinaryOperatorTransform(
channel_idx=list(range(-len(foreground_labels), 0)), strel_size=(1, 8), p_per_label=1
),
apply_probability=0.4,
)
)
transforms.append(
RandomTransform(
RemoveRandomConnectedComponentFromOneHotEncodingTransform(
channel_idx=list(range(-len(foreground_labels), 0)),
fill_with_other_class_p=0,
dont_do_if_covers_more_than_x_percent=0.15,
p_per_label=1,
),
apply_probability=0.2,
)
)

if regions is not None:
# the ignore label must also be converted
transforms.append(
ConvertSegmentationToRegionsTransform(
regions=[*list(regions), ignore_label] if ignore_label is not None else regions, channel_in_seg=0
)
)

if deep_supervision_scales is not None:
transforms.append(DownsampleSegForDSTransform(ds_scales=deep_supervision_scales))

)
return ComposeTransforms(transforms)


Expand Down Expand Up @@ -175,55 +129,16 @@ def get_training_transforms(
if do_dummy_2d_data_aug:
transforms.append(Convert2DTo3DTransform())

if use_mask_for_norm is not None and any(use_mask_for_norm):
transforms.append(
MaskImageTransform(
apply_to_channels=[i for i in range(len(use_mask_for_norm)) if use_mask_for_norm[i]],
channel_idx_in_seg=0,
set_outside_to=0,
)
transforms.extend(
nnunet_tail_transforms(
use_mask_for_norm=use_mask_for_norm,
deep_supervision_scales=None,
is_cascaded=is_cascaded,
foreground_labels=foreground_labels,
regions=regions,
ignore_label=ignore_label,
)

transforms.append(RemoveLabelTansform(-1, 0))

# The following augmentations are related to special nnunet executions
if is_cascaded:
assert foreground_labels is not None, "We need foreground_labels for cascade augmentations"
transforms.append(
MoveSegAsOneHotToDataTransform(source_channel_idx=1, all_labels=foreground_labels, remove_channel_from_source=True)
)
transforms.append(
RandomTransform(
ApplyRandomBinaryOperatorTransform(
channel_idx=list(range(-len(foreground_labels), 0)), strel_size=(1, 8), p_per_label=1
),
apply_probability=0.4,
)
)
transforms.append(
RandomTransform(
RemoveRandomConnectedComponentFromOneHotEncodingTransform(
channel_idx=list(range(-len(foreground_labels), 0)),
fill_with_other_class_p=0,
dont_do_if_covers_more_than_x_percent=0.15,
p_per_label=1,
),
apply_probability=0.2,
)
)

if regions is not None:
# the ignore label must also be converted
transforms.append(
ConvertSegmentationToRegionsTransform(
regions=[*list(regions), ignore_label] if ignore_label is not None else regions, channel_in_seg=0
)
)

# NOTE: DownsampleSegForDSTransform is now handled in train_step for GPU augmentations
# if deep_supervision_scales is not None:
# transforms.append(DownsampleSegForDSTransform(ds_scales=deep_supervision_scales))

)
return ComposeTransforms(transforms)

def train_step(self, batch: dict) -> dict:
Expand Down
Loading
Loading