具有 3D 循环卷积自动编码器的扩散 MRI 中的角超分辨率

该项目通过使用循环 CNN 来提高 dMRI 数据的角分辨率。

目录

• 安装
• 推理
• 训练
• 码头工人
• 球面谐波基线

安装

dMRI-RCNN可以通过 pip 安装：

pip install dmri-rcnn

要求

dMRI-RCNN使用TensorFlow作为深度学习架构。要在 TensorFlow中启用GPU ，您应该确保安装了适当的先决条件。

下面列出了此软件包的要求。

• tensorflow>=2.6.0
• numpy
• einops
• nibabel
• tqdm

推理

安装后，用于run_dmri_rcnn.py执行新 dMRI 卷的推断。下面列出了使用脚本的数据要求，以及可用于推理的命令行参数。

数据

要运行此脚本，需要以下格式的 dMRI 数据：

• 上下文 dMRI 文件。dMRI 数据用作模型中的上下文以推断其他体积
  • 文件格式：NIfTI
  • 单壳：仅包含一个 b 值。
  • 尺寸：(i, j, k, q_in).
    • (i, j, k)是数据的空间维度
    • q_inq 空间维度内的样本数。这可以是6、10或30，并且会影响使用哪个训练模型。
• 上下文 b 向量文件。上下文 dMRI 文件的相应 b 向量。
  • 文件格式：文本文件，空格分隔。
  • 3对应于x, y, zq 空间坐标的行
  • q_in对应于采样的 q 空间方向的列。q_in必须是6,10或30.
• 目标 b 向量文件。推断的 dMRI 数据的相应 b 向量。
  • 文件格式：文本文件，空格分隔。
  • 3对应于x, y, zq 空间坐标的行
  • q_out对应于采样的 q 空间方向的列。
• 脑膜文件。dMRI 数据的二进制脑掩码文件。
  • 文件格式：NIfTI
  • 尺寸：(i, j, k). 与 dMRI 数据中使用的空间维度相同。

该脚本将创建以下数据：

• 推断的 dMRI 文件。根据目标 b 向量定义的模型推断的 dMRI 体积。
  • 文件格式：NIfTI
  • 尺寸：(i, j, k, q_out).
    • q_outq 空间维度内的样本数。这可以是任何数字，但如果使用更高的数字将需要更多的 GPU 内存。

命令行

通过 调出以下帮助信息run_dmri_rcnn.py -h：

usage: `run_dmri_rcnn.py` [-h] -dmri_in DMRI_IN -bvec_in BVEC_IN -bvec_out BVEC_OUT -mask MASK -dmri_out DMRI_OUT -s {1000,2000,3000} [-m {1,3}] [-c] [-b BATCH_SIZE]

optional arguments:
  -h, --help            show this help message and exit
  -dmri_in DMRI_IN      Context dMRI NIfTI volume. Must be single-shell and contain q_in 3D volumes
  -bvec_in BVEC_IN      Context b-vector text file. Whitespace delimited with 3 rows and q_in columns
  -bvec_out BVEC_OUT    Target b-vector text file. Whitespace delimited with 3 rows and q_out columns
  -mask MASK            Brain mask NIfTI volume. Must have space spatial dimensions as dmri_in.
  -dmri_out DMRI_OUT    Inferred dMRI NIfTI volume. This will contain q_out inferred volumes.
  -s {1000,2000,3000}, --shell {1000,2000,3000}
                        Shell to perform inference with. Must be same shell as context/target dMRI and b-vectors
  -m {1,3}, --model-dim {1,3}
                        Model dimensionality, choose either 1 or 3. Default: 3.
  -c, --combined        Use combined shell model. Currently only applicable with 3D model and 10 q_in.
  -n, --norm            Perform normalisation using 99 percentile of data. Only implemented with --combined flag, and only for q_in = 10
  -b BATCH_SIZE, --batch-size BATCH_SIZE
                        Batch size to run model inference with.

NB权重默认下载并存储在其中~/.dmri_rcnn。要将权重存储在不同的目录中，请设置环境变量DMRI_RCNN_DIR="/your/custom/directory"

例子

以下示例b = 1000使用 3D dMRI RCNN 对HCP 数据执行推理。

$ run_dmri_rcnn.py -dmri_in context_dmri.nii.gz -bvec_in context_bvecs -bvec_out target_bvecs -mask brain_mask.nii.gz -dmri_out inferred_dmri.nii.gz -s 1000 -m 3

这个例子需要大约 2 分钟来推断NVIDIA RTX 3080.

要对 HCP 数据集之外的数据执行推理，请使用标志-c和-n. 这目前仅适用于 $q_{in} = 10$。

训练

以下是有关如何训练给定模型以及所涉及的预处理步骤的详细信息。

数据预处理

训练数据集通常太大而无法一次全部放入内存。为了克服这个问题，这个项目使用了 TensorFlow 的.tfrecord文件格式和 tf.data.Dataset API。因此，训练数据应该在开始之前以这种格式保存。以下是有关如何使用该dMRI-RCNN项目执行此操作的示例。

import numpy as np

from dmri_rcnn.core import io
from dmri_rcnn.core.processing import save_tfrecord_data

# First load a subject into memory
dmri, _ = io.load_nifti('/path/to/dmri/data.nii.gz')
mask, _ = io.load_nifti('/path/to/brain/mask.nii.gz', dtype=np.int8)
bvecs = io.load_bvec('/path/to/dmri/bvecs') # bvecs & bvals should be in FSL format
bvals = io.load_bval('/path/to/dmri/bvals')

# Optionally crop image data
dmri, mask = io.autocrop_dmri(dmri, mask)

# .tfrecord format uses a maximum filesize of 2 GiB, therefore for high
# resolution dMRI data, the image may need to be split into smaller parts
# to do this use the function below. It is recommended to first try to save
# each subject as a whole before splitting the image into separate files.
dmri_list = io.split_image_to_octants(dmri)
mask_list = io.split_image_to_octants(mask)

# Now save data in .tfrecord format
save_tfrecord_data(dmri, bvecs, bvals, mask, '/path/to/saved/data.tfrecord')

# Alternatively save the list of image parts if dmri is too large
for i in range(len(dmri_list)):
  save_tfrecord_data(dmri_list[i], bvecs, bvals, mask_list[i], '/path/to/saved/data' + str(i) + '.tfrecord')

训练模型

预处理完成后，您就可以训练模型了。

from dmri_rcnn.core.weights import get_weights
from dmri_rcnn.core.model import get_1d_autoencoder, get_3d_autoencoder
from dmri_rcnn.core.processing import TrainingProcessor, TrainingProcessorNorm

# If we want to fine-tune the model we can load the previously obtained weights.
# In this example we'll load the weights for the 3D RCNN trained on the b = 1000
# shell and 6 q-space samples per input.
weights = get_weights(model_dim=3, shell=1000, q_in=6)

# Now we can instantiate the pre-compiled 3D model
model = get_3d_autoencoder(weights) # Omit the weights argument to load without pre-trained weights

# Instantiate the training processor
processor = TrainingProcessor(shells=[1000], q_in=6)

# If using non-HCP data, the TrainingProcessorNorm should be used instead.
processor = TrainingProcessorNorm(shells=[1000], q_in=6)

# Important: Here our q_in = 6, and the processor uses a default q_out = 10, therefore our dmri data must
# contain at least 16 volumes.

# Load dataset mapping
train_data = processor.load_data(['/path/to/train_data0.tfrecord', '/path/to/train_data1.tfrecord'])
validation_data = processor.load_data(['/path/to/val_data0.tfrecord'], validation=True)

# Begin training
model.fit(train_data, epochs=10, validation_data=validation_data)

码头工人

您也可以dMRI-RCNN直接通过Docker使用。该项目的 CPU 和 GPU 版本均可用。

中央处理器

要dMRI-RCNN仅与 CPU 一起使用，请使用：

sudo docker run -v /absolute/path/to/my/data/directory:/data -it -t mlyon93/dmri-rcnn-cpu:latest

图形处理器

要dMRI-RCNN与 GPU 一起使用，首先确保已安装适当的 NVIDIA 先决条件。然后使用：

sudo docker run --gpus all -v /absolute/path/to/my/data/directory:/data -it -t mlyon93/dmri-rcnn-gpu:latest

球面谐波基线

要运行本文中使用的 Spherical Harmonic 基线模型，首先dipy要安装确保。您可以通过以下提示dipy直接安装或安装此项目。pip

pip install dmri-rcnn[sh]

命令行

通过 调出以下帮助信息dmri_sh_baseline.py -h：

usage: dMRI Spherical Harmonic Baseline Inference [-h] -dmri_in DMRI_IN -bvec_in BVEC_IN -bvec_out BVEC_OUT -dmri_out DMRI_OUT -s SHELL

optional arguments:
  -h, --help            show this help message and exit
  -dmri_in DMRI_IN      Context dMRI NIfTI volume. Must be single-shell and contain q_in 3D volumes
  -bvec_in BVEC_IN      Context b-vectory text file. Whitespace delimited with 3 rows and q_in columns
  -bvec_out BVEC_OUT    Target b-vector text file. Whitespace delimited with 3 rows and q_out columns
  -dmri_out DMRI_OUT    Inferred dMRI NIfTI volume. This will contain q_out inferred volumes.
  -s SHELL, --shell SHELL
                        Shell to perform inference on. Must be same shell as context/target dMRI and b-vecs

例子

以下示例执行b = 1000球谐函数推断。

$ dmri_sh_baseline.py -dmri_in context_dmri.nii.gz -bvec_in context_bvecs -bvec_out target_bvecs -dmri_out inferred_dmri.nii.gz -s 1000

使用或检查球谐模型，代码可以在dmri_rcnn.core.processing.sph_harmonic.

路线图

未来的补充和改进：

• 绘图功能