5ttgen deep_atropos: Multiple changes

- Add capability take as input the concatenated tissue probability images as an alternative to the segmentation label image. - Change default allocation of brain stem to be the 5th volume, in line with 5ttgen hsvs behaviour; this can be overridden using the -white-stem option. - Change allocation of cerebellum to be subcortical grey matter. - Import updated script test data, and add tests for the new command. - Add RS to command author list.
MRtrix3 · Jan 8, 2025 · b90c14f · b90c14f
1 parent c9f746a
commit b90c14f
Show file tree

Hide file tree

Showing 8 changed files with 85 additions and 45 deletions.
diff --git a/python/mrtrix3/commands/5ttgen/__init__.py b/python/mrtrix3/commands/5ttgen/__init__.py
@@ -14,4 +14,4 @@
 # For more details, see http://www.mrtrix.org/.
 
 # pylint: disable=unused-variable
-ALGORITHMS = ['freesurfer', 'fsl', 'gif', 'hsvs', 'deep_atropos']
+ALGORITHMS = ['deep_atropos', 'freesurfer', 'fsl', 'gif', 'hsvs']
diff --git a/python/mrtrix3/commands/5ttgen/deep_atropos.py b/python/mrtrix3/commands/5ttgen/deep_atropos.py
@@ -1,6 +1,6 @@
-import os.path, shutil
+import os
 from mrtrix3 import MRtrixError
-from mrtrix3 import app, image, path, run
+from mrtrix3 import app, image, run
 
 def usage(base_parser, subparsers):
     parser = subparsers.add_parser('deep_atropos', parents=[base_parser])
@@ -12,57 +12,67 @@ def usage(base_parser, subparsers):
     parser.add_argument('output',
                        type=app.Parser.ImageOut(),
                        help='The output 5TT image')
-
-def check_deep_atropos_input(image_path):
-    dim = image.Header(image_path).size()
-    if len(dim) != 3:
-        raise MRtrixError(f'Image \'{str(image_path)}\' does not look like Deep Atropos segmentation (number of spatial dimensions is not 3)')
+    parser.add_argument('-white_stem',
+                        action='store_true',
+                        default=None,
+                        help='Classify the brainstem as white matter')
 
 def execute():
-    check_deep_atropos_input(app.ARGS.input)
-    run.command(['mrconvert', app.ARGS.input, 'input.mif'])
 
-    # Generate initial tissue-specific maps
-    run.command('mrcalc input.mif 1 -eq CSF.mif')
-    run.command('mrcalc input.mif 2 -eq cGM.mif')
-    run.command('mrcalc input.mif 3 -eq WM1.mif')
-    run.command('mrcalc input.mif 5 -eq WM2.mif')
-    run.command('mrcalc input.mif 6 -eq WM3.mif')
-    run.command('mrmath WM1.mif WM2.mif WM3.mif sum WM.mif')
-    run.command('mrcalc input.mif 4 -eq sGM.mif')
+    if app.ARGS.sgm_amyg_hipp:
+        app.warn('Option -sgm_amyg_hipp has no effect on deep_atropos algorithm')
 
-    # Run connected components on WM for cleanup
-    run.command('mrthreshold WM.mif - -abs 0.001 | '
-                'maskfilter - connect - -connectivity | '
-                'mrcalc 1 - 1 -gt -sub remove_unconnected_wm_mask.mif -datatype bit')
+    dim = image.Header(app.ARGS.input).size()
+    if not(len(dim) == 3 or (len(dim) == 4 and dim[3] == 7)):
+        raise MRtrixError(f'Image \'{str(app.ARGS.input)}\' does not look like Deep Atropos segmentation'
+                          f' (expected either a 3D image, or a 4D image with 7 volumes; input image is {dim})')
 
-    # Preserve CSF and handle volume fractions
-    run.command('mrcalc CSF.mif remove_unconnected_wm_mask.mif -mult csf_clean.mif')
-    run.command('mrcalc 1.0 csf_clean.mif -sub sGM.mif -min sgm_clean.mif')
-
-    # Calculate multiplier for volume fraction correction
-    run.command('mrcalc 1.0 csf_clean.mif sgm_clean.mif -add -sub cGM.mif WM.mif -add -div multiplier.mif')
-    run.command('mrcalc multiplier.mif -finite multiplier.mif 0.0 -if multiplier_noNAN.mif')
-
-    # Apply corrections
-    run.command('mrcalc cGM.mif multiplier_noNAN.mif -mult remove_unconnected_wm_mask.mif -mult cgm_clean.mif')
-    run.command('mrcalc WM.mif multiplier_noNAN.mif -mult remove_unconnected_wm_mask.mif -mult wm_clean.mif')
-
-    # Create empty pathological tissue map
-    run.command('mrcalc wm_clean.mif 0 -mul path.mif')
+    run.command(['mrconvert', app.ARGS.input, 'input.mif'])
+
+    if len(dim) == 3:
+        # Generate tissue-specific masks
+        run.command('mrcalc input.mif 2 -eq cGM.mif')
+        run.command('mrcalc input.mif 4 -eq input.mif 6 -eq -add sGM.mif')
+        run.command(f'mrcalc input.mif 3 -eq{" input.mif 5 -eq -add" if app.ARGS.white_stem else ""} WM.mif')
+        run.command('mrcalc input.mif 1 -eq CSF.mif')
+        run.command(f'mrcalc input.mif {"0 -mult" if app.ARGS.white_stem else "5 -eq"} path.mif')
+    else:
+        # Brain mask = non-brain probability is <50%
+        run.command('mrconvert input.mif -coord 3 0 -axes 0,1,2 - | '
+                    'mrthreshold - -abs 0.5 -comparison le mask.mif')
+        # Need to rescale model probabilities so that, excluding the non-brain component,
+        #   the sum across all tissues will be 1.0
+        run.command('mrconvert input.mif -coord 3 1:end - | '
+                    'mrmath - sum -axis 3 - | '
+                    'mrcalc 1 - -div multiplier.mif')
+        # Generate tissue-specific probability maps
+        run.command('mrconvert input.mif -coord 3 2 -axes 0,1,2 cGM.mif')
+        run.command('mrconvert input.mif -coord 3 4,6 - | '
+                    'mrmath - sum -axis 3 sGM.mif')
+        if app.ARGS.white_stem:
+            run.command(f'mrconvert input.mif -coord 3 3,5 - | '
+                        'mrmath - sum -axis 3 WM.mif')
+        else:
+            run.command('mrconvert input.mif -coord 3 3 -axes 0,1,2 WM.mif')
+        run.command('mrconvert input.mif -coord 3 1 -axes 0,1,2 CSF.mif')
+        if app.ARGS.white_stem:
+            run.command('mrcalc cGM.mif 0 -mult path.mif')
+        else:
+            run.command('mrconvert input.mif -coord 3 5 -axes 0,1,2 path.mif')
 
-    # Combine into 5TT format
-    run.command('mrcat cgm_clean.mif sgm_clean.mif wm_clean.mif csf_clean.mif path.mif - -axis 3 | '
+    # Concatenate into the 5TT image
+    run.command('mrcat cGM.mif sGM.mif WM.mif CSF.mif path.mif - -axis 3 | '
+                f'{"mrcalc - multiplier.mif -mult mask.mif -mult - | " if len(dim) == 4 else ""}'
                 'mrconvert - combined_precrop.mif -strides +2,+3,+4,+1')
 
     # Apply cropping unless disabled
     if app.ARGS.nocrop:
         run.function(os.rename, 'combined_precrop.mif', 'result.mif')
     else:
         run.command('mrmath combined_precrop.mif sum - -axis 3 | '
-                   'mrthreshold - - -abs 0.5 | '
-                   'mrgrid combined_precrop.mif crop result.mif -mask -')
+                    'mrthreshold - - -abs 0.5 | '
+                    'mrgrid combined_precrop.mif crop result.mif -mask -')
 
     run.command(['mrconvert', 'result.mif', app.ARGS.output],
                 mrconvert_keyval=app.ARGS.input,
-                force=app.FORCE_OVERWRITE)
+                force=app.FORCE_OVERWRITE)
diff --git a/testing/CMakeLists.txt b/testing/CMakeLists.txt
@@ -16,7 +16,7 @@ include(ExternalProject)
 ExternalProject_Add(BinariesTestData
     PREFIX ${CMAKE_CURRENT_BINARY_DIR}/binaries_data
     GIT_REPOSITORY ${mrtrix_binaries_data_url}
-    GIT_TAG 2169ebc06040a0b1380017f5f2a11d6380c69922 
+    GIT_TAG 2169ebc06040a0b1380017f5f2a11d6380c69922
     GIT_PROGRESS TRUE
     CONFIGURE_COMMAND ""
     BUILD_COMMAND ""
@@ -27,19 +27,19 @@ ExternalProject_Add(BinariesTestData
 ExternalProject_Add(ScriptsTestData
     PREFIX ${CMAKE_CURRENT_BINARY_DIR}/scripts_data
     GIT_REPOSITORY ${mrtrix_scripts_data_url}
-    GIT_TAG 76f47633cd0a37e901c42320f4540ecaffd51367
+    GIT_TAG 7f3dae1e1bbbb383d710c0db66f469b5f812a298
     GIT_PROGRESS TRUE
     CONFIGURE_COMMAND ""
     BUILD_COMMAND ""
     INSTALL_COMMAND ""
     LOG_DOWNLOAD ON
 )
 
-set(BINARY_DATA_DIR 
+set(BINARY_DATA_DIR
     ${CMAKE_CURRENT_BINARY_DIR}/binaries_data/src/BinariesTestData
 )
 
-set(SCRIPT_DATA_DIR 
+set(SCRIPT_DATA_DIR
     ${CMAKE_CURRENT_BINARY_DIR}/scripts_data/src/ScriptsTestData
 )
 

diff --git a/testing/scripts/CMakeLists.txt b/testing/scripts/CMakeLists.txt
@@ -27,6 +27,10 @@ function(add_bash_script_test file_path)
     )
 endfunction()
 
+add_bash_script_test(5ttgen/deepatropos_fromseg_default "pythonci")
+add_bash_script_test(5ttgen/deepatropos_fromseg_whitestem)
+add_bash_script_test(5ttgen/deepatropos_fromprob_default "pythonci")
+add_bash_script_test(5ttgen/deepatropos_fromprob_whitestem)
 add_bash_script_test(5ttgen/freesurfer_default "pythonci")
 add_bash_script_test(5ttgen/freesurfer_nocrop)
 add_bash_script_test(5ttgen/freesurfer_piping)

diff --git a/testing/scripts/tests/5ttgen/deepatropos_fromprob_default b/testing/scripts/tests/5ttgen/deepatropos_fromprob_default
@@ -0,0 +1,6 @@
+#!/bin/bash
+# Verify default operation of "5ttgen deep_atropos"
+#   where the input is the concatenation of tissue probability images
+# Outcome is compared to that generated using a prior software version
+5ttgen deep_atropos 5ttgen/deep_atropos/probability_images.nii.gz tmp.mif -force
+testing_diff_image tmp.mif 5ttgen/deep_atropos/fromprob_default.mif.gz
diff --git a/testing/scripts/tests/5ttgen/deepatropos_fromprob_whitestem b/testing/scripts/tests/5ttgen/deepatropos_fromprob_whitestem
@@ -0,0 +1,7 @@
+#!/bin/bash
+# Verify default operation of "5ttgen deep_atropos"
+#   where the input is the concatenation of tissue probability images
+#   and the brain stem is allocated to WM rather than 5th volume
+# Outcome is compared to that generated using a prior software version
+5ttgen deep_atropos 5ttgen/deep_atropos/probability_images.nii.gz tmp.mif -white_stem -force
+testing_diff_image tmp.mif 5ttgen/deep_atropos/fromprob_whitestem.mif.gz
diff --git a/testing/scripts/tests/5ttgen/deepatropos_fromseg_default b/testing/scripts/tests/5ttgen/deepatropos_fromseg_default
@@ -0,0 +1,6 @@
+#!/bin/bash
+# Verify default operation of "5ttgen deep_atropos"
+#   where input image is the segmentation label image
+# Outcome is compared to that generated using a prior software version
+5ttgen deep_atropos 5ttgen/deep_atropos/segmentation_image.nii.gz tmp.mif -force
+testing_diff_image tmp.mif 5ttgen/deep_atropos/fromseg_default.mif.gz
diff --git a/testing/scripts/tests/5ttgen/deepatropos_fromseg_whitestem b/testing/scripts/tests/5ttgen/deepatropos_fromseg_whitestem
@@ -0,0 +1,7 @@
+#!/bin/bash
+# Verify default operation of "5ttgen deep_atropos"
+#   where input image is the segmentation label image
+#   and the brain stem is allocated to WM rather than 5th volume
+# Outcome is compared to that generated using a prior software version
+5ttgen deep_atropos 5ttgen/deep_atropos/segmentation_image.nii.gz tmp.mif -white_stem -force
+testing_diff_image tmp.mif 5ttgen/deep_atropos/fromseg_whitestem.mif.gz