gitbook/PyTorch深度学习实战/docs/455415.md
2022-09-03 22:05:03 +08:00

524 lines
24 KiB
Markdown
Raw Blame History

This file contains invisible Unicode characters

This file contains invisible Unicode characters that are indistinguishable to humans but may be processed differently by a computer. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# 20 | 图像分割(下):如何构建一个图像分割模型?
你好,我是方远。
在上一节课中,我们掌握了图像分割的理论知识,你是不是已经迫不及待要上手体验一下,找找手感了呢?
今天我们就从头开始来完成一个图像分割项目。项目的内容是对图片中的小猫进行语义分割。为了实现这个项目我会引入一个简单但实用的网络结构UNet。通过这节课的学习你不但能再次体验一下完整机器学习的模型实现过程还能实际训练一个语义分割模型。
课程代码你可以从[这里](https://github.com/syuu1987/geekTime-semantic-segmentation/tree/main)下载。
## 数据部分
我们还是从机器学习开发三件套:数据、训练、评估说起。首先是数据准备部分,我们先对训练数据进行标记,然后完成数据读取工作。
### 分割图像的标记
之前也提到过图像分割的准备相比图像分类的准备更加复杂。那我们如何标记语义分割所需要的图片呢在图像分割中我们使用的每张图片都要有一张与之对应的Mask如下所示
![图片](https://static001.geekbang.org/resource/image/c2/db/c258c4f2ffd1f819c662aa1e9f6a8cdb.jpg?wh=1024x640)
![图片](https://static001.geekbang.org/resource/image/1a/a0/1a35623ceccb0750cd8058568d847fa0.png?wh=1024x640)
上节课我们说过Mask就是含有像素类别的特征图。结合这里的示例图片我们可以看到Mask就是原图所对应的一张图片它的每个位置都记录着原图每个位置对应的像素类别。对于Mask的标记我们需要使用到Labelme工具。
标记的方法一共包括七步,我们挨个看一下。
第一步,下载安装[Labelme](https://github.com/wkentaro/labelme)。我们按照Github中的安装方式进行安装即可。如果安装比较慢的话你可以使用国内的镜像例如清华的进行安装。
第二步我们要将需要标记的图⽚放到⼀个⽂件夹中。这里我是将所有猫的图片放入到cats文件夹中了。
![图片](https://static001.geekbang.org/resource/image/f3/04/f3c8cc99959c74f363ec290558a51d04.png?wh=1722x882)
第三步我们事先准备好⼀个label.txt的⽂件⾥⾯每⼀⾏写好的需要标记的类别。我的label.txt如下
```python
__ignore__
_background_
cat
```
这里我要提醒你的是前两行最好这么写。不这样写的话使用label2voc.py转换就会报错但label2voc.py不是唯一的数据转换方式还可以使用labelme\_json\_to\_dataset但推荐你使用label2voc.py。从第三行开始表示要标记的类别。
第四步执行后面的这条命令就会自动启动Labelme。
```python
labelme --labels labels.txt --nodata
```
第五步点我们击左侧的Open Dir选择第二步中的文件夹就会自动导入需要标记的图片。在右下角选择需要标记的文件后会自动显示出来如下图所示。
![图片](https://static001.geekbang.org/resource/image/0d/81/0d37591417e44d51a21bac11f409c381.png?wh=1780x1470)
第六步点击左侧的Create Polygons。就可以开始标注了。标记的方式就是将小猫沿着它的边界给圈出来当形成一个闭环的时候Labelme会自动提示你输入类别我们选择cat类即可。
标记成功后,结果如下图所示。
![图片](https://static001.geekbang.org/resource/image/88/y6/888dc5b576ccc9629cd1f3fd2d9cbyy6.png?wh=1676x1438)
当标记完成后我们需要保存一下保存之后会生成标记好的json文件。如下所示
```python
fangyuan@geektime data $ ls cats
1.jpeg  1.json  10.jpeg 10.json 2.jpeg  3.jpeg  4.jpeg  4.json
```
第七步执行下面的代码将标记好的数据转换成Mask。
```python
python label2voc.py cats cats_output --label label.txt 
```
上面代码里用到的label2voc.py你可以通过后面这个链接获取它[https://github.com/wkentaro/labelme/blob/main/examples/semantic\_segmentation/labelme2voc.py](https://github.com/wkentaro/labelme/blob/main/examples/semantic_segmentation/labelme2voc.py)。
其中cats为标记好的数据cats\_output为输出文件夹。在cats\_output下会自动生成4个文件夹我们只需要两个文件夹分别是JPEGImages训练原图与SegmentationClassPNG转换后的Mask
到此为止我们的数据就准备好了。我一共标记了8张图片如下所示。当然了在实际的项目中需要大量标记好的图片这里主要是为了方便演示。
![图片](https://static001.geekbang.org/resource/image/7c/1d/7ca1ecafd3ce893610c5eb89yy7ca51d.png?wh=686x285)
![图片](https://static001.geekbang.org/resource/image/53/67/53e088956a41de56ea1010af8a2a6d67.png?wh=703x324)
到此为止,标记工作宣告完成。
### 数据读取
完成了标记工作之后我们就要用PyTorch把这些数据给读入进来了我们把数据相关的写在dataset.py中。具体还是和之前讲的一样要继承Dataset类然后实现\_\_init\_\_、\_\_len\_\_和\_\_getitem\_\_方法。
dataset.py的代码如下所示我已经在代码中写好注释了相信结合注释你很容易就能领会意思。
```python
import os
import torch
import numpy as np
from torch.utils.data import Dataset
from PIL import Image
class CatSegmentationDataset(Dataset):
# 模型输入是3通道数据
in_channels = 3
# 模型输出是1通道数据
out_channels = 1
def __init__(
self,
images_dir,
image_size=256,
):
print("Reading images...")
# 原图所在的位置
image_root_path = images_dir + os.sep + 'JPEGImages'
# Mask所在的位置
mask_root_path = images_dir + os.sep + 'SegmentationClassPNG'
# 将图片与Mask读入后分别存在image_slices与mask_slices中
self.image_slices = []
self.mask_slices = []
for im_name in os.listdir(image_root_path):
# 原图与mask的名字是相同的只不过是后缀不一样
mask_name = im_name.split('.')[0] + '.png'
image_path = image_root_path + os.sep + im_name
mask_path = mask_root_path + os.sep + mask_name
im = np.asarray(Image.open(image_path).resize((image_size, image_size)))
mask = np.asarray(Image.open(mask_path).resize((image_size, image_size)))
self.image_slices.append(im / 255.)
self.mask_slices.append(mask)
def __len__(self):
return len(self.image_slices)
def __getitem__(self, idx):
image = self.image_slices[idx]
mask = self.mask_slices[idx]
# tensor的顺序是Batch_size, 通道而numpy读入后的顺序是(高,宽,通道)
image = image.transpose(2, 0, 1)
# Mask是单通道数据所以要再加一个维度
mask = mask[np.newaxis, :, :]
image = image.astype(np.float32)
mask = mask.astype(np.float32)
return image, mask
```
然后我们的训练代码写在train.py中train.py中的main函数为主函数在main中我们会调用data\_loaders来加载数据。代码如下所示
```python
import torch
from torch.utils.data import DataLoader
from torch.utils.data import DataLoader
from dataset import CatSegmentationDataset as Dataset
def data_loaders(args):
dataset_train = Dataset(
images_dir=args.images,
image_size=args.image_size,
)
loader_train = DataLoader(
dataset_train,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers,
)
return loader_train
# args是传入的参数
def main(args):
loader_train = data_loaders(args)
```
以上就是数据处理的全部内容了。接下来,我们再来看看模型训练部分的内容。
## 模型训练
我们先来回忆一下,模型训练的老三样,分别是网络结构、损失函数和优化方法。
先从网络结构说起今天我要为你介绍一个叫做UNet的语义分割网络。
### 网络结构UNet
[UNet](https://arxiv.org/pdf/1505.04597.pdf)是一个非常实用的网络。它是一个典型的Encoder-Decoder类型的分割网络网络结构非常简单如下图所示。
![图片](https://static001.geekbang.org/resource/image/11/b9/1196c6fcff2fe8c601f608b01bf82ab9.jpg?wh=1920x1130 "图片来自论文https://arxiv.org/pdf/1505.04597.pdf")
它的网络结构虽然简单但是效果并不“简单”我在很多项目中都用它与一些主流的语义分割做对比而UNet都取得了非常好的效果。
整体网络结构跟论文中给出的示意图一样,我们重点去关注几个实现细节。
第一点图中横向蓝色的箭头它们都是重复的相同结构都是由两个3x3的卷积层组合而成的在每层卷积之后会跟随一个BN层与ReLU的激活层。按照[第14节课](https://time.geekbang.org/column/article/442442)讲的这一部分重复的组织是可以单独提取出来的。我们先来创建一个unet.py文件用来定义网络结构。
现在unet.py中创建Block类它是用来定义刚才所说的重复的卷积块
```python
class Block(nn.Module):
    def __init__(self, in_channels, features):
        super(Block, self).__init__()
        self.features = features
        self.conv1 = nn.Conv2d(
                            in_channels=in_channels,
                            out_channels=features,
                            kernel_size=3,
                            padding='same',
                        )
        self.conv2 = nn.Conv2d(
                            in_channels=features,
                            out_channels=features,
                            kernel_size=3,
                            padding='same',
                        )
    def forward(self, input):
        x = self.conv1(input)
        x = nn.BatchNorm2d(num_features=self.features)(x)
        x = nn.ReLU(inplace=True)(x)
        x = self.conv2(x)
        x = nn.BatchNorm2d(num_features=self.features)(x)
        x = nn.ReLU(inplace=True)(x)
        return x
```
这里需要注意的是同一个块内特征图的尺寸是不变的所以padding为same。
第二点,就是绿色向上的箭头,也就是上采样的过程。这块的实现就是采用上一节课所讲的转置卷积来实现的。
最后一点我们现在是要对小猫进行分割也就是说一共有两个类别——猫与背景。对于二分类的问题我们可以直接输出一张特征图然后通过概率来进行判断是正例还是负例背景也就是下面代码中的第71行。同时下述代码也补全了unet.py中的所有代码。
```python
import torch
import torch.nn as nn
class Block(nn.Module):
...
class UNet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1, init_features=32):
        super(UNet, self).__init__()
        features = init_features
        self.conv_encoder_1 = Block(in_channels, features)
        self.conv_encoder_2 = Block(features, features * 2)
        self.conv_encoder_3 = Block(features * 2, features * 4)
        self.conv_encoder_4 = Block(features * 4, features * 8)
        self.bottleneck = Block(features * 8, features * 16)
        self.upconv4 = nn.ConvTranspose2d(
            features * 16, features * 8, kernel_size=2, stride=2
        )
        self.conv_decoder_4 = Block((features * 8) * 2, features * 8)
        self.upconv3 = nn.ConvTranspose2d(
            features * 8, features * 4, kernel_size=2, stride=2
        )
        self.conv_decoder_3 = Block((features * 4) * 2, features * 4)
        self.upconv2 = nn.ConvTranspose2d(
            features * 4, features * 2, kernel_size=2, stride=2
        )
        self.conv_decoder_2 = Block((features * 2) * 2, features * 2)
        self.upconv1 = nn.ConvTranspose2d(
            features * 2, features, kernel_size=2, stride=2
        )
        self.decoder1 = Block(features * 2, features)
        self.conv = nn.Conv2d(
            in_channels=features, out_channels=out_channels, kernel_size=1
        )
    def forward(self, x):
        conv_encoder_1_1 = self.conv_encoder_1(x)
        conv_encoder_1_2 = nn.MaxPool2d(kernel_size=2, stride=2)(conv_encoder_1_1)
        conv_encoder_2_1 = self.conv_encoder_2(conv_encoder_1_2)
        conv_encoder_2_2 = nn.MaxPool2d(kernel_size=2, stride=2)(conv_encoder_2_1)
        conv_encoder_3_1 = self.conv_encoder_3(conv_encoder_2_2)
        conv_encoder_3_2 = nn.MaxPool2d(kernel_size=2, stride=2)(conv_encoder_3_1)
        conv_encoder_4_1 = self.conv_encoder_4(conv_encoder_3_2)
        conv_encoder_4_2 = nn.MaxPool2d(kernel_size=2, stride=2)(conv_encoder_4_1)
        bottleneck = self.bottleneck(conv_encoder_4_2)
        conv_decoder_4_1 = self.upconv4(bottleneck)
        conv_decoder_4_2 = torch.cat((conv_decoder_4_1, conv_encoder_4_1), dim=1)
        conv_decoder_4_3 = self.conv_decoder_4(conv_decoder_4_2)
        conv_decoder_3_1 = self.upconv3(conv_decoder_4_3)
        conv_decoder_3_2 = torch.cat((conv_decoder_3_1, conv_encoder_3_1), dim=1)
        conv_decoder_3_3 = self.conv_decoder_3(conv_decoder_3_2)
        conv_decoder_2_1 = self.upconv2(conv_decoder_3_3)
        conv_decoder_2_2 = torch.cat((conv_decoder_2_1, conv_encoder_2_1), dim=1)
        conv_decoder_2_3 = self.conv_decoder_2(conv_decoder_2_2)
        conv_decoder_1_1 = self.upconv1(conv_decoder_2_3)
        conv_decoder_1_2 = torch.cat((conv_decoder_1_1, conv_encoder_1_1), dim=1)
        conv_decoder_1_3 = self.decoder1(conv_decoder_1_2)
        return torch.sigmoid(self.conv(conv_decoder_1_3))
```
到这里,网络结构我们就搭建好了,然后我们来我看看损失函数。
### 损失函数Dice Loss
这里我们来看一下语义分割中常用的损失函数Dice Loss。
想要知道这个损失函数如何生成你需要先了解一个语义分割的评价指标但更常用的还是后面要讲的的mIoU它就是Dice系数常用于计算两个集合的相似度取值范围在0-1之间。
Dice系数的公式如下。
$$Dice=\\frac{2|P\\cap G|}{|P|+|G|}$$
其中,$|P\\cap G|$是集合P与集合G之间交集元素的个数$|P|$和$|G|$分别表示集合P和G的元素个数。分子的系数2这是为了抵消分母中P和G之间的共同元素。对语义分割任务而言集合P就是预测值的Mask集合G就是真实值的Mask。
根据Dice系数我们就能设计出一种损失函数也就是Dice Loss。它的计算公式非常简单如下所示。
$$Dice Loss=1-\\frac{2|P\\cap G|}{|P|+|G|}$$
从公式中可以看出当预测值的Mask与GT越相似损失就越小当预测值的Mask与GT差异度越大损失就越大。
对于二分类问题GT只有0和1两个值。当我们直接使用模型输出的预测概率而不是使用阈值将它们转换为二值Mask时这种损失函数就被称为Soft Dice Loss。此时$|P\\cap G|$的值近似为GT与预测概率矩阵的点乘。
定义损失函数的代码如下。
```python
import torch.nn as nn
class DiceLoss(nn.Module):
    def __init__(self):
        super(DiceLoss, self).__init__()
        self.smooth = 1.0
    def forward(self, y_pred, y_true):
        assert y_pred.size() == y_true.size()
        y_pred = y_pred[:, 0].contiguous().view(-1)
        y_true = y_true[:, 0].contiguous().view(-1)
        intersection = (y_pred * y_true).sum()
        dsc = (2. * intersection + self.smooth) / (
            y_pred.sum() + y_true.sum() + self.smooth
        )
        return 1. - dsc
```
其中self.smooth是一个平滑值这是为了防止分子和分母为0的情况。
### 训练流程
最后,我们将模型、损失函数和优化方法串起来,看下整体的训练流程,训练的代码如下。
```python
def main(args):
    makedirs(args)
# 根据cuda可用情况选择使用cpu或gpu
    device = torch.device("cpu" if not torch.cuda.is_available() else args.device)
# 加载训练数据
    loader_train = data_loaders(args)
# 实例化UNet网络模型
    unet = UNet(in_channels=Dataset.in_channels, out_channels=Dataset.out_channels)
    # 将模型送入gpu或cpu中
unet.to(device)
# 损失函数
    dsc_loss = DiceLoss()
# 优化方法
    optimizer = optim.Adam(unet.parameters(), lr=args.lr)
    loss_train = []
    step = 0
# 训练n个Epoch
    for epoch in tqdm(range(args.epochs), total=args.epochs):
        unet.train()
        for i, data in enumerate(loader_train):
            step += 1
            x, y_true = data
            x, y_true = x.to(device), y_true.to(device)
            y_pred = unet(x)
            optimizer.zero_grad()
            loss = dsc_loss(y_pred, y_true)
            loss_train.append(loss.item())
            loss.backward()
            optimizer.step()
            if (step + 1) % 10 == 0:
                print('Step ', step, 'Loss', np.mean(loss_train))
                loss_train = []
        torch.save(unet, args.weights + '/unet_epoch_{}.pth'.format(epoch))
```
需要注意的点,我都在注释中进行了说明,你可以自己看一看。其实就是我们一直说的模型训练的那几件事情:数据加载、构建网络以及迭代更新网络参数。
我用训练数据训练了若干个Epoch同时也保存了若干个模型保存为pth格式。到这里就完成了模型训练的整个环节我们可以使用保存好的模型进行预测来看看分割效果如何。
## 模型预测
现在我们要用训练生成的模型来进行语义分割,看看结果是什么样子的。
模型预测的代码如下。
```python
import torch
import numpy as np
from PIL import Image
img_size = (256, 256)
# 加载模型
unet = torch.load('./weights/unet_epoch_51.pth')
unet.eval()
# 加载并处理输入图片
ori_image = Image.open('data/JPEGImages/6.jpg')
im = np.asarray(ori_image.resize(img_size))
im = im / 255.
im = im.transpose(2, 0, 1)
im = im[np.newaxis, :, :]
im = im.astype('float32')
# 模型预测
output = unet(torch.from_numpy(im)).detach().numpy()
# 模型输出转化为Mask图片
output = np.squeeze(output)
output = np.where(output>0.5, 1, 0).astype(np.uint8)
mask = Image.fromarray(output, mode='P')
mask.putpalette([0,0,0, 0,128,0])
mask = mask.resize(ori_image.size)
mask.save('output.png')
```
这段代码也很好理解。首先用torch.load函数加载模型。接着加载一张待分割的图片并进行数据预处理。然后将处理好的数据送入模型中得到预测值output。最后将预测值转化为可视化的Mask图片进行保存。
输入图片也就是待分割的图片如下左图所示。最终的输出即可视化的Mask图片如下右图所示。
![图片](https://static001.geekbang.org/resource/image/c2/db/c258c4f2ffd1f819c662aa1e9f6a8cdb.jpeg?wh=1024x640)
![图片](https://static001.geekbang.org/resource/image/fb/61/fbfecd56d8c31589890fcd05c7995461.png?wh=1024x640)
在将预测值转化为Mask图片的过程中最终预测值的概率卡了0.5的阈值超过阈值的像素点在output矩阵中的值为1表示猫的区域没有超过阈值的像素点在output矩阵中的值为0表示背景区域。
为了将output矩阵输出为可视化的图像我们使用Image.fromarray函数将Numpy的array转化为Image格式并将模式设置为“P”即调色板模式。然后用putpalette函数来给Image对象上色。
其中putpalette函数的参数是一个列表\[0, 0, 0, 0, 128, 0\]列表前三个数表示值为0的像素的RGB\[0, 0, 0\]表示黑色列表后三个数表示值为1的像素的RGB\[0, 128, 0\]表示绿色。这样我们保存的Mask图片黑色部分即为背景区域绿色部分即为猫的区域。
不过这样分开的轮廓图可能无法让我们很直观地看出语义分割的效果。所以我们将原图和Mask合成一张图片来看看效果。具体的代码如下。
```python
image = ori_image.convert('RGBA')
mask = mask.convert('RGBA')
# 合成
image_mask = Image.blend(image, mask, 0.3)
image_mask.save("output_mask.png")
```
首先我们将原图image和Mask图片都转换为RGBA带透明度的模式。然后使用Image.blend函数将两张图片合成一张图片最后一个参数0.3表示Mask图片透明度为30%原图的透明度为70%。
最终的结果如下图所示。
![图片](https://static001.geekbang.org/resource/image/4d/7b/4d804527a87cc92aab8173da85f0ff7b.png?wh=1024x640)
这样我们就可以直观地看出哪些地方预测得不准确了。
## 模型评估
在语义分割中常用的评价指标是mIoU。mIoU全称为mean Intersection over Union即平均交并比。交并比是真实值和预测值的交集和并集之比。
真实值就是我们刚刚用labelme标注的Mask也是Ground TruthGT。如下左图所示。
预测值就是模型预测出的Mask用Prediction表示。如后面右图所示。
![图片](https://static001.geekbang.org/resource/image/61/0b/61afb79172dfa0bd652f237fd1c5bd0b.png?wh=1024x640)
![图片](https://static001.geekbang.org/resource/image/c3/9d/c31ec50a2a67262728a9fd8e84a1729d.png?wh=1024x640)
交集是指真实值与预测值的交集,如下图黄色区域所示。并集是指真实值与预测值的并集,如下图蓝色区域所示。
![图片](https://static001.geekbang.org/resource/image/53/29/53e81fae1ceb5f21a269c3a461c6b129.png?wh=1024x640)
通过上面几个图我们很容易就能理解mIoU了。mIoU的公式如下所示。
$$mIoU=\\frac{1}{k}\\sum\_{i=1}^{k}{\\frac{P\\cap G}{P\\cup G}}$$
其中k为所有类别数在我们的例子中只有“cat”一类因此k为1我们通常不将背景计算到mIoU中P为预测值G是真实值。
## 小结
恭喜你,完成了今天的学习任务。这节课我们一起完成了一个图像分割项目的实践。
首先我带你了解了图像分割的数据准备需要使用Labelme工具为图像做标记。数据质量的好坏决定了最终模型的质量所以你要对数据的标注好好把握。在使用Labelme标记完成之后我们可以使用label2voc.py将json转换为Mask。
之后我们学习了一种非常高效且实用的模型UNet并使用PyTorch实现了其网络结构。
然后我为你讲解了图像分割的评估指标mIoU和损失函数Dice Loss。
mIoU的公式如下
$$mIoU=\\frac{1}{k}\\sum\_{i=1}^{k}{\\frac{P\\cap G}{P\\cup G}}$$
mIoU主要是从预测结果与GT的重合度这一角度来衡量分割模型的好与坏的它是图像分割中经常使用的评价指标。
最后,我们使用训练好的模型进行预测,并对分割结果进行了可视化绘制。相信通过之前学习的图像分类项目与今天学习的图像分割项目,对于图像处理,你会获得更深层次的理解。
## 每课一练
你可以根据今天的内容,自己动手试试建立一个图像分割模型,然后用一张图片来测一下效果如何。
欢迎你在留言区跟我交流讨论,也推荐你把今天的内容分享给更多同事、朋友,我们下节课见。