医学图像算法之基于Unet的视网膜血管分割

第一步：准备数据

视网膜血管分割数据比较少，但效果好，总共40张

第二步：搭建模型

UNet主要贡献是在U型结构上，该结构可以使它使用更少的训练图片的同时，且分割的准确度也不会差，UNet的网络结构如下图：

（1）UNet采用全卷积神经网络。
（2）左边网络为特征提取网络：使用conv和pooling
（3）右边网络为特征融合网络：使用上采样产生的特征图与左侧特征图进行concatenate操作。（pooling层会丢失图像信息和降低图像分辨率且是永久性的，对于图像分割任务有一些影响，对图像分类任务的影响不大，为什么要做上采样呢？上采样可以让包含高级抽象特征低分辨率图片在保留高级抽象特征的同时变为高分辨率，然后再与左边低级表层特征高分辨率图片进行concatenate操作）
（4）最后再经过两次卷积操作，生成特征图，再用两个卷积核大小为1*1的卷积做分类得到最后的两张heatmap，例如第一张表示第一类的得分，第二张表示第二类的得分heatmap，然后作为softmax函数的输入，算出概率比较大的softmax，然后再进行loss，反向传播计算。

第三步：代码

1）损失函数为：MSE损失函数

2）网络代码（将输出层增加一个sigmoid函数，将输出值约束到0-1之间）：

class U_Net(nn.Module):
    def __init__(self,img_ch=3,output_ch=1):
        super(U_Net,self).__init__()
        
        self.Maxpool = nn.MaxPool2d(kernel_size=2,stride=2)

        self.Conv1 = conv_block(ch_in=img_ch,ch_out=64)
        self.Conv2 = conv_block(ch_in=64,ch_out=128)
        self.Conv3 = conv_block(ch_in=128,ch_out=256)
        self.Conv4 = conv_block(ch_in=256,ch_out=512)
        self.Conv5 = conv_block(ch_in=512,ch_out=1024)

        self.Up5 = up_conv(ch_in=1024,ch_out=512)
        self.Up_conv5 = conv_block(ch_in=1024, ch_out=512)

        self.Up4 = up_conv(ch_in=512,ch_out=256)
        self.Up_conv4 = conv_block(ch_in=512, ch_out=256)
        
        self.Up3 = up_conv(ch_in=256,ch_out=128)
        self.Up_conv3 = conv_block(ch_in=256, ch_out=128)
        
        self.Up2 = up_conv(ch_in=128,ch_out=64)
        self.Up_conv2 = conv_block(ch_in=128, ch_out=64)

        self.Conv_1x1 = nn.Conv2d(64,output_ch,kernel_size=1,stride=1,padding=0)


    def forward(self,x):
        # encoding path
        x1 = self.Conv1(x)

        x2 = self.Maxpool(x1)
        x2 = self.Conv2(x2)
        
        x3 = self.Maxpool(x2)
        x3 = self.Conv3(x3)

        x4 = self.Maxpool(x3)
        x4 = self.Conv4(x4)

        x5 = self.Maxpool(x4)
        x5 = self.Conv5(x5)

        # decoding + concat path
        d5 = self.Up5(x5)
        d5 = torch.cat((x4,d5),dim=1)
        
        d5 = self.Up_conv5(d5)
        
        d4 = self.Up4(d5)
        d4 = torch.cat((x3,d4),dim=1)
        d4 = self.Up_conv4(d4)

        d3 = self.Up3(d4)
        d3 = torch.cat((x2,d3),dim=1)
        d3 = self.Up_conv3(d3)

        d2 = self.Up2(d3)
        d2 = torch.cat((x1,d2),dim=1)
        d2 = self.Up_conv2(d2)

        d1 = self.Conv_1x1(d2)

        d1 = F.sigmoid(d1)

        return d1

第四步：统计一些指标（训练过程中的loss）