feat: add adaptive sampling, motion registration, MEDL, and e2e ADS configs

direct/checkpointer.py

@@ -177,7 +177,9 @@ def _load_model(self, obj, state_dict):
         # Link has more elaborate checking for incompatibles in _log_incompatible_keys
         incompatible = obj.load_state_dict(state_dict, strict=False)
         if incompatible.missing_keys:
-            raise NotImplementedError
+            raise RuntimeError(
+                f"Missing keys when loading checkpoint into model {type(obj).__name__}: {incompatible.missing_keys}"
+            )
         if incompatible.unexpected_keys:
             self.logger.warning("Unexpected keys provided which cannot be loaded: %s.", incompatible.unexpected_keys)
 

direct/common/subsample.py

@@ -1418,10 +1418,12 @@ def __load_masks(self, acceleration: int) -> dict[tuple[int, int], tuple[np.ndar
             If the download fails.
         """
         masks_path = DIRECT_CACHE_DIR / "calgary_campinas_masks"
+        # Accelerations are stored as floats in BaseMaskFunc; filenames use integer R.
+        accel = int(acceleration)
         paths = [
-            f"R{acceleration}_218x170.npy",
-            f"R{acceleration}_218x174.npy",
-            f"R{acceleration}_218x180.npy",
+            f"R{accel}_218x170.npy",
+            f"R{accel}_218x174.npy",
+            f"R{accel}_218x180.npy",
         ]
 
         downloaded = [download_url(self.BASE_URL + _, masks_path, md5=self.MASK_MD5S[_]) is None for _ in paths]

direct/config/defaults.py

@@ -28,13 +28,20 @@ class TensorboardConfig(BaseConfig):
 @dataclass
 class LoggingConfig(BaseConfig):
     log_as_image: list[str] | None = None
+    # How often (in iterations) to flush scalars / write TensorBoard. Default: 20.
+    log_interval: int = 20
     tensorboard: TensorboardConfig = field(default_factory=TensorboardConfig)
 
 
 @dataclass
 class FunctionConfig(BaseConfig):
     function: str = MISSING
     multiplier: float = 1.0
+    # Optional tensor keys for loss comparison. When omitted, defaults are inferred
+    # from ``function`` (image → output_image/target, kspace → output_kspace/kspace,
+    # displacement_field → displacement_field/displacement_field).
+    source_key: str | None = None
+    target_key: str | None = None
 
 
 @dataclass
@@ -106,6 +113,7 @@ class ValidationConfig(BaseConfig):
 class InferenceConfig(BaseConfig):
     dataset: DatasetConfig = field(default_factory=DatasetConfig)
     batch_size: int = 1
+    metrics: list[str] = field(default_factory=list)
     crop: str | None = None
 
 
@@ -130,8 +138,10 @@ class DefaultConfig(BaseConfig):
 
     physics: PhysicsConfig = field(default_factory=PhysicsConfig)
 
-    training: TrainingConfig = field(default_factory=TrainingConfig)  # This should be optional.
-    validation: ValidationConfig = field(default_factory=ValidationConfig)  # This should be optional.
+    # Optional so inference-only YAMLs need not declare training/validation.
+    training: TrainingConfig | None = None
+    validation: ValidationConfig | None = None
+
     inference: InferenceConfig | None = None
 
     logging: LoggingConfig = field(default_factory=LoggingConfig)

direct/data/datasets.py

@@ -122,8 +122,11 @@ class FakeMRIBlobsDataset(Dataset):
         Pass the attributes of the generated sample.
     text_description: str
         Description of dataset, can be useful for logging.
-    kspace_context: bool
-        If true corresponds to 3D reconstruction, else reconstruction is 2D.
+    kspace_context: bool or str
+        If set (e.g. ``True`` or ``"time"``), each item is a full 3D / 2D+time
+        volume of shape ``(time_or_slice, coils, height, width)``. Otherwise
+        reconstruction is 2D (with optional per-slice indexing when
+        ``spatial_shape`` is 3D).
     """
 
     def __init__(
@@ -136,7 +139,7 @@ def __init__(
         filenames: list[str] | str | None = None,
         pass_attrs: bool | None = None,
         text_description: str | None = None,
-        kspace_context: bool | None = None,
+        kspace_context: bool | str | int | None = None,
         **kwargs,
     ) -> None:
         """Inits :class:`FakeMRIBlobsDataset`."""
@@ -157,8 +160,18 @@ def __init__(
         if self.text_description:
             self.logger.info("Dataset description: %s.", self.text_description)
 
+        # Volume mode: return full (T/S, coils, H, W) tensors for dynamic / registration.
+        self.kspace_context = kspace_context if kspace_context not in (None, 0, False, "") else 0
+        self.volume_mode = bool(self.kspace_context)
+        if self.volume_mode and len(spatial_shape) != 3:
+            raise ValueError(
+                "FakeMRIBlobsDataset volume mode (kspace_context set) requires "
+                f"spatial_shape (time_or_slice, height, width). Got {spatial_shape}."
+            )
+        self.ndim = 3 if self.volume_mode else 2
+
         self.fake_data: Callable = FakeMRIData(
-            ndim=len(self.spatial_shape),
+            ndim=3 if (self.volume_mode or len(self.spatial_shape) == 3) else 2,
             blobs_n_samples=kwargs.get("blobs_n_samples", None),
             blobs_cluster_std=kwargs.get("blobs_cluster_std", None),
         )
@@ -167,20 +180,18 @@ def __init__(
         self.rng = np.random.RandomState()
 
         with temp_seed(self.rng, seed):
-            # size = sample_size * num_slices if data is 3D
-            self.data = [
-                (filename, slice_no, seed)
-                for (filename, seed) in zip(
-                    self.parse_filenames_data(filenames),
-                    list(self.rng.choice(a=range(int(1e5)), size=self.sample_size, replace=False)),
-                )  # ensure reproducibility
-                for slice_no in range(self.spatial_shape[0] if len(spatial_shape) == 3 else 1)
-            ]
-        self.kspace_context = kspace_context if kspace_context else 0
-        self.ndim = 2 if self.kspace_context == 0 else 3
-
-        if self.kspace_context != 0:
-            raise NotImplementedError("3D reconstruction is not yet supported with FakeMRIBlobsDataset.")
+            filenames_parsed = self.parse_filenames_data(filenames)
+            seeds = list(self.rng.choice(a=range(int(1e5)), size=self.sample_size, replace=False))
+            if self.volume_mode:
+                # One dataset item per volume.
+                self.data = list(zip(filenames_parsed, seeds))
+            else:
+                # size = sample_size * num_slices if data is 3D (slice-wise 2D)
+                self.data = [
+                    (filename, slice_no, sample_seed)
+                    for (filename, sample_seed) in zip(filenames_parsed, seeds)
+                    for slice_no in range(self.spatial_shape[0] if len(spatial_shape) == 3 else 1)
+                ]
 
     def parse_filenames_data(self, filenames):
         if filenames is None:
@@ -198,7 +209,10 @@ def parse_filenames_data(self, filenames):
                 # pylint: disable=logging-fstring-interpolation
                 self.logger.info(f"Parsing: {(idx + 1) / len(filenames) * 100:.2f}%.")
 
-            num_slices = self.spatial_shape[0] if len(self.spatial_shape) == 3 else 1
+            if self.volume_mode:
+                num_slices = 1
+            else:
+                num_slices = self.spatial_shape[0] if len(self.spatial_shape) == 3 else 1
             self.volume_indices[pathlib.PosixPath(filename)] = range(
                 current_slice_number, current_slice_number + num_slices
             )
@@ -221,7 +235,11 @@ def __len__(self):
 
     def __getitem__(self, index: int) -> dict[str, Any]:
         """Get a sample from the dataset."""
-        filename, slice_no, sample_seed = self.data[index]
+        if self.volume_mode:
+            filename, sample_seed = self.data[index]  # ty: ignore[invalid-assignment]
+            slice_no = 0
+        else:
+            filename, slice_no, sample_seed = self.data[index]  # ty: ignore[invalid-assignment]
 
         sample = self.fake_data(
             sample_size=1,
@@ -230,7 +248,11 @@ def __getitem__(self, index: int) -> dict[str, Any]:
             name=[filename],
             seed=sample_seed,
         )[0]
-        sample["kspace"] = sample["kspace"][slice_no]
+        if self.volume_mode:
+            # FakeMRIData returns (time/slice, coils, H, W); pipeline expects (coils, time/slice, H, W).
+            sample["kspace"] = np.swapaxes(sample["kspace"], 0, 1)
+        else:
+            sample["kspace"] = sample["kspace"][slice_no]
 
         if "attrs" in sample:
             metadata = self._get_metadata(sample["attrs"])

direct/data/datasets_config.py

@@ -20,14 +20,17 @@
 from omegaconf import MISSING
 
 from direct.common.subsample_config import MaskingConfig
-from direct.config.defaults import BaseConfig
+from direct.config import BaseConfig
 from direct.data.mri_transforms import (
+    DemonsFilterType,
     HalfSplitType,
     MaskSplitterType,
     RandomFlipType,
     ReconstructionType,
+    RegistrationSimulateReferenceType,
     RescaleMode,
     SensitivityMapType,
+    TransformKey,
     TransformsType,
 )
 
@@ -73,6 +76,25 @@ class NormalizationTransformConfig(BaseConfig):
     scale_percentile: float | None = 0.99
 
 
+@dataclass
+class RegistrationTransformConfig(BaseConfig):
+    registration: bool = False
+    registration_simulate_reference: RegistrationSimulateReferenceType | None = None
+    registration_simulate_elastic_sigma: float = 3.0
+    registration_simulate_elastic_points: int = 3
+    registration_simulate_elastic_rotate: float = 0.0
+    registration_simulate_elastic_zoom: float = 0.0
+    registration_estimate_displacement: bool = True
+    registration_simulate_reference_from_key_index: int = 0
+    registration_moving_key: TransformKey = TransformKey.TARGET
+    demons_filter_type: DemonsFilterType = DemonsFilterType.SYMMETRIC_FORCES
+    demons_num_iterations: int = 100
+    demons_smooth_displacement_field: bool = True
+    demons_standard_deviations: float = 1.5
+    demons_intensity_difference_threshold: float | None = None
+    demons_maximum_rms_error: float | None = None
+
+
 @dataclass
 class TransformsConfig(BaseConfig):
     """Configuration for the transforms.
@@ -81,6 +103,8 @@ class TransformsConfig(BaseConfig):
     ----------
     masking : MaskingConfig
         Configuration for the masking.
+    target_acceleration : float, optional
+        Target acceleration to override the sampled acceleration with. Default is None.
     cropping : CropTransformConfig
         Configuration for the cropping.
     augmentation : AugmentationTransformConfig
@@ -95,6 +119,8 @@ class TransformsConfig(BaseConfig):
         Configuration for the sensitivity map estimation.
     normalization : NormalizationTransformConfig
         Configuration for the normalization.
+    use_acs_as_mask : bool
+        Use the ACS mask as the sampling mask. Default is False.
     delete_acs_mask : bool
         Delete ACS mask after its use. Default is True.
     delete_kspace : bool
@@ -107,6 +133,8 @@ class TransformsConfig(BaseConfig):
         Default is None.
     pad_coils : int, optional
         Pad coils. Default is None.
+    registration : RegistrationTransformConfig
+        Configuration for the registration transforms.
     use_seed : bool
         Use seed for the transforms. Typically this should be set to True for reproducibility (e.g. inference),
         and False for training. Default is True.
@@ -136,6 +164,9 @@ class TransformsConfig(BaseConfig):
     """
 
     masking: MaskingConfig | None = field(default_factory=MaskingConfig)
+    target_acceleration: float | None = None
+    # Paper adaptive DYNAMIC sampling: independent init/ACS mask per time/slice frame.
+    dynamic_mask: bool = False
     cropping: CropTransformConfig = field(default_factory=CropTransformConfig)
     augmentation: AugmentationTransformConfig = field(default_factory=AugmentationTransformConfig)
     random_augmentations: RandomAugmentationTransformsConfig = field(default_factory=RandomAugmentationTransformsConfig)
@@ -145,11 +176,13 @@ class TransformsConfig(BaseConfig):
         default_factory=SensitivityMapEstimationTransformConfig
     )
     normalization: NormalizationTransformConfig = field(default_factory=NormalizationTransformConfig)
+    use_acs_as_mask: bool = False
     delete_acs_mask: bool = True
     delete_kspace: bool = True
     image_recon_type: ReconstructionType = ReconstructionType.RSS
     compress_coils: int | None = None
     pad_coils: int | None = None
+    registration: RegistrationTransformConfig = field(default_factory=RegistrationTransformConfig)
     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
@@ -183,7 +216,6 @@ class H5SliceConfig(DatasetConfig):
 
 @dataclass
 class CMRxReconConfig(DatasetConfig):
-    regex_filter: str | None = None
     data_root: str | None = None
     filenames_filter: list[str] | None = None
     filenames_lists: list[str] | None = None
@@ -207,6 +239,8 @@ class CalgaryCampinasConfig(H5SliceConfig):
 @dataclass
 class FakeMRIBlobsConfig(DatasetConfig):
     pass_attrs: bool = True
+    # If set (e.g. True / "time"), each sample is a full volume (T/S, coils, H, W).
+    kspace_context: bool | str | int | None = None
 
 
 @dataclass