PyTorch12—Unet图像语义分割

语义分割简介

图像语义分割是计算机视觉中十分重要的领域。它是指像素级地识别图像，即标注出图像中每个像素所属的对象类别。下图为语义分割的一个实例，其目标是预测出图像中每一个像素的类标签。

图像语义分割是图像处理和是计算机视觉技术中关于图像理解的重要一环，也是 AI 领域中一个重要的分支。 语义分割对图像中每一个像素点进行分类，确定每个点的类别（如属于背景、边缘或身体等）。

这里需要和实例分割区分开来。它没有分离同一类的实例；我们关心的只是每个像素的类别，如果输入对象中有两个相同类别的对象，则分割本身不将他们区分为单独的对象。

---

图像语义分割应用：

•  自动驾驶汽车：我们需要为汽车增加必要的感知，以了解他们所处的环境，以便自动
• 驾驶的汽车可以安全行驶；
•  医学图像诊断：机器可以增强放射医生进行的分析，大大减少了运行诊断测试所需的时间；
• 无人机着陆点判断等

下图是对街景的语义分割；

图像语义分割实质：一般是将一张RGB图像（height*width*3）或是灰度图（height*width*1）作为输入，输出的是分割图，其中每一个像素包含了其类别的标签（height*width*1）。

图像语义分割的实现 

目前在图像分割领域比较成功的算法，有很大一部分都来自于同一个先驱：Long等人提出的Fully Convolutional Network（FCN），或者叫全卷积网络。FCN将分类网络转换成用于分割任务的网络结构，并证明了在分割问题上，可以实现端到端的网络训练。FCN成为了深度学习解决分割问题的奠基石。

分类网络结构尽管表面上来看可以接受任意尺寸的图片作为输入，但是由于网络结构最后全连接层的存在，使其丢失了输入的空间信息，因此，这些网络并没有办法直接用于解决诸如分割等稠密估计的问题。考虑到这一点，FCN用卷积层和池化层替代了分类网络中的全连接层，从而使得网络结构可以适应像素级的稠密估计任务。

语义分割的UNET网络结构 

Unet是2015年诞生的模型，它几乎是当前segmentation项目中应用最广的模型。 Unet能从更少的训练图像中进行学习。当它在少于 40 张图的生物医学数据集上训练时，IOU 值仍能达到 92%。

Unet的左侧是convolution layers，右侧则是upsamping layers，convolutions layers中每个pooling layer前输出值会concatenate到对应的upsamping层的输入值中。注意是concatenate，而FCN是add。

Unet已经成为大多做医疗影像语义分割任务的最基础的网络结构。也启发了大量研究者去思考U型语义分割网络。 即使在自然影像理解方面，也有越来越多的语义分割和目标检测模型
开始关注和使用U型结构。

U-net网络非常简单，前半部分作用是特征提取，后半部分是上采样。在一些文献中也把这样的结构叫做编码器-解码器结构。由于此网络整体结构类似于大写的英文字母U，故得名U-net。

U-net与其他常见的分割网络有一点非常不同的地方：U-net采用了完全不同的特征融合方式：拼接（tf.concat）。

U-net采用将特征在channel维度拼接在一起，形成更厚的特征。而FCN融合时使用的对应点相加，并不形成更厚的特征。

语义分割网络在特征融合时有两种办法：

• 1. FCN式的对应点相加，对应于TensorFlow中的tf.add()函数；
• 2. U-net式的channel维度拼接融合，对应于TensorFlow的tf.concat()函数，比较占显存。

UNET网络结构特点：

• 1、网络对图像特征的多尺度特征识别。
• 2、上采样部分会融合特征提取部分的输出，这样做实际上是将多尺度特征融合在了一起，以最后一个上采样为例，它的特征既来自第一个卷积block的输出(同尺度特征)，也来自上采样的输出(大尺度特征)。

代码实战

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils import data
import numpy as np
import matplotlib.pyplot as plt
import torchvision
from torchvision import transforms
import os
import glob
from PIL import Image

BATCH_SIZE = 8

plt.figure('测试一张图片')
pil_img = Image.open(r'hk\training\00001.png') # 读取一张图片，使之成为PIL对象
np_img = np.array(pil_img) # 将PIL对象转换为ndarray形式
plt.imshow(np_img)
plt.show()

plt.figure('测试一张蒙版图片')
pil_img = Image.open(r'hk\training\00001_matte.png') # 对应的蒙版图片
np_img = np.array(pil_img)
plt.imshow(np_img)
plt.show()

 

---

print('np.unique(np_img)：\n',np.unique(np_img)) # 目标图是一个二分类问题。我们希望它的唯一值只有两类，但是事与愿违
print(f'np_img.max():{np_img.max()}, np_img.min():{np_img.min()}') # 不符合我们预期。我们希望它只有两类
"""
因为在实际的标注过程中，边缘的地方的像素取值都包含了，我们要人为的将其转换为两类就可以了。
np_img[np_img>0]=1
0代表人，1代表背景
"""
print(f'np_img.shape:{np_img.shape}')   # (800, 600)

 

---

创建Datasset 

all_pics = glob.glob(r'hk\training\*.png') # 获取所有图片路径
print(f'all_pics[:5]:{all_pics[:5]}')
images = [p for p in all_pics if 'matte' not in p] # 由于原图和目标图像混合在一起的。我们希望原图与蒙版图分开
annotations = [p for p in all_pics if 'matte' in p] # annotations是蒙版图
print(f'len(images):{len(images)}，len(annotations):{len(annotations)}') # 1700 1700

# 确保图片与蒙版图片是一一对应的
images[:5]
annotations[:5]

# 打乱顺序
np.random.seed(2021)
index = np.random.permutation(len(images))
images = np.array(images)[index] # 转换成array，方便索引
anno = np.array(annotations)[index]
# 测试打乱顺序后原图与蒙版图是否一一对应
print(f'images[:5]:\n{images[:5]}')
print(f'anno[:5]:\n{anno[:5]}')
print(f'images[-5:]:\n{images[-5:]}')
print(f'anno[-5:]:\n{anno[-5:]}')


all_test_pics = glob.glob(r'hk\testing\*.png')
test_images = [p for p in all_test_pics if 'matte' not in p]
test_anno = [p for p in all_test_pics if 'matte' in p]



# 输入的图片进行转换的方法
transform = transforms.Compose([
                    transforms.Resize((256, 256)), # 将图片转换成 256x256。
                    transforms.ToTensor(),
])

class Portrait_dataset(data.Dataset):
    def __init__(self, img_paths, anno_paths):
        self.imgs = img_paths # 原图路径
        self.annos = anno_paths # 蒙版图片路径

    def __getitem__(self, index):
        img = self.imgs[index]
        anno = self.annos[index]

        pil_img = Image.open(img)
        img_tensor = transform(pil_img)  # 转换成tensor，对原图的处理。

        pil_anno = Image.open(anno)
        anno_tensor = transform(pil_anno)

        # 由于蒙版图都是黑白图，会产生channel为1的维度。经过转换后，256x256x1，这个1并不是我们需要的。
        anno_tensor = torch.squeeze(anno_tensor).type(torch.long)  # 将channel去掉。squeeze方法会把维度为1的去掉，比如你的维度是3x4x1，squeeze之后就是3x4。
        # torch.long代表整形，目标值必须是0、1整数。
        anno_tensor[anno_tensor > 0] = 1  # 语义分割。二分类。tensor也可以向量运算

        return img_tensor, anno_tensor

    def __len__(self):
        return len(self.imgs)

train_dataset = Portrait_dataset(images, anno)
test_dataset = Portrait_dataset(test_images, test_anno)

# 创建dataloader
train_dl = data.DataLoader(
                           train_dataset,
                           batch_size=BATCH_SIZE,
                           shuffle=True,
)
test_dl = data.DataLoader(
                          test_dataset,
                          batch_size=BATCH_SIZE,
)

# 取出一个批次的数据玩玩
imgs_batch, annos_batch = next(iter(train_dl))
print(f'imgs_batch.shape:{imgs_batch.shape},annos_batch.shape:{annos_batch.shape}')
# imgs_batch.shape:torch.Size([8, 3, 256, 256])
# annos_batch.shape:torch.Size([8, 256, 256])

img = imgs_batch[0].permute(1,2,0).numpy()  # permute方法转换维度，将channel放在最后一个维度
anno = annos_batch[0].numpy()
plt.subplot(1,2,1)
plt.imshow(img)
plt.subplot(1,2,2)
plt.imshow(anno)
plt.show()

---

创建下采样模型、上采样模型和Unet结构

class Downsample(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(Downsample, self).__init__()

        # 两层卷积加激活
        self.conv_relu = nn.Sequential(
                            nn.Conv2d(in_channels, out_channels,
                                      kernel_size=3, padding=1),    # padding=1，希望图像经过卷积之后大小不变。
                            nn.ReLU(inplace=True),  # inplace=True，就地改变

                            # 第二层卷积，输入是多少个channel，输出仍然是多少个channel。
                            nn.Conv2d(out_channels, out_channels,
                                      kernel_size=3, padding=1),
                            nn.ReLU(inplace=True)
            )
        # 下采样
        self.pool = nn.MaxPool2d(kernel_size=2)
    def forward(self, x, is_pool=True):
        if is_pool: # 是否需要进行下采样
            x = self.pool(x)    # 先下采样，再卷积。
        x = self.conv_relu(x)
        return x


# 上采样模型。卷积、卷积、上采样（反卷积实现上采样）
class Upsample(nn.Module):
    # 上采样中，输出channel会变成输入channel的一半。所以并不需要给两个channel值。看unet网络模型就知道了。
    def __init__(self, channels):
        super(Upsample, self).__init__()
        self.conv_relu = nn.Sequential(
            nn.Conv2d(2 * channels, channels,
                      kernel_size=3, padding=1),
            nn.ReLU(inplace=True),
            # 第二层卷积，channel不变。
            nn.Conv2d(channels, channels,
                      kernel_size=3, padding=1),
            nn.ReLU(inplace=True)
        )
        # 上采样，也需要激活。
        self.upconv_relu = nn.Sequential(
            # 反卷积层
            nn.ConvTranspose2d(channels,    # 输入channel
                               channels // 2,   # 输出channel
                               kernel_size=3,
                               stride=2,    # 跨度必须为2，才能放大，长和宽会变为原来的两倍
                               padding=1,   # 输入的kernel所在的位置，起始位置在图像的第一个像素做反卷积。
                               output_padding=1),   # 反卷积之后，在最外层（周边）做的填充。边缘填充为1，
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = self.conv_relu(x)
        x = self.upconv_relu(x)
        return x

# 创建Unet。我们要初始化上、下采样层，还有其他的一些层
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.down1 = Downsample(3, 64)  # 输入的是一张图片，channel是3。
        self.down2 = Downsample(64, 128)
        self.down3 = Downsample(128, 256)
        self.down4 = Downsample(256, 512)
        self.down5 = Downsample(512, 1024)

        self.up = nn.Sequential(
            nn.ConvTranspose2d(1024,
                               512,
                               kernel_size=3,
                               stride=2,
                               padding=1,
                               output_padding=1),
            nn.ReLU(inplace=True)
        )

        self.up1 = Upsample(512)
        self.up2 = Upsample(256)
        self.up3 = Upsample(128)

        self.conv_2 = Downsample(128, 64)   # 最后的两层卷积
        self.last = nn.Conv2d(64, 2, kernel_size=1) # 输出层。

    # 前向传播。
    def forward(self, x):
        x1 = self.down1(x, is_pool=False)
        x2 = self.down2(x1)
        x3 = self.down3(x2)
        x4 = self.down4(x3)
        x5 = self.down5(x4)

        x5 = self.up(x5)

        # 我们需要将x4的输出将x5（up上采样的输出）做一个合并。使得channel变厚。
        # 将前面下采样的特征合并过来,有利于重复利用这些特征，有利于模型的每一个像素进行分类。

        x5 = torch.cat([x4, x5], dim=1)  # 32*32*1024。沿着channel这个维度进行合并。[batch channel height weight]
        x5 = self.up1(x5)  # 64*64*256) # 上采样
        x5 = torch.cat([x3, x5], dim=1)  # 64*64*512
        x5 = self.up2(x5)  # 128*128*128
        x5 = torch.cat([x2, x5], dim=1)  # 128*128*256
        x5 = self.up3(x5)  # 256*256*64
        x5 = torch.cat([x1, x5], dim=1)  # 256*256*128

        x5 = self.conv_2(x5, is_pool=False)  # 256*256*64。最后的两层卷积

        x5 = self.last(x5)  # 256*256*3 # 最后的输出层。每一个像素点都输出为长度为2的向量。
        return x5

print('\n创建模型')
model = Net()
if torch.cuda.is_available():
    model.to('cuda')

---

开始训练

loss_fn = nn.CrossEntropyLoss() # 二分类。输出长度为2的一个张量，损失是CrossEntropyLoss。看着两个值哪个大
from torch.optim import lr_scheduler
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)  # 优化函数
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)   # 学习速率衰减


def fit(epoch, model, trainloader, testloader):
    correct = 0
    total = 0
    running_loss = 0

    model.train()
    for x, y in trainloader:
        if torch.cuda.is_available():
            x, y = x.to('cuda'), y.to('cuda')
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        with torch.no_grad():
            y_pred = torch.argmax(y_pred, dim=1)
            correct += (y_pred == y).sum().item()
            total += y.size(0)
            running_loss += loss.item()
    exp_lr_scheduler.step()
    epoch_loss = running_loss / len(trainloader.dataset)
    epoch_acc = correct / (total * 256 * 256)

    test_correct = 0
    test_total = 0
    test_running_loss = 0

    model.eval()
    with torch.no_grad():
        for x, y in testloader:
            if torch.cuda.is_available():
                x, y = x.to('cuda'), y.to('cuda')
            y_pred = model(x)
            loss = loss_fn(y_pred, y)
            y_pred = torch.argmax(y_pred, dim=1)
            test_correct += (y_pred == y).sum().item()
            test_total += y.size(0)
            test_running_loss += loss.item()

    epoch_test_loss = test_running_loss / len(testloader.dataset)
    epoch_test_acc = test_correct / (test_total * 256 * 256)

    print('epoch: ', epoch,
          'loss： ', round(epoch_loss, 3),
          'accuracy:', round(epoch_acc, 3),
          'test_loss： ', round(epoch_test_loss, 3),
          'test_accuracy:', round(epoch_test_acc, 3)
          )

    return epoch_loss, epoch_acc, epoch_test_loss, epoch_test_acc


epochs = 10
train_loss = []
train_acc = []
test_loss = []
test_acc = []

for epoch in range(epochs):
    epoch_loss, epoch_acc, epoch_test_loss, epoch_test_acc = fit(epoch,
                                                                 model,
                                                                 train_dl,
                                                                 test_dl)
    train_loss.append(epoch_loss)
    train_acc.append(epoch_acc)
    test_loss.append(epoch_test_loss)
    test_acc.append(epoch_test_acc)

# 保存模型
PATH = 'unet_model.pth'
torch.save(model.state_dict(), PATH)    # 保存的是模型的参数

# 测试模型
my_model = Net()
my_model.load_state_dict(torch.load(PATH))

# 在测试数据集上预测
num=3   # 绘制3张图片
image, mask = next(iter(test_dl))
pred_mask = my_model(image)

plt.figure(figsize=(10, 10))
for i in range(num):
    plt.subplot(num, 3, i*num+1)
    plt.imshow(image[i].permute(1,2,0).cpu().numpy())   # 原图。由于是在显卡上运行，所以需要将它放到cpu上
    plt.subplot(num, 3, i*num+2)
    plt.imshow(mask[i].cpu().numpy())   # 实际的分割图
    plt.subplot(num, 3, i*num+3)

    # 由于进行了permute，维度放在了最后一个。所以是在最后一个维度上进行argmax
    # detach取出实际的data，没有梯度的data。
    plt.imshow(torch.argmax(pred_mask[i].permute(1,2,0), axis=-1).detach().numpy()) # 预测的分割图
    plt.show()

# 在train数据集上测试
image, mask = next(iter(train_dl))
pred_mask = my_model(image)

plt.figure(figsize=(10, 10))
for i in range(num):
    plt.subplot(num, 3, i*num+1)
    plt.imshow(image[i].permute(1,2,0).cpu().numpy())
    plt.subplot(num, 3, i*num+2)
    plt.imshow(mask[i].cpu().numpy())
    plt.subplot(num, 3, i*num+3)
    plt.imshow(torch.argmax(pred_mask[i].permute(1,2,0), axis=-1).detach().numpy())

---

用模型进行预测 

# 预测应用
path = 'z1.jpg'
pil_img = Image.open(path)
img_tensor = transform(pil_img)

# 添加维度（添加batch维度）
img_tensor_batch = torch.unsqueeze(img_tensor,0)    # 在第0个维度添加 维度值为1 的维度（1,3,256,256）
pred = my_model(img_tensor_batch)
print(pred.shape) # torch.Size([1,2,256,256])，1是批次大小
# pred[0] 就是取第一个元素，这个元素的形状就是[2,256,256]
pred_img = torch.argmax(pred[0].permute(1,2,0),axis=-1).numpy()
plt.subplot(1,2,1)
plt.imshow(img_tensor.permute(1,2,0 ).numpy())
plt.subplot(1,2,2)
plt.imshow(pred_img)
plt.show()