图像分割之——U-Net实现

语义分割(Semantic Segmentation):是图像处理和机器视觉一个重要分支。与分类任务不同，语义分割需要判断图像每个像素点的类别，进行精确分割。语义分割目前在自动驾驶、自动抠图、医疗影像等领域有着比较广泛的应用。——语义分割也是一个分类问题！

Unet可以说是最常用、最简单的一种分割模型了，它简单、高效、易懂、容易构建、可以从小数据集中训练。UNet主要贡献是在U型结构上，该结构可以使它使用更少的训练图片的同时，且分割的准确度也不会差，UNet的网络结构如下图：

unet网络非常的简单，前半部分就是特征提取，后半部分是上采样。在一些文献中把这种结构叫做编码器-解码器结构(自编码：标签是数据,编解码结构:标签是分割图),由于网络的整体结构是一个大些的英文字母U，所以叫做U-net。

• Encoder：左半部分，由两个3x3的卷积层（RELU）再加上一个2x2的maxpooling层组成一个下采样的模块（后面代码可以看出）；
• Decoder：右半部分，由一个上采样的卷积层（去卷积层）+特征拼接concat+两个3x3的卷积层（ReLU）反复构成（代码中可以看出来）；

在当时，Unet相比更早提出的FCN网络，使用拼接来作为特征图的融合方式。

• FCN是通过特征图对应像素值的相加来融合特征的；
• U-net通过通道数的拼接，这样可以形成更厚的特征，当然这样会更佳消耗显存；

Unet的好处：网络层越深得到的特征图，有着更大的视野域，浅层卷积关注纹理特征，深层网络关注本质的那种特征，所以深层浅层特征都是有各自的意义的；另外一点是通过反卷积得到的更大的尺寸的特征图的边缘，是缺少信息的，毕竟每一次下采样提炼特征的同时，也必然会损失一些边缘特征，而失去的特征并不能从上采样中找回，因此通过特征的拼接，来实现边缘特征的一个找回。

这个结构就是先对图片进行卷积和池化，在Unet论文中是池化4次，比方说一开始的图片224x224的，那么就会变成112x112，56x56,28x28,14x14四个不同尺寸的特征。然后我们对14x14的特征图做上采样或者反卷积，得到28x28的特征图，这个28x28的特征图与之前的28x28的特征图进行通道上的拼接concat，然后再对拼接之后的特征图做卷积和上采样，得到56x56的特征图，再与之前的56x56的特征拼接、卷积，再上采样，经过四次上采样可以得到一个与输入图像尺寸相同的224x224的预测结果。

代码实例：

import torch
import torch.nn as nn
from torch.nn.functional import interpolate

#unet
class CNNlayer(nn.Module):
    def __init__(self,c_in,c_Out):
        super(CNNlayer, self).__init__()
        self.layer=nn.Sequential(
            nn.Conv2d(c_in,c_Out,3,1,padding=1,padding_mode="reflect",bias=False),#"reflect"镜像翻转
            nn.BatchNorm2d(c_Out),
            nn.LeakyReLU(),#原模型为Rule(可能会造成网络退化，负半轴为0）
            #LeakyReLU不能防止网络过拟合
            nn.Dropout2d(0.3),#随机抑制%30的神经元
            #图像分割容易出现过拟合（数据量比较小）
            nn.Conv2d(c_Out, c_Out, 3, 1, 1,padding_mode="reflect",bias=False),
            nn.BatchNorm2d(c_Out),
            nn.LeakyReLU(),
            nn.Dropout2d(0.4)
        )
    def forward(self,x):
        return self.layer(x)

#下采样(使用最大池化）---降噪能力较强
class DownSampling(nn.Module):
    def __init__(self):
        super(DownSampling, self).__init__()
        self.layer=nn.Sequential(
            nn.MaxPool2d(2)
        )
    def forward(self,x):
        return self.layer(x)
# #2:使用步长为2的卷积做下采样
# class DownSampling(nn.Module):
#     def __init__(self,C):
#         super(DownSampling, self).__init__()
#         self.layer=nn.Sequential(
#             nn.Conv2d(C,C,3,2,1,padding_mode="reflect"),
#             nn.LeakyReLU(),
#             nn.BatchNorm2d(C)
#         )
#     def forward(self,x):
#         return self.layer(x)
#上采样
class UpSampling(nn.Module):
    def __init__(self,c):
        super(UpSampling, self).__init__()
        #特征图大小扩大两倍，通道数减半
        self.layer=nn.Sequential(
        nn.Conv2d(c, c //2, 3, 1, 1, padding_mode="reflect",bias=False),
        nn.BatchNorm2d(c//2),
        nn.LeakyReLU(),
        )
    def forward(self,x,r):
        #使用临近插值法进行上采样
        up=interpolate(x,scale_factor=2,mode="nearest")#特征图放大两倍
        x=self.layer(up)#通道数减半
        #通道拼接（cat)
        return torch.cat((x,r),dim=1)#nchw
class UNet(nn.Module):
    def __init__(self):
        super(UNet, self).__init__()
        self.C1=CNNlayer(3,64)
        #4次下采样
        self.D1=DownSampling()
        self.C2=CNNlayer(64,128)
        self.D2 = DownSampling()
        self.C3 = CNNlayer(128, 256)
        self.D3 = DownSampling()
        self.C4 = CNNlayer(256, 512)
        self.D4 = DownSampling()
        self.C5 = CNNlayer(512,1024)
        #4次上采用
        self.U1=UpSampling(1024)
        self.C6=CNNlayer(1024,512)
        self.U2 = UpSampling(512)
        self.C7 = CNNlayer(512, 256)
        self.U3 = UpSampling(256)
        self.C8 = CNNlayer(256, 128)
        self.U4 = UpSampling(128)
        self.C9 = CNNlayer(128, 64)

        self.pre=nn.Conv2d(64,3,3,1,1)#64-->3
    def forward(self,x):
        #下采样部分
        R1=self.C1(x)
        R2=self.C2(self.D1(R1))
        R3 = self.C3(self.D2(R2))
        R4 = self.C4(self.D3(R3))
        R5 = self.C5(self.D4(R4))
        #上采用部分
        O1 = self.C6(self.U1(R5, R4))
        O2 = self.C7(self.U2(O1, R3))
        O3 = self.C8(self.U3(O2, R2))
        O4 = self.C9(self.U4(O3, R1))
        return self.pre(O4)

if __name__ == '__main__':
    x=torch.randn(1,3,256,256)
    net=UNet()
    out=net(x)
    print(out.shape)