计算机视觉语义分割——U-Net(Convolutional Networks for Biomedical Image Segmentation)

摘要

本周的工作聚焦于U-Net网络在语义分割任务中的研究与实现。U-Net作为一种经典的全卷积神经网络，因其编码器-解码器的对称结构和跳跃连接机制，广泛应用于生物医学图像分割。本次工作详细解析了U-Net的网络结构，包括收缩路径、扩展路径，以及其与FCN的关系。此外，还探讨了U-Net在生成模型中的改进及其在Stable Diffusion中的应用。通过搭建基于PyTorch的U-Net网络，完成了数据预处理、模型训练及预测，并在医学图像分割任务中取得了良好的效果。实验表明，U-Net可以有效减少信息丢失，提升分割精度，为后续生成模型中的应用奠定了基础。

Abstract

This week’s work focused on the study and implementation of the U-Net network in semantic segmentation tasks. U-Net, as a classic fully convolutional neural network, is widely used in biomedical image segmentation due to its symmetrical encoder-decoder structure and skip connection mechanism. This report provides a detailed analysis of the U-Net architecture, including its contracting path, expanding path, and its relationship with FCN. Additionally, improvements to U-Net in generative models and its application in Stable Diffusion were explored. By building a PyTorch-based U-Net network, tasks such as data preprocessing, model training, and prediction were completed, achieving promising results in medical image segmentation. The experiments demonstrate that U-Net effectively reduces information loss and improves segmentation accuracy, laying a solid foundation for its application in generative models.

一、U-Net

1. 基本思想

U-Net由 Olaf Ronneberger 等人为生物医学图像分割而开发。该架构包含两条路径。第一条路径是收缩路径（也称为编码器），用于捕获图像中的上下文。编码器只是卷积层和最大池化层的传统堆栈。第二条路径是扩展路径（也称为解码器），用于使用转置卷积实现精确定位。因此，它是一个端到端的全卷积网络 (FCN)，即它只包含卷积层，不包含任何密集层，因此它可以接受任何大小的图像。

2. 收缩路径(Contracting path)

收缩路径使用卷积层和池化层的组合来提取和捕获图像中的特征，同时减少其空间维度。现在让我们来看看下面收缩路径中的5 个block中的每一个。

Block 1:

1. 一张尺寸为572² 的输入图像被输入到 U-Net 中。该输入图像仅包含1 个灰度通道。
2. 然后将两个3x3 卷积层 （无填充）应用于输入图像，每个层后跟一个ReLU层。同时将通道数增加到**64 个，**以便捕获更高级别的特征。
3. 然后应用步长为2 的2x2 最大池化层。这会将特征图下采样为其大小的一半，即284²。

Block 2:

1. 与Block 1 一样，两个3x3 卷积层（无填充）应用于Block 1 的输出，每个层后又跟有一个ReLU层。在每个新块中，特征通道数量加倍，现在为128。
2. 接下来，再次将2x2 最大池化层应用于生成的特征图，将空间尺寸减少一半至140²。

Block 3:

1. 与Block 1和 Block 2中的步骤相似，最后将特征图下采样至68²。

Block 4:

1. 与Block 1和 Block 2中的步骤相似，最后将特征图下采样至32²。

Block 5:

1. 在收缩路径的最后一个块中，特征通道数在每个块中翻倍，达到1024个。
2. 此块还包含两个3x3 卷积层（未填充），每个层后面都有一个ReLU层。但是，出于对称目的，我只包含一个层，并将第二层包含在扩展路径中。

在提取出复杂的特征之后，特征图就会进入到扩展路径的运算阶段。

3. 扩展路径(Expanding path)

扩展路径使用卷积和转置卷积操作来组合学习到的特征并对输入特征图进行上采样，直到生成分割图。扩展路径与收缩路径非常相似，下面将逐步了解每个Block。

注意：上图的灰色箭头代表“跳跃连接”，跳跃连接用于将图像直接从收缩路径发送到扩展路径，而无需经过所有块。这样可以保留和学习高级和低级特征，从而减少收缩路径期间发生的任何信息丢失。

Block 5：

1. 从收缩路径继续，应用第二个3x3 卷积（未填充），其后接着一个ReLU层。
2. 然后应用2x2 卷积（上卷积）层，将空间维度上采样两倍，并将通道数减半为512。

Block 4：

1. 然后使用跳跃连接将收缩路径中的相应特征图连接起来，将特征通道数加倍至1024。请注意，必须裁剪此连接以匹配扩展路径的尺寸。
2. 应用两个3x3 卷积层（未填充），每个层后面都有一个ReLU层，将通道数减少到512。
3. 之后，应用2x2 卷积（上卷积）层，将空间维度上采样两倍，并将通道数减半为256。

Block 3：

1. 与Block 5和 Block 4中的步骤相似，最后将特征图上采样两倍至200²，并将通道数减半为128。

Block 2：

1. 与Block 5和 Block 4中的步骤相似，最后将特征图上采样两倍至392²，并将通道数减半为64。

Block 1：

1. 在扩展路径的最后一块中，连接跳跃连接后共有128 个通道。
2. 接下来，在特征图上应用两个3x3 卷积层（未填充），中间的ReLU层将特征通道数减少到64。
3. 最后，使用1x1 卷积层，然后是激活层（二元分类中的Sigmoid函数），将通道数减少到所需的类别数。在本例中，为 2 个类别，因为二元分类通常用于医学成像。

在扩展路径上对特征图进行上采样后，应该生成分割图，其中每个像素都被单独分类。

4. U-Net代码实践

4.1 数据集选择

收集的30张果蝇的电镜细胞图，分辨率为512x512，上图第一张为原图，第二张则为标签图。

训练集图片：

训练集标签：

4.2 加载数据集

数据加载要做哪些处理，是根据任务和数据集而决定的，对于我们的分割任务，不用做太多处理，但由于数据量很少，仅30张，我们可以使用一些数据增强方法，来扩大我们的数据集。

Pytorch 给我们提供了一个方法，方便我们加载数据，我们可以使用这个框架，去加载我们的数据。看下伪代码：

# ================================================================== #
#                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)

这是一个标准的模板，我们就使用这个模板，来加载数据，定义标签，以及进行数据增强。

创建一个dataset.py文件，编写代码如下：

import torch
import cv2
import os
import glob
from torch.utils.data import Dataset
import random

class ISBI_Loader(Dataset):
    def __init__(self, data_path):
        # 初始化函数，读取所有data_path下的图片
        self.data_path = data_path
        self.imgs_path = glob.glob(os.path.join(data_path, 'image/*.png'))

    def augment(self, image, flipCode):
        # 使用cv2.flip进行数据增强，filpCode为1水平翻转，0垂直翻转，-1水平+垂直翻转
        flip = cv2.flip(image, flipCode)
        return flip
        
    def __getitem__(self, index):
        # 根据index读取图片
        image_path = self.imgs_path[index]
        # 根据image_path生成label_path
        label_path = image_path.replace('image', 'label')
        # 读取训练图片和标签图片
        image = cv2.imread(image_path)
        label = cv2.imread(label_path)
        # 将数据转为单通道的图片
        image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        label = cv2.cvtColor(label, cv2.COLOR_BGR2GRAY)
        image = image.reshape(1, image.shape[0], image.shape[1])
        label = label.reshape(1, label.shape[0], label.shape[1])
        # 处理标签，将像素值为255的改为1
        if label.max() > 1:
            label = label / 255
        # 随机进行数据增强，为2时不做处理
        flipCode = random.choice([-1, 0, 1, 2])
        if flipCode != 2:
            image = self.augment(image, flipCode)
            label = self.augment(label, flipCode)
        return image, label

    def __len__(self):
        # 返回训练集大小
        return len(self.imgs_path)

    
if __name__ == "__main__":
    isbi_dataset = ISBI_Loader("data/train/")
    print("数据个数：", len(isbi_dataset))
    train_loader = torch.utils.data.DataLoader(dataset=isbi_dataset,
                                               batch_size=2, 
                                               shuffle=True)
    for image, label in train_loader:
        print(image.shape)

使用dataset.py文件加载数据的输出结果为：

4.3 U-Net模型搭建

如果完全按照论文的结构，模型输出的尺寸会稍微小于图片输入的尺寸，如果使用论文的网络结构需要在结果输出后，做一个 resize 操作。为了省去这一步，我们可以修改网络，使网络的输出尺寸正好等于图片的输入尺寸。

创建unet_parts.py文件，编写如下代码：

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

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)

创建unet_model.py文件，编写如下代码：

"""利用unet_parts.py中的部件组装形成完整的Unet网络"""

import torch.nn.functional as F
from .unet_parts import *

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

if __name__ == '__main__':
    net = UNet(n_channels=3, n_classes=1)
    print(net)

这样调整过后，网络的输出尺寸就与图片的输入尺寸相同了。

4.4 训练U-Net

我们今天的任务，只需要分割出细胞边缘，也就是一个很简单的二分类任务，所以我们可以使用BCEWithLogitsLoss。

BCEWithLogitsLoss是 Pytorch 提供的用来计算二分类交叉熵的函数：

这个公式就是 Logistic 回归的损失函数，它利用的是 Sigmoid 函数阈值在[0,1]这个特性来进行分类的。

from model.unet_model import UNet
from utils.dataset import ISBI_Loader
from torch import optim
import torch.nn as nn
import torch

def train_net(net, device, data_path, epochs=40, batch_size=1, lr=0.00001):
    # 加载训练集
    isbi_dataset = ISBI_Loader(data_path)
    train_loader = torch.utils.data.DataLoader(dataset=isbi_dataset,
                                               batch_size=batch_size, 
                                               shuffle=True)
    # 定义RMSprop算法
    optimizer = optim.RMSprop(net.parameters(), lr=lr, weight_decay=1e-8, momentum=0.9)
    # 定义Loss算法
    criterion = nn.BCEWithLogitsLoss()
    # best_loss统计，初始化为正无穷
    best_loss = float('inf')
    # 训练epochs次
    for epoch in range(epochs):
        # 训练模式
        net.train()
        # 按照batch_size开始训练
        for image, label in train_loader:
            optimizer.zero_grad()
            # 将数据拷贝到device中
            image = image.to(device=device, dtype=torch.float32)
            label = label.to(device=device, dtype=torch.float32)
            # 使用网络参数，输出预测结果
            pred = net(image)
            # 计算loss
            loss = criterion(pred, label)
            print('Loss/train', loss.item())
            # 保存loss值最小的网络参数
            if loss < best_loss:
                best_loss = loss
                torch.save(net.state_dict(), 'best_model.pth')
            # 更新参数
            loss.backward()
            optimizer.step()

if __name__ == "__main__":
    # 选择设备，有cuda用cuda，没有就用cpu
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 加载网络，图片单通道1，分类为1。
    net = UNet(n_channels=1, n_classes=1)
    # 将网络拷贝到deivce中
    net.to(device=device)
    # 指定训练集地址，开始训练
    data_path = "data/train/"
    train_net(net, device, data_path)

训练的输出LOSS为：

D:\Pytorch\Anaconda3\envs\d2l\python.exe E:/Unet/train.py
Loss/train 0.7607325911521912
Loss/train 0.7042417526245117
Loss/train 0.6359086036682129
Loss/train 0.5970470905303955
Loss/train 0.5956597328186035
Loss/train 0.5168007016181946
Loss/train 0.509236216545105
Loss/train 0.4616139233112335
Loss/train 0.4775354862213135
Loss/train 0.44973015785217285
Loss/train 0.4584605395793915
Loss/train 0.4683246314525604
Loss/train 0.42488470673561096
Loss/train 0.4471128582954407
Loss/train 0.4211571216583252
..............................
Loss/train 0.08856159448623657
Loss/train 0.09392605721950531
Loss/train 0.10895076394081116
Loss/train 0.10066412389278412
Loss/train 0.10104762017726898
Loss/train 0.09606537222862244
Loss/train 0.09012659639120102
Loss/train 0.10549496114253998
Loss/train 0.0879889726638794
Loss/train 0.10525484383106232
Loss/train 0.10839355736970901

Process finished with exit code 0

4.5 预测

模型训练好了，我们可以用它在测试集上看下效果。在工程根目录创建 predict.py 文件，编写如下代码：

import glob
import numpy as np
import torch
import os
import cv2
from model.unet_model import UNet

if __name__ == "__main__":
    # 选择设备，有cuda用cuda，没有就用cpu
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    # 加载网络，图片单通道，分类为1。
    net = UNet(n_channels=1, n_classes=1)
    # 将网络拷贝到deivce中
    net.to(device=device)
    # 加载模型参数
    net.load_state_dict(torch.load('best_model.pth', map_location=device))
    # 测试模式
    net.eval()
    # 读取所有图片路径
    tests_path = glob.glob('data/test/*.png')
    # 遍历所有图片
    for test_path in tests_path:
        # 保存结果地址
        save_res_path = test_path.split('.')[0] + '_res.png'
        # 读取图片
        img = cv2.imread(test_path)
        # 转为灰度图
        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
        # 转为batch为1，通道为1，大小为512*512的数组
        img = img.reshape(1, 1, img