Refactor dwi2response brain mask derivation with shell subsets

docs/reference/commands/dwi2response.rst

@@ -30,6 +30,8 @@ https://mrtrix.readthedocs.io/en/3.0.8/constrained_spherical_deconvolution/respo
 Note that if the -mask command-line option is not specified, the MRtrix3 command dwi2mask will automatically be called to derive an initial voxel exclusion mask. More information on mask derivation from DWI data can be found at: 
 https://mrtrix.readthedocs.io/en/3.0.8/dwi_preprocessing/masking.html
 
+In the absence of a user-specified mask (option -mask), the whole DWI series will be used for derivation of the brain mask, even where only a subset of the DWI volumes is used for response function estimation (whether because the -shells option has been specified, or because the nominated algorithm operates on only a single b-value shell). If it is instead desired that the same subset of shells used for response function estimation also be used for brain mask derivation, then the user has two alternatives: either generate that subset of shells themselves---e.g. using the dwiextract command---and provide the result as the input to dwi2response; or generate a brain mask from that subset of shells and provide that mask via the -mask option.
+
 Options
 -------
 

python/mrtrix3/commands/dwi2response/dwi2response.py

@@ -38,6 +38,19 @@ def usage(cmdline): #pylint: disable=unused-variable
                           ' derive an initial voxel exclusion mask.'
                           ' More information on mask derivation from DWI data can be found at: \n'
                           f'https://mrtrix.readthedocs.io/en/{version.TAG}/dwi_preprocessing/masking.html')
+  cmdline.add_description('In the absence of a user-specified mask (option -mask),'
+                          ' the whole DWI series will be used for derivation of the brain mask,'
+                          ' even where only a subset of the DWI volumes is used for response function estimation'
+                          ' (whether because the -shells option has been specified,'
+                          ' or because the nominated algorithm operates on only a single b-value shell).'
+                          ' If it is instead desired that the same subset of shells used for response function estimation'
+                          ' also be used for brain mask derivation,'
+                          ' then the user has two alternatives:'
+                          ' either generate that subset of shells themselves---'
+                          'e.g. using the dwiextract command---'
+                          'and provide the result as the input to dwi2response;'
+                          ' or generate a brain mask from that subset of shells'
+                          ' and provide that mask via the -mask option.')
 
   # General options
   common_options = cmdline.add_argument_group('General dwi2response options')
@@ -95,14 +108,48 @@ def execute(): #pylint: disable=unused-variable
                       'either in image header, or using -grad / -fslgrad option')
 
   app.activate_scratch_dir()
+
+  # Determine whether only a subset of the DWI volumes is to be used for response function estimation;
+  #   this is the case if the user has explicitly requested a subset of shells,
+  #   or if the nominated algorithm operates on only a single b-value shell.
+  need_to_extract = bool(alg.NEEDS_SINGLE_SHELL or shells_option)
+  extract_option = shells_option + singleshell_option
+
+  # Import the user-provided mask, if one was given
+  if app.ARGS.mask:
+    app.console(f'Importing mask ({app.ARGS.mask})...')
+    run.command(['mrconvert', app.ARGS.mask, 'mask.mif', '-datatype', 'bit'],
+                show=False,
+                preserve_pipes=True)
+
   # Get standard input data into the scratch directory
-  if alg.NEEDS_SINGLE_SHELL or shells_option:
+  if need_to_extract and alg.SUPPORTS_MASK and not app.ARGS.mask:
+    # The user has requested that only a subset of the data be used for response function estimation
+    #   (or the nominated algorithm operates on only a single shell),
+    #   but no mask has been provided and therefore one must be derived.
+    # Import the whole DWI series so that brain mask derivation can make use of all available data,
+    #   derive the mask, and only then extract the requested subset of shells
+    #   for the purpose of response function estimation.
+    app.console(f'Importing DWI data ({app.ARGS.input})...')
+    run.command(['mrconvert', app.ARGS.input, 'dwi_full.mif', '-strides', '0,0,0,1']
+                + grad_import_option,
+                show=False,
+                preserve_pipes=True)
+    dwi2mask_algo = CONFIG['Dwi2maskAlgorithm']
+    app.console(f'Computing brain mask (dwi2mask {dwi2mask_algo})...')
+    run.command(f'dwi2mask {dwi2mask_algo} dwi_full.mif mask.mif', show=False)
+    app.console('Extracting requested b-value shells for response function estimation...')
+    run.command(['dwiextract', 'dwi_full.mif', 'dwi.mif'] + extract_option,
+                show=False)
+    app.cleanup('dwi_full.mif')
+  elif need_to_extract:
+    # Either a mask has been provided by the user, or the nominated algorithm makes no use of a mask;
+    #   the requested subset of shells can therefore be extracted immediately upon import.
     app.console(f'Importing DWI data ({app.ARGS.input}) and selecting b-values...')
     run.command(['mrconvert', app.ARGS.input, '-', '-strides', '0,0,0,1']
                 + grad_import_option
                 + ['|', 'dwiextract', '-', 'dwi.mif']
-                + shells_option
-                + singleshell_option,
+                + extract_option,
                 show=False,
                 preserve_pipes=True)
   else: # Don't discard b=0 in multi-shell algorithms
@@ -111,11 +158,10 @@ def execute(): #pylint: disable=unused-variable
                 + grad_import_option,
                 show=False,
                 preserve_pipes=True)
-  if app.ARGS.mask:
-    app.console(f'Importing mask ({app.ARGS.mask})...')
-    run.command(['mrconvert', app.ARGS.mask, 'mask.mif', '-datatype', 'bit'],
-                show=False,
-                preserve_pipes=True)
+    if alg.SUPPORTS_MASK and not app.ARGS.mask:
+      dwi2mask_algo = CONFIG['Dwi2maskAlgorithm']
+      app.console(f'Computing brain mask (dwi2mask {dwi2mask_algo})...')
+      run.command(f'dwi2mask {dwi2mask_algo} dwi.mif mask.mif', show=False)
 
   if alg.SUPPORTS_MASK:
     if app.ARGS.mask:
@@ -125,10 +171,6 @@ def execute(): #pylint: disable=unused-variable
         raise MRtrixError('Dimensions of provided mask image do not match DWI')
       if not (len(mask_header.size()) == 3 or (len(mask_header.size()) == 4 and mask_header.size()[3] == 1)):
         raise MRtrixError('Provided mask image needs to be a 3D image')
-    else:
-      dwi2mask_algo = CONFIG['Dwi2maskAlgorithm']
-      app.console(f'Computing brain mask (dwi2mask {dwi2mask_algo})...')
-      run.command(f'dwi2mask {dwi2mask_algo} dwi.mif mask.mif', show=False)
 
     if not image.statistics('mask.mif', mask='mask.mif').count:
       raise MRtrixError(('Provided' if app.ARGS.mask else 'Generated') + ' mask image does not contain any voxels')