Diffusion Tensor Maps Enhancement from Reduced-Direction Acquisition

Code for DTI map enhancement preprocessing (FSL/HD-BET/MATLAB) and network training accompanying the M.Sc. thesis and manuscript.

Overview

This repository contains the implementation used to enhance scalar DTI-derived parametric maps derived from 16-direction (ND16) — fractional anisotropy (FA), mean diffusivity (MD), axial diffusivity (AD), and radial diffusivity (RD) — to approximate the quality of 32-direction (ND32) reference maps, using a modified 3D denoising convolutional neural network (3D-DnCNN).

The pipeline has two main parts:

• preprocessing/ — DWI preprocessing pipeline (distortion correction, motion correction, brain extraction, registration, intensity normalization, low-angular-resolution generation, tissue segmentation, diffusion tensor estimation, and data augmentation).
• model/ — Network architecture, custom loss function, and training configuration for the 3D-DnCNN model.

Preprocessing (preprocessing/)

Implemented in Python using FSL and HD-BET, with intensity normalization, low-angular-resolution generation, and direction selection implemented as in-house MATLAB functions invoked through the MATLAB Engine API for Python.

Step	Description	File
Distortion correction	Susceptibility-induced distortion correction using FSL TOPUP on merged AP/PA volumes	preprocessing/data_preprocess.py
Motion correction	Rigid-body (12 DOF) motion correction using FSL MCFLIRT	preprocessing/data_preprocess.py
Brain extraction	Skull-stripping of T1w and DWI volumes using HD-BET	preprocessing/data_preprocess.py
Registration	Affine (12 DOF) registration of DWI to T1w using FSL FLIRT, applied to full 4D data via MCFLIRT	preprocessing/data_preprocess.py
Intensity normalization	Volume-wise mean-intensity normalization (MATLAB)	preprocessing/normwithmask.m
AP/PA averaging	Voxel-wise averaging of normalized AP and PA datasets (MATLAB)	preprocessing/average.m
Direction selection (calibration)	One-time identification of the 16 directions in the 32-direction scheme with minimum angular difference to a reference 16-direction scheme (MATLAB)	preprocessing/loc_select.m
Low-angular-resolution generation	Subject-level subsampling of the 32-direction dataset to 16 directions using the calibrated indices (MATLAB)	preprocessing/subsample.m
Tissue segmentation & input assembly	CSF/GM/WM partial-volume segmentation of T1w using FSL FAST; WM-masking of T1w (wm_masked.m); registration and assembly of final network input/label volumes (inputall.m)	preprocessing/input_preparation.py, preprocessing/wm_masked.m, preprocessing/inputall.m
Tensor estimation	Diffusion tensor fitting (WLS) using FSL DTIFIT	preprocessing/data_preprocess.py
Data augmentation	Random rigid spatial transformations (rotation + translation) applied identically to input and label volumes (MATLAB)	preprocessing/augmentation.m

Network (model/)

Implemented in TensorFlow using the Keras functional API (model/8layers-DnCNN-with-skipconn.py).

• Architecture: 8-layer 3D convolutional network (3×3×3 kernels, 64 filters), with batch normalization, ReLU activation, and dropout, plus three skip connections (conv1→conv7, conv2→conv6, conv3→conv5) to stabilize training and preserve low-level features. Input channels: FA, MD, AD, RD, T1w, WM; output channels: enhanced FA, MD, AD, RD.
• Loss function: Weighted normalized root mean square error (NRMSE), computed over non-background voxels, with per-channel weights [1.0, 3.0, 2.0, 3.0] for FA, MD, AD, RD respectively.
• Training: Adam optimizer (lr = 1e-4), up to 200 epochs, with TensorBoard logging and checkpointing on best validation loss.

Requirements

External neuroimaging tools

• FSL (TOPUP, MCFLIRT, FLIRT, FAST, DTIFIT) — used via command-line calls in preprocessing/data_preprocess.py and preprocessing/input_preparation.py
• HD-BET — skull-stripping, used via command-line calls in the same scripts

MATLAB

• MATLAB R2017b or later (for native niftiinfo/niftiread/niftiwrite support)
• Image Processing Toolbox (for imrotate3, imtranslate used in augmentation.m)
• MATLAB Engine API for Python — required to call the .m functions from the Python preprocessing scripts

Python

• Python 3.8+
• See requirements.txt:

  tensorflow
  keras
  numpy
  nibabel
  mat73
  matplotlib
  scikit-image
  scikit-learn
  matlabengine
  

  Install with:

  pip install -r requirements.txt
  

Implementation

Run the preprocessing pipeline in the following order before training:

1. preprocessing/data_preprocess.py — distortion correction, motion correction, brain extraction, registration, intensity normalization, low-angular-resolution generation, and tensor estimation. Calls normwithmask.m, average.m, and subsample.m via the MATLAB Engine API.
2. preprocessing/input_preparation.py — tissue segmentation and assembly of the final network input/label volumes. Calls wm_masked.m and inputall.m via the MATLAB Engine API.
3. preprocessing/augmentation.m — applies random rigid spatial augmentation to the paired input/label volumes produced in step 2.

Once augmented training pairs are generated, train the network with:

4. model/8layers-DnCNN-with-skipconn.py — builds, compiles, and trains the 3D-DnCNN model on the prepared and augmented data.

Citation

If you use this code, please cite the accompanying thesis and manuscript (citation details to be added upon publication).