update hydra import, suppress nnunet prints

nndet/io/augmentation/bg_aug.py

@@ -18,6 +18,7 @@ from typing import Sequence, List
 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,
@@ -38,11 +39,12 @@ from batchgenerators.transforms import (
 from batchgenerators.transforms.color_transforms import BrightnessTransform
 from batchgenerators.transforms.utility_transforms import RemoveLabelTransform
 
-from nnunet.training.data_augmentation.custom_transforms import (
-    Convert3DTo2DTransform,
-    Convert2DTo3DTransform,
-    MaskTransform,
-    )
+with SuppressPrint():
+    from nnunet.training.data_augmentation.custom_transforms import (
+        Convert3DTo2DTransform,
+        Convert2DTo3DTransform,
+        MaskTransform,
+        )
 
 from nndet.io.augmentation import AUGMENTATION_REGISTRY
 

nndet/planning/architecture/boxes/base.py

@@ -9,7 +9,10 @@ from tqdm import tqdm
 from loguru import logger
 from torchvision.models.detection.rpn import AnchorGenerator
 
-from nnunet.experiment_planning.common_utils import get_pool_and_conv_props
+from nndet.utils.info import SuppressPrint
+
+with SuppressPrint():
+    from nnunet.experiment_planning.common_utils import get_pool_and_conv_props
 
 from nndet.io.load import load_pickle
 from nndet.arch.abstract import AbstractModel

nndet/preprocessing/resampling.py

@@ -16,7 +16,10 @@ limitations under the License.
 
 import numpy as np
 
-import nnunet.preprocessing.preprocessing as nn_preprocessing
+from nndet.utils.info import SuppressPrint
+
+with SuppressPrint():
+    import nnunet.preprocessing.preprocessing as nn_preprocessing
 
 
 def resize_segmentation(segmentation, new_shape, order=3, cval=0):

nndet/utils/config.py

@@ -20,7 +20,7 @@ from pathlib import Path
 
 import yaml
 from omegaconf import OmegaConf
-from hydra.experimental import compose as hydra_compose
+from hydra import compose as hydra_compose
 
 from nndet.io.paths import Pathlike, get_task
 

nndet/utils/info.py

@@ -15,6 +15,7 @@ limitations under the License.
 """
 
 import os
+import sys
 import copy
 import pathlib
 import warnings
@@ -33,6 +34,16 @@ from pathlib import Path
 from git import Repo, InvalidGitRepositoryError
 
 
+class SuppressPrint:
+    def __enter__(self):
+        self._original_stdout = sys.stdout
+        sys.stdout = open(os.devnull, 'w')
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        sys.stdout.close()
+        sys.stdout = self._original_stdout
+
+
 def get_requirements():
     """
     Get all installed packages from currently active environment

scripts/convert_cls2fg.py

@@ -20,7 +20,7 @@ import sys
 from pathlib import Path
 from tokenize import PseudoExtras
 
-from hydra.experimental import initialize_config_module
+from hydra import initialize_config_module
 from loguru import logger
 
 from nndet.io import get_task, load_json, save_json

scripts/convert_seg2det.py

@@ -24,7 +24,7 @@ from typing import Sequence
 
 import numpy as np
 import SimpleITK as sitk
-from hydra.experimental import initialize_config_module
+from hydra import initialize_config_module
 from loguru import logger
 from scipy import ndimage
 from scipy.ndimage import label

scripts/convert_seg2det_restore.py

@@ -18,7 +18,7 @@ import argparse
 import os
 from pathlib import Path
 
-from hydra.experimental import initialize_config_module
+from hydra import initialize_config_module
 
 from nndet.utils.config import compose
 

scripts/nnunet/nnunet_export.py

@@ -23,7 +23,7 @@ from pathlib import Path
 
 from loguru import logger
 from omegaconf import OmegaConf
-from hydra.experimental import initialize_config_module
+from hydra import initialize_config_module
 
 from nnunet.paths import nnUNet_raw_data
 

scripts/nnunet/nnunet_import.py

@@ -28,7 +28,7 @@ from typing import Union, Sequence, Optional
 
 import numpy as np
 
-from hydra.experimental import initialize_config_module
+from hydra import initialize_config_module
 from loguru import logger
 
 from nndet.evaluator.registry import evaluate_box_dir

scripts/preprocess.py

@@ -28,7 +28,7 @@ from itertools import repeat
 from typing import Dict, Sequence, Tuple, List
 from pathlib import Path
 from multiprocessing import Pool
-from hydra.experimental import initialize_config_module
+from hydra import initialize_config_module
 from omegaconf import OmegaConf
 
 from nndet.utils.config import compose

scripts/train.py

@@ -28,7 +28,7 @@ from pytorch_lightning.loggers import MLFlowLogger
 from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
 
 from loguru import logger
-from hydra.experimental import initialize_config_module
+from hydra import initialize_config_module
 from omegaconf.omegaconf import OmegaConf
 
 import nndet