VoxelMorph-无监督医学图像配准-代码讲解

VoxelMorph官方Github地址：https://github.com/voxelmorph/voxelmorph，本文按照官方的Tutorial提供的路线进行讲解。

原文：Visit the VoxelMorph tutorial to learn about VoxelMorph and Learning-based Registration. Here’s an additional small tutorial on warping annotations together with images, and another on template (atlas) construction with VoxelMorph.

主要分为：

• Additional small tutorial 根据注释对图像进行变换
• VoxelMorph tutorial VoxelMorph教程
• Template (atlas) construction 模版搭建教程
  本文Github地址：https://github.com/MaybeRichard/VoxelMorph-explain

第一部分：Additional small tutorial 根据注释对图像进行变换

环境及背景介绍：

本部分的官方代码地址：https://colab.research.google.com/drive/1V0CutSIfmtgDJg1XIkEnGteJuw0u7qT-#scrollTo=h1KXYz-Nauwn
这一部分主要介绍的是如何使用vxm库里的方法对图像进行变换，代码中的方法是随机生成一个矩阵，然后根据该矩阵对图像进行仿射变换。
环境要求：tensorflow2.4,VoxelMorph
完整代码如下：

# 安装和导包
!pip install voxelmorph
import voxelmorph as vxm
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
 
# 对输入图像进行适当的预处理
pad_amt = 10
(x_train,_),_ = tf.keras.datasets.mnist.load_data()
# float64占用的内存是float32的两倍，是float16的4倍；比如对于CIFAR10数据集，如果采用float64来表示，需要60000323238/1024**3=1.4G，光把数据集调入内存就需要1.4G；如果采用float32，只需要0.7G，如果采用float16，只需要0.35G左右；占用内存的多少，会对系统运行效率有严重影响；（因此数据集文件都是采用uint8来存在数据，保持文件最小）
im = x_train[0,...].astype('float')/255
# np.pad(需要填充的array，((上，下),(左，右))，mode=constant...),这一步是为了增加边缘，累死padding，作用是防止后面的平移导致其超出范围
im = np.pad(im,((pad_amt,pad_amt),(pad_amt,pad_amt)))

# 手工创建变换矩阵
aff = np.eye(3) # 创建主对角矩阵
aff[:2,:2]+=np.random.randn(2,2)*0.1 # 在上半部分的2*2区域加入随机噪声
aff[:2, 2] = np.random.uniform(-10, 10, (2, )) # 均匀分布，(low,high,size) aff[:2, 2]的尺寸是(2,)
aff_inv = np.linalg.inv(aff)

margin=10
nb_annotations = 5
annotations = [np.random.uniform(margin,f-margin,nb_annotations) for f in im.shape] # 创建两个注释，(48,48)表示两个
annotations = np.stack(annotations,1)

# np.newaxis 的功能是增加新的维度，但是要注意 np.newaxis 放的位置不同，产生的矩阵形状也不同。放在第一个，给行上增加维度，放在最后一个，给列上增加维度
im_keras = im[np.newaxis,...,np.newaxis]
aff_keras = aff[np.newaxis,:2,:]
annotations_keras = annotations[np.newaxis,...]

# 进行仿射变换
im_warped = vxm.layers.SpatialTransformer()([im_keras, aff_keras])
im_warped = im_warped[0, ..., 0]

# get dense field of inverse affine
field_inv = vxm.utils.affine_to_dense_shift(aff_inv[:-1,:], im.shape, shift_center=True)[np.newaxis, ...]

# warp annotations
data = [tf.convert_to_tensor(f, dtype=tf.float32) for f in [annotations_keras, field_inv]]
annotations_warped = vxm.utils.point_spatial_transformer(data)[0, ...].numpy()

# 结果可视化
plt.figure()
# note that x and y need to be flipped due to xy indexing in matplotlib
plt.subplot(1, 2, 1)
plt.imshow(im, cmap='gray')
plt.plot(*[annotations[:, f] for f in [1, 0]], 'o')  
plt.axis('off')

plt.subplot(1, 2, 2)
plt.imshow(im_warped, cmap='gray')
plt.plot(*[annotations_warped[:, f] for f in [1, 0]], 'o')
plt.axis('off');
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54

代码分析与讲解：

1. 库的导入：

!pip install voxelmorph
import voxelmorph as vxm
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
1
2
3
4
5

2. 图像的输入与预处理：

教程中使用的是mnist的数据集，数据集的预处理步骤包括:

• 对原始图像进行边缘填充
• 灰度归一化（具体作用参考：https://blog.csdn.net/qq_41383956/article/details/88593538）

# 加载mnist数据集
# 其标准输出应为： (x_train, y_train), (x_test, y_test)，但是只需要x_train数据展示，所以其他的丢掉
(x_train,_),_ = tf.keras.datasets.mnist.load_data()
# 灰度归一化，从0-255压缩到0-1
im = x_train[0,...].astype('float')/255
# 边缘填充
# 这一步的目的是，在后面对图像进行变换时，原本的Mnist数据集的28*28在变换后，
# 数字可能会移出图像区域，所以扩大原始数据的大小，也就是空白部分，方便展示变换的效果。
# pad_amt设置为10，及补充的区域为10个pixel
pad_amt = 10
# np.pad(需要填充的array，((上，下),(左，右))，mode=constant...),这一步是为了增加边缘，可以理解为padding
# 原始数据28*28，填补大小，上下左右各10，处理后数据48*48
im = np.pad(im,((pad_amt,pad_amt),(pad_amt,pad_amt)))
1
2
3
4
5
6
7
8
9
10
11
12
13

数据处理前后效果：

3. 手动创建变换矩阵：

# 手动生成仿射变换矩阵，方便后面affine操作

# 创建主对角矩阵
aff = np.eye(3) 
 # 在左上半部分的2*2区域加入随机噪声
aff[:2,:2]+=np.random.randn(2,2)*0.1
# 前两行的第三列的内容使用(-10,10)之间的均匀随机采样数字来替换
#  np.random.uniform(low,high,size)，使用(2,)的原因是aff[:2,2]数组就是一个两行一列的值
aff[:2, 2] = np.random.uniform(-10, 10, (2, ))
# 对上面计算后的矩阵求逆
aff_inv = np.linalg.inv(aff)

# 手动生成annotation变换矩阵，方便后面warp操作
margin=10
nb_annotations = 5
 # 创建一个列表，其中包含两个annotations，每个中包含nb_annotations个随机数字，范围在(margin,f-margin)之间
annotations = [np.random.uniform(margin,f-margin,nb_annotations) for f in im.shape]
# np.stack的简单用法在我的notion中有说明：
# https://sandy-property-d5e.notion.site/np-stack-48a69e31be084aa98cd15ce7d093c2ec
annotations = np.stack(annotations,1)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

处理后的数据分别为：
aff_inv:
[ 1.01 − 3.137 − 7.602 5.857 8.561 3.138 0 0 1 ]

$$\begin{bmatrix} 1.01& -3.137 & -7.602\\ 5.857& 8.561& 3.138\\ 0& 0 &1 \end{bmatrix}$$ ⎣ ⎡​1.015.8570​−3.1378.5610​−7.6023.1381​⎦ ⎤​
annotations:
$$\left[\begin{array}{cc}23.46& -3.3\\ 2.34& 7.33\\ 6.42& 14.34\\ 5.32& 37.34\\ 7.14& 17.34\end{array}\right]$$

4. Warp Data

# np.newaxis 的功能是增加新的维度。放在第一个，给行上增加维度，放在最后一个，给列上增加维度
im_keras = im[np.newaxis,...,np.newaxis] # (1, 48, 48, 1)
aff_keras = aff[np.newaxis,:2,:] #(1, 2, 3)
annotations_keras = annotations[np.newaxis,...] # (1,5,2)

# warp image
# 调用vxm库里的SpatialTransformer类，([im_keras, aff_keras])放在后面是什么用法暂时还没搞懂
# ([im_keras, aff_keras])分别代表的是图像数据和形变数据，通过空间变换将形变数据作用到图像数据中
im_warped = vxm.layers.SpatialTransformer()([im_keras, aff_keras])
im_warped = im_warped[0, ..., 0] # 取im_warped中的第0行第0列

# 获取取逆操作后的仿射矩阵的密集场Dense field
# 此处的affine_to_dense_shift和官方的教程不同，因为新版的vxm已经更新为此方法，此处已通过Issue询问过开发者
# vxm.utils.affine_to_dense_shift(array,shape,shift_center=True)
# 最后[np.newaxis, ...]的作用等价于field_inv = field_inv[np.newaxis, ...]，即给输出的结果的第一个位置增加一个维度
field_inv = vxm.utils.affine_to_dense_shift(aff_inv[:-1,:], im.shape, shift_center=True)[np.newaxis, ...]

# warp annotations
# 我的理解是：annotation是一些随机生成的点，在变换前后的图像中都是存在的
# 其作用是，帮助更明显的看出图像变化的方向和形式（涉及形变、整体移动的方向等信息）
# data为长度为2的列表，存储的分别是annotations_keras, field_inv，且两个都被转换为tf.Tensor形式，用于输入到vxm.utils.point_spatial_transformer中
data = [tf.convert_to_tensor(f, dtype=tf.float32) for f in [annotations_keras, field_inv]]
# 将辅助点和形变场都放入 vxm.utils.point_spatial_transformer，获取辅助点在该形变场下的变换信息
# [0,...]：从[1,5,2]中获取第0维度的信息=>[5,2]
annotations_warped = vxm.utils.point_spatial_transformer(data)[0, ...].numpy()
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

5.展示结果

plt.figure()
# 分别展示初始的图像和生成的辅助点
plt.subplot(1, 2, 1)
plt.imshow(im, cmap='gray')
plt.plot(*[annotations[:, f] for f in [1, 0]], 'o')  
plt.axis('off')

# 分别展示变换后的图像和变换后的辅助点
plt.subplot(1, 2, 2)
plt.imshow(im_warped, cmap='gray')
plt.plot(*[annotations_warped[:, f] for f in [1, 0]], 'o')
plt.axis('off');
1
2
3
4
5
6
7
8
9
10
11
12

第二部分：VoxelMorph tutorial VoxelMorph模型和训练教程

本部分的官方代码地址：https://colab.research.google.com/drive/1WiqyF7dCdnNBIANEY80Pxw_mVz4fyV-S?usp=sharing#scrollTo=joVczQLTPXMZ
这一部分主要介绍VoxelMorph基于深度学习的配准的实现，主要介绍以下四部分：

• MNIST数据集的介绍和使用
  如何处理数据集，建立模型，训练，配准和一般化的使用
• 现实的使用场景：颅脑MRI（2维切片）
  展示这些模型是如何在2d的颅脑数据上工作的，并展示更复杂场景下的使用
• 3D颅脑数据的使用
  展示完整3D图像的配准
• 高级功能
  使用更高级的功能，包括差分形态和微调模型

代码分析与讲解：

一、MNIST数据集的介绍和使用：

1.库的导入：

# 库的安装和导入
!pip install voxelmorph 
import os, sys
import numpy as np
import tensorflow as tf
assert tf.__version__.startswith('2.'), 'This tutorial assumes Tensorflow 2.0+'
import voxelmorph as vxm
import neurite as ne
1
2
3
4
5
6
7
8

2.数据的准备：

在这部分代码中，主要介绍2D MNIST数据的配准，在之后会尝试配准2维医学图像数据。如果数据量很小，可以将其加载到内存中，因为这样测试和训练起来更快。但是如果数据量很大的话，则需要按需扫描加载到内存中，这一点会在后续谈到。

# 导入MNIST数据集，需要使用到tensorflow.keras库
from tensorflow.keras.datasets import mnist
# 分别存储训练和测试数据
(x_train_load,y_train_load),(x_test_load,y_test_load) = mnist.load_data()
# 本文以数字5的图像配准为例
digit_sel = 5

# 分别获取下标为5的数据，分别存储为训练和测试集
x_train = x_train_load[y_train_load==digit_sel,...]
y_train = y_train_load[y_train_load==digit_sel]
x_test = x_test_load[y_test_load==digit_sel, ...]
y_test = y_test_load[y_test_load==digit_sel]

# 输出数据的尺寸以供检查
print('shape of x_train:{},y_train:{}'.format(x_train.shape,y_train.shape))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

shape of x_train: (5421, 28, 28), y_train: (5421,)

测试/验证集的划分

ML的弯路：把数据只分在训练/测试中往往会导致问题的出现
反复（A）建立一个模型，（B）在训练数据上训练，（C）在测试数据上测试
这样做会导致过拟合（因为你会根据测试数据调整你的算法）。这在深度学习中是一个常见的错误。我们将把 "训练 "分成 "训练/验证 "数据，并把测试集留待以后使用。而只有在最后才会看测试数据。

# 抽出1000个作为验证集
nb_val = 1000
x_val = x_train[-nb_val:,...]
y_val = y_train[-nb_val:]
x_train = x_train[:-nb_val,...]
y_train = y_train[:-nb_val]
1
2
3
4
5
6

可视化数据

# numebr of visualize展示的数据的个数
nb_vis = 5
# nb.random.choice(需要抽取的数组，抽取的个数，是否允许重复)
idx = np.random.choice(x_train.shape[0],nb_vis,replace=False)
# example_digits是一个列表，存储的是分别是随机选择的数字的灰度值的矩阵
example_digits = [f for f in x_train[idx,...]]
# ne.plot.slices工具用来可视化数据
ne.plot.slices(example_digits,cmaps=['gray'],do_colorbars=True);
1
2
3
4
5
6
7
8

对数据进行归一化并重新显示：

x_train = x_train.astype('float')/255
x_val = x_val.astype('float')/255
x_test = x_test.astype('float')/255

example_digits = [f for f in x_train[idx,...]]
ne.plot.slices(example_digits,camps=['gray'],do_colorbars=True];
1
2
3
4
5
6

扩展图像

# 从28*28拓展到32*32
# 第一维是图像个数，无需处理，后两维是长度和宽度
pad_amount = ((0,0),(2,2),(2,2))

x_train = np.pad(x_train,pad_amount,'constant')
x_val = np.pad(x_val,pad_amount,'constant')
x_text = np.pad(x_text,pad_amount,'constant')
print('shape of training data', x_train.shape)
1
2
3
4
5
6
7
8

shape of training data (4421, 32, 32)

3. CNN模型

提供参考和浮动图像，配准的目标是找到二者之间的变形矩阵。在基于学习的方法中，VoxelMorph选择两幅图像作为输入（参考和浮动图像，3232的MNIST数据），输出为密集形变场