Minor change

NKI-AI · Feb 1, 2024 · ef6f359 · ef6f359
1 parent d983ff0
commit ef6f359
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/direct/data/mri_transforms.py b/direct/data/mri_transforms.py
@@ -2266,14 +2266,22 @@ def build_mri_transforms(
             target_kspace_key=KspaceKey.MASKED_KSPACE,
         ),
     ]
-
+    if scaling_key == TransformKey.ACS_KSPACE:
+        mri_transforms += [
+            ApplyMask(
+                sampling_mask_key="acs_mask",
+                input_kspace_key=KspaceKey.KSPACE,
+                target_kspace_key=KspaceKey.ACS_KSPACE,
+            ),
+        ]
+
     mri_transforms += [
         ComputeScalingFactor(
             normalize_key=scaling_key,
             percentile=scale_percentile,
             scaling_factor_key=TransformKey.SCALING_FACTOR,
         ),
-        Normalize(scaling_factor_key=TransformKey.SCALING_FACTOR),
+        Normalize(scaling_factor_key=None if scale_percentile is None else TransformKey.SCALING_FACTOR),
     ]
 
     mri_transforms += [

diff --git a/direct/types.py b/direct/types.py
@@ -47,12 +47,14 @@ def __hash__(self) -> int:
 class KspaceKey(DirectEnum):
     KSPACE = "kspace"
     MASKED_KSPACE = "masked_kspace"
+    ACS_KSPACE = "acs_kspace"
 
 
 class TransformKey(DirectEnum):
     SENSITIVITY_MAP = "sensitivity_map"
     TARGET = "target"
     KSPACE = "kspace"
+    ACS_KSPACE = "acs_kspace"
     MASKED_KSPACE = "masked_kspace"
     SAMPLING_MASK = "sampling_mask"
     ACS_MASK = "acs_mask"