图像分割之Unet解析及实现代码_overlap-tile代码实现

论文连接: https://arxiv.org/pdf/1505.04597.pdf

源码连接: https://github.com/FENGShuanglang/unet

图像分割

图像分割就是把图像分成若干个特定的、具賄独特性质的区域并提出感兴趣目标的技术和过程。它是由图像处理到图像分析的关键步骤。现有的图像分割方法主要分以下几类:基于阈值的分割方法、基于区域的分割方法、基于边缘的分割方法以及基于特定理论的分割方法等。从数学角度来看，图像分割是将数字图像划分成互不相交的区域的过程。图像分割的过程也是一个标记过程，即把属于同一区域的像素赋予相同的编号。

Unet网络结构详解

Unet网络结构：

该网络结构主要分为三部分:下采样， 上采样以及 跳跃连接。首先将该网络分为左右部分来分析，左边是压缩的过程，即Encoder。通过卷积和下采样来降低图像尺寸,提取一些浅显的特征。 右边Decoder是解码的过程。 通过卷积和上采样来获取一些深层次的特征。其中卷积采用的无填充方式来保证结果都是基于没有缺失上下文特征得到的,因此每次经过卷积后,图像的大小会减小。中间通过concat的方式，将编码阶段获得的feature map同解码阶段获得的feature map结合在-起,结合深层次和浅层次的特征,得到更加精细的细节，根据得到的feature map进行预测分割。要注意的是这里两层的feature map大小不同,因此需要经过切割。最后通过1x1的卷积做分类。

overlap-tile 策略

该策略的思想是:对图像的某一块像点(黄框内部分)进行预测时，要该图像块周围的像素点(色框内)提供上下文信息
(context) ,以获得更准确的预测。

医学图像是一般相当大， 但是分割时候不可能将原图太小输入网络,所以必须切成一张一 张的小patch, 在切成小patch的时候,Unet由于网络结构原因适合有overlap的切图，可以看图,红框是要分割区域，但是在切图时要包含周围区域，overlap另- 个重要原因是周围overlap部分可以为分割区域边缘部分提供文理等信息。可以看黄框的边缘,分割结果并没有受到切成小patch而造成分割情况不好。

损失函数

ak(x)表示像素x在特征图中的第k层的激活值，
k表示是第几个特征通道，
x表示像素点，
K表示类别的个数。

l: S2→{…,K.} 每个像素的真实标签

表示训练构成中像素点的重要权重
W0 表示平衡类别频率的权重图，
d1 表示此像素点到离它最近cel边界的距离，
d2 表示此像素点到离他第二近ell边界的距离。

实现代码(python)

环境：pytorch 1.0以上。可直接运行

""" Full assembly of the parts to form the complete network """
"""Refer https://github.com/milesial/Pytorch-UNet/blob/master/unet/unet_model.py"""

import torch
import torch.nn as nn
import torch.nn.functional as F


class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=True):
        super(UNet, self).__init__()
        self.n_channels = n_channels
        self.n_classes = n_classes
        self.bilinear = bilinear

        self.inc = DoubleConv(n_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, bilinear)
        self.up2 = Up(512, 128, bilinear)
        self.up3 = Up(256, 64, bilinear)
        self.up4 = Up(128, 64, bilinear)
        self.outc = OutConv(64, n_classes)

    def forward(self, x):
        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)
        logits = self.outc(x)
        return logits
class DoubleConv(nn.Module):
    """(convolution => [BN] => ReLU) * 2"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)


class Down(nn.Module):
    """Downscaling with maxpool then double conv"""

    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.maxpool_conv = nn.Sequential(
            nn.MaxPool2d(2),
            DoubleConv(in_channels, out_channels)
        )

    def forward(self, x):
        return self.maxpool_conv(x)


class Up(nn.Module):
    """Upscaling then double conv"""

    def __init__(self, in_channels, out_channels, bilinear=True):
        super().__init__()

        # if bilinear, use the normal convolutions to reduce the number of channels
        if bilinear:
            self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        else:
            self.up = nn.ConvTranspose2d(in_channels // 2, in_channels // 2, kernel_size=2, stride=2)

        self.conv = DoubleConv(in_channels, out_channels)

    def forward(self, x1, x2):
        x1 = self.up(x1)
        # input is CHW
        diffY = torch.tensor([x2.size()[2] - x1.size()[2]])
        diffX = torch.tensor([x2.size()[3] - x1.size()[3]])

        x1 = F.pad(x1, [diffX // 2, diffX - diffX // 2,
                        diffY // 2, diffY - diffY // 2])

        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)


class OutConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(OutConv, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)

    def forward(self, x):
        return self.conv(x)
if __name__ == '__main__':
    net = UNet(n_channels=3, n_classes=1)
    print(net)

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
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108

本人在读研一，如果有想要学习的神经网络模型，可以私信我，会尽快更新。喜欢点个赞吧，多谢支持！