autodmri.scripts

  1import numpy as np
  2import nibabel as nib
  3
  4import os
  5import argparse
  6import logging
  7
  8from autodmri.estimator import estimate_from_dwis, estimate_from_nmaps
  9
 10
 11DESCRIPTION = """
 12Main script for automatically characterizing noise distributions.
 13"""
 14
 15EPILOG = """
 16Reference:
 17
 18MICCAI version
 19---------------
 20St-Jean S., De Luca A., Viergever M.A., Leemans A. (2018)
 21Automatic, Fast and Robust Characterization of Noise Distributions for Diffusion MRI, MICCAI 2018.
 22Springer International Publishing, pp. 304-312. doi: 10.1007/978-3-030-00928-1_35.
 23Available at: https://arxiv.org/abs/1805.12071
 24
 25Journal version
 26-----------------
 27Samuel St-Jean, Alberto De Luca, Chantal M.W. Tax, Max A. Viergever, Alexander Leemans,
 28Automated characterization of noise distributions in diffusion MRI data,
 29Medical Image Analysis, 2020, 101758, ISSN 1361-8415,
 30doi: 10.1016/j.media.2020.101758
 31Available at: http://www.sciencedirect.com/science/article/pii/S1361841520301225
 32"""
 33
 34
 35class CustomFormatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawTextHelpFormatter):
 36    pass
 37
 38
 39def buildArgsParser():
 40
 41    p = argparse.ArgumentParser(description=DESCRIPTION,
 42                                epilog=EPILOG,
 43                                formatter_class=CustomFormatter)
 44
 45    p.add_argument('data', metavar='input',
 46                   help='Path of the input data.')
 47
 48    p.add_argument('sigma', metavar='sigma',
 49                   help='Path of the output sigma volume.')
 50
 51    p.add_argument('N', metavar='N',
 52                   help='Path of the output N volume.')
 53
 54    p.add_argument('mask', metavar='mask',
 55                   help='Path of the output mask for voxels identified as noise.')
 56
 57    p.add_argument('-a', '--axis', type=int, default=-2, choices=[0, 1, 2],
 58                   help='Axis (0, 1 or 2 typically) which is assumed to contain uniform noise.')
 59
 60    p.add_argument('-m', '--method', default='moments', choices=['moments', 'maxlk'], metavar='string',
 61                   help='Method to use for estimating parameters, either "moments" or "maxlk".')
 62
 63    p.add_argument('--ncores', metavar='int', type=int, default=-1,
 64                   help='Number of cores to use for multithreading.')
 65
 66    p.add_argument('--exclude', metavar='file',
 67                   help='Mask indicating which voxels to exclude from the computation. Useful to remove gross artifacts.')
 68
 69    p.add_argument('--noise_maps', action='store_true',
 70                   help='Estimate in small windows instead of whole slices over the input volume. Only valid in theory for noise maps.')
 71
 72    p.add_argument('--subsample', action='store_true',
 73                   help='If supplied, estimate in non-overlapping windows with option --noise_maps.')
 74
 75    p.add_argument('--fast_median', action='store_true',
 76                   help='If supplied, computes the median of medians from each volume instead of one median value.\n'
 77                      'Useful for large datasets with many volumes (e.g. HCP) since the median requires a full copy of the data and sorting.')
 78
 79    p.add_argument('--size', metavar='int', type=int, default=5,
 80                   help='Size of the window for local noise maps estimation.')
 81
 82    p.add_argument('-f', '--force', action='store_true', dest='overwrite',
 83                   help='If set, overwrites the output text file if it already exists.')
 84
 85    p.add_argument('-v', '--verbose', action='store_true', dest='verbose',
 86                   help='If set, shows a progress bar during processing and prints useful information.')
 87
 88    p.add_argument('-l', '--log', dest='logfile', metavar='file',
 89                   help='Save the logging output to this file. Implies verbose output.')
 90
 91    return p
 92
 93
 94def main():
 95    parser = buildArgsParser()
 96    args = parser.parse_args()
 97
 98    logger = logging.getLogger('autodmri')
 99
100    if args.logfile is not None:
101        handler = logging.FileHandler(args.logfile)
102        args.verbose = True
103    else:
104        handler = logging.StreamHandler(args.logfile)
105
106    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s', "%Y-%m-%d %H:%M:%S")
107    handler.setFormatter(formatter)
108    logger.addHandler(handler)
109
110    if args.verbose:
111        logger.setLevel(logging.INFO)
112        logger.info('Verbosity is on')
113
114    overwritable_files = [args.sigma,
115                          args.N,
116                          args.mask]
117
118    for f in overwritable_files:
119        if f is not None and os.path.isfile(f):
120            if args.overwrite:
121                logger.warning(f'Overwriting {os.path.realpath(f)}')
122            else:
123                parser.error(f'{f} already exists! Use -f or --force to overwrite it.')
124
125    vol = nib.load(args.data)
126    data = vol.get_fdata(dtype=np.float32)
127    aff = vol.affine
128    hdr = vol.header
129
130    ncores = args.ncores
131    method = args.method
132    axis = args.axis
133    full = not args.subsample
134    size = args.size
135    noise_maps = args.noise_maps
136
137    if args.exclude is not None:
138        exclude_mask = nib.load(args.exclude).get_fdata().astype(bool)
139        logger.info(f'Excluding voxels from file {args.exclude}')
140    else:
141        exclude_mask = None
142
143    logger.info(f'Now estimating over file {args.data} with method = {method} and axis = {axis}')
144
145    if noise_maps:
146        if full:
147            overlap = 'overlapping windows'
148        else:
149            overlap = 'non-overlapping windows'
150
151        if data.ndim == 3:
152            data = data[..., None]
153
154        logger.info(f'Estimation will be done over noise maps with a window of size {size} and {overlap}')
155        sigma, N, mask = estimate_from_nmaps(data, size=size, return_mask=True, method=method, full=full, ncores=ncores, use_rejection=False, verbose=args.verbose)
156
157    else:
158        if axis < 0:
159            axis = data.ndim + axis
160
161        if args.fast_median:
162            logger.info('Estimation of the medians will be done over each volume, then on the median of the medians.')
163        elif data.shape[-1] > 100:
164            logger.warning(f'Estimation of the median will be done over the whole volume, but you have {data.shape[-1]} volumes.\n' +
165                           '\tConsider the option --fast_median if memory usage is high and startup time is too long.')
166
167        sigma, N, mask = estimate_from_dwis(data, axis=axis, return_mask=True, exclude_mask=exclude_mask, ncores=ncores,
168                                            method=method, verbose=args.verbose, fast_median=args.fast_median)
169
170        # Broadcast the 1D arrays to full 3D
171        if axis == 0:
172            sigma = sigma[:, None, None]
173            N = N[:, None, None]
174        elif axis == 1:
175            sigma = sigma[None, :, None]
176            N = N[None, :, None]
177        elif axis == 2:
178            sigma = sigma[None, None, :]
179            N = N[None, None, :]
180        else:
181            raise ValueError(f'axis = {axis} is not 0, 1 or 2, but that should never happen!')
182
183        sigma = np.ones(mask.shape) * sigma
184        N = np.ones(mask.shape) * N
185
186    # Save the data
187    logger.info(f'Output files are {args.sigma}, {args.N} and {args.mask}')
188    mask = mask.astype(np.int16)
189    sigma = sigma.astype(np.float32)
190    N = N.astype(np.float32)
191
192    nib.Nifti1Image(sigma, aff).to_filename(args.sigma)
193    nib.Nifti1Image(N, aff).to_filename(args.N)
194    nib.Nifti1Image(mask, aff).to_filename(args.mask)