Add augmentation mri transforms

.github/workflows/black.yml

@@ -7,5 +7,5 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
-      - uses: actions/setup-python@v2
+      - uses: actions/setup-python@v4
       - uses: psf/black@22.12.0

.github/workflows/coverage.yml

@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.8]
+        python-version: ['3.10']
     steps:
     - uses: actions/checkout@v2
     - name: Set up Python ${{ matrix.python-version }}

.github/workflows/pylint.yml

@@ -7,10 +7,10 @@ jobs:
     runs-on: ubuntu-latest
     steps:
     - uses: actions/checkout@v2
-    - name: Set up Python 3.8
+    - name: Set up Python 3.10
       uses: actions/setup-python@v1
       with:
-        python-version: 3.8
+        python-version: '3.10'
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip

.github/workflows/tox.yml

@@ -9,12 +9,12 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: [3.8, 3.9]
+        python-version: ['3.9', '3.10']
 
     steps:
     - uses: actions/checkout@v2
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v2
+      uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies

contributing.rst

@@ -110,7 +110,7 @@ Before you submit a pull request, check that it meets these guidelines:
 #. If the pull request adds functionality, the docs should be updated. Put
    your new functionality into a function with a docstring, and add the
    feature to the list in README.rst.
-#. The pull request should work for Python 3.8 and 3.9 and for PyPy. Check Github actions and see that all tests pass.
+#. The pull request should work for Python 3.10 and for PyPy. Check Github actions and see that all tests pass.
 
 Tests
 ^^^^^

direct/common/_gaussian.pyx

@@ -57,7 +57,7 @@ def gaussian_mask_1d(
     int n,
     int center,
     double std,
-    cnp.ndarray[cnp.int_t, ndim=1, mode='c'] mask,
+    cnp.ndarray[cnp.int64_t, ndim=1, mode='c'] mask,
     int seed,
 ):
     cdef int count, ind
@@ -82,7 +82,7 @@ def gaussian_mask_2d(
     int center_x,
     int center_y,
     cnp.ndarray[cnp.float_t, ndim=1, mode='c'] std,
-    cnp.ndarray[cnp.int_t, ndim=2, mode='c'] mask,
+    cnp.ndarray[cnp.int64_t, ndim=2, mode='c'] mask,
     int seed,
 ):
     cdef int count, indx, indy

direct/common/_poisson.pyx

@@ -38,7 +38,7 @@ def poisson(
     int nx,
     int ny,
     int max_attempts,
-    cnp.ndarray[cnp.int_t, ndim=2, mode='c'] mask,
+    cnp.ndarray[cnp.int64_t, ndim=2, mode='c'] mask,
     cnp.ndarray[cnp.float64_t, ndim=2, mode='c'] radius_x,
     cnp.ndarray[cnp.float64_t, ndim=2, mode='c'] radius_y,
     int seed
@@ -62,8 +62,8 @@ def poisson(
     cdef Py_ssize_t startx, endx, starty, endy, px, py
 
     # initialize active list
-    cdef cnp.ndarray[cnp.int_t, ndim=1, mode='c'] pxs = np.empty(nx * ny, dtype=int)
-    cdef cnp.ndarray[cnp.int_t, ndim=1, mode='c'] pys = np.empty(nx * ny, dtype=int)
+    cdef cnp.ndarray[cnp.int64_t, ndim=1, mode='c'] pxs = np.empty(nx * ny, dtype=int)
+    cdef cnp.ndarray[cnp.int64_t, ndim=1, mode='c'] pys = np.empty(nx * ny, dtype=int)
 
     srand(seed)
 

direct/data/datasets_config.py

@@ -2,8 +2,10 @@
 
 """Classes holding the typed configurations for the datasets."""
 
+from __future__ import annotations
+
 from dataclasses import dataclass
-from typing import List, Optional, Tuple
+from typing import Optional
 
 from omegaconf import MISSING
 
@@ -14,6 +16,7 @@
     MaskSplitterType,
     RandomFlipType,
     ReconstructionType,
+    RescaleMode,
     SensitivityMapType,
     TransformsType,
 )
@@ -37,9 +40,17 @@ class SensitivityMapEstimationTransformConfig(BaseConfig):
     sensitivity_maps_gaussian: Optional[float] = 0.7
 
 
+@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
+
+
 @dataclass
 class RandomAugmentationTransformsConfig(BaseConfig):
-    random_rotation_degrees: Tuple[int, ...] = (-90, 90)
+    random_rotation_degrees: tuple[int, ...] = (-90, 90)
     random_rotation_probability: float = 0.0
     random_flip_type: Optional[RandomFlipType] = RandomFlipType.RANDOM
     random_flip_probability: float = 0.0
@@ -62,8 +73,10 @@ class TransformsConfig(BaseConfig):
         Configuration for the masking.
     cropping : CropTransformConfig
         Configuration for the cropping.
+    augmentation : AugmentationTransformConfig
+        Configuration for the augmentation. Currently only rescale and pad are supported.
     random_augmentations : RandomAugmentationTransformsConfig
-        Configuration for the random augmentations.
+        Configuration for the random augmentations. Currently only random rotation, flip and reverse are supported.
     padding_eps : float
         Padding epsilon. Default is 0.001.
     estimate_body_coil_image : bool
@@ -79,6 +92,9 @@ class TransformsConfig(BaseConfig):
         Default is True.
     image_recon_type : ReconstructionType
         Image reconstruction type. Default is ReconstructionType.RSS.
+    compress_coils : int, optional
+        Number of coils to compress input k-space. It is not recommended to be used in combination with `pad_coils`.
+        Default is None.
     pad_coils : int, optional
         Pad coils. Default is None.
     use_seed : bool
@@ -89,10 +105,10 @@ class TransformsConfig(BaseConfig):
         To use SSL transforms, set transforms_type to `SSL_SSDU`. This will require additional parameters to be set:
         mask_split_ratio, mask_split_acs_region, mask_split_keep_acs, mask_split_type, mask_split_gaussian_std.
         Default is `TransformsType.SUPERVISED`.
-    mask_split_ratio : Tuple[float, ...]
+    mask_split_ratio : tuple[float, ...]
         Ratio of the mask to split into input and target mask. Ignored if transforms_type is not `SSL_SSDU`.
         Default is (0.4,).
-    mask_split_acs_region : Tuple[int, int]
+    mask_split_acs_region : tuple[int, int]
         Region of the ACS k-space to keep in the input mask. Ignored if transforms_type is not `SSL_SSDU`.
         Default is (0, 0).
     mask_split_keep_acs : bool, optional
@@ -111,6 +127,7 @@ class TransformsConfig(BaseConfig):
 
     masking: Optional[MaskingConfig] = MaskingConfig()
     cropping: CropTransformConfig = CropTransformConfig()
+    augmentation: AugmentationTransformConfig = AugmentationTransformConfig()
     random_augmentations: RandomAugmentationTransformsConfig = RandomAugmentationTransformsConfig()
     padding_eps: float = 0.001
     estimate_body_coil_image: bool = False
@@ -119,12 +136,13 @@ class TransformsConfig(BaseConfig):
     delete_acs_mask: bool = True
     delete_kspace: bool = True
     image_recon_type: ReconstructionType = ReconstructionType.RSS
+    compress_coils: Optional[int] = None
     pad_coils: Optional[int] = None
     use_seed: bool = True
     transforms_type: TransformsType = TransformsType.SUPERVISED
     # Next attributes are for the mask splitter in case of transforms_type is set to SSL_SSDU
-    mask_split_ratio: Tuple[float, ...] = (0.4,)
-    mask_split_acs_region: Tuple[int, int] = (0, 0)
+    mask_split_ratio: tuple[float, ...] = (0.4,)
+    mask_split_acs_region: tuple[int, int] = (0, 0)
     mask_split_keep_acs: Optional[bool] = False
     mask_split_type: MaskSplitterType = MaskSplitterType.GAUSSIAN
     mask_split_gaussian_std: float = 3.0
@@ -146,21 +164,21 @@ class H5SliceConfig(DatasetConfig):
     kspace_context: int = 0
     pass_mask: bool = False
     data_root: Optional[str] = None
-    filenames_filter: Optional[List[str]] = None
-    filenames_lists: Optional[List[str]] = None
+    filenames_filter: Optional[list[str]] = None
+    filenames_lists: Optional[list[str]] = None
     filenames_lists_root: Optional[str] = None
 
 
 @dataclass
 class CMRxReconConfig(DatasetConfig):
     regex_filter: Optional[str] = None
     data_root: Optional[str] = None
-    filenames_filter: Optional[List[str]] = None
-    filenames_lists: Optional[List[str]] = None
+    filenames_filter: Optional[list[str]] = None
+    filenames_lists: Optional[list[str]] = None
     filenames_lists_root: Optional[str] = None
     kspace_key: str = "kspace_full"
     compute_mask: bool = False
-    extra_keys: Optional[List[str]] = None
+    extra_keys: Optional[list[str]] = None
     kspace_context: Optional[str] = None
 
 
@@ -181,11 +199,11 @@ class FakeMRIBlobsConfig(DatasetConfig):
 
 @dataclass
 class SheppLoganDatasetConfig(DatasetConfig):
-    shape: Tuple[int, int, int] = (100, 100, 30)
+    shape: tuple[int, int, int] = (100, 100, 30)
     num_coils: int = 12
     seed: Optional[int] = None
     B0: float = 3.0
-    zlimits: Tuple[float, float] = (-0.929, 0.929)
+    zlimits: tuple[float, float] = (-0.929, 0.929)
 
 
 @dataclass