Requested changes

NKI-AI · Jan 22, 2024 · dcdd5c3 · dcdd5c3
1 parent 9149732
commit dcdd5c3
Show file tree

Hide file tree

Showing 9 changed files with 85 additions and 98 deletions.
diff --git a/direct/data/datasets_config.py b/direct/data/datasets_config.py
@@ -32,14 +32,11 @@ class SensitivityMapEstimationTransformConfig(BaseConfig):
 
 @dataclass
 class RandomAugmentationTransformsConfig(BaseConfig):
-    random_rotation: bool = False
     random_rotation_degrees: Tuple[int, ...] = (-90, 90)
-    random_rotation_probability: Optional[float] = 0.5
-    random_flip: bool = False
+    random_rotation_probability: float = 0.0
     random_flip_type: Optional[str] = "random"
-    random_flip_probability: Optional[float] = 0.5
-    random_reverse: bool = False
-    random_reverse_probability: Optional[float] = 0.5
+    random_flip_probability: float = 0.0
+    random_reverse_probability: float = 0.0
 
 
 @dataclass

diff --git a/direct/data/mri_transforms.py b/direct/data/mri_transforms.py
@@ -544,7 +544,7 @@ def __init__(
         self,
         kspace_key: KspaceKey = KspaceKey.KSPACE,
         padding_key: str = "padding",
-        eps: Union[float, None] = 0.0001,
+        eps: Optional[float] = 0.0001,
     ) -> None:
         """Inits :class:`ComputeZeroPadding`.
 
@@ -1428,12 +1428,10 @@ def build_pre_mri_transforms(
     crop: Optional[Union[Tuple[int, int], str]] = None,
     crop_type: Optional[str] = "uniform",
     image_center_crop: bool = True,
-    random_rotation: bool = False,
     random_rotation_degrees: Optional[Sequence[int]] = (-90, 90),
-    random_rotation_probability: Optional[float] = 0.5,
-    random_flip: bool = False,
+    random_rotation_probability: float = 0.0,
     random_flip_type: Optional[RandomFlipType] = RandomFlipType.RANDOM,
-    random_flip_probability: Optional[float] = 0.5,
+    random_flip_probability: float = 0.0,
     padding_eps: float = 0.0001,
     estimate_body_coil_image: bool = False,
     use_seed: bool = True,
@@ -1467,19 +1465,16 @@ def build_pre_mri_transforms(
     image_center_crop : bool
         If True the backprojected kspace will be cropped around the center, otherwise randomly.
         This will be ignored if `crop` is None. Default: True.
-    random_rotation : bool
-        If True, random rotations will be applied of `random_rotation_degrees` degrees, with probability
-        `random_rotation_probability`. Default: False.
     random_rotation_degrees : Sequence[int], optional
         Default: (-90, 90).
     random_rotation_probability : float, optional
-        Default: 0.5.
-    random_flip : bool
-        If True, random rotation of `random_flip_type` type, with probability `random_flip_probability`. Default: False.
+        If greater than 0.0, random rotations will be applied of `random_rotation_degrees` degrees, with probability
+        `random_rotation_probability`. Default: 0.0.
     random_flip_type : RandomFlipType, optional
         Default: RandomFlipType.RANDOM.
     random_flip_probability : float, optional
-        Default: 0.5.
+        If greater than 0.0, random rotation of `random_flip_type` type, with probability `random_flip_probability`.
+        Default: 0.0.
     padding_eps: float
         Padding epsilon. Default: 0.0001.
     estimate_body_coil_image : bool
@@ -1506,15 +1501,15 @@ def build_pre_mri_transforms(
                 random_crop_sampler_use_seed=use_seed,
             )
         ]
-    if random_rotation:
+    if random_rotation_probability > 0.0:
         mri_transforms += [
             RandomRotation(
                 degrees=random_rotation_degrees,
                 p=random_rotation_probability,
                 keys_to_rotate=(TransformKey.KSPACE, TransformKey.SENSITIVITY_MAP),
             )
         ]
-    if random_flip:
+    if random_flip_probability > 0.0:
         mri_transforms += [
             RandomFlip(
                 flip=random_flip_type,
@@ -1655,14 +1650,11 @@ def build_mri_transforms(
     crop: Optional[Union[Tuple[int, int], str]] = None,
     crop_type: Optional[str] = "uniform",
     image_center_crop: bool = True,
-    random_rotation: bool = False,
     random_rotation_degrees: Optional[Sequence[int]] = (-90, 90),
-    random_rotation_probability: Optional[float] = 0.5,
-    random_flip: bool = False,
+    random_rotation_probability: float = 0.0,
     random_flip_type: Optional[RandomFlipType] = RandomFlipType.RANDOM,
-    random_flip_probability: Optional[float] = 0.5,
-    random_reverse: bool = False,
-    random_reverse_probability: float = 0.5,
+    random_flip_probability: float = 0.0,
+    random_reverse_probability: float = 0.0,
     padding_eps: float = 0.0001,
     estimate_body_coil_image: bool = False,
     estimate_sensitivity_maps: bool = True,
@@ -1708,23 +1700,19 @@ def build_mri_transforms(
     image_center_crop : bool
         If True the backprojected kspace will be cropped around the center, otherwise randomly.
         This will be ignored if `crop` is None. Default: True.
-    random_rotation : bool
-        If True, random rotations will be applied of `random_rotation_degrees` degrees, with probability
-        `random_rotation_probability`. Default: False.
     random_rotation_degrees : Sequence[int], optional
         Default: (-90, 90).
     random_rotation_probability : float, optional
-        Default: 0.5.
-    random_flip : bool
-        If True, random rotation of `random_flip_type` type, with probability `random_flip_probability`. Default: False.
+        If greater than 0.0, random rotations will be applied of `random_rotation_degrees` degrees, with probability
+        `random_rotation_probability`. Default: 0.0.
     random_flip_type : RandomFlipType, optional
         Default: RandomFlipType.RANDOM.
     random_flip_probability : float, optional
-        Default: 0.5.
-    random_reverse : bool
-        If True will perform random reversion along the time or slice dimension (2). Default: False.
+        If greater than 0.0, random rotation of `random_flip_type` type, with probability `random_flip_probability`.
+        Default: 0.0.
     random_reverse_probability : float
-        Default: 0.5.
+        If greater than 0.0, will perform random reversion along the time or slice dimension (2) with probability
+        `random_reverse_probability`. Default: 0.0.
     padding_eps: float
         Padding epsilon. Default: 0.0001.
     estimate_body_coil_image : bool
@@ -1779,23 +1767,23 @@ def build_mri_transforms(
                 random_crop_sampler_use_seed=use_seed,
             )
         ]
-    if random_rotation:
+    if random_rotation_probability > 0.0:
         mri_transforms += [
             RandomRotation(
                 degrees=random_rotation_degrees,
                 p=random_rotation_probability,
                 keys_to_rotate=(TransformKey.KSPACE, TransformKey.SENSITIVITY_MAP),
             )
         ]
-    if random_flip:
+    if random_flip_probability > 0.0:
         mri_transforms += [
             RandomFlip(
                 flip=random_flip_type,
                 p=random_flip_probability,
                 keys_to_flip=(TransformKey.KSPACE, TransformKey.SENSITIVITY_MAP),
             )
         ]
-    if random_reverse:
+    if random_reverse_probability > 0.0:
         mri_transforms += [
             RandomReverse(
                 p=random_reverse_probability,

diff --git a/direct/nn/unet/unet_3d.py b/direct/nn/unet/unet_3d.py
@@ -51,7 +51,11 @@ def forward(self, input_data: torch.Tensor) -> torch.Tensor:
 
 
 class UnetModel3d(nn.Module):
-    """PyTorch implementation of a 3D U-Net model."""
+    """PyTorch implementation of a 3D U-Net model.
+
+    This class defines a 3D U-Net architecture consisting of down-sampling and up-sampling layers with 3D convolutional
+    blocks. This is an extension of Unet2dModel, but for volumes.
+    """
 
     def __init__(
         self,

diff --git a/direct/nn/vsharp/__init__.py b/direct/nn/vsharp/__init__.py
@@ -1,2 +1 @@
-# coding=utf-8
 # Copyright (c) DIRECT Contributors
diff --git a/projects/vSHARP/fastmri_prostate/configs/base_recurrentvarnet.yaml b/projects/vSHARP/fastmri_prostate/configs/base_recurrentvarnet.yaml
@@ -11,8 +11,8 @@ training:
                 estimate_sensitivity_maps: true
                 scaling_key: masked_kspace  # Compute the image normalization based on the masked_kspace maximum
                 image_center_crop: false
-                random_flip: true
-                random_rotation: true
+                random_flip_probability: 0.5
+                random_rotation_probability: 0.5
                 masking:
                     name: FastMRIEquispaced
                     accelerations: [4, 8, 16]
@@ -28,8 +28,8 @@ training:
                 estimate_sensitivity_maps: true
                 scaling_key: masked_kspace  # Compute the image normalization based on the masked_kspace maximum
                 image_center_crop: false
-                random_flip: true
-                random_rotation: true
+                random_flip_probability: 0.5
+                random_rotation_probability: 0.5
                 masking:
                     name: FastMRIEquispaced
                     accelerations: [4, 8, 16]
@@ -45,8 +45,8 @@ training:
                 estimate_sensitivity_maps: true
                 scaling_key: masked_kspace  # Compute the image normalization based on the masked_kspace maximum
                 image_center_crop: false
-                random_flip: true
-                random_rotation: true
+                random_flip_probability: 0.5
+                random_rotation_probability: 0.5
                 masking:
                     name: FastMRIEquispaced
                     accelerations: [4, 8, 16]
@@ -62,8 +62,8 @@ training:
                 estimate_sensitivity_maps: true
                 scaling_key: masked_kspace  # Compute the image normalization based on the masked_kspace maximum
                 image_center_crop: false
-                random_flip: true
-                random_rotation: true
+                random_flip_probability: 0.5
+                random_rotation_probability: 0.5
                 masking:
                     name: FastMRIEquispaced
                     accelerations: [4, 8, 16]
@@ -79,8 +79,8 @@ training:
                 estimate_sensitivity_maps: true
                 scaling_key: masked_kspace  # Compute the image normalization based on the masked_kspace maximum
                 image_center_crop: false
-                random_flip: true
-                random_rotation: true
+                random_flip_probability: 0.5
+                random_rotation_probability: 0.5
                 masking:
                     name: FastMRIEquispaced
                     accelerations: [4, 8, 16]
@@ -96,8 +96,8 @@ training:
                 estimate_sensitivity_maps: true
                 scaling_key: masked_kspace  # Compute the image normalization based on the masked_kspace maximum
                 image_center_crop: false
-                random_flip: true
-                random_rotation: true
+                random_flip_probability: 0.5
+                random_rotation_probability: 0.5
                 masking:
                     name: FastMRIEquispaced
                     accelerations: [4, 8, 16]
@@ -113,8 +113,8 @@ training:
                 estimate_sensitivity_maps: true
                 scaling_key: masked_kspace  # Compute the image normalization based on the masked_kspace maximum
                 image_center_crop: false
-                random_flip: true
-                random_rotation: true
+                random_flip_probability: 0.5
+                random_rotation_probability: 0.5
                 masking:
                     name: FastMRIEquispaced
                     accelerations: [4, 8, 16]

diff --git a/projects/vSHARP/fastmri_prostate/configs/base_unet.yaml b/projects/vSHARP/fastmri_prostate/configs/base_unet.yaml
@@ -11,8 +11,8 @@ training:
                 estimate_sensitivity_maps: true
                 scaling_key: masked_kspace  # Compute the image normalization based on the masked_kspace maximum
                 image_center_crop: false
-                random_flip: true
-                random_rotation: true
+                random_flip_probability: 0.5
+                random_rotation_probability: 0.5
                 masking:
                     name: FastMRIEquispaced
                     accelerations: [4, 8, 16]
@@ -28,8 +28,8 @@ training:
                 estimate_sensitivity_maps: true
                 scaling_key: masked_kspace  # Compute the image normalization based on the masked_kspace maximum
                 image_center_crop: false
-                random_flip: true
-                random_rotation: true
+                random_flip_probability: 0.5
+                random_rotation_probability: 0.5
                 masking:
                     name: FastMRIEquispaced
                     accelerations: [4, 8, 16]
@@ -45,8 +45,8 @@ training:
                 estimate_sensitivity_maps: true
                 scaling_key: masked_kspace  # Compute the image normalization based on the masked_kspace maximum
                 image_center_crop: false
-                random_flip: true
-                random_rotation: true
+                random_flip_probability: 0.5
+                random_rotation_probability: 0.5
                 masking:
                     name: FastMRIEquispaced
                     accelerations: [4, 8, 16]
@@ -62,8 +62,8 @@ training:
                 estimate_sensitivity_maps: true
                 scaling_key: masked_kspace  # Compute the image normalization based on the masked_kspace maximum
                 image_center_crop: false
-                random_flip: true
-                random_rotation: true
+                random_flip_probability: 0.5
+                random_rotation_probability: 0.5
                 masking:
                     name: FastMRIEquispaced
                     accelerations: [4, 8, 16]
@@ -79,8 +79,8 @@ training:
                 estimate_sensitivity_maps: true
                 scaling_key: masked_kspace  # Compute the image normalization based on the masked_kspace maximum
                 image_center_crop: false
-                random_flip: true
-                random_rotation: true
+                random_flip_probability: 0.5
+                random_rotation_probability: 0.5
                 masking:
                     name: FastMRIEquispaced
                     accelerations: [4, 8, 16]
@@ -96,8 +96,8 @@ training:
                 estimate_sensitivity_maps: true
                 scaling_key: masked_kspace  # Compute the image normalization based on the masked_kspace maximum
                 image_center_crop: false
-                random_flip: true
-                random_rotation: true
+                random_flip_probability: 0.5
+                random_rotation_probability: 0.5
                 masking:
                     name: FastMRIEquispaced
                     accelerations: [4, 8, 16]
@@ -113,8 +113,8 @@ training:
                 estimate_sensitivity_maps: true
                 scaling_key: masked_kspace  # Compute the image normalization based on the masked_kspace maximum
                 image_center_crop: false
-                random_flip: true
-                random_rotation: true
+                random_flip_probability: 0.5
+                random_rotation_probability: 0.5
                 masking:
                     name: FastMRIEquispaced
                     accelerations: [4, 8, 16]
Original file line number	Diff line number	Diff line change
		@@ -1,2 +1 @@
		# coding=utf-8
		# Copyright (c) DIRECT Contributors