2D医学影像分割任务-Medical Image Segmentation Task
2D医学影像分割任务-Medical Image Segmentation Task
本文档主要基于Pytorch-Medical-Segmentation-master进行2D医学图像分割任务的开展
我自己已跑通的代码
提取码：wwu2
目录
├──背景介绍
├──算法原理
├──Unet
├──Unet++
├──核心代码解析
├──基于pytorch进行U-net网络的搭建
├──loss函数的构建
├──train函数的构建
├──数据处理
├──环境要求
├──准备您的数据集
├──调试代码
├──训练
├──测试
├──结果展示
一、背景介绍
医学影像是研究借助于某种介质(如X射线、电磁场、超声波等)与人体相互作用， 从而用图像的方式将人体内部的器官组织结构、密度表现出来，作为相关人员进行诊断的重要依据之一。由于医学影像除了人眼最为敏感的形状结构特征以外，还包含了大量其他维度的容易被人眼所忽略的信息(如标准方差、能量、复杂灰度、共生矩阵特征等)，利用图像处理技术对图像进行分析和处理，实现对人体器官、软组织和病变体的位置检测、分割提取、三维重建和三维显示，可以对感兴趣区域（Region Of Interest, ROI)进行定性甚至定量的分析，从而大大提高临床诊断的效率、准确性和可靠性。
医学图像分割处理的对象主要是各种细胞、组织、器官的图像，根据区域间的相似或不同，把图像分割成若干区域。医学图像分割的早期方法通常依赖于边缘检测、模式匹配技术、统计形状模型、活动轮廓和传统机器学习技术。
这些方法在一定程度上取得了不错的效果，但由于特征表示的困难，图像分割仍然是计算机视觉领域中最具挑战性的课题之一。特别是医学图像的特征提取比普通RGB图像更难，因为前者往往存在模糊、噪声、对比度低等问题。由于深度学习技术的快速发展，卷积神经网络(CNN)成功实现了图像的层次特征表示，从而成为图像处理和计算机视觉领域最热门的研究课题。由于用于特征学习的CNN对图像噪声、模糊、对比度等不敏感，对医学图像提供了优秀的分割结果。
医学影像的分割方法主要包括监督学习和弱监督学习。监督学习的优点在于可以通过精确标记的数据来进行模型训练，主要包括对Backbone、Network Block、Loss Function进行设计。通常采用编码器-解码器结构的全卷积网络方法，编码器通常用于提取图像特征，而解码器通常用于将提取的特征恢复到原始图像大小并输出最终分割结果，典型的有FCN，U-Net，Deeplab等。
由于在医学影像领域很难获取大量的高质量的标记数据集，因此提出基于使用不完整或不完备的数据集进行训练，也就是弱监督学习，其数据集中只包含一小部分具有标签，大部分数据都是没有标签的。在缺乏大量标注的情况下，数据增强(Data Augmentation)可以有效的改善训练网络的泛化性能。传统的数据增强方法主要包括噪声抑制、改变图像强度、旋转、失真、裁剪等。这些操作因为不需要计算成本，因此一般放在训练前进行。生成对抗网络(GAN)也可以作为数据增强的一种方法，它通过构建一个生成器和一个判别器，生成器可以看作一位画家，判别器可以看作一位鉴定家，画家进行作画，将真画与假画一起交给鉴定家进行鉴别，这个过程重复一次又一次，最后画家作画的能力和鉴定家鉴别的能力一起提升，我们就可以利用此时训练完成的生成器进行数据增强。除数据增强的方法外，迁移学习也可以使标签有限的数据实现快速的模型训练，一种方法是针对目标医学图像分析任务微调ImageNet上的预训练模型，而另一种方法是将训练的数据从多个领域迁移。
与普通的图像分割不同，医学图像通常含有噪声且边界模糊。因此，仅仅依靠图像的底层特征很难对医学图像中的目标进行检测和识别。同时，由于缺乏图像细节信息，仅靠图像语义特征无法获得准确的边界。而U-Net通过跳跃连接，将低分辨率和高分辨率的特征图结合起来，有效地融合了低分辨率和高分辨率的图像特征，是医学图像分割任务的完美解决方案。目前，U-Net已经成为大多数医学图像分割任务的基准，并激发了许多有意义的改进。
二、算法原理
主要对Unet、Unet++两种分割算法进行分析
2.1 Unet
目前在医学图像分割任务中，Unet是最为经典的网络之一，仍然使用编码-解码的结构，核心思想就是输入是一幅图，输出是目标的分割结果。继续简化就是，一幅图，编码，或者说降采样，然后解码，也就是升采样，然后输出一个分割结果。根据结果和真实分割的差异，反向传播来训练这个分割网络。
有这样一个问题：既然输入和输出都是相同大小的图，为什么要折腾去降采样一下再升采样呢？按论文中所讲，这可以增加对输入图像的一些小扰动的鲁棒性，比如图像平移，旋转等，减少过拟合的风险，降低运算量，和增加感受野的大小。升采样的最大的作用其实就是把抽象的特征再还原解码到原图的尺寸，最终得到分割结果。
Unet是基于全卷积网络FCN提出的，其中比较重要的特点在于它使用了跳跃连接即图中的灰色线条，将低分辨率与高分辨率特征图进行融合，提高了网络的泛化性能，与FCN不同的是，FCN采用的是sum相加操作，而Unet使用的是cat拼接操作。这与残差网络类似，同时它在解码部分采用的上采样与编码部分完全对称。目的是经过下采样的编码，得到一串比原先图像更小的特征，相当于压缩，然后再经过一个解码，理想状况下还原到原来的图像。解码过程与编码过程完全对称，对底层特征进行上采样后，与调整了size的高分辨率图像进行连接，同样经过两次卷积后重复上采样操作4次。最后使用1x1的卷积用于调整通道数得到分割结果。
从UNet网络中可以看出，不管是下采样过程还是上采样过程，每一层都会连续进行两次卷积操作,蓝色箭头代表3x3的卷积操作，并且stride是1，padding策略是vaild，因此，每个该操作以后，featuremap的大小会减2，同时扩张通道数为64。
红色箭头代表2x2的maxpooling操作，需要注意的是，此时的padding策略也是vaild，这就会导致如果pooling之前featuremap的大小是奇数，会损失一些信息 。所以要选取合适的输入大小，因为2*2的max-pooling算子适用于偶数像素点的图像长宽。
Upsample模块中除了常规的上采样操作，还有进行特征的融合。绿色箭头代表2x2的反卷积操作,灰色箭头代表跳跃连接。对frature map进行上采样有两种方法：Upsample和ConvTranspose2d，也就是双线性插值和反卷积。前者其实就是在像素点中间补点，补的点的值是多少，是由相邻像素点的值决定的，后者其实就是先对小的feature map进行padding后再卷积使其尺寸增加一倍。
从左至右分析其网络结构图，蓝色方框为feature map，也就是一个输入到网络中的Tensor(张量)，左下角数字代表width,height，上方数字代表channel。
可以看到输入的图像为(572,572,1)的灰度图，然后经过一次3x3x64的卷积与ReLu激活函数，进行特征提取，并调整通道数，得到(570,570,64)的feature map，这里没有加入BN标准化可能是因为论文所使用的数据量较小，所以没有考虑。宽高减少了2是因为在原作者进行卷积前没有进行padding填充。重复上述操作后会得到一个(568,568,64)的feature map。
红色的箭头表示max pooling操作，使用2x2的池化核进行最大池化，步长为1，提取2x2窗口中的最大像素值，会将feature map长宽降低为原来的一半。这里的尺寸变换为从（568，568，64）到（284，284，64），随后的两个卷积层将feature map通道数扩张为128。
重复池化和2次卷积操作四次为Unet的左半编码部分，会得到(28,28,1024)的高度压缩的feature map，对其进行Upsample(上采样)，论文中这里所使用的上采样方法为反卷积，进行feature map宽高的扩张。得到(54,54,1024)的feature map。
此时需要将下采样第四层中的浅层特征即size为(64,64,512)的feature map通过copy and crop，剪裁其尺寸为(54,54,512)后在通道维度与上采样第四层的抽象特征进行拼接，获得一个(54,54,1536)的feature map。
同理进行两次卷积压缩通道数和尺寸减2后重复上采样与跳跃连接操作，在最上层跳跃连接并进行两次卷积后获得一个(388,388,64)的feature map，最后通过一个1x1的卷积调整其通道数为2后输出，代表二分类任务,这些过程属于Unet的右半解码部分。
2.2 Unet++
Unet++是Unet的改进版本，其作者对Unet的跳跃连接以及网络结构的深度发出疑问：既然跳跃连接使得浅层特征与抽象特征进行融合，使得分割效果大大提升，那么为什么必须要在上采样的过程中才进行跳跃连接呢？为什么Unet一定是四次下采样呢？所以他改进了Unet的网络结构，使其在下采样的过程中就开始了上采样和跳跃连接，在每一个深度上都有结果输出，其好处是经过实验，能够比对哪一层的性能最佳，从而在预测网络中进行剪枝，大大减少预测网络参数量，使其提高预测速度。
那为什么不在训练网络中剪枝呢？因为其每一深度所得网络的性能本就是依据完全的网络结构进行训练的，如果剪去其中权重较小的部分，那么训练时的反向传播将与原来不同，其重新得到的性能结果也不会不同，故只能在预测网络中进行剪枝。Unet++所采用的卷积，上下采样操作与Unet相似，关键在于跳跃连接部分变得更加复杂，增加了训练时的参数量，但它考虑的更加全面，不同层次的特征所带来的精度的提升，同时灵活的网络结构配合深监督，让参数量巨大的深度网络在可接受的精度范围内大幅度的缩减参数量。
三、核心代码解析
主要包括对U-net网络结构搭建，loss函数构造，train函数构造三部分核心代码进行解析注释。
3.1 基于pytorch进行U-net网络的搭建
- 主要的Unet结构代码块，该部分代码中将Unet中集成了输入卷积(InConv)、下采样(Down)、上采样(Up)、输出卷积(OutConv)四个主要部分。从forward()方法开始，实例化U-net类后使用该方法，x表示输入到U-net中对应原始图像的张量，经过一次输入卷积，四次下采样，四次上采样，一次输出卷积后得到2通道的张量。
class Unet(nn.Module): #定义Unet类
def __init__(self, in_channels, classes): #self为实例变量；in_channels代表输入通道数；classses代表种类数目，该文档进行的分割任务为简单的二分类的任务
super(Unet, self).__init__() #子类Unet继承父类nn.Module中的字段和方法
self.n_channels = in_channels #实例化后传入以下各个参数
self.n_classes = classes
self.inc = InConv(in_channels, 64)
self.down1 = Down(64, 128)
self.down2 = Down(128, 256)
self.down3 = Down(256, 512)
self.down4 = Down(512, 512)
self.up1 = Up(1024, 256)
self.up2 = Up(512, 128)
self.up3 = Up(256, 64)
self.up4 = Up(128, 64)
self.outc = OutConv(64, classes)
def forward(self, x): #x为的输入tensor，之后进入U-net中训练后返回2通道张量
x1 = self.inc(x)
x2 = self.down1(x1)
x3 = self.down2(x2)
x4 = self.down3(x3)
x5 = self.down4(x4)
x = self.up1(x5, x4)
x = self.up2(x, x3)
x = self.up3(x, x2)
x = self.up4(x, x1)
x = self.outc(x)
return x
- 输入卷积(InConv)：利用pytorch，padding=1使输入图像边界填充一圈为0的白边后，经过一次步长为1的3x3卷积,调整通道数由out_ch决定，经过标准化使特征图满足均值为0，方差为1的规律，再原地经过ReLU激活函数，大于0的值保留，小于等于0的值，返回为0。重复一次上述过程即为DoubleConv,也是输入卷积。
class DoubleConv(nn.Module):
def __init__(self, in_ch, out_ch):
super(DoubleConv, self).__init__() #两次2D卷积+标准化+Relu激活函数
self.conv = nn.Sequential(
nn.Conv2d(in_ch, out_ch, 3, padding=1), #(输入通道数；输出通道数；卷积核size；添加白边，提高对边界信息的提取)
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True),
nn.Conv2d(out_ch, out_ch, 3, padding=1),
nn.BatchNorm2d(out_ch),
nn.ReLU(inplace=True))
def forward(self, x):
x = self.conv(x)
return x
class InConv(nn.Module): #定义输入卷积类
def __init__(self, in_ch, out_ch):
super(InConv, self).__init__()
self.conv = DoubleConv(in_ch, out_ch)
def forward(self, x):
x = self.conv(x)
return x
- 下采样(Down)：对经过输入卷积后的张量先进行2x2最大池化，步长默认为1，相当于每一次在2x2窗口的Tensor中提取其最大值；后进行DoubleConv如上文所讲。相当于在DoubleConv中加入池化操作即为下采样，重复四次。
class Down(nn.Module):
def __init__(self, in_ch, out_ch):
super(Down, self).__init__()
self.mpconv = nn.Sequential(
nn.MaxPool2d(2),
DoubleConv(in_ch, out_ch)
) #nn.Sequential相当于模块容器，表明依次执行MaxPool和DoubleConv
def forward(self, x):
x = self.mpconv(x)
return x
- 上采样(Up)：看需求设置是否采用插值方法，pytorch中UPsample支持多种插值方法：‘linear’、‘bilinear’、‘bicubic’、‘trilibear’、default:‘nearest’。根据像素值与像素点的坐标进行计算得到四个像素值与坐标完成宽高扩张，拟合理想逆卷积过程；同时该Up方法中也要进行跳跃连接，融合浅层特征与抽象特征，2D图像中一般神经网络中训练的维度为[bacth_size,channels,width,heigh]，为保证跳跃连接中两个特征图(feature_map)的size一致，对其中一个张量进行调整后拼接。最后进行DoubleConv输出。
class Up(nn.Module):
def __init__(self, in_ch, out_ch, bilinear=True): #bilinear是否使用双线性插值
super(Up, self).__init__()
#使用双线性插值进行上采样，优点是上采样过程没有需要训练的参数；或使用ConvTranspose2d逆卷积进行上采样
if bilinear:
self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
else:
self.up = nn.ConvTranspose2d(in_ch // 2, in_ch // 2, 2, stride=2)
self.conv = DoubleConv(in_ch, out_ch)
def forward(self, x1, x2):
x1 = self.up(x1)
diffY = x2.size()[2] - x1.size()[2] #跳跃连接中两个张量的宽差
diffX = x2.size()[3] - x1.size()[3] #高差
x1 = F.pad(x1, (diffX // 2, diffX - diffX // 2,
diffY // 2, diffY - diffY // 2)) #扩张x1的w,h使其与拼接张量size一致
x = torch.cat([x2, x1], dim=1) #张量拼接在[1]，即通道维度
x = self.conv(x)
return x
- 输出卷积(OutConv)：使用1x1卷积调整通道数，输出2通道的张量
class OutConv(nn.Module):
def __init__(self, in_ch, out_ch):
super(OutConv, self).__init__()
self.conv = nn.Conv2d(in_ch, out_ch, 1)
def forward(self, x):
x = self.conv(x)
return x
3.2 loss函数的构建
- 该任务使用了Loss=DiceLoss+CrossEntropyLoss(交叉熵损失)，Dice损失是由V-net一文中提出，常用于图像二值分割中，loss_function.py中进行DiceLoss损失函数的构建。从forward方法开始解析，可以从main.py中分析得知inputs为图像经过U-net后得到的2通道的张量，target为标签使用argmax(dim=1)后得到的索引one-hot向量，weight表示是否使用了预训练权重进行训练，softmax表示是否使用该激活函数对输入tensor进行归一化处理。target首先经过_one_hot_encoder方法得到与inputs相同size的张量，然后使用_dice_loss方法计算两者损失，将权重参数保存在class_wise_dice当中
class DiceLoss(nn.Module):
def __init__(self, n_classes):
super(DiceLoss, self).__init__()
self.n_classes = n_classes
def _one_hot_encoder(self, input_tensor):
tensor_list = []
for i in range(self.n_classes):
temp_prob = input_tensor == i * torch.ones_like(input_tensor)
tensor_list.append(temp_prob)
output_tensor = torch.stack(tensor_list, dim=1)
return output_tensor.float()
def _dice_loss(self, score, target):
target = target.float()
smooth = 1e-5
intersect = torch.sum(score * target)
y_sum = torch.sum(target * target)
z_sum = torch.sum(score * score)
loss = (2 * intersect + smooth) / (z_sum + y_sum + smooth)
loss = 1 - loss
return loss
def forward(self, inputs, target, weight=None, softmax=False):
if softmax:
inputs = torch.softmax(inputs, dim=1)
target = self._one_hot_encoder(target)
if weight is None:
weight = [1] * self.n_classes
assert inputs.size() == target.size(), 'predict & target shape do not match'
class_wise_dice = []
loss = 0.0
for i in range(0, self.n_classes):
dice = self._dice_loss(inputs[:, i], target[:, i])
class_wise_dice.append(1.0 - dice.item())
loss += dice * weight[i]
return loss / self.n_classes
- 可以看到对两种损失实例化后进行计算并相加，每次迭代运行。
criterion_dice = DiceLoss(2).cuda() #实例化DiceLoss
criterion_ce = CrossEntropyLoss().cuda() #torch中的交叉熵方法
loss = criterion_ce(outputs, y.argmax(dim=1)) + criterion_dice(outputs, y.argmax(dim=1))
num_iters += 1
loss.backward() #反向传播更新权重参数
optimizer.step() #更新学习率
iteration += 1 #迭代次数计数器
3.2 train函数的构建
- 在hparm.py中配置训练参数
class hparams:
train_or_test = 'train' #选择main,py的运行模式
output_dir = 'logs' #存放权重文件的文件夹路径
aug = None #是否启用数据增强
latest_checkpoint_file = 'checkpoint_latest.pt' #使用预训练权重进行训练时权重的路径
total_epochs = 9999999 #设置的训练轮次
epochs_per_checkpoint = 10 #每10次保存一次权重
batch_size = 2 #设置批量大小
ckpt = None #是否具有预训练权重
init_lr = 0.002 #学习率
scheduer_step_size = 20
scheduer_gamma = 0.8
debug = False
mode = '2d' # '2d or '3d'
#这里的in_class、out_class、channel并不代表输入输出U-net中的通道数，原作者为方便使用各种网络结构进行训练而初步设置为1，方便在调用不同网络结构时通过+1、+2的方式来获得真正的通道数。需要通过分析main.py中的相关参数调用进行设置。
in_class = 1
out_class = 1
channel = 3
#数据增强设置的裁剪size
crop_or_pad_size = 256,256,1 # if 2D: 256,256,1 3D:512,512,32
#随机局部特征进行train和loss计算
patch_size = 256,256,1 # if 2D: 128,128,1 3D:32,32,32
# for test
patch_overlap = 4,4,0 # if 2D: 4,4,0 3D:4,4,4
fold_arch = '*.jpg' #图像格式
save_arch = '.jpg'
source_train_dir = 'train/CHASEDB1/image'
label_train_dir = 'train/CHASEDB1/labels'
source_test_dir = 'train/CHASEDB1/test_image'
label_test_dir = 'train/CHASEDB1/test_label'
output_dir_test = 'results'
- 在data_function.py中进行对数据与标签的增强和提取
- 使用torchio进行图像与标签的提取
if (hp.in_class == 1) and (hp.out_class == 1) :
images_dir = Path(images_dir) #传入图像路径
self.image_paths = sorted(images_dir.glob(hp.fold_arch)) #sorted()排序函数对所有可迭代的对像进行排序操作，返回列表
labels_dir = Path(labels_dir) #同上
self.label_paths = sorted(labels_dir.glob(hp.fold_arch))
#获取图像与标签
for (image_path, label_path) in zip(self.image_paths, self.label_paths):
subject = tio.Subject(
source=tio.ScalarImage(image_path),
label=tio.LabelMap(label_path),
)
self.subjects.append(subject) #存储图像与标签信息到subjects中
- 在训练数据集较少的情况下可以使用数据增强提高网络鲁棒性。
if hp.aug:
training_transform = Compose([
CropOrPad((hp.crop_or_pad_size), padding_mode='reflect'), #切割与扩充
# RandomMotion(), #运动伪影
RandomBiasField(), #随机MRI偏置场伪影
ZNormalization(), #正则化
RandomNoise(), #高斯噪声
RandomFlip(axes=(0,)), #反转图像中的元素顺序
OneOf({
RandomAffine(): 0.8, #倾斜
RandomElasticDeformation(): 0.2, #随机弹性变形
}),])
四、数据处理
这部分主要介绍数据集的预处理，以及代码调试。
4.1 环境要求
- torch=1.7.1
- torchio=0.18.79
- python=3.7
- 深度学习环境的搭建可以参考这里
4.2准备您的数据集
train
├──image
├── source_1.jpg
├── source_2.jpg
├── source_3.jpg
├── source_4.jpg
└── ...
├──label
├── label_1.jpg
├── label_2.jpg
├── label_3.jpg
├── label_4.jpg
└── ...
- 保证进行2D医学图像分割所使用数据集的输入与标签格式一致，包括不限于".jpg"格式，只需要设置U-net中输入图片的通道数与您所使用的图片通道数一致即可。
#可以查看所使用图片的通道数
img = Image.open('image_path')
print(len(img.split()))
4.3调试代码
- 在main.py中调整输入通道数(line 94)
#根据您的数据集通道数进行修改
model = Unet(in_channels=hp.in_class, classes=hp.out_class+1) #表明输入U-net中的通道数为1
model = Unet(in_channels=hp.in_class+2, classes=hp.out_class+1) #表明输入U-net中的通道数为3
- 也可以预处理训练图片将’RGB’模式转换为’L’灰度图模式，通道数变为1，则不用修改line94。
from PIL import Image
import os
pic_path = 'train/label/'#需要修改的图片路径
save_path = 'train/labels/'#保存图片路径
pic_name = os.listdir(pic_path)#获取原路径下的图片
for i in range(len(pic_name)):
name = pic_name.pop()
print(name)
img = Image.open(pic_path + name).convert('1')#转换图片
img.save(save_path + name)
- train模式下，在hparam.py中设置2D模式相关参数，fold_arch与source_train_dir的设置格式要如下，才能正确提取到训练队列中
fold_arch = '*.jpg'
save_arch = '.jpg'
source_train_dir = 'train/image'
label_train_dir = 'train/label'
- train模式下，在main.py文件中修改(可能目标语句所在位置在所给参考附近，找到相同语句修改其中的不同点即可)
y[y!=0] = 1 --> y[y!=0] = 255 #line248，二值化保存ground-truth图像
model_output_one_hot = torch.nn.functional.one_hot(labels, num_classes=hp.out_class+1).permute(0,3,1,2) #line258，2D训练维数为4，(batch,通道，w,h)
y_one_hot = torch.nn.functional.one_hot(y_argmax, num_classes=hp.out_class+1).permute(0,3,1,2) #line283
outputs = outputs.unsqueeze(4)
model_output_one_hot = model_output_one_hot.unsqueeze(4) #line343，与上方代码同缩进下添加该行代码，生成4D预测权重
model_output_one_hot = np.expand_dims(model_output_one_hot, axis=1)
model_output_one_hot[model_output_one_hot==1] = 255 #line356，与上方代码同缩进下添加该行代码，二值化处理预测输出图像
- train模式下，在loss_function.py中修改
#output_tensor = torch.cat(tensor_list, dim=1)
output_tensor = torch.stack(tensor_list, dim=1) #line112，cat为宽高拼接张量，stack，dim=1为拼接通道
4.4 训练
- 不使用预训练模型
set hparam.train_or_test to 'train'
set hparam.mode to '2D'
python main.py
- 使用预训练模型
set hparam.ckpt to 'True'
python main.py
4.5 测试
set hparam.train_or_test to 'test'
set hparam.mode to '2D'
python main.py
五、结果展示
source |
ground-truth |
predict |
参考资料
- 浅谈医学图像处理
- 医学影像到底是什么？
- 医疗是什么？AI医疗概念解析AI
- AI医学影像技术展望
- 研习U-net
- U-Net网络原理分析与pytorch实现
[1] Risheng Wang, Tao Lei, Ruixia Cui, Bingtao Zhang, Hongying Meng, Asoke K. Nandi. Medical Image Segmentation Using Deep Learning: A Survey[J]. IET Image Processing, 2022.
[2] Chen C, Zhou K. An Effective Deep Neural Network for Lung Lesions Segmentation from COVID-19 CT Images[J]. IEEE Transactions on Industrial Informatics, 2021.
[3] Chen C, Zhang T, et al. Pathological lung segmentation in chest CT images based on improved random walker[J]. Computer Methods and Programs in Biomedicine, 2021, 200: 105864.