前言

研究一个深度学习算法，可以先看网络结构，看懂网络结构后，再Loss计算方法、训练方法等。本文主要针对UNet的网络结构进行讲解

卷积神经网络被大规模的应用在分类任务中，输出的结果是整个图像的类标签。但是UNet是像素级分类，输出的则是每个像素点的类别，且不同类别的像素会显示不同颜色，UNet常常用在生物医学图像上，而该任务中图片数据往往较少。所以，Ciresan等人训练了一个卷积神经网络，用滑动窗口提供像素的周围区域（patch）作为输入来预测每个像素的类标签。

  • 优点
    • 输出结果可以定位出目标类别的位置；
    • 由于输入的训练数据是patches，这样就相当于进行了数据增强，从而解决了生物医学图像数量少的问题，数据增强有利于模型的训练
  • 缺点
    • 训练过程较慢，网络必须训练每个patches，由于每个patches具有较多的重叠部分，这样持续训练patches，就会导致相当多的图片特征被多次训练，造成资源的浪费，导致训练时间加长且效率会低下。但是也会认为网络对这个特征进行多次训练，会对这个特征影响十分深刻，从而准确率得到改进。但是这里你拿一张图片复制100次去训练，很可能会出现过拟合的现象，对于这张图片确实十分敏感，但是拿另外一张图片来就可能识别不出了啦
    • 定位准确性和获取上下文信息不可兼得，大的patches需要更多的max-pooling，这样会减少定位准确性，因为最大池化会丢失目标像素和周围像素之间的空间关系，而小patches只能看到很小的局部信息，包含的背景信息不够。

网络结构原理

UNet网络结构，最主要的两个特点是：U型网络结构和Skip Connection跳层连接。

UNet网络结构分为三个部分，原理图如下：

  • 第一部分是主干特征提取部分，我们可以利用主干部分获得一个又一个的特征层，Unet的主干特征提取部分与VGG相似，为卷积和最大池化的堆叠。利用主干特征提取部分我们可以获得五个初步有效特征层，在第二步中，我们会利用这五个有效特征层可以进行特征融合。

    • 下采样

    • 左边特征提取网络：使用conv和pooling，就是每次向下采样之前都会进行两次的卷积操作，然后向下采样，然后再进行两次卷积操作，以此往复，向下连续采样五次

  • 第二部分是加强特征提取部分，我们可以利用主干部分获取到的五个初步有效特征层进行上采样，并且进行特征融合，获得一个最终的，融合了所有特征的有效特征层。

    • 上采样

    • 右边网络为特征融合网络：使用上采样产生的特征图与左侧特征图进行concatenate操作

    • Skip Connection中间四条灰色的平行线，Skip Connection就是在上采样的过程中，融合下采样过过程中的feature map。Skip Connection用到的融合的操作也很简单，就是将feature map的通道进行叠加，俗称Concat。

    • Concat操作也很好理解，举个例子：一本大小为10cm10cm，厚度为3cm的书A，和一本大小为10cm10cm，厚度为4cm的书B。将书A和书B，边缘对齐地摞在一起。这样就得到了，大小为10cm*10cm厚度为7cm的一摞书（就是直接把书叠起来的意思）

    • 对于feature map，一个大小为256 256 64的feature map，即feature map的w（宽）为256，h（高）为256，c（通道数）为64。和一个大小为256 256 32的feature map进行Concat融合，就会得到一个大小为256 256 96的feature map。

      在实际使用中，Concat融合的两个feature map的大小不一定相同，例如256 256 64的feature map和240 240 32的feature map进行Concat。

      这种时候，就有两种办法：

      • 第一种：将大256 256 64的feature map进行裁剪，裁剪为240 240 64的feature map，比如上下左右，各舍弃8 pixel，裁剪后再进行Concat，得到240 240 96的feature map。

      • 第二种：将小240 240 32的feature map进行padding操作，padding为256 256 32的feature map，比如上下左右，各补8 pixel，padding后再进行Concat，得到256 256 96的feature map。

      UNet采用的Concat方案就是第二种，将小的feature map进行padding，padding的方式是补0，一种常规的常量填充。

  • 第三部分是预测部分，我们会利用最终获得的最后一个有效特征层对每一个特征点进行分类，相当于对每一个像素点进行分类。（将最后特征层调整通道数，也就是我们要分类个数）

    • 最后再经过两次卷积操作，生成特征图，再用两个卷积核大小为1*1的卷积做分类得到最后的两张heatmap，例如第一张表示第一类的得分，第二张表示第二类的得分heatmap，然后作为softmax函数的输入，算出概率比较大的softmax，然后再进行loss，反向传播计算。

网络代码实现

按照UNet的网络结构分parts去实现Unet结构，采取一种搭积木的方式，先定义各个独立的模块，最后组合拼接就可以！

DoubleConv模块

如下图所示模块，连续的两个卷积的操作，在整个UNet网络中，主干特征提取网络和加强特征网络中各自使用了五次，每一层都会采取这个操作，故可以提取出来：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=0),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=0),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )
    def forward(self, x):
        return self.double_conv(x)

  • nn.Sequential 是一个时许的容器，会将里面的 modle 逐一执行，执行顺序为：卷积->BN->ReLU->卷积->BN->ReLU。

  • in_channels, out_channels，输入输出通道定义为参数，增强扩展使用

  • 卷积 nn.Conv2d 的输出：

    • nn. Conv2d(in_channels, out_channels, kernel_size, stride=1, padding=0,dilation=1, groups=1, bias=True, padding_mode= ‘zeros’ )

      • in_channels:输入的四维张量[N, C, H, W]中的C，也就是说输入张量的channels数。这个形参是确定权重等可学习参数的shape所必需的。
      • out_channels:也很好理解，即期望的四维输出张量的channels数，不再多说。
      • kernel_size:卷积核的大小，一般我们会使用5x5、3x3这种左右两个数相同的卷积核，因此这种情况只需要写kernel_size = 5这样的就行了。如果左右两个数不同，比如3x5的卷积核，那么写作kernel_size = (3, 5)，注意需要写一个tuple，而不能写一个列表（list）。
        stride = 1:卷积核在图像窗口上每次平移的间隔，即所谓的步长。这个概念和Tensorflow等其他框架没什么区别，不再多言。
      • padding:这是Pytorch与Tensorflow在卷积层实现上最大的差别，padding也就是指图像填充，后面的int型常数代表填充的多少（行数、列数），默认为0。需要注意的是这里的填充包括图像的上下左右，以padding=1为例，若原始图像大小为32 32，那么padding后的图像大小就变成了34 34，而不是33*33。
        Pytorch不同于Tensorflow的地方在于，Tensorflow提供的是padding的模式，比如same、valid，且不同模式对应了不同的输出图像尺寸计算公式。而Pytorch则需要手动输入padding的数量，当然，Pytorch这种实现好处就在于输出图像尺寸计算公式是唯一的，
      • dilation:这个参数决定了是否采用空洞卷积，默认为1（不采用）。从中文上来讲，这个参数的意义从卷积核上的一个参数到另一个参数需要走过的距离，那当然默认是1了，毕竟不可能两个不同的参数占同一个地方吧（为0）。更形象和直观的图示可以观察Github上的Dilated convolution animations，展示了dilation=2的情况。
      • groups:决定了是否采用分组卷积，groups参数可以参考groups参数详解
      • bias:即是否要添加偏置参数作为可学习参数的一个，默认为True。
      • padding_mode:即padding的模式，默认采用零填充。

    • 输出通道就是 out_channels

    • 输出的 X * X 计算公式：

- I 为输入feature map的大小，O为输出feature map的大小，K为卷积核的大小，P为padding的大小，S为步长
Down（下采样模块）

UNet的下采样模块有着4次的下采样过程，过程如下

1
2
3
4
5
6
7
8
9
10
class Down(nn.Module):
    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)
  • 代码很简单，就是一个maxpool池化层，进行下采样，然后接一个DoubleConv模块。
  • 到这里，左边的网络完成！！
Up（上采样模块）

上采样模块就是出来常规的上采样操作以外，还需要进行特征融合，

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
class Up(nn.Module):
   
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.up = nn.ConvTranspose2d(in_channels, 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 = 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])
     
        x = torch.cat([x2, x1], dim=1)
        return self.conv(x)

  • 初始化函数里定义的上采样方法（反卷积）以及卷积采用DoubleConv

  • 反卷积，顾名思义，就是反着卷积。卷积是让featuer map越来越小，反卷积就是让feature map越来越大，

    下面蓝色为原始图片，周围白色的虚线方块为padding结果，通常为0，上面绿色为卷积后的图片。

    这个示意图，就是一个从 2 * 2的feature map —-> 4 * 4 的feature map过程。

    在forward前向传播函数中，x1接收的是上采样的数据，x2接收的是特征融合的数据。特征融合方法就是，上文提到的，先对小的feature map进行padding，再进行concat。

OutConv模块

用上述的DoubleConv模块、Down模块、Up模块就可以拼出UNet的主体网络结构了。UNet网络的输出需要根据分割数量，整合输出通道。

利用前面的模块，我们可以获取输入进来的图片的特征，此时，我们需要利用特征获得预测结果

利用特征获得预测结果的过程为：

  • 利用一个1x1卷积进行通道调整，将最终特征层的通道数调整成num_classes。 （即对每一个像素点进行分类）

这个过程简单，顺便也包装一下吧

1
2
3
4
5
6
7
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)

到这里，所有的积木已经完成了，接下来就是搭建的过程了。

UNet模块

到这里，按照UNet网络结构，设置每个模块的输入输出通道个数以及调用顺序，代码如下：

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
import torch.nn.functional as F
import torch.nn as nn
from nets.net_of_me.unet_parts import *
class UNet(nn.Module):
    def __init__(self, n_channels, n_classes, bilinear=False):
        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, 1024)
        self.up1 = Up(1024, 512, bilinear)
        self.up2 = Up(512, 256, bilinear)
        self.up3 = Up(256, 128, 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

训练网络

训练网络的代码：

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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
import os
import time
import numpy as np
import torch
import torch.backends.cudnn as cudnn
import torch.optim as optim
from torch.utils.data import DataLoader
from tqdm import tqdm
from nets.net_of_me.unet_model import UNet
from nets.unet_training import CE_Loss, Dice_loss, LossHistory
from utils.dataloader import DeeplabDataset, deeplab_dataset_collate
from utils.metrics import f_score

def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']

def fit_one_epoch(net,epoch,epoch_size,epoch_size_val,gen,genval,Epoch,cuda):
    net = net.train()
    total_loss = 0
    total_f_score = 0

    val_toal_loss = 0
    val_total_f_score = 0
    start_time = time.time()

    print('Start Training')
    for iteration, batch in enumerate(gen):
        if iteration >= epoch_size:
            break
        imgs, pngs, labels = batch
        with torch.no_grad():
            imgs = torch.from_numpy(imgs).type(torch.FloatTensor)
            pngs = torch.from_numpy(pngs).type(torch.FloatTensor).long()
            labels = torch.from_numpy(labels).type(torch.FloatTensor)
            if cuda:
                imgs = imgs.cuda()
                pngs = pngs.cuda()
                labels = labels.cuda()
        optimizer.zero_grad()
        #进行训练
        outputs = net(imgs)
        loss    = CE_Loss(outputs, pngs, num_classes = NUM_CLASSES)
        if dice_loss:
            main_dice = Dice_loss(outputs, labels)
            loss      = loss + main_dice

        with torch.no_grad():
            #-------------------------------#
            #   计算f_score
            #-------------------------------#
            _f_score = f_score(outputs, labels)

        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        total_f_score += _f_score.item()

        waste_time = time.time() - start_time #训练epoch需要的时间
        start_time = time.time()

        if (iteration % 50 == 0):
            print("epoch = {} and loss = {} and waste_time = {}".format(epoch,loss.item(),waste_time))
            #写入日志文件
            with open("log/train_logs.txt", "a") as f:  # 格式化字符串还能这么用！
                f.write("epoch = {} and loss = {}".format(epoch,loss.item()) + "\n")

    print('Finish Training')

    print('Start Validation')
    for iteration, batch in enumerate(genval):
        if iteration >= epoch_size_val:
            break
        imgs, pngs, labels = batch
        with torch.no_grad():
            imgs = torch.from_numpy(imgs).type(torch.FloatTensor)
            pngs = torch.from_numpy(pngs).type(torch.FloatTensor).long()
            labels = torch.from_numpy(labels).type(torch.FloatTensor)
            if cuda:
                imgs = imgs.cuda()
                pngs = pngs.cuda()
                labels = labels.cuda()
			# 开始训练
            outputs = net(imgs)
            #计算损失函数
            val_loss = CE_Loss(outputs, pngs, num_classes=NUM_CLASSES)
            if dice_loss:
                main_dice = Dice_loss(outputs, labels)
                val_loss = val_loss + main_dice
            # -------------------------------#
            #   计算f_score
            # -------------------------------#
            _f_score = f_score(outputs, labels)

            val_toal_loss += val_loss.item()
            val_total_f_score += _f_score.item()

            if (iteration % 50 == 0):
                print("epoch = {} and val_loss = {} ".format(epoch, val_loss.item()))
                # 写入日志文件
                with open("log/val_logs.txt", "a") as f:  # 格式化字符串还能这么用！
                    f.write("epoch = {} and loss = {}".format(epoch, val_loss.item()) + "\n")

    print('Finish Validation')
    print('Epoch:' + str(epoch + 1) + '/' + str(Epoch))
    print('Total Loss: %.4f || Val Loss: %.4f ' % (total_loss / (epoch_size + 1), val_toal_loss / (epoch_size_val + 1)))

    print('Saving state, iter:', str(epoch + 1))
    torch.save(model.state_dict(), 'model/Epoch%d-Total_Loss%.4f-%.4f.pth' % ((epoch + 1), total_loss / (epoch_size + 1), val_toal_loss / (epoch_size_val + 1)))



if __name__ == "__main__":
    #------------------------------#
    #   输入图片的大小
    #------------------------------#
    inputs_size = [512,512,3]
    #---------------------#
    #   分类个数+1
    #   2+1
    #---------------------#
    NUM_CLASSES = 21
    #   Cuda的使用
    #-------------------------------#
    Cuda = True
    #linux服务器
    dataset_path = "/data/xwd/pro_datas/VOCdevkit/VOC2007"

    #网络
    model = UNet(n_channels=inputs_size[-1], n_classes=NUM_CLASSES).train()

    if Cuda:
        net = torch.nn.DataParallel(model)
        cudnn.benchmark = True
        net = net.cuda()

    # 打开训练数据集的txt
    with open(os.path.join(dataset_path, "ImageSets/Segmentation/train.txt"),"r") as f:
        train_lines = f.readlines()

    # 打开验证数据集的txt
    with open(os.path.join(dataset_path, "ImageSets/Segmentation/val.txt"),"r") as f:
        val_lines = f.readlines()

    if True:
        lr = 1e-4
        Init_Epoch = 0
        Interval_Epoch = 5
        Batch_size = 4

        optimizer = optim.Adam(model.parameters(), lr)  #优化器
        lr_scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.92) #学习率的调整
        
        #封装数据
        train_dataset = DeeplabDataset(train_lines, inputs_size, NUM_CLASSES, True, dataset_path)
        val_dataset = DeeplabDataset(val_lines, inputs_size, NUM_CLASSES, False, dataset_path)
        gen = DataLoader(train_dataset, batch_size=Batch_size, num_workers=4, pin_memory=True,
                         drop_last=True, collate_fn=deeplab_dataset_collate)
        gen_val = DataLoader(val_dataset, batch_size=Batch_size, num_workers=4, pin_memory=True,
                             drop_last=True, collate_fn=deeplab_dataset_collate)

        epoch_size = len(train_lines) // Batch_size
        epoch_size_val = len(val_lines) // Batch_size

        if epoch_size == 0 or epoch_size_val == 0:
            raise ValueError("数据集过小，无法进行训练，请扩充数据集。")

        for epoch in range(Init_Epoch, Interval_Epoch):
            fit_one_epoch(model, epoch, epoch_size, epoch_size_val, gen, gen_val, Interval_Epoch, Cuda)
            lr_scheduler.step()

封装数据集的办法主要采用：自定义类继承Dataset，下面展示的是他的伪代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# ================================================================== #
#                Input pipeline for custom dataset                 #
# ================================================================== 
# You should build your custom dataset as below.
class CustomDataset(torch.utils.data.Dataset):
    def __init__(self):
        # TODO
        # 1. Initialize file paths or a list of file names. 
        pass
    def __getitem__(self, index):
        # TODO
        # 1. Read one data from file (e.g. using numpy.fromfile, PIL.Image.open).
        # 2. Preprocess the data (e.g. torchvision.Transform).
        # 3. Return a data pair (e.g. image and label).
        pass
    def __len__(self):
        # You should change 0 to the total size of your dataset.
        return 0 
# You can then use the prebuilt data loader. 
custom_dataset = CustomDataset()
train_loader = torch.utils.data.DataLoader(dataset=custom_dataset,
                                           batch_size=64, 
                                           shuffle=True)
  • init函数是这个类的初始化函数，根据指定的图片路径，读取所有图片数据，
  • len函数可以返回数据的多少，这个类实例化后，通过len()函数调用。
  • getitem函数是数据获取函数，在这个函数里你可以写数据怎么读，怎么处理，并且可以一些数据预处理、数据增强都可以在这里进行

下面的是自定义的这个方法：

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
class DeeplabDataset(Dataset):
    def __init__(self,train_lines,image_size,num_classes,random_data,dataset_path):
        super(DeeplabDataset, self).__init__()

        self.train_lines    = train_lines
        self.train_batches  = len(train_lines)
        self.image_size     = image_size
        self.num_classes    = num_classes
        self.random_data    = random_data
        self.dataset_path   = dataset_path

    def __len__(self):
        return self.train_batches

    def rand(self, a=0, b=1):
        return np.random.rand() * (b - a) + a

    def get_random_data(self, image, label, input_shape, jitter=.3, hue=.1, sat=1.5, val=1.5):
        label = Image.fromarray(np.array(label))

        h, w = input_shape
        # resize image
        rand_jit1 = rand(1-jitter,1+jitter)
        rand_jit2 = rand(1-jitter,1+jitter)
        new_ar = w/h * rand_jit1/rand_jit2

        scale = rand(0.5,1.5)
        if new_ar < 1:
            nh = int(scale*h)
            nw = int(nh*new_ar)
        else:
            nw = int(scale*w)
            nh = int(nw/new_ar)
        image = image.resize((nw,nh), Image.BICUBIC)
        label = label.resize((nw,nh), Image.NEAREST)
        label = label.convert("L")
        
        # flip image or not
        flip = rand()<.5
        if flip: 
            image = image.transpose(Image.FLIP_LEFT_RIGHT)
            label = label.transpose(Image.FLIP_LEFT_RIGHT)
        
        # place image
        dx = int(rand(0, w-nw))
        dy = int(rand(0, h-nh))
        new_image = Image.new('RGB', (w,h), (128,128,128))
        new_label = Image.new('L', (w,h), (0))
        new_image.paste(image, (dx, dy))
        new_label.paste(label, (dx, dy))
        image = new_image
        label = new_label

        # distort image
        hue = rand(-hue, hue)
        sat = rand(1, sat) if rand()<.5 else 1/rand(1, sat)
        val = rand(1, val) if rand()<.5 else 1/rand(1, val)
        x = cv2.cvtColor(np.array(image,np.float32)/255, cv2.COLOR_RGB2HSV)
        x[..., 0] += hue*360
        x[..., 0][x[..., 0]>1] -= 1
        x[..., 0][x[..., 0]<0] += 1
        x[..., 1] *= sat
        x[..., 2] *= val
        x[x[:,:, 0]>360, 0] = 360
        x[:, :, 1:][x[:, :, 1:]>1] = 1
        x[x<0] = 0
        image_data = cv2.cvtColor(x, cv2.COLOR_HSV2RGB)*255
        return image_data,label
    def __getitem__(self, index):
        if index == 0:
            shuffle(self.train_lines)  
        annotation_line = self.train_lines[index]
        name = annotation_line.split()[0]
        # 从文件中读取图像
        jpg = Image.open(os.path.join(os.path.join(self.dataset_path, "JPEGImages"), name + ".jpg"))
        png = Image.open(os.path.join(os.path.join(self.dataset_path, "SegmentationClass"), name + ".png"))

        if self.random_data:
            jpg, png = self.get_random_data(jpg,png,(int(self.image_size[1]),int(self.image_size[0])))
        else:
            jpg, png = letterbox_image(jpg, png, (int(self.image_size[1]),int(self.image_size[0])))
        png = np.array(png)
        png[png >= self.num_classes] = self.num_classes
        #-------------------------------------------------------#
        #   转化成one_hot的形式
        #   在这里需要+1是因为voc数据集有些标签具有白边部分
        #   我们需要将白边部分进行忽略，+1的目的是方便忽略。
        #-------------------------------------------------------#
        seg_labels = np.eye(self.num_classes+1)[png.reshape([-1])]
        seg_labels = seg_labels.reshape((int(self.image_size[1]),int(self.image_size[0]),self.num_classes+1))
        jpg = np.transpose(np.array(jpg),[2,0,1])/255
        return jpg, png, seg_labels

我这边设置的epoch并不算很大，采用3090的显卡也是运行了一段时间是时间，可以看到网络，loss实在逐渐在收敛的：

采用训练好的模型进行预测，看看结果如何：

这边采用的是在网络上copy的图片预处理和后续处理的代码，本人目前对图片处理还是比较菜，把别人的代码贴在这里，最后给出自己的预测结果：

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
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
import colorsys
import copy
import time
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn
from nets.net_of_me.unet_model import UNet
#-------------------------------------------#
#   使用自己训练好的模型预测需要修改2个参数
#   model_path和num_classes都需要修改！
#   如果出现shape不匹配
#   一定要注意训练时的model_path和num_classes数的修改
#--------------------------------------------#
class Unet(object):
    _defaults = {
        "model_path"        : 'model\Epoch2-Total_Loss1.0039-0.8573.pth', #保存的训练模型的路径
        "model_image_size"  : (512, 512, 3), #输入图片的大小
        "num_classes"       : 21, 
        "cuda"              : True,
        #--------------------------------#
        #   blend参数用于控制是否
        #   让识别结果和原图混合
        #--------------------------------#
        "blend"             : True
    }

    #---------------------------------------------------#
    #   初始化UNET
    #---------------------------------------------------#
    def __init__(self, **kwargs):
        self.__dict__.update(self._defaults)
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.generate()

    #---------------------------------------------------#
    #   获得所有的分类
    #---------------------------------------------------#
    def generate(self):
        self.net = UNet(n_channels=self.model_image_size[-1],n_classes=self.num_classes).eval()

        # 加载本地的模型参数
        state_dict = torch.load(self.model_path,self.device)
        self.net.load_state_dict(state_dict)

        if self.cuda:
            self.net = nn.DataParallel(self.net) #可以调用多个GPU，帮助加速训练
            self.net = self.net.to(self.device)

        print('{} model loaded.'.format(self.model_path))

        if self.num_classes <= 21:
            self.colors = [(0, 0, 0), (128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128), (128, 0, 128),                              (0, 128, 128),  (128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0), 
                           (192, 128, 0), (64, 0, 128), (192, 0, 128), (64, 128, 128), (192, 128, 128), 
                           (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64,128),(128, 64, 12)]
        else:
            # 画框设置不同的颜色
            hsv_tuples = [(x / len(self.class_names), 1., 1.)
                        for x in range(len(self.class_names))]
            self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
            self.colors = list(
                map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)),
                    self.colors))

    def letterbox_image(self ,image, size):
        image = image.convert("RGB")
        iw, ih = image.size
        w, h = size
        scale = min(w/iw, h/ih)
        nw = int(iw*scale)
        nh = int(ih*scale)

        image = image.resize((nw,nh), Image.BICUBIC)
        new_image = Image.new('RGB', size, (128,128,128))
        new_image.paste(image, ((w-nw)//2, (h-nh)//2))
        return new_image,nw,nh

    #---------------------------------------------------#
    #   检测图片，处理图片
    #---------------------------------------------------#
    def detect_image(self, image):
        #---------------------------------------------------------#
        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
        #---------------------------------------------------------#
        image = image.convert('RGB')
        
        #---------------------------------------------------#
        #   对输入图像进行一个备份，后面用于绘图
        #---------------------------------------------------#
        old_img = copy.deepcopy(image)

        orininal_h = np.array(image).shape[0]
        orininal_w = np.array(image).shape[1]


        #---------------------------------------------------#
        #   进行不失真的resize，添加灰条，进行图像归一化
        #---------------------------------------------------#
        image, nw, nh = self.letterbox_image(image,(self.model_image_size[1],self.model_image_size[0]))
        a = np.array(image).shape
        images = [np.array(image)/255]

        images = np.transpose(images,(0,3,1,2))

        #---------------------------------------------------#
        #   图片传入网络进行预测
        #---------------------------------------------------#
        with torch.no_grad():
            images = torch.from_numpy(images).type(torch.FloatTensor) #转化为tensor
            if self.cuda:
                #images =images.cuda()
                images = images.cpu()

            pr = self.net(images)

            pr = pr[0]
            pr1 = pr[1]
            #---------------------------------------------------#
            #   取出每一个像素点的种类
            #---------------------------------------------------#
            pr = F.softmax(pr.permute(1,2,0),dim = -1).cpu().numpy().argmax(axis=-1)
            #--------------------------------------#
            #   将灰条部分截取掉
            #--------------------------------------#
            pr = pr[int((self.model_image_size[0]-nh)//2):int((self.model_image_size[0]-nh)//2+nh), 
                    int((self.model_image_size[1]-nw)//2):int((self.model_image_size[1]-nw)//2+nw)]

        #------------------------------------------------#
        #   创建一副新图，并根据每个像素点的种类赋予颜色
        #------------------------------------------------#
        seg_img = np.zeros((np.shape(pr)[0],np.shape(pr)[1],3))
        for c in range(self.num_classes):
            seg_img[:,:,0] += ((pr[:,: ] == c )*( self.colors[c][0] )).astype('uint8')
            seg_img[:,:,1] += ((pr[:,: ] == c )*( self.colors[c][1] )).astype('uint8')
            seg_img[:,:,2] += ((pr[:,: ] == c )*( self.colors[c][2] )).astype('uint8')

        #------------------------------------------------#
        #   将新图片转换成Image的形式
        #------------------------------------------------#
        image = Image.fromarray(np.uint8(seg_img)).resize((orininal_w,orininal_h))
        #------------------------------------------------#
        #   将新图片和原图片混合
        #------------------------------------------------#
        if self.blend:
             image = Image.blend(old_img,image,0.7)      
        return image,  old_img


if __name__ == "__main__":
    unet = Unet()
    while True:
        img = input('Input image filename:')
        try:
            image = Image.open(img)
        except:
            print('Open Error! Try again!')
            continue
        else:
            r_image,old_image= unet.detect_image(image)
            old_image.show()
            r_image.show()

调用这个函数，得到的预测结果如下：

语义分割的MIOU指标

语义分割的标准度量。其计算所有类别交集和并集之比的平均值.，语义分割说到底也还是一个分割任务，既然是一个分割的任务，预测的结果往往就是四种情况：

  • true positive（TP）：预测正确, 预测结果是正类, 真实是正类

  • false positive（FP）：预测错误, 预测结果是正类, 真实是负类

  • true negative（TN）：预测错误, 预测结果是负类, 真实是正类

  • false negative（FN）：预测正确, 预测结果是负类, 真实是负类

mIOU 的定义：计算真实值和预测值两个集合的交集和并集之比。这个比例可以变形为TP（交集）比上TP、FP、FN之和（并集）。即：mIOU=TP/(FP+FN+TP)。

计算公式：

等价于：

mIOU一般都是基于类进行计算的，将每一类的IOU计算之后累加，再进行平均，得到的就是基于全局的评价。

MIoU：计算两圆交集（橙色部分）与两圆并集（红色+橙色+黄色）之间的比例，理想情况下两圆重合，比例为1。

计算本网络的MIoU可以采样训练好的模型进行计算，计算的结果比例越接近1效果越好。

代码实现后续把，hhhhhhhhhhhhhhhhhhhhh。。。。。。。