Fix inference YAMLs for joint registration models

Enable registration transforms (FROM_KEY reference frame) and DYNAMIC
masking for phase/LOUPE policies so `direct predict` matches training.

README.md

@@ -83,6 +83,17 @@ direct predict ./predictions \
 
 The first argument is the **prediction output directory**.
 
+These models include a `registration_model`, so inference YAMLs enable registration transforms that build a `reference_image` (default: drop frame index `6` via `FROM_KEY`). Volumes must have enough temporal frames for that index. To use a different reference frame, edit:
+
+```yaml
+transforms:
+  registration:
+    registration: true
+    registration_simulate_reference: FROM_KEY
+    registration_simulate_reference_from_key_index: 6
+    registration_estimate_displacement: false
+```
+
 ## Citation
 
 If you use these models or [DIRECT](https://github.com/NKI-AI/direct), please cite the DIRECT toolkit and the method paper(s) below.

varnet_ads_1d_phase_reg.yaml

@@ -14,8 +14,8 @@ additional_models:
         slope: 10.0
         fix_sign_leakage: true
         kspace_shape:
-            - 512
-            - 246
+        - 512
+        - 246
         num_layers: 1
         kspace_sampler: false
         sampler_detach_mask: false
@@ -57,9 +57,10 @@ inference:
             masking:
                 name: FastMRIEquispaced
                 accelerations:
-                    - 4.0327
+                - 4.0327
                 center_fractions:
-                    - 0.04
+                - 0.04
+                mode: DYNAMIC
             cropping:
                 image_center_crop: false
             sensitivity_map_estimation:
@@ -72,6 +73,11 @@ inference:
             padding_eps: 0
             augmentation:
                 pad:
-                    - 512
-                    - 246
+                - 512
+                - 246
+            registration:
+                registration: true
+                registration_simulate_reference: FROM_KEY
+                registration_simulate_reference_from_key_index: 6
+                registration_estimate_displacement: false
         kspace_context: time

varnet_ads_1d_reg.yaml

@@ -14,9 +14,9 @@ additional_models:
         slope: 10.0
         fix_sign_leakage: true
         kspace_shape:
-            - 11
-            - 512
-            - 246
+        - 11
+        - 512
+        - 246
         num_layers: 1
         kspace_sampler: false
         sampler_detach_mask: false
@@ -57,9 +57,9 @@ inference:
             masking:
                 name: FastMRIEquispaced
                 accelerations:
-                    - 4.0327
+                - 4.0327
                 center_fractions:
-                    - 0.04
+                - 0.04
             cropping:
                 image_center_crop: false
             sensitivity_map_estimation:
@@ -72,6 +72,11 @@ inference:
             padding_eps: 0
             augmentation:
                 pad:
-                    - 512
-                    - 246
+                - 512
+                - 246
+            registration:
+                registration: true
+                registration_simulate_reference: FROM_KEY
+                registration_simulate_reference_from_key_index: 6
+                registration_estimate_displacement: false
         kspace_context: time

vsharp_ads_1d_init_reg.yaml

@@ -6,15 +6,15 @@ model:
     no_parameter_sharing: true
     auxiliary_steps: -1
     initializer_channels:
-        - 32
-        - 32
-        - 64
-        - 64
+    - 32
+    - 32
+    - 64
+    - 64
     initializer_dilations:
-        - 1
-        - 1
-        - 2
-        - 4
+    - 1
+    - 1
+    - 2
+    - 4
     initializer_multiscale: 1
     initializer_activation: PRELU
     unet_num_filters: 16
@@ -31,9 +31,9 @@ additional_models:
         slope: 10.0
         fix_sign_leakage: true
         kspace_shape:
-            - 11
-            - 512
-            - 246
+        - 11
+        - 512
+        - 246
         num_layers: 1
         kspace_sampler: false
         sampler_detach_mask: false
@@ -72,9 +72,9 @@ inference:
             masking:
                 name: FastMRIEquispaced
                 accelerations:
-                    - 8.2
+                - 8.2
                 center_fractions:
-                    - 0.04
+                - 0.04
             cropping:
                 image_center_crop: false
             sensitivity_map_estimation:
@@ -87,6 +87,11 @@ inference:
             padding_eps: 0
             augmentation:
                 pad:
-                    - 512
-                    - 246
+                - 512
+                - 246
+            registration:
+                registration: true
+                registration_simulate_reference: FROM_KEY
+                registration_simulate_reference_from_key_index: 6
+                registration_estimate_displacement: false
         kspace_context: time

vsharp_ads_1d_init_reg_disjoint.yaml

@@ -6,15 +6,15 @@ model:
     no_parameter_sharing: true
     auxiliary_steps: -1
     initializer_channels:
-        - 32
-        - 32
-        - 64
-        - 64
+    - 32
+    - 32
+    - 64
+    - 64
     initializer_dilations:
-        - 1
-        - 1
-        - 2
-        - 4
+    - 1
+    - 1
+    - 2
+    - 4
     initializer_multiscale: 1
     initializer_activation: PRELU
     unet_num_filters: 16
@@ -31,9 +31,9 @@ additional_models:
         slope: 10.0
         fix_sign_leakage: true
         kspace_shape:
-            - 11
-            - 512
-            - 246
+        - 11
+        - 512
+        - 246
         num_layers: 1
         kspace_sampler: false
         sampler_detach_mask: false
@@ -73,9 +73,9 @@ inference:
             masking:
                 name: FastMRIEquispaced
                 accelerations:
-                    - 8.2
+                - 8.2
                 center_fractions:
-                    - 0.04
+                - 0.04
             cropping:
                 image_center_crop: false
             sensitivity_map_estimation:
@@ -88,6 +88,11 @@ inference:
             padding_eps: 0
             augmentation:
                 pad:
-                    - 512
-                    - 246
+                - 512
+                - 246
+            registration:
+                registration: true
+                registration_simulate_reference: FROM_KEY
+                registration_simulate_reference_from_key_index: 6
+                registration_estimate_displacement: false
         kspace_context: time

vsharp_ads_1d_phase_init_reg.yaml

@@ -6,15 +6,15 @@ model:
     no_parameter_sharing: true
     auxiliary_steps: -1
     initializer_channels:
-        - 32
-        - 32
-        - 64
-        - 64
+    - 32
+    - 32
+    - 64
+    - 64
     initializer_dilations:
-        - 1
-        - 1
-        - 2
-        - 4
+    - 1
+    - 1
+    - 2
+    - 4
     initializer_multiscale: 1
     initializer_activation: PRELU
     unet_num_filters: 16
@@ -31,8 +31,8 @@ additional_models:
         slope: 10.0
         fix_sign_leakage: true
         kspace_shape:
-            - 512
-            - 246
+        - 512
+        - 246
         num_layers: 1
         kspace_sampler: false
         sampler_detach_mask: false
@@ -72,9 +72,10 @@ inference:
             masking:
                 name: FastMRIEquispaced
                 accelerations:
-                    - 8.2
+                - 8.2
                 center_fractions:
-                    - 0.04
+                - 0.04
+                mode: DYNAMIC
             cropping:
                 image_center_crop: false
             sensitivity_map_estimation:
@@ -87,6 +88,11 @@ inference:
             padding_eps: 0
             augmentation:
                 pad:
-                    - 512
-                    - 246
+                - 512
+                - 246
+            registration:
+                registration: true
+                registration_simulate_reference: FROM_KEY
+                registration_simulate_reference_from_key_index: 6
+                registration_estimate_displacement: false
         kspace_context: time

vsharp_ads_1d_phase_init_reg_disjoint.yaml

@@ -6,15 +6,15 @@ model:
     no_parameter_sharing: true
     auxiliary_steps: -1
     initializer_channels:
-        - 32
-        - 32
-        - 64
-        - 64
+    - 32
+    - 32
+    - 64
+    - 64
     initializer_dilations:
-        - 1
-        - 1
-        - 2
-        - 4
+    - 1
+    - 1
+    - 2
+    - 4
     initializer_multiscale: 1
     initializer_activation: PRELU
     unet_num_filters: 16
@@ -31,8 +31,8 @@ additional_models:
         slope: 10.0
         fix_sign_leakage: true
         kspace_shape:
-            - 512
-            - 246
+        - 512
+        - 246
         num_layers: 1
         kspace_sampler: false
         sampler_detach_mask: false
@@ -73,9 +73,10 @@ inference:
             masking:
                 name: FastMRIEquispaced
                 accelerations:
-                    - 8.2
+                - 8.2
                 center_fractions:
-                    - 0.04
+                - 0.04
+                mode: DYNAMIC
             cropping:
                 image_center_crop: false
             sensitivity_map_estimation:
@@ -88,6 +89,11 @@ inference:
             padding_eps: 0
             augmentation:
                 pad:
-                    - 512
-                    - 246
+                - 512
+                - 246
+            registration:
+                registration: true
+                registration_simulate_reference: FROM_KEY
+                registration_simulate_reference_from_key_index: 6
+                registration_estimate_displacement: false
         kspace_context: time

vsharp_ads_1d_phase_reg.yaml

@@ -6,15 +6,15 @@ model:
     no_parameter_sharing: true
     auxiliary_steps: -1
     initializer_channels:
-        - 32
-        - 32
-        - 64
-        - 64
+    - 32
+    - 32
+    - 64
+    - 64
     initializer_dilations:
-        - 1
-        - 1
-        - 2
-        - 4
+    - 1
+    - 1
+    - 2
+    - 4
     initializer_multiscale: 1
     initializer_activation: PRELU
     unet_num_filters: 16
@@ -31,8 +31,8 @@ additional_models:
         slope: 10.0
         fix_sign_leakage: true
         kspace_shape:
-            - 512
-            - 246
+        - 512
+        - 246
         num_layers: 1
         kspace_sampler: false
         sampler_detach_mask: false
@@ -72,9 +72,9 @@ inference:
             masking:
                 name: KtUniform
                 accelerations:
-                    - 4.0327
+                - 4.0327
                 center_fractions:
-                    - 0.04
+                - 0.04
             cropping:
                 image_center_crop: false
             sensitivity_map_estimation:
@@ -87,6 +87,11 @@ inference:
             padding_eps: 0
             augmentation:
                 pad:
-                    - 512
-                    - 246
+                - 512
+                - 246
+            registration:
+                registration: true
+                registration_simulate_reference: FROM_KEY
+                registration_simulate_reference_from_key_index: 6
+                registration_estimate_displacement: false
         kspace_context: time

vsharp_ads_1d_phase_reg_disjoint.yaml

@@ -6,15 +6,15 @@ model:
     no_parameter_sharing: true
     auxiliary_steps: -1
     initializer_channels:
-        - 32
-        - 32
-        - 64
-        - 64
+    - 32
+    - 32
+    - 64
+    - 64
     initializer_dilations:
-        - 1
-        - 1
-        - 2
-        - 4
+    - 1
+    - 1
+    - 2
+    - 4
     initializer_multiscale: 1
     initializer_activation: PRELU
     unet_num_filters: 16
@@ -31,8 +31,8 @@ additional_models:
         slope: 10.0
         fix_sign_leakage: true
         kspace_shape:
-            - 512
-            - 246
+        - 512
+        - 246
         num_layers: 1
         kspace_sampler: false
         sampler_detach_mask: false
@@ -73,9 +73,10 @@ inference:
             masking:
                 name: FastMRIEquispaced
                 accelerations:
-                    - 4.0327
+                - 4.0327
                 center_fractions:
-                    - 0.04
+                - 0.04
+                mode: DYNAMIC
             cropping:
                 image_center_crop: false
             sensitivity_map_estimation:
@@ -88,6 +89,11 @@ inference:
             padding_eps: 0
             augmentation:
                 pad:
-                    - 512
-                    - 246
+                - 512
+                - 246
+            registration:
+                registration: true
+                registration_simulate_reference: FROM_KEY
+                registration_simulate_reference_from_key_index: 6
+                registration_estimate_displacement: false
         kspace_context: time

vsharp_ads_1d_reg.yaml

@@ -6,15 +6,15 @@ model:
     no_parameter_sharing: true
     auxiliary_steps: -1
     initializer_channels:
-        - 32
-        - 32
-        - 64
-        - 64
+    - 32
+    - 32
+    - 64
+    - 64
     initializer_dilations:
-        - 1
-        - 1
-        - 2
-        - 4
+    - 1
+    - 1
+    - 2
+    - 4
     initializer_multiscale: 1
     initializer_activation: PRELU
     unet_num_filters: 16
@@ -31,9 +31,9 @@ additional_models:
         slope: 10.0
         fix_sign_leakage: true
         kspace_shape:
-            - 11
-            - 512
-            - 246
+        - 11
+        - 512
+        - 246
         num_layers: 1
         kspace_sampler: false
         sampler_detach_mask: false
@@ -72,9 +72,9 @@ inference:
             masking:
                 name: FastMRIEquispaced
                 accelerations:
-                    - 4.0327
+                - 4.0327
                 center_fractions:
-                    - 0.04
+                - 0.04
             cropping:
                 image_center_crop: false
             sensitivity_map_estimation:
@@ -87,6 +87,11 @@ inference:
             padding_eps: 0
             augmentation:
                 pad:
-                    - 512
-                    - 246
+                - 512
+                - 246
+            registration:
+                registration: true
+                registration_simulate_reference: FROM_KEY
+                registration_simulate_reference_from_key_index: 6
+                registration_estimate_displacement: false
         kspace_context: time

vsharp_ads_1d_reg_disjoint.yaml

@@ -6,15 +6,15 @@ model:
     no_parameter_sharing: true
     auxiliary_steps: -1
     initializer_channels:
-        - 32
-        - 32
-        - 64
-        - 64
+    - 32
+    - 32
+    - 64
+    - 64
     initializer_dilations:
-        - 1
-        - 1
-        - 2
-        - 4
+    - 1
+    - 1
+    - 2
+    - 4
     initializer_multiscale: 1
     initializer_activation: PRELU
     unet_num_filters: 16
@@ -31,9 +31,9 @@ additional_models:
         slope: 10.0
         fix_sign_leakage: true
         kspace_shape:
-            - 11
-            - 512
-            - 246
+        - 11
+        - 512
+        - 246
         num_layers: 1
         kspace_sampler: false
         sampler_detach_mask: false
@@ -73,9 +73,9 @@ inference:
             masking:
                 name: FastMRIEquispaced
                 accelerations:
-                    - 4.0327
+                - 4.0327
                 center_fractions:
-                    - 0.04
+                - 0.04
             cropping:
                 image_center_crop: false
             sensitivity_map_estimation:
@@ -88,6 +88,11 @@ inference:
             padding_eps: 0
             augmentation:
                 pad:
-                    - 512
-                    - 246
+                - 512
+                - 246
+            registration:
+                registration: true
+                registration_simulate_reference: FROM_KEY
+                registration_simulate_reference_from_key_index: 6
+                registration_estimate_displacement: false
         kspace_context: time

vsharp_fixed_1d_phase_reg.yaml

@@ -6,15 +6,15 @@ model:
     no_parameter_sharing: true
     auxiliary_steps: -1
     initializer_channels:
-        - 32
-        - 32
-        - 64
-        - 64
+    - 32
+    - 32
+    - 64
+    - 64
     initializer_dilations:
-        - 1
-        - 1
-        - 2
-        - 4
+    - 1
+    - 1
+    - 2
+    - 4
     initializer_multiscale: 1
     initializer_activation: PRELU
     unet_num_filters: 16
@@ -50,9 +50,10 @@ inference:
             masking:
                 name: FastMRIEquispaced
                 accelerations:
-                    - 4.0327
+                - 4.0327
                 center_fractions:
-                    - 0.04
+                - 0.04
+                mode: DYNAMIC
             cropping:
                 image_center_crop: false
             sensitivity_map_estimation:
@@ -65,6 +66,11 @@ inference:
             padding_eps: 0
             augmentation:
                 pad:
-                    - 512
-                    - 246
+                - 512
+                - 246
+            registration:
+                registration: true
+                registration_simulate_reference: FROM_KEY
+                registration_simulate_reference_from_key_index: 6
+                registration_estimate_displacement: false
         kspace_context: time

vsharp_fixed_1d_reg.yaml

@@ -6,15 +6,15 @@ model:
     no_parameter_sharing: true
     auxiliary_steps: -1
     initializer_channels:
-        - 32
-        - 32
-        - 64
-        - 64
+    - 32
+    - 32
+    - 64
+    - 64
     initializer_dilations:
-        - 1
-        - 1
-        - 2
-        - 4
+    - 1
+    - 1
+    - 2
+    - 4
     initializer_multiscale: 1
     initializer_activation: PRELU
     unet_num_filters: 16
@@ -50,9 +50,9 @@ inference:
             masking:
                 name: FastMRIEquispaced
                 accelerations:
-                    - 4.0327
+                - 4.0327
                 center_fractions:
-                    - 0.04
+                - 0.04
             cropping:
                 image_center_crop: false
             sensitivity_map_estimation:
@@ -65,6 +65,11 @@ inference:
             padding_eps: 0
             augmentation:
                 pad:
-                    - 512
-                    - 246
+                - 512
+                - 246
+            registration:
+                registration: true
+                registration_simulate_reference: FROM_KEY
+                registration_simulate_reference_from_key_index: 6
+                registration_estimate_displacement: false
         kspace_context: time

vsharp_loupe_1d_phase_reg.yaml

@@ -6,15 +6,15 @@ model:
     no_parameter_sharing: true
     auxiliary_steps: -1
     initializer_channels:
-        - 32
-        - 32
-        - 64
-        - 64
+    - 32
+    - 32
+    - 64
+    - 64
     initializer_dilations:
-        - 1
-        - 1
-        - 2
-        - 4
+    - 1
+    - 1
+    - 2
+    - 4
     initializer_multiscale: 1
     initializer_activation: PRELU
     unet_num_filters: 16
@@ -31,7 +31,7 @@ additional_models:
         slope: 10.0
         fix_sign_leakage: true
         kspace_shape:
-            - 246
+        - 246
         num_time_steps: 11
         non_uniform: false
     sensitivity_model:
@@ -64,9 +64,10 @@ inference:
             masking:
                 name: FastMRIEquispaced
                 accelerations:
-                    - 4.0327
+                - 4.0327
                 center_fractions:
-                    - 0.04
+                - 0.04
+                mode: DYNAMIC
             cropping:
                 image_center_crop: false
             sensitivity_map_estimation:
@@ -79,6 +80,11 @@ inference:
             padding_eps: 0
             augmentation:
                 pad:
-                    - 512
-                    - 246
+                - 512
+                - 246
+            registration:
+                registration: true
+                registration_simulate_reference: FROM_KEY
+                registration_simulate_reference_from_key_index: 6
+                registration_estimate_displacement: false
         kspace_context: time

vsharp_loupe_1d_reg.yaml

@@ -6,15 +6,15 @@ model:
     no_parameter_sharing: true
     auxiliary_steps: -1
     initializer_channels:
-        - 32
-        - 32
-        - 64
-        - 64
+    - 32
+    - 32
+    - 64
+    - 64
     initializer_dilations:
-        - 1
-        - 1
-        - 2
-        - 4
+    - 1
+    - 1
+    - 2
+    - 4
     initializer_multiscale: 1
     initializer_activation: PRELU
     unet_num_filters: 16
@@ -31,7 +31,7 @@ additional_models:
         slope: 10.0
         fix_sign_leakage: true
         kspace_shape:
-            - 246
+        - 246
     sensitivity_model:
         model_name: unet.unet_2d.UnetModel2d
         in_channels: 2
@@ -62,9 +62,10 @@ inference:
             masking:
                 name: FastMRIEquispaced
                 accelerations:
-                    - 4.0327
+                - 4.0327
                 center_fractions:
-                    - 0.04
+                - 0.04
+                mode: DYNAMIC
             cropping:
                 image_center_crop: false
             sensitivity_map_estimation:
@@ -77,6 +78,11 @@ inference:
             padding_eps: 0
             augmentation:
                 pad:
-                    - 512
-                    - 246
+                - 512
+                - 246
+            registration:
+                registration: true
+                registration_simulate_reference: FROM_KEY
+                registration_simulate_reference_from_key_index: 6
+                registration_estimate_displacement: false
         kspace_context: time