batchgen 0.23 support

nndet/io/augmentation/bg_aug.py

@@ -20,24 +20,36 @@ from loguru import logger
 from nndet.io.augmentation.base import AugmentationSetup, get_patch_size
 from nndet.utils.info import SuppressPrint
 
-from batchgenerators.transforms import (
-    DataChannelSelectionTransform,
-    SegChannelSelectionTransform,
+from batchgenerators.transforms.abstract_transforms import Compose
+from batchgenerators.transforms.spatial_transforms import (
     SpatialTransform,
-    GammaTransform,
     MirrorTransform,
-    Compose,
+)
+from batchgenerators.transforms.resample_transforms import (
+    SimulateLowResolutionTransform,
+)
+from batchgenerators.transforms.crop_and_pad_transforms import (
+    CenterCropTransform,
+)
+from batchgenerators.transforms.color_transforms import (
+    GammaTransform,
     BrightnessMultiplicativeTransform,
+    BrightnessTransform,
     ContrastAugmentationTransform,
-    GaussianNoiseTransform,
+)
+from batchgenerators.transforms.noise_transforms import (
     GaussianBlurTransform,
-    SimulateLowResolutionTransform,
+    GaussianNoiseTransform,
+)
+from batchgenerators.transforms.channel_selection_transforms import (
+    DataChannelSelectionTransform,
+    SegChannelSelectionTransform,
+)
+from batchgenerators.transforms.utility_transforms import (
+    RemoveLabelTransform,
     RenameTransform,
     NumpyToTensor,
-    CenterCropTransform,
     )
-from batchgenerators.transforms.color_transforms import BrightnessTransform
-from batchgenerators.transforms.utility_transforms import RemoveLabelTransform
 
 with SuppressPrint():
     from nnunet.training.data_augmentation.custom_transforms import (

nndet/io/datamodule/bg_loader.py

@@ -21,7 +21,7 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 import numpy as np
 from loguru import logger
-from batchgenerators.dataloading import SlimDataLoaderBase
+from batchgenerators.dataloading.data_loader import SlimDataLoaderBase
 
 from nndet.io.datamodule import DATALOADER_REGISTRY
 from nndet.io.load import load_pickle

nndet/io/datamodule/bg_module.py

@@ -20,10 +20,8 @@ from typing import Iterable, Optional, List, Sequence, Type
 import numpy as np
 from loguru import logger
 
-from batchgenerators.dataloading import (
-    MultiThreadedAugmenter,
-    SingleThreadedAugmenter,
-    )
+from batchgenerators.dataloading.single_threaded_augmenter import SingleThreadedAugmenter
+from batchgenerators.dataloading.multi_threaded_augmenter import MultiThreadedAugmenter
 
 from nndet.io.augmentation import AUGMENTATION_REGISTRY
 from nndet.io.datamodule import DATALOADER_REGISTRY