[ENH] Fix registration + use mask apply in tractoflow

modules/nf-neuro/image/applymask/main.nf

@@ -16,12 +16,14 @@ process IMAGE_APPLYMASK {
 
     script:
     def prefix = task.ext.prefix ?: "${meta.id}"
+    def suffix = task.ext.first_suffix ? task.ext.first_suffix + "_masked"  : "masked"
     def nthreads_mrtrix = task.ext.single_thread ? "-nthreads 0" : "-nthreads ${task.cpus}"
+    def data_type = task.ext.data_type ? "-datatype ${task.ext.data_type}" : "-datatype float32"
 
     """
     export OMP_NUM_THREADS=${task.ext.single_thread ? 1 : task.cpus}
 
-    mrcalc $image $mask -mult ${prefix}_masked.nii.gz -force -quiet ${nthreads_mrtrix}
+    mrcalc $image $mask -mult ${prefix}_${suffix}.nii.gz -force -quiet ${nthreads_mrtrix} ${data_type}
 
     cat <<-END_VERSIONS > versions.yml
     "${task.process}":
@@ -31,9 +33,9 @@ process IMAGE_APPLYMASK {
 
     stub:
     def prefix = task.ext.prefix ?: "${meta.id}"
-
+    def suffix = task.ext.first_suffix ? task.ext.first_suffix + "_masked"  : "masked"
     """
-    touch ${prefix}_masked.nii.gz
+    touch ${prefix}_${suffix}.nii.gz
 
     mrcalc -h
 

modules/nf-neuro/image/applymask/meta.yml

@@ -32,10 +32,20 @@ input:
         ontologies:
           - edam: http://edamontology.org/format_3989 # GZIP format
 args:
+  - first_suffix:
+      type: string
+      description: Suffix for the output image file. This will be used as the first part of the suffix, followed by "_masked".
+      default: ""
   - single_thread:
       type: boolean
       description: If true, the command will be run in single-threaded mode. By default, the command will use multiple threads based on the number of CPUs allocated to the task.
       default: false
+  - data_type:
+      type: string
+      description: |
+        Data type for the output image. Use the following format:
+        uint8, int16, int32, int64, float16, float32 or float64.
+      default: "float32"
 output:
   image:
     - - meta:

modules/nf-neuro/registration/anattodwi/main.nf

@@ -5,10 +5,11 @@ process REGISTRATION_ANATTODWI {
     container "scilus/scilus:2.2.2"
 
     input:
-        tuple val(meta), path(fixed_reference), path(moving_anat), path(metric)
+        tuple val(meta), path(fixed_reference), path(moving_anat), path(metric), path(fixed_mask), path(moving_mask)
 
     output:
         tuple val(meta), path("*_warped.nii.gz")                            , emit: anat_warped
+        tuple val(meta), path("*_warped_reference.nii.gz")                  , emit: fixed_warped
         tuple val(meta), path("*_forward1_affine.mat")                      , emit: forward_affine
         tuple val(meta), path("*_forward0_warp.nii.gz")                     , emit: forward_warp
         tuple val(meta), path("*_backward1_warp.nii.gz")                    , emit: backward_warp
@@ -25,15 +26,18 @@ process REGISTRATION_ANATTODWI {
 
     script:
     def prefix = task.ext.prefix ?: "${meta.id}"
+    def suffix = task.ext.suffix ? "${task.ext.suffix}_warped" : "warped"
     def run_qc = task.ext.run_qc as Boolean || false
+    def args = task.ext.args ?: ''
 
+    if (( task.ext.masking_strategy == "both" || task.ext.masking_strategy == "internal" ) && (fixed_mask || moving_mask)) args += " -x \"${fixed_mask ?: 'NULL'},${moving_mask ?: 'NULL'}\""
     """
     export ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS=${task.ext.single_thread ? 1 : task.cpus}
     export OMP_NUM_THREADS=${task.ext.single_thread ? 1 : task.cpus}
     export ANTS_RANDOM_SEED=${task.ext.ants_rng_seed ? task.ext.ants_rng_seed : "1234"}
 
     antsRegistration --dimensionality 3 --float 0\
-        --output [forward,warped.nii.gz]\
+        --output [forward,warped.nii.gz,InverseWarped.nii.gz]\
         --interpolation Linear --use-histogram-matching 0\
         --winsorize-image-intensities [0.005,0.995]\
         --initial-moving-transform [$fixed_reference,$moving_anat,1]\
@@ -49,14 +53,16 @@ process REGISTRATION_ANATTODWI {
         --metric MI[$fixed_reference,$moving_anat,1,32]\
         --metric CC[$metric,$moving_anat,1,4]\
         --convergence [50x25x10,1e-6,10] --shrink-factors 4x2x1\
-        --smoothing-sigmas 3x2x1
+        --smoothing-sigmas 3x2x1\
+        $args
 
-    moving_base=\$(basename "${moving_anat}")
+    moving_base=\$(basename $moving_anat .nii.gz)
     ext=\${moving_base#*.}
     moving_id=\${moving_base%.\${ext}}
     moving_id=\${moving_id#${prefix}_*}
 
-    mv warped.nii.gz ${prefix}_\${moving_id}_warped.nii.gz
+    mv warped.nii.gz ${prefix}_\${moving_id}_${suffix}.nii.gz
+    mv InverseWarped.nii.gz ${prefix}_warped_reference.nii.gz
     mv forward0GenericAffine.mat ${prefix}_forward1_affine.mat
     mv forward1Warp.nii.gz ${prefix}_forward0_warp.nii.gz
     mv forward1InverseWarp.nii.gz ${prefix}_backward1_warp.nii.gz
@@ -67,7 +73,7 @@ process REGISTRATION_ANATTODWI {
     ### ** QC ** ###
     if $run_qc; then
         # Extract dimensions.
-        dim=\$(mrinfo ${prefix}_\${moving_id}_warped.nii.gz -size)
+        dim=\$(mrinfo ${prefix}_\${moving_id}_${suffix}.nii.gz -size)
         read sagittal_dim coronal_dim axial_dim <<< "\${dim}"
 
         # Get middle slices.
@@ -83,7 +89,7 @@ process REGISTRATION_ANATTODWI {
         fixed_id=\${fixed_id#${prefix}_*}
 
         # Iterate over images.
-        for image in \${moving_id}_warped \${fixed_id}; do
+        for image in \${moving_id}_${suffix} \${fixed_id}; do
             mrconvert *\${image}.nii.gz \${image}_viz.nii.gz -stride -1,2,3
             scil_viz_volume_screenshot \${image}_viz.nii.gz \${image}_coronal.png \
                 --slices \$coronal_mid --axis coronal \$viz_params
@@ -109,11 +115,11 @@ process REGISTRATION_ANATTODWI {
 
         # Create GIF.
         convert -delay 10 -loop 0 -morph 10 \
-            \${moving_id}_warped_mosaic.png \${fixed_id}_mosaic.png \${moving_id}_warped_mosaic.png \
+            \${moving_id}_${suffix}_mosaic.png \${fixed_id}_mosaic.png \${moving_id}_${suffix}_mosaic.png \
             ${prefix}_registration_anattodwi_mqc.gif
 
         # Clean up.
-        rm \${moving_id}_warped_mosaic.png \${fixed_id}_mosaic.png
+        rm \${moving_id}_${suffix}_mosaic.png \${fixed_id}_mosaic.png
     fi
 
     cat <<-END_VERSIONS > versions.yml
@@ -128,18 +134,21 @@ process REGISTRATION_ANATTODWI {
     stub:
     def prefix = task.ext.prefix ?: "${meta.id}"
     def run_qc = task.ext.run_qc as Boolean || false
-
+    def suffix = task.ext.suffix ? "${task.ext.suffix}_warped" : "warped"
     """
     antsRegistration -h
     antsApplyTransforms -h
     mrconvert -h
     scil_viz_volume_screenshot -h
     convert -help .
 
-    moving_id=\$(basename $moving_anat .nii.gz)
+    moving_base=\$(basename $moving_anat .nii.gz)
+    ext=\${moving_base#*.}
+    moving_id=\${moving_base%.\${ext}}
     moving_id=\${moving_id#${prefix}_*}
 
-    touch ${prefix}_\${moving_id}_warped.nii.gz
+    touch ${prefix}_\${moving_id}_${suffix}.nii.gz
+    touch ${prefix}_warped_reference.nii.gz
     touch ${prefix}_forward1_affine.mat
     touch ${prefix}_forward0_warp.nii.gz
     touch ${prefix}_backward1_warp.nii.gz

modules/nf-neuro/registration/anattodwi/meta.yml

@@ -43,6 +43,15 @@ args:
       type: integer
       description: Random seed for ANTs registration. Setting a fixed seed can help with reproducibility of results.
       default: 1234
+  - masking_strategy:
+      type: string
+      choices:
+        - none
+        - apriori
+        - internal
+        - both
+      description: Masking strategy to use for registration. Options are 'none', 'apriori', 'internal' or 'both'.
+      default: "none"
 input:
   - - meta:
         type: map
@@ -70,6 +79,20 @@ input:
         mandatory: true
         ontologies:
           - edam: http://edamontology.org/format_4001 # NIFTI format
+    - fixed_mask:
+        type: file
+        description: Fixed image mask
+        pattern: "*.{nii,nii.gz}"
+        mandatory: false
+        ontologies:
+          - edam: http://edamontology.org/format_4001 # NIFTI format
+    - moving_mask:
+        type: file
+        description: Moving image mask
+        pattern: "*.{nii,nii.gz}"
+        mandatory: false
+        ontologies:
+          - edam: http://edamontology.org/format_4001 # NIFTI format
 output:
   anat_warped:
     - - meta:
@@ -83,6 +106,18 @@ output:
           pattern: "*_warped.nii.gz"
           ontologies:
             - edam: http://edamontology.org/format_4001 # NIFTI format
+  fixed_warped:
+    - - meta:
+          type: map
+          description: |
+            Groovy Map containing sample information
+            e.g. `[ id:'test', single_end:false ]`
+      - "*_warped_reference.nii.gz":
+          type: file
+          description: B0 warped to T1 space
+          pattern: "*_warped_reference.nii.gz"
+          ontologies:
+            - edam: http://edamontology.org/format_4001 # NIFTI format
   forward_affine:
     - - meta:
           type: map

modules/nf-neuro/registration/anattodwi/tests/main.nf.test

@@ -56,6 +56,59 @@ nextflow_process {
                 input[0] = ch_b0
                     .join(ch_t1w)
                     .join(ch_fa)
+                    .map{ meta, fixed, moving, fixed_metric -> [meta, fixed, moving, fixed_metric, [], []] }
+                """
+            }
+        }
+
+        then {
+            assertAll(
+                { assert process.success },
+                { assert snapshot(process.out).match() }
+            )
+        }
+    }
+
+    test("registration - anattodwi - fixed mask") {
+        config "./nextflow.config"
+        when {
+            process {
+                """
+                ch_split_test_data = LOAD_DATA.out.test_data_directory
+                    .branch{
+                        t1w: it.simpleName == "T1w"
+                        b0: it.simpleName == "b0"
+                        dti: it.simpleName == "dti"
+                    }
+                ch_t1w = ch_split_test_data.t1w.map{
+                    test_data_directory -> [
+                        [ id:'test', single_end:false ],
+                        file("\${test_data_directory}/T1w.nii.gz")
+                    ]
+                }
+                ch_b0 = ch_split_test_data.b0.map{
+                    test_data_directory -> [
+                        [ id:'test', single_end:false ],
+                        file("\${test_data_directory}/b0.nii.gz")
+                    ]
+                }
+                ch_fa = ch_split_test_data.dti.map{
+                    test_data_directory -> [
+                        [ id:'test', single_end:false ],
+                        file("\${test_data_directory}/fa.nii.gz")
+                    ]
+                }
+                ch_fixed_mask = ch_split_test_data.t1w.map{
+                    test_data_directory -> [
+                        [ id:'test', single_end:false ],
+                        file("\${test_data_directory}/T1w_mask.nii.gz")
+                    ]
+                }
+                input[0] = ch_b0
+                    .join(ch_t1w)
+                    .join(ch_fa)
+                    .join(ch_fixed_mask)
+                    .map{ meta, fixed, moving, fixed_metric, fixed_mask -> [meta, fixed, moving, fixed_metric, fixed_mask, []] }
                 """
             }
         }
@@ -101,6 +154,7 @@ nextflow_process {
                 input[0] = ch_b0
                     .join(ch_t1w)
                     .join(ch_fa)
+                    .map{ meta, fixed, moving, fixed_metric -> [meta, fixed, moving, fixed_metric, [], []] }
                 """
             }
         }