Support more dynamic configurations for datamodules (#138)

* Rename folder + Add resolvers

* Add Hydra configs for data augmentation in datamodule

* Add `flatten_dict` function and refactor datamodule

* Add resolvers for RandArtfclAliasingd transform

Also cast list returned by resolver to ListConfig so that it is resolvable by OmegaConf

* Rename ArtfclAliasing to RandArtfclAliasing + Small refactoring

* Add RandArtfclAliasing transform in augmentation config group

* Add dealias_2d augmentation config + Uncomment configs in default_2d

* Finish refactoring all datamodules

* Update datamodule configs for dealiasing + Add nnunet_reg config

* Add Hydra configs for data loading transforms

* Add `data_keys` parameter in datamodule + Update augmentation configs

* Refactor datamodules

- `nnunet_dealias_datamodule` now can handle all experiments for dealiasing
- `nnunet_datamodule` now can also handle data loading for regression

* Slightly clean up `transforms` script

* Improve override syntax to prevent creation of additional config files

* Rename augmentations by adding `rand` to the prefix

* Fix wrong target in rand_spatial_crop + Add `rand` prefix

* Update dealiasreg_2d datamodule config

* Override `augmentation` with `default_2d` in case of 2D experiments

* Final refactor of datamodules for cleaner composition data augmentations

* Add docstring + Change private method to public

* Add documentation to explain datamodule configs overriding

ascent/configs/datamodule/augmentation/aliasing/rand_artificial_aliasing.yaml

@@ -0,0 +1,6 @@
+rand_artificial_aliasing:
+  _target_: ascent.utils.transforms.RandArtfclAliasingd
+  keys: ${datamodule.data_keys.all_keys}
+  prob: 0.5
+  wrap_range: [0.6, 0.9]
+  seg_key: ${get_seg_key:${datamodule.augmentation.aliasing.rand_artificial_aliasing.keys}}

ascent/configs/datamodule/augmentation/dealias_2d.yaml

@@ -0,0 +1,16 @@
+defaults:
+  - default_2d
+  - aliasing:
+      - rand_artificial_aliasing
+
+noise:
+  rand_gaussian_noise:
+    keys: ${get_noise_and_intensity_transform_key:${datamodule.separate_transform},${datamodule.data_keys.image_key}}
+  rand_gaussian_smooth:
+    keys: ${get_noise_and_intensity_transform_key:${datamodule.separate_transform},${datamodule.data_keys.image_key}}
+
+intensity:
+  rand_scale_intensity:
+    keys: ${get_noise_and_intensity_transform_key:${datamodule.separate_transform},${datamodule.data_keys.image_key}}
+  rand_adjust_contrast:
+    keys: ${get_noise_and_intensity_transform_key:${datamodule.separate_transform},${datamodule.data_keys.image_key}}

ascent/configs/datamodule/augmentation/default_2d.yaml

@@ -0,0 +1,14 @@
+defaults:
+  - rotation:
+      - rand_rotate
+  - zoom:
+      - rand_zoom
+  - flip:
+      - rand_flip_x
+      - rand_flip_y
+  - noise:
+      - rand_gaussian_noise
+      - rand_gaussian_smooth
+  - intensity:
+      - rand_scale_intensity
+      - rand_adjust_contrast

ascent/configs/datamodule/augmentation/default_3d.yaml

@@ -0,0 +1,15 @@
+defaults:
+  - rotation:
+      - rand_rotate
+  - zoom:
+      - rand_zoom
+  - flip:
+      - rand_flip_x
+      - rand_flip_y
+      - rand_flip_z
+  - noise:
+      - rand_gaussian_noise
+      - rand_gaussian_smooth
+  - intensity:
+      - rand_scale_intensity
+      - rand_adjust_contrast

ascent/configs/datamodule/augmentation/flip/rand_flip_x.yaml

@@ -0,0 +1,5 @@
+rand_flip_x:
+  _target_: monai.transforms.RandFlipd
+  keys: ${datamodule.data_keys.all_keys}
+  spatial_axis: 0
+  prob: 0.5

ascent/configs/datamodule/augmentation/flip/rand_flip_y.yaml

@@ -0,0 +1,5 @@
+rand_flip_y:
+  _target_: monai.transforms.RandFlipd
+  keys: ${datamodule.data_keys.all_keys}
+  spatial_axis: 1
+  prob: 0.5

ascent/configs/datamodule/augmentation/flip/rand_flip_z.yaml

@@ -0,0 +1,5 @@
+rand_flip_z:
+  _target_: monai.transforms.RandFlipd
+  keys: ${datamodule.data_keys.all_keys}
+  spatial_axis: 2
+  prob: 0.5

ascent/configs/datamodule/augmentation/intensity/rand_adjust_contrast.yaml

@@ -0,0 +1,5 @@
+rand_adjust_contrast:
+  _target_: monai.transforms.RandAdjustContrastd
+  keys: ${datamodule.data_keys.image_key}
+  gamma: [0.7, 1.5]
+  prob: 0.3

ascent/configs/datamodule/augmentation/intensity/rand_scale_intensity.yaml

@@ -0,0 +1,5 @@
+rand_scale_intensity:
+  _target_: monai.transforms.RandScaleIntensityd
+  keys: ${datamodule.data_keys.image_key}
+  factors: 0.3
+  prob: 0.15

ascent/configs/datamodule/augmentation/noise/rand_gaussian_noise.yaml

@@ -0,0 +1,5 @@
+rand_gaussian_noise:
+  _target_: monai.transforms.RandGaussianNoised
+  keys: ${datamodule.data_keys.image_key}
+  std: 0.01
+  prob: 0.15

ascent/configs/datamodule/augmentation/noise/rand_gaussian_smooth.yaml

@@ -0,0 +1,6 @@
+rand_gaussian_smooth:
+  _target_: monai.transforms.RandGaussianSmoothd
+  keys: ${datamodule.data_keys.image_key}
+  sigma_x: [0.5, 1.15]
+  sigma_y: [0.5, 1.15]
+  prob: 0.15

ascent/configs/datamodule/augmentation/rotation/rand_rotate.yaml

@@ -0,0 +1,11 @@
+rand_rotate:
+  _target_: monai.transforms.RandRotated
+  keys: ${datamodule.data_keys.all_keys}
+  range_x: ${get_rot_range:${datamodule.patch_size},${datamodule.do_dummy_2D_data_aug},"x"}
+  range_y: ${get_rot_range:${datamodule.patch_size},${datamodule.do_dummy_2D_data_aug},"y"}
+  range_z: ${get_rot_range:${datamodule.patch_size},${datamodule.do_dummy_2D_data_aug},"z"}
+  mode:
+    - ${get_interp_mode:${datamodule.patch_size},${datamodule.do_dummy_2D_data_aug},"rotation"}
+    - "nearest"
+  padding_mode: "zeros"
+  prob: 0.2

ascent/configs/datamodule/augmentation/zoom/rand_zoom.yaml

@@ -0,0 +1,13 @@
+rand_zoom:
+  _target_: monai.transforms.RandZoomd
+  keys: ${datamodule.data_keys.all_keys}
+  min_zoom: 0.7
+  max_zoom: 1.4
+  mode:
+    - ${get_interp_mode:${datamodule.patch_size},${datamodule.do_dummy_2D_data_aug},"zoom"}
+    - "nearest"
+  padding_mode: "constant"
+  align_corners:
+    - True
+    - null # null in yaml is None in python
+  prob: 0.2

ascent/configs/datamodule/camus_challenge_2d.yaml

@@ -1,5 +1,6 @@
 defaults:
   - nnunet
+  - override augmentation: default_2d
 
 batch_size: 2
 dataset_name: CAMUS_challenge

ascent/configs/datamodule/cardinal_2d.yaml

@@ -1,5 +1,6 @@
 defaults:
   - nnunet
+  - override augmentation: default_2d
 
 batch_size: 2
 dataset_name: CARDINAL

ascent/configs/datamodule/dealias_2d.yaml

@@ -1,9 +1,9 @@
 defaults:
   - nnunet
+  - override augmentation: dealias_2d
 
 _target_: ascent.datamodules.nnunet_dealias_datamodule.nnUNetDealiasDataModule
-alias_transform: True
-separate_transform: True
+separate_transform: ${do_separate_transform:${datamodule.in_channels}}
 exclude_Dpower: True
 batch_size: 4
 dataset_name: DEALIASC

ascent/configs/datamodule/dealiasc_2d.yaml

@@ -1,11 +1,4 @@
 defaults:
-  - nnunet
+  - dealias_2d
 
-_target_: ascent.datamodules.nnunet_dealias_datamodule.nnUNetDealiasDataModule
-alias_transform: True
-separate_transform: True
 exclude_Dpower: False
-batch_size: 4
-dataset_name: DEALIASC
-do_dummy_2D_data_aug: False
-num_workers: 12

ascent/configs/datamodule/dealiasm_2d.yaml

@@ -1,11 +1,6 @@
 defaults:
-  - nnunet
+  - dealias_2d
 
-_target_: ascent.datamodules.nnunet_dealias_datamodule.nnUNetDealiasDataModule
-alias_transform: True
 separate_transform: False
 exclude_Dpower: False
-batch_size: 4
 dataset_name: DEALIASM
-do_dummy_2D_data_aug: False
-num_workers: 12

ascent/configs/datamodule/dealiasreg_2d.yaml

@@ -1,12 +1,6 @@
 defaults:
-  - nnunet
+  - dealiasv_2d
 
-_target_: ascent.datamodules.dealias_datamodule.DealiasDataModule
-alias_transform: True
-batch_size: 4
 dataset_name: UNWRAPV2
-do_dummy_2D_data_aug: False
-num_workers: 12
-seg_label: False
-in_channels: ${model.net.in_channels}
-patch_size: ${model.net.patch_size}
+separate_transform: True
+exclude_Dpower: False

ascent/configs/datamodule/dealiasv_2d.yaml

@@ -1,12 +1,54 @@
 defaults:
-  - nnunet
+  - dealiasm_2d
+  - override loading/[email protected]_loading: unfolded_loadnpy
+  - override loading/[email protected]_loading: unfolded_loadnpy
 
-_target_: ascent.datamodules.dealias_datamodule.DealiasDataModule
-alias_transform: False
-batch_size: 4
 dataset_name: DEALIASV
-do_dummy_2D_data_aug: False
-num_workers: 12
+separate_transform: False
 seg_label: False
-in_channels: 1
-patch_size: ${model.net.patch_size}
+data_keys:
+  all_keys:
+    - ${datamodule.data_keys.image_key}
+    - ${datamodule.data_keys.label_key}
+    - "seg"
+
+loading:
+  train:
+    crop:
+      label_key: "seg"
+
+augmentation:
+  rotation:
+    rand_rotate:
+      mode:
+        - ${get_interp_mode:${datamodule.patch_size},${datamodule.do_dummy_2D_data_aug},"rotation"}
+        - ${get_interp_mode:${datamodule.patch_size},${datamodule.do_dummy_2D_data_aug},"rotation"}
+        - "nearest"
+  zoom:
+    rand_zoom:
+      mode:
+        - ${get_interp_mode:${datamodule.patch_size},${datamodule.do_dummy_2D_data_aug},"zoom"}
+        - ${get_interp_mode:${datamodule.patch_size},${datamodule.do_dummy_2D_data_aug},"zoom"}
+        - "nearest"
+      align_corners:
+        - True
+        - True
+        - null
+  noise:
+    rand_gaussian_noise:
+      keys:
+        - ${get_noise_and_intensity_transform_key:${datamodule.separate_transform},${datamodule.data_keys.image_key}}
+        - ${datamodule.data_keys.label_key}
+    rand_gaussian_smooth:
+      keys:
+        - ${get_noise_and_intensity_transform_key:${datamodule.separate_transform},${datamodule.data_keys.image_key}}
+        - ${datamodule.data_keys.label_key}
+  intensity:
+    rand_scale_intensity:
+      keys:
+        - ${get_noise_and_intensity_transform_key:${datamodule.separate_transform},${datamodule.data_keys.image_key}}
+        - ${datamodule.data_keys.label_key}
+    rand_adjust_contrast:
+      keys:
+        - ${get_noise_and_intensity_transform_key:${datamodule.separate_transform},${datamodule.data_keys.image_key}}
+        - ${datamodule.data_keys.label_key}

ascent/configs/datamodule/loading/crop/rand_posneg_crop.yaml

@@ -0,0 +1,17 @@
+# Crop random fixed sized regions with the center being a foreground or background voxel based on
+# the Pos Neg Ratio. Refer to https://docs.monai.io/en/latest/transforms.html#RandCropByPosNegLabeld
+# for more details.
+_target_: monai.transforms.RandCropByPosNegLabeld
+keys: ${datamodule.data_keys.all_keys}
+label_key: ${datamodule.data_keys.label_key}
+spatial_size: ${get_crop_size:${datamodule.patch_size}}
+pos: 0.33
+neg: 0.67
+num_samples: 1
+image_key: null
+image_threshold: 0
+fg_indices_key: null
+bg_indices_key: null
+allow_smaller: False
+allow_missing_keys: False
+lazy: False

ascent/configs/datamodule/loading/crop/rand_spatial_crop.yaml

@@ -0,0 +1,9 @@
+# General purpose cropper to produce sub-volume region of interest. Refer to
+# https://docs.monai.io/en/latest/transforms.html#RandSpatialCropd for more details.
+_target_: monai.transforms.RandSpatialCropd
+keys: ${datamodule.data_keys.all_keys}
+roi_size: ${get_crop_size:${datamodule.patch_size}}
+max_roi_size: null
+random_center: True
+random_size: False
+lazy: False

ascent/configs/datamodule/loading/data_loading/loadimage.yaml

@@ -0,0 +1,14 @@
+# Generic loader from Monai. Refer to https://docs.monai.io/en/latest/transforms.html#monai.transforms.LoadImaged
+# for more details.
+_target_: monai.transforms.LoadImaged
+keys: ${datamodule.data_keys.all_keys}
+reader: null
+meta_keys: null
+overwriting: False
+image_only: True
+ensure_channel_first: False
+simple_keys: False
+prune_meta_pattern: null
+prune_meta_sep: "."
+allow_missing_keys: False
+expanduser: True

ascent/configs/datamodule/loading/data_loading/loadnpy.yaml

@@ -0,0 +1,5 @@
+_target_: ascent.utils.transforms.LoadNpyd
+keys:
+  - "data" # Load ["data"] or ["data", "image_meta_dict"]
+allow_missing_keys: False
+seg_label: ${datamodule.seg_label}

ascent/configs/datamodule/loading/data_loading/unfolded_loadnpy.yaml

@@ -0,0 +1,4 @@
+_target_: ascent.utils.transforms.DealiasLoadNpyd
+keys:
+  - "data" # Load ["data"] or ["data", "image_meta_dict"]
+allow_missing_keys: False

ascent/configs/datamodule/loading/default_test.yaml

@@ -0,0 +1,7 @@
+defaults:
+  - data_loading: loadnpy
+  - utility@channel_first: channel_first
+  - utility@maybe_squeeze: maybe_squeeze
+
+data_loading:
+  keys: ["data", "image_meta_dict"]

ascent/configs/datamodule/loading/default_train.yaml

@@ -0,0 +1,6 @@
+defaults:
+  - data_loading: loadnpy
+  - utility@channel_first: channel_first
+  - pad: spatial_pad
+  - crop: rand_posneg_crop
+  - utility@maybe_squeeze: maybe_squeeze

ascent/configs/datamodule/loading/pad/spatial_pad.yaml

@@ -0,0 +1,9 @@
+# Performs padding to the data, symmetric for all sides or all on one side for each dimension.
+# Refer to https://docs.monai.io/en/latest/transforms.html#SpatialPadd for more details.
+_target_: monai.transforms.SpatialPadd
+keys: ${datamodule.data_keys.all_keys}
+spatial_size: ${get_crop_size:${datamodule.patch_size}}
+method: "symmetric"
+mode: "constant"
+allow_missing_keys: False
+lazy: False

ascent/configs/datamodule/loading/utility/channel_first.yaml

@@ -0,0 +1,8 @@
+# Adjust or add the channel dimension of input data to ensure channel_first shape. Refer to
+# https://docs.monai.io/en/latest/transforms.html#monai.transforms.EnsureChannelFirstd for more
+# details.
+_target_: monai.transforms.EnsureChannelFirstd
+keys: ${datamodule.data_keys.all_keys}
+strict_check: True
+allow_missing_keys: False
+channel_dim: null

ascent/configs/datamodule/loading/utility/maybe_squeeze.yaml

@@ -0,0 +1,4 @@
+_target_: ascent.utils.transforms.MayBeSqueezed
+keys: ${datamodule.data_keys.all_keys}
+dim: -1
+allow_missing_keys: False

ascent/configs/datamodule/myosaiq+emidec_2d.yaml

@@ -1,5 +1,6 @@
 defaults:
   - nnunet
+  - override augmentation: default_2d
 
 batch_size: 32
 dataset_name: MYOSAIQ+EMIDEC

ascent/configs/datamodule/myosaiq_2d.yaml

@@ -1,5 +1,6 @@
 defaults:
   - nnunet
+  - override augmentation: default_2d
 
 batch_size: 32
 dataset_name: MYOSAIQ

ascent/configs/datamodule/nnunet.yaml

@@ -1,3 +1,8 @@
+defaults:
+  - augmentation: default_3d
+  - [email protected]: default_train
+  - [email protected]: default_test
+
 _target_: ascent.datamodules.nnunet_datamodule.nnUNetDataModule
 data_dir: ${paths.data_dir}
 batch_size: ???
@@ -11,3 +16,9 @@ num_workers: 12
 pin_memory: True
 test_splits: True
 seg_label: True
+data_keys:
+  image_key: "image"
+  label_key: "label"
+  all_keys:
+    - ${datamodule.data_keys.image_key}
+    - ${datamodule.data_keys.label_key}

ascent/configs/datamodule/nnunet_reg.yaml

@@ -0,0 +1,18 @@
+defaults:
+  - nnunet
+  - override loading/[email protected]: rand_spatial_crop
+
+seg_label: False
+
+augmentation:
+  rotation:
+    rand_rotate:
+      mode:
+        - ${get_interp_mode:${model.net.patch_size},${datamodule.do_dummy_2D_data_aug},"rotation"}
+        - ${get_interp_mode:${model.net.patch_size},${datamodule.do_dummy_2D_data_aug},"rotation"}
+  zoom:
+    rand_zoom:
+      mode:
+        - ${get_interp_mode:${model.net.patch_size},${datamodule.do_dummy_2D_data_aug},"zoom"}
+        - ${get_interp_mode:${model.net.patch_size},${datamodule.do_dummy_2D_data_aug},"zoom"}
+      align_corners: True