Пример #1
0
def test_tensor_algebra():
    # Test that the computation of tensor determinant and norm is correct
    test_arr = np.random.rand(10, 3, 3)
    t_det = dti.determinant(test_arr)
    t_norm = dti.norm(test_arr)
    for i, x in enumerate(test_arr):
        npt.assert_almost_equal(np.linalg.det(x), t_det[i])
        npt.assert_almost_equal(np.linalg.norm(x), t_norm[i])
Пример #2
0
def test_tensor_algebra():
    """
    Test that the computation of tensor determinant and norm is correct
    """
    test_arr = np.random.rand(10, 3, 3)
    t_det = dti.determinant(test_arr)
    t_norm = dti.norm(test_arr)
    for i, x in enumerate(test_arr):
        assert_almost_equal(np.linalg.det(x), t_det[i])
        assert_almost_equal(np.linalg.norm(x), t_norm[i])
Пример #3
0
def main():
    parser = _build_args_parser()
    args = parser.parse_args()

    if not args.not_all:
        args.fa = args.fa or 'fa.nii.gz'
        args.ga = args.ga or 'ga.nii.gz'
        args.rgb = args.rgb or 'rgb.nii.gz'
        args.md = args.md or 'md.nii.gz'
        args.ad = args.ad or 'ad.nii.gz'
        args.rd = args.rd or 'rd.nii.gz'
        args.mode = args.mode or 'mode.nii.gz'
        args.norm = args.norm or 'tensor_norm.nii.gz'
        args.tensor = args.tensor or 'tensor.nii.gz'
        args.evecs = args.evecs or 'tensor_evecs.nii.gz'
        args.evals = args.evals or 'tensor_evals.nii.gz'
        args.residual = args.residual or 'dti_residual.nii.gz'
        args.p_i_signal =\
            args.p_i_signal or 'physically_implausible_signals_mask.nii.gz'
        args.pulsation = args.pulsation or 'pulsation_and_misalignment.nii.gz'

    outputs = [args.fa, args.ga, args.rgb, args.md, args.ad, args.rd,
               args.mode, args.norm, args.tensor, args.evecs, args.evals,
               args.residual, args.p_i_signal, args.pulsation]
    if args.not_all and not any(outputs):
        parser.error('When using --not_all, you need to specify at least ' +
                     'one metric to output.')

    assert_inputs_exist(
        parser, [args.input, args.bvals, args.bvecs], args.mask)
    assert_outputs_exist(parser, args, outputs)

    img = nib.load(args.input)
    data = img.get_data()
    affine = img.get_affine()
    if args.mask is None:
        mask = None
    else:
        mask = nib.load(args.mask).get_data().astype(np.bool)

    # Validate bvals and bvecs
    logging.info('Tensor estimation with the %s method...', args.method)
    bvals, bvecs = read_bvals_bvecs(args.bvals, args.bvecs)

    if not is_normalized_bvecs(bvecs):
        logging.warning('Your b-vectors do not seem normalized...')
        bvecs = normalize_bvecs(bvecs)

    check_b0_threshold(args, bvals.min())
    gtab = gradient_table(bvals, bvecs, b0_threshold=bvals.min())

    # Get tensors
    if args.method == 'restore':
        sigma = ne.estimate_sigma(data)
        tenmodel = TensorModel(gtab, fit_method=args.method, sigma=sigma,
                               min_signal=_get_min_nonzero_signal(data))
    else:
        tenmodel = TensorModel(gtab, fit_method=args.method,
                               min_signal=_get_min_nonzero_signal(data))

    tenfit = tenmodel.fit(data, mask)

    FA = fractional_anisotropy(tenfit.evals)
    FA[np.isnan(FA)] = 0
    FA = np.clip(FA, 0, 1)

    if args.tensor:
        # Get the Tensor values and format them for visualisation
        # in the Fibernavigator.
        tensor_vals = lower_triangular(tenfit.quadratic_form)
        correct_order = [0, 1, 3, 2, 4, 5]
        tensor_vals_reordered = tensor_vals[..., correct_order]
        fiber_tensors = nib.Nifti1Image(
            tensor_vals_reordered.astype(np.float32), affine)
        nib.save(fiber_tensors, args.tensor)

    if args.fa:
        fa_img = nib.Nifti1Image(FA.astype(np.float32), affine)
        nib.save(fa_img, args.fa)

    if args.ga:
        GA = geodesic_anisotropy(tenfit.evals)
        GA[np.isnan(GA)] = 0

        ga_img = nib.Nifti1Image(GA.astype(np.float32), affine)
        nib.save(ga_img, args.ga)

    if args.rgb:
        RGB = color_fa(FA, tenfit.evecs)
        rgb_img = nib.Nifti1Image(np.array(255 * RGB, 'uint8'), affine)
        nib.save(rgb_img, args.rgb)

    if args.md:
        MD = mean_diffusivity(tenfit.evals)
        md_img = nib.Nifti1Image(MD.astype(np.float32), affine)
        nib.save(md_img, args.md)

    if args.ad:
        AD = axial_diffusivity(tenfit.evals)
        ad_img = nib.Nifti1Image(AD.astype(np.float32), affine)
        nib.save(ad_img, args.ad)

    if args.rd:
        RD = radial_diffusivity(tenfit.evals)
        rd_img = nib.Nifti1Image(RD.astype(np.float32), affine)
        nib.save(rd_img, args.rd)

    if args.mode:
        # Compute tensor mode
        inter_mode = dipy_mode(tenfit.quadratic_form)

        # Since the mode computation can generate NANs when not masked,
        # we need to remove them.
        non_nan_indices = np.isfinite(inter_mode)
        mode = np.zeros(inter_mode.shape)
        mode[non_nan_indices] = inter_mode[non_nan_indices]

        mode_img = nib.Nifti1Image(mode.astype(np.float32), affine)
        nib.save(mode_img, args.mode)

    if args.norm:
        NORM = norm(tenfit.quadratic_form)
        norm_img = nib.Nifti1Image(NORM.astype(np.float32), affine)
        nib.save(norm_img, args.norm)

    if args.evecs:
        evecs = tenfit.evecs.astype(np.float32)
        evecs_img = nib.Nifti1Image(evecs, affine)
        nib.save(evecs_img, args.evecs)

        # save individual e-vectors also
        e1_img = nib.Nifti1Image(evecs[..., 0], affine)
        e2_img = nib.Nifti1Image(evecs[..., 1], affine)
        e3_img = nib.Nifti1Image(evecs[..., 2], affine)

        nib.save(e1_img, add_filename_suffix(args.evecs, '_v1'))
        nib.save(e2_img, add_filename_suffix(args.evecs, '_v2'))
        nib.save(e3_img, add_filename_suffix(args.evecs, '_v3'))

    if args.evals:
        evals = tenfit.evals.astype(np.float32)
        evals_img = nib.Nifti1Image(evals, affine)
        nib.save(evals_img, args.evals)

        # save individual e-values also
        e1_img = nib.Nifti1Image(evals[..., 0], affine)
        e2_img = nib.Nifti1Image(evals[..., 1], affine)
        e3_img = nib.Nifti1Image(evals[..., 2], affine)

        nib.save(e1_img, add_filename_suffix(args.evals, '_e1'))
        nib.save(e2_img, add_filename_suffix(args.evals, '_e2'))
        nib.save(e3_img, add_filename_suffix(args.evals, '_e3'))

    if args.p_i_signal:
        S0 = np.mean(data[..., gtab.b0s_mask], axis=-1, keepdims=True)
        DWI = data[..., ~gtab.b0s_mask]
        pis_mask = np.max(S0 < DWI, axis=-1)

        if args.mask is not None:
            pis_mask *= mask

        pis_img = nib.Nifti1Image(pis_mask.astype(np.int16), affine)
        nib.save(pis_img, args.p_i_signal)

    if args.pulsation:
        STD = np.std(data[..., ~gtab.b0s_mask], axis=-1)

        if args.mask is not None:
            STD *= mask

        std_img = nib.Nifti1Image(STD.astype(np.float32), affine)
        nib.save(std_img, add_filename_suffix(args.pulsation, '_std_dwi'))

        if np.sum(gtab.b0s_mask) <= 1:
            logger.info('Not enough b=0 images to output standard '
                        'deviation map')
        else:
            if len(np.where(gtab.b0s_mask)) == 2:
                logger.info('Only two b=0 images. Be careful with the '
                            'interpretation of this std map')

            STD = np.std(data[..., gtab.b0s_mask], axis=-1)

            if args.mask is not None:
                STD *= mask

            std_img = nib.Nifti1Image(STD.astype(np.float32), affine)
            nib.save(std_img, add_filename_suffix(args.pulsation, '_std_b0'))

    if args.residual:
        # Mean residual image
        S0 = np.mean(data[..., gtab.b0s_mask], axis=-1)
        data_p = tenfit.predict(gtab, S0)
        R = np.mean(np.abs(data_p[..., ~gtab.b0s_mask] -
                           data[..., ~gtab.b0s_mask]), axis=-1)

        if args.mask is not None:
            R *= mask

        R_img = nib.Nifti1Image(R.astype(np.float32), affine)
        nib.save(R_img, args.residual)

        # Each volume's residual statistics
        if args.mask is None:
            logger.info("Outlier detection will not be performed, since no "
                        "mask was provided.")
        stats = [dict.fromkeys(['label', 'mean', 'iqr', 'cilo', 'cihi', 'whishi',
                                'whislo', 'fliers', 'q1', 'med', 'q3'], [])
                 for i in range(data.shape[-1])]  # stats with format for boxplots
        # Note that stats will be computed manually and plotted using bxp
        # but could be computed using stats = cbook.boxplot_stats
        # or pyplot.boxplot(x)
        R_k = np.zeros(data.shape[-1])    # mean residual per DWI
        std = np.zeros(data.shape[-1])  # std residual per DWI
        q1 = np.zeros(data.shape[-1])   # first quartile per DWI
        q3 = np.zeros(data.shape[-1])   # third quartile per DWI
        iqr = np.zeros(data.shape[-1])  # interquartile per DWI
        percent_outliers = np.zeros(data.shape[-1])
        nb_voxels = np.count_nonzero(mask)
        for k in range(data.shape[-1]):
            x = np.abs(data_p[..., k] - data[..., k])[mask]
            R_k[k] = np.mean(x)
            std[k] = np.std(x)
            q3[k], q1[k] = np.percentile(x, [75, 25])
            iqr[k] = q3[k] - q1[k]
            stats[k]['med'] = (q1[k] + q3[k]) / 2
            stats[k]['mean'] = R_k[k]
            stats[k]['q1'] = q1[k]
            stats[k]['q3'] = q3[k]
            stats[k]['whislo'] = q1[k] - 1.5 * iqr[k]
            stats[k]['whishi'] = q3[k] + 1.5 * iqr[k]
            stats[k]['label'] = k

            # Outliers are observations that fall below Q1 - 1.5(IQR) or
            # above Q3 + 1.5(IQR) We check if a voxel is an outlier only if
            # we have a mask, else we are biased.
            if args.mask is not None:
                outliers = (x < stats[k]['whislo']) | (x > stats[k]['whishi'])
                percent_outliers[k] = np.sum(outliers)/nb_voxels*100
                # What would be our definition of too many outliers?
                # Maybe mean(all_means)+-3SD?
                # Or we let people choose based on the figure.
                # if percent_outliers[k] > ???? :
                #    logger.warning('   Careful! Diffusion-Weighted Image'
                #                   ' i=%s has %s %% outlier voxels',
                #                   k, percent_outliers[k])

        # Saving all statistics as npy values
        residual_basename, _ = split_name_with_nii(args.residual)
        res_stats_basename = residual_basename + ".npy"
        np.save(add_filename_suffix(
            res_stats_basename, "_mean_residuals"), R_k)
        np.save(add_filename_suffix(res_stats_basename, "_q1_residuals"), q1)
        np.save(add_filename_suffix(res_stats_basename, "_q3_residuals"), q3)
        np.save(add_filename_suffix(res_stats_basename, "_iqr_residuals"), iqr)
        np.save(add_filename_suffix(res_stats_basename, "_std_residuals"), std)

        # Showing results in graph
        if args.mask is None:
            fig, axe = plt.subplots(nrows=1, ncols=1, squeeze=False)
        else:
            fig, axe = plt.subplots(nrows=1, ncols=2, squeeze=False,
                                    figsize=[10, 4.8])
            # Default is [6.4, 4.8]. Increasing width to see better.

        medianprops = dict(linestyle='-', linewidth=2.5, color='firebrick')
        meanprops = dict(linestyle='-', linewidth=2.5, color='green')
        axe[0, 0].bxp(stats, showmeans=True, meanline=True, showfliers=False,
                      medianprops=medianprops, meanprops=meanprops)
        axe[0, 0].set_xlabel('DW image')
        axe[0, 0].set_ylabel('Residuals per DWI volume. Red is median,\n'
                             'green is mean. Whiskers are 1.5*interquartile')
        axe[0, 0].set_title('Residuals')
        axe[0, 0].set_xticks(range(0, q1.shape[0], 5))
        axe[0, 0].set_xticklabels(range(0, q1.shape[0], 5))

        if args.mask is not None:
            axe[0, 1].plot(range(data.shape[-1]), percent_outliers)
            axe[0, 1].set_xticks(range(0, q1.shape[0], 5))
            axe[0, 1].set_xticklabels(range(0, q1.shape[0], 5))
            axe[0, 1].set_xlabel('DW image')
            axe[0, 1].set_ylabel('Percentage of outlier voxels')
            axe[0, 1].set_title('Outliers')
        plt.savefig(residual_basename + '_residuals_stats.png')
Пример #4
0
def main():
    parser = _build_args_parser()
    args = parser.parse_args()

    if not args.not_all:
        args.fa = args.fa or 'fa.nii.gz'
        args.ga = args.ga or 'ga.nii.gz'
        args.rgb = args.rgb or 'rgb.nii.gz'
        args.md = args.md or 'md.nii.gz'
        args.ad = args.ad or 'ad.nii.gz'
        args.rd = args.rd or 'rd.nii.gz'
        args.mode = args.mode or 'mode.nii.gz'
        args.norm = args.norm or 'tensor_norm.nii.gz'
        args.tensor = args.tensor or 'tensor.nii.gz'
        args.evecs = args.evecs or 'tensor_evecs.nii.gz'
        args.evals = args.evals or 'tensor_evals.nii.gz'
        args.residual = args.residual or 'dti_residual.nii.gz'
        args.p_i_signal =\
            args.p_i_signal or 'physically_implausible_signals_mask.nii.gz'
        args.pulsation = args.pulsation or 'pulsation_and_misalignment.nii.gz'

    outputs = [args.fa, args.ga, args.rgb, args.md, args.ad, args.rd,
               args.mode, args.norm, args.tensor, args.evecs, args.evals,
               args.residual, args.p_i_signal, args.pulsation]
    if args.not_all and not any(outputs):
        parser.error('When using --not_all, you need to specify at least ' +
                     'one metric to output.')

    assert_inputs_exist(
        parser, [args.input, args.bvals, args.bvecs], [args.mask])
    assert_outputs_exists(parser, args, outputs)

    img = nib.load(args.input)
    data = img.get_data()
    affine = img.get_affine()
    if args.mask is None:
        mask = None
    else:
        mask = nib.load(args.mask).get_data().astype(np.bool)

    # Validate bvals and bvecs
    logging.info('Tensor estimation with the %s method...', args.method)
    bvals, bvecs = read_bvals_bvecs(args.bvals, args.bvecs)

    if not is_normalized_bvecs(bvecs):
        logging.warning('Your b-vectors do not seem normalized...')
        bvecs = normalize_bvecs(bvecs)

    check_b0_threshold(args, bvals.min())
    gtab = gradient_table(bvals, bvecs, b0_threshold=bvals.min())

    # Get tensors
    if args.method == 'restore':
        sigma = ne.estimate_sigma(data)
        tenmodel = TensorModel(gtab, fit_method=args.method, sigma=sigma,
                               min_signal=_get_min_nonzero_signal(data))
    else:
        tenmodel = TensorModel(gtab, fit_method=args.method,
                               min_signal=_get_min_nonzero_signal(data))

    tenfit = tenmodel.fit(data, mask)

    FA = fractional_anisotropy(tenfit.evals)
    FA[np.isnan(FA)] = 0
    FA = np.clip(FA, 0, 1)

    if args.tensor:
        # Get the Tensor values and format them for visualisation
        # in the Fibernavigator.
        tensor_vals = lower_triangular(tenfit.quadratic_form)
        correct_order = [0, 1, 3, 2, 4, 5]
        tensor_vals_reordered = tensor_vals[..., correct_order]
        fiber_tensors = nib.Nifti1Image(
            tensor_vals_reordered.astype(np.float32), affine)
        nib.save(fiber_tensors, args.tensor)

    if args.fa:
        fa_img = nib.Nifti1Image(FA.astype(np.float32), affine)
        nib.save(fa_img, args.fa)

    if args.ga:
        GA = geodesic_anisotropy(tenfit.evals)
        GA[np.isnan(GA)] = 0

        ga_img = nib.Nifti1Image(GA.astype(np.float32), affine)
        nib.save(ga_img, args.ga)

    if args.rgb:
        RGB = color_fa(FA, tenfit.evecs)
        rgb_img = nib.Nifti1Image(np.array(255 * RGB, 'uint8'), affine)
        nib.save(rgb_img, args.rgb)

    if args.md:
        MD = mean_diffusivity(tenfit.evals)
        md_img = nib.Nifti1Image(MD.astype(np.float32), affine)
        nib.save(md_img, args.md)

    if args.ad:
        AD = axial_diffusivity(tenfit.evals)
        ad_img = nib.Nifti1Image(AD.astype(np.float32), affine)
        nib.save(ad_img, args.ad)

    if args.rd:
        RD = radial_diffusivity(tenfit.evals)
        rd_img = nib.Nifti1Image(RD.astype(np.float32), affine)
        nib.save(rd_img, args.rd)

    if args.mode:
        # Compute tensor mode
        inter_mode = dipy_mode(tenfit.quadratic_form)

        # Since the mode computation can generate NANs when not masked,
        # we need to remove them.
        non_nan_indices = np.isfinite(inter_mode)
        mode = np.zeros(inter_mode.shape)
        mode[non_nan_indices] = inter_mode[non_nan_indices]

        mode_img = nib.Nifti1Image(mode.astype(np.float32), affine)
        nib.save(mode_img, args.mode)

    if args.norm:
        NORM = norm(tenfit.quadratic_form)
        norm_img = nib.Nifti1Image(NORM.astype(np.float32), affine)
        nib.save(norm_img, args.norm)

    if args.evecs:
        evecs = tenfit.evecs.astype(np.float32)
        evecs_img = nib.Nifti1Image(evecs, affine)
        nib.save(evecs_img, args.evecs)

        # save individual e-vectors also
        e1_img = nib.Nifti1Image(evecs[..., 0], affine)
        e2_img = nib.Nifti1Image(evecs[..., 1], affine)
        e3_img = nib.Nifti1Image(evecs[..., 2], affine)

        nib.save(e1_img, add_filename_suffix(args.evecs, '_v1'))
        nib.save(e2_img, add_filename_suffix(args.evecs, '_v2'))
        nib.save(e3_img, add_filename_suffix(args.evecs, '_v3'))

    if args.evals:
        evals = tenfit.evals.astype(np.float32)
        evals_img = nib.Nifti1Image(evals, affine)
        nib.save(evals_img, args.evals)

        # save individual e-values also
        e1_img = nib.Nifti1Image(evals[..., 0], affine)
        e2_img = nib.Nifti1Image(evals[..., 1], affine)
        e3_img = nib.Nifti1Image(evals[..., 2], affine)

        nib.save(e1_img, add_filename_suffix(args.evals, '_e1'))
        nib.save(e2_img, add_filename_suffix(args.evals, '_e2'))
        nib.save(e3_img, add_filename_suffix(args.evals, '_e3'))

    if args.p_i_signal:
        S0 = np.mean(data[..., gtab.b0s_mask], axis=-1, keepdims=True)
        DWI = data[..., ~gtab.b0s_mask]
        pis_mask = np.max(S0 < DWI, axis=-1)

        if args.mask is not None:
            pis_mask *= mask

        pis_img = nib.Nifti1Image(pis_mask.astype(np.int16), affine)
        nib.save(pis_img, args.p_i_signal)

    if args.pulsation:
        STD = np.std(data[..., ~gtab.b0s_mask], axis=-1)

        if args.mask is not None:
            STD *= mask

        std_img = nib.Nifti1Image(STD.astype(np.float32), affine)
        nib.save(std_img, add_filename_suffix(args.pulsation, '_std_dwi'))

        if np.sum(gtab.b0s_mask) <= 1:
            logger.info('Not enough b=0 images to output standard '
                        'deviation map')
        else:
            if len(np.where(gtab.b0s_mask)) == 2:
                logger.info('Only two b=0 images. Be careful with the '
                            'interpretation of this std map')

            STD = np.std(data[..., gtab.b0s_mask], axis=-1)

            if args.mask is not None:
                STD *= mask

            std_img = nib.Nifti1Image(STD.astype(np.float32), affine)
            nib.save(std_img, add_filename_suffix(args.pulsation, '_std_b0'))

    if args.residual:
        if args.mask is None:
            logger.info("Outlier detection will not be performed, since no "
                        "mask was provided.")
        S0 = np.mean(data[..., gtab.b0s_mask], axis=-1)
        data_p = tenfit.predict(gtab, S0)
        R = np.mean(np.abs(data_p[..., ~gtab.b0s_mask] -
                           data[..., ~gtab.b0s_mask]), axis=-1)

        if args.mask is not None:
            R *= mask

        R_img = nib.Nifti1Image(R.astype(np.float32), affine)
        nib.save(R_img, args.residual)

        R_k = np.zeros(data.shape[-1])  # mean residual per DWI
        std = np.zeros(data.shape[-1])  # std residual per DWI
        q1 = np.zeros(data.shape[-1])   # first quartile
        q3 = np.zeros(data.shape[-1])   # third quartile
        iqr = np.zeros(data.shape[-1])  # interquartile
        for i in range(data.shape[-1]):
            x = np.abs(data_p[..., i] - data[..., i])[mask]
            R_k[i] = np.mean(x)
            std[i] = np.std(x)
            q3[i], q1[i] = np.percentile(x, [75, 25])
            iqr[i] = q3[i] - q1[i]

            # Outliers are observations that fall below Q1 - 1.5(IQR) or
            # above Q3 + 1.5(IQR) We check if a volume is an outlier only if
            # we have a mask, else we are biased.
            if args.mask is not None and R_k[i] < (q1[i] - 1.5 * iqr[i]) \
                    or R_k[i] > (q3[i] + 1.5 * iqr[i]):
                logger.warning('WARNING: Diffusion-Weighted Image i=%s is an '
                               'outlier', i)

        residual_basename, _ = split_name_with_nii(args.residual)
        res_stats_basename = residual_basename + ".npy"
        np.save(add_filename_suffix(
            res_stats_basename, "_mean_residuals"), R_k)
        np.save(add_filename_suffix(res_stats_basename, "_q1_residuals"), q1)
        np.save(add_filename_suffix(res_stats_basename, "_q3_residuals"), q3)
        np.save(add_filename_suffix(res_stats_basename, "_iqr_residuals"), iqr)
        np.save(add_filename_suffix(res_stats_basename, "_std_residuals"), std)

        # To do: I would like to have an error bar with q1 and q3.
        # Now, q1 acts as a std
        dwi = np.arange(R_k[~gtab.b0s_mask].shape[0])
        plt.bar(dwi, R_k[~gtab.b0s_mask], 0.75,
                color='y', yerr=q1[~gtab.b0s_mask])
        plt.xlabel('DW image')
        plt.ylabel('Mean residuals +- q1')
        plt.title('Residuals')
        plt.savefig(residual_basename + '_residuals_stats.png')
Пример #5
0
def nonlinfit_fn(dwi, bvecs, bvals, base_name):
    import nibabel as nb
    import numpy as np
    import os.path as op
    import dipy.reconst.dti as dti
    from dipy.core.gradients import GradientTable

    dwi_img = nb.load(dwi)
    dwi_data = dwi_img.get_data()
    dwi_affine = dwi_img.get_affine()
    
    from dipy.segment.mask import median_otsu
    b0_mask, mask = median_otsu(dwi_data, 2, 4)
    # Mask the data so that tensors are not fit for
    # unnecessary voxels
    mask_img = nb.Nifti1Image(mask.astype(np.float32), dwi_affine)
    b0_imgs = nb.Nifti1Image(b0_mask.astype(np.float32), dwi_affine)
    b0_img = nb.four_to_three(b0_imgs)[0]

    out_mask_name = op.abspath(base_name + '_binary_mask.nii.gz')
    out_b0_name = op.abspath(base_name + '_b0_mask.nii.gz')
    nb.save(mask_img, out_mask_name)
    nb.save(b0_img, out_b0_name)

    # Load the gradient strengths and directions
    bvals = np.loadtxt(bvals)
    gradients = np.loadtxt(bvecs)

    # Dipy wants Nx3 arrays
    if gradients.shape[0] == 3:
        gradients = gradients.T
        assert(gradients.shape[1] == 3)

    # Place in Dipy's preferred format
    gtab = GradientTable(gradients)
    gtab.bvals = bvals

    # Fit the tensors to the data
    tenmodel = dti.TensorModel(gtab, fit_method="NLLS")
    tenfit = tenmodel.fit(dwi_data, mask)

    # Calculate the fit, fa, and md of each voxel's tensor
    tensor_data = tenfit.lower_triangular()
    print('Computing anisotropy measures (FA, MD, RGB)')
    from dipy.reconst.dti import fractional_anisotropy, color_fa

    evals = tenfit.evals.astype(np.float32)
    FA = fractional_anisotropy(np.abs(evals))
    FA = np.clip(FA, 0, 1)

    MD = dti.mean_diffusivity(np.abs(evals))
    norm = dti.norm(tenfit.quadratic_form)

    RGB = color_fa(FA, tenfit.evecs)

    evecs = tenfit.evecs.astype(np.float32)
    mode = tenfit.mode.astype(np.float32)
    mode = np.nan_to_num(mode)


    # Write tensor as a 4D Nifti image with the original affine
    tensor_fit_img = nb.Nifti1Image(tensor_data.astype(np.float32), dwi_affine)
    mode_img = nb.Nifti1Image(mode.astype(np.float32), dwi_affine)
    norm_img = nb.Nifti1Image(norm.astype(np.float32), dwi_affine)
    FA_img = nb.Nifti1Image(FA.astype(np.float32), dwi_affine)
    evecs_img = nb.Nifti1Image(evecs, dwi_affine)
    evals_img = nb.Nifti1Image(evals, dwi_affine)
    rgb_img = nb.Nifti1Image(np.array(255 * RGB, 'uint8'), dwi_affine)
    MD_img = nb.Nifti1Image(MD.astype(np.float32), dwi_affine)

    out_tensor_file = op.abspath(base_name + "_tensor.nii.gz")
    out_mode_file = op.abspath(base_name + "_mode.nii.gz")
    out_fa_file = op.abspath(base_name + "_fa.nii.gz")
    out_norm_file = op.abspath(base_name + "_norm.nii.gz")
    out_evals_file = op.abspath(base_name + "_evals.nii.gz")
    out_evecs_file = op.abspath(base_name + "_evecs.nii.gz")
    out_rgb_fa_file = op.abspath(base_name + "_rgb_fa.nii.gz")
    out_md_file = op.abspath(base_name + "_md.nii.gz")

    nb.save(rgb_img, out_rgb_fa_file)
    nb.save(norm_img, out_norm_file)
    nb.save(mode_img, out_mode_file)
    nb.save(tensor_fit_img, out_tensor_file)
    nb.save(evecs_img, out_evecs_file)
    nb.save(evals_img, out_evals_file)
    nb.save(FA_img, out_fa_file)
    nb.save(MD_img, out_md_file)
    print('Tensor fit image saved as {i}'.format(i=out_tensor_file))
    print('FA image saved as {i}'.format(i=out_fa_file))
    print('MD image saved as {i}'.format(i=out_md_file))
    return out_tensor_file, out_fa_file, out_md_file, \
        out_evecs_file, out_evals_file, out_rgb_fa_file, out_norm_file, \
        out_mode_file, out_mask_name, out_b0_name
def nonlinfit_fn(dwi, bvecs, bvals, base_name):
    import nibabel as nb
    import numpy as np
    import os.path as op
    import dipy.reconst.dti as dti
    from dipy.core.gradients import GradientTable

    dwi_img = nb.load(dwi)
    dwi_data = dwi_img.get_data()
    dwi_affine = dwi_img.get_affine()
    
    from dipy.segment.mask import median_otsu
    b0_mask, mask = median_otsu(dwi_data, 2, 4)
    # Mask the data so that tensors are not fit for
    # unnecessary voxels
    mask_img = nb.Nifti1Image(mask.astype(np.float32), dwi_affine)
    b0_imgs = nb.Nifti1Image(b0_mask.astype(np.float32), dwi_affine)
    b0_img = nb.four_to_three(b0_imgs)[0]

    out_mask_name = op.abspath(base_name + '_binary_mask.nii.gz')
    out_b0_name = op.abspath(base_name + '_b0_mask.nii.gz')
    nb.save(mask_img, out_mask_name)
    nb.save(b0_img, out_b0_name)

    # Load the gradient strengths and directions
    bvals = np.loadtxt(bvals)
    gradients = np.loadtxt(bvecs).T

    # Place in Dipy's preferred format
    gtab = GradientTable(gradients)
    gtab.bvals = bvals

    # Fit the tensors to the data
    tenmodel = dti.TensorModel(gtab, fit_method="NLLS")
    tenfit = tenmodel.fit(dwi_data, mask)

    # Calculate the fit, fa, and md of each voxel's tensor
    tensor_data = tenfit.lower_triangular()
    print('Computing anisotropy measures (FA, MD, RGB)')
    from dipy.reconst.dti import fractional_anisotropy, color_fa

    evals = tenfit.evals.astype(np.float32)
    FA = fractional_anisotropy(np.abs(evals))
    FA = np.clip(FA, 0, 1)

    MD = dti.mean_diffusivity(np.abs(evals))
    norm = dti.norm(tenfit.quadratic_form)

    RGB = color_fa(FA, tenfit.evecs)

    evecs = tenfit.evecs.astype(np.float32)
    mode = tenfit.mode.astype(np.float32)

    # Write tensor as a 4D Nifti image with the original affine
    tensor_fit_img = nb.Nifti1Image(tensor_data.astype(np.float32), dwi_affine)
    mode_img = nb.Nifti1Image(mode.astype(np.float32), dwi_affine)
    norm_img = nb.Nifti1Image(norm.astype(np.float32), dwi_affine)
    FA_img = nb.Nifti1Image(FA.astype(np.float32), dwi_affine)
    evecs_img = nb.Nifti1Image(evecs, dwi_affine)
    evals_img = nb.Nifti1Image(evals, dwi_affine)
    rgb_img = nb.Nifti1Image(np.array(255 * RGB, 'uint8'), dwi_affine)
    MD_img = nb.Nifti1Image(MD.astype(np.float32), dwi_affine)

    out_tensor_file = op.abspath(base_name + "_tensor.nii.gz")
    out_mode_file = op.abspath(base_name + "_mode.nii.gz")
    out_fa_file = op.abspath(base_name + "_fa.nii.gz")
    out_norm_file = op.abspath(base_name + "_norm.nii.gz")
    out_evals_file = op.abspath(base_name + "_evals.nii.gz")
    out_evecs_file = op.abspath(base_name + "_evecs.nii.gz")
    out_rgb_fa_file = op.abspath(base_name + "_rgb_fa.nii.gz")
    out_md_file = op.abspath(base_name + "_md.nii.gz")

    nb.save(rgb_img, out_rgb_fa_file)
    nb.save(norm_img, out_norm_file)
    nb.save(mode_img, out_mode_file)
    nb.save(tensor_fit_img, out_tensor_file)
    nb.save(evecs_img, out_evecs_file)
    nb.save(evals_img, out_evals_file)
    nb.save(FA_img, out_fa_file)
    nb.save(MD_img, out_md_file)
    print('Tensor fit image saved as {i}'.format(i=out_tensor_file))
    print('FA image saved as {i}'.format(i=out_fa_file))
    print('MD image saved as {i}'.format(i=out_md_file))
    return out_tensor_file, out_fa_file, out_md_file, \
        out_evecs_file, out_evals_file, out_rgb_fa_file, out_norm_file, \
        out_mode_file, out_mask_name, out_b0_name