Minor fixes

NKI-AI · Jun 5, 2024 · bcd58fa · bcd58fa
1 parent 1cbbb45
commit bcd58fa
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/direct/data/datasets_config.py b/direct/data/datasets_config.py
@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import List, Optional
+from typing import Optional
 
 from omegaconf import MISSING
 
@@ -42,10 +42,10 @@ class SensitivityMapEstimationTransformConfig(BaseConfig):
 
 @dataclass
 class AugmentationTransformConfig(BaseConfig):
-    rescale: Optional[tuple[int, ...]] = (None,)
-    rescale_mode: Optional[RescaleMode] = (RescaleMode.NEAREST,)
-    rescale_2d_if_3d: Optional[bool] = (False,)
-    pad: Optional[tuple[int, ...]] = (None,)
+    rescale: Optional[tuple[int, ...]] = None
+    rescale_mode: Optional[RescaleMode] = RescaleMode.NEAREST
+    rescale_2d_if_3d: Optional[bool] = False
+    pad: Optional[tuple[int, ...]] = None
 
 
 @dataclass

diff --git a/tests/test_train.py b/tests/test_train.py
@@ -25,12 +25,13 @@
 )
 from direct.launch import launch
 from direct.train import setup_train
+from direct.types import MaskFuncMode
 
 
 def create_test_transform_cfg(transforms_type):
     transforms_config = TransformsConfig(
         normalization=NormalizationTransformConfig(scaling_key="masked_kspace"),
-        masking=MaskingConfig(name="FastMRIRandom"),
+        masking=MaskingConfig(name="FastMRIRandom", mode=MaskFuncMode.STATIC),
         cropping=CropTransformConfig(crop=(32, 32)),
         sensitivity_map_estimation=SensitivityMapEstimationTransformConfig(estimate_sensitivity_maps=True),
         transforms_type=transforms_type,