基于pytorch简单实现u-net

前言

​ 最近在看经典的卷积网络架构，打算自己尝试复现一下，在此系列文章中，会参考很多文章，有些已经忘记了出处，所以就不贴链接了，希望大家理解。

​ 完整的代码在最后。

本系列必须的基础

​ python基础知识、CNN原理知识、pytorch基础知识

本系列的目的

​ 一是帮助自己巩固知识点；

​ 二是自己实现一次，可以发现很多之前的不足；

​ 三是希望可以给大家一个参考。

参考资料

​ 来自b站大佬的项目库：

b站链接： https://space.bilibili.com/18161609
GitHub链接：https://github.com/WZMIAOMIAO/deep-learning-for-image-processing

目录结构

1. 前言：

​ 本篇文章打算简单实现一下图像分割中的u-net网络，并讲解一下实现的流程和思路。其中，代码有很多的参考别人的地方，有些方法甚至是直接拷贝过来用的（有现成的你不用—狗头保命），特此申明一下。当然，虽然有些代码是参考的，但是自己把主要的数据加载器、网络、训练等自己实现了一次。

​ 另外，这个实现不像之前实现分类网络，这个实现更复杂，因此代码讲解上会偏少，主要把思路理清。

2. 数据集介绍与下载：

下载

​ u_net主要应用于医学数据上的分割，这里使用的数据集名为DRIVE，可以从下面的链接中下载：

链接：https://pan.baidu.com/s/1YjW9wOm-sLC5oDs8T6a9bw 
提取码：0b0w 

介绍

​ 图像目录结构：

├─test
│  ├─1st_manual
│  ├─2nd_manual
│  ├─images
│  └─mask
└─training
    ├─1st_manual
    ├─images
    └─mask

​ 其中：

• manual是手工划分的真实图像分隔，黑色（0）为背景，白色（255）为对象
• images是真实图像，大小都为565*584*3，彩色的
• mask为一个蒙版，即黑色区域为不感兴趣的区域，白色区域为感兴趣的区域

3. 项目流程说明：

​ 整体流程如下：

• 首先，定义数据加载器，加载数据，并返回image和mask（这里的mask与数据集中的mask不同，不同点在下面的4小节中介绍）
• 接着，定义网络结构、损失函数，并开始训练
• 保存训练的权重参数，用于后期的预测

​ **提前说明：**当你把我写的代码下载后，需要修改路径参数，因为我的数据集是放在了和项目不同的文件夹下的。

​ 另外，我的整个项目文件夹结构为：

├─network_files
│  └─u_net.py
├─save_weights
├─utils
│  └─Loss.py
|  └─transforms.py
└─My_Dataset.py
└─train,py
└─predict.py

4. 数据加载器：

​ 这次的数据加载器实现起来很简单，主要思路如下：

​ 其中，我们需要特别对mask进行一定的处理：

​ 根据上述思路，可以完成代码（看注释）：

import os
import torch
from torch import nn
from torch.utils.data import Dataset
from PIL import Image
import numpy as np

class My_Dataset(Dataset):
    def __init__(self,root,train=True,transforms=None):
        '''
        :param root: 路径，比如.\data
        :param train: 加载的数据集为训练还是测试集，默认为训练集
        :param transforms:  预处理方法
        '''
        super(My_Dataset, self).__init__()
        # 获取基础路径
        self.flag = 'training' if train else 'test'
        self.root = os.path.join(root,'DRIVE',self.flag) # .\data\DRIVE\training
        # 获取所需的路径
        self.img_name = [i for i in os.listdir(os.path.join(self.root,'images'))] # '21_training.tif'
        self.img_list = [os.path.join(self.root,'images',i) for i in self.img_name] # .\data\DRIVE\training\images\21_training.tif
        self.manual = [os.path.join(self.root,'1st_manual',i.split('_')[0])+'_manual1.gif' for i in self.img_name ] # .\data\DRIVE\training\1st_manual\21_manual1.gif
        self.roi_mask = [os.path.join(self.root,'mask',i.split('_')[0]+'_'+self.flag+'_mask.gif') for i in self.img_name] # .\data\DRIVE\training\mask\21_training_mask.gif
        # 初始化其它变量
        self.transforms = transforms

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

    def __getitem__(self, idx):
        # 打开图像(RGB模式)、manual(L-灰度图)、roi_mask（L-灰度图）
        image = Image.open(self.img_list[idx])
        manual = Image.open(self.manual[idx]).convert('L')
        roi_mask = Image.open(self.roi_mask[idx]).convert('L')
        # 对图像进行处理
        # 将对象设置为1，背景设置为0，不感兴趣的区域设置为255
        manual = np.array(manual) / 255  # 白色（对象）为255，变为了1
        roi_mask = 255 - np.array(roi_mask) # 不感兴趣区域为黑色（0），变为了255 ； 感兴趣为白色（255），变为了0
        mask = np.clip(manual+roi_mask,a_min=0,a_max=255)
        mask = Image.fromarray(mask) # 转为PIL，方便做预处理操作
        if self.transforms is not None:
            image,mask = self.transforms(image,mask)
        return image,mask

5. 网络架构：

​ 这里用b站up主的图来描述网络架构：

​ 这里要注意几点：

• 原论文中每层双卷积后，图像尺寸都会减小，导致最后输出图像大小与输入大小不一致，这样一来，想要得到最终的分割图需要做一定的后处理。而我们这里双卷积不改变图像尺寸大小，只由下采样和上采样改变，如此一来，最终输出大小和输入大小一致，方便最后出结果。
• 下采样最后一层的双卷积，不改变图像通道数，仍然为512。是因为我们上采样用双线性插值，不改变通道数，如此一来，上采样后仍然为512，可以方便的和之前的进行拼接。
• 最终输出的通道数为2，是因为这里进行二分类，一个为背景，一个为分割对象。

​ 基于上述结构图，可以理清实现网络架构的思路图：

​ 可以轻松写出下面的代码：

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

# 上采样+拼接
class Up(nn.Module):
    def __init__(self,in_channels,out_channels,bilinear=True):
        '''
        :param in_channels: 输入通道数
        :param out_channels:  输出通道数
        :param bilinear: 是否采用双线性插值，默认采用
        '''
        super(Up, self).__init__()
        if bilinear:
            # 双线性差值
            self.up = nn.Upsample(scale_factor=2,mode='bilinear',align_corners=True)
            self.conv = doubleConv(in_channels,out_channels,in_channels//2) # 拼接后为1024，经历第一个卷积后512
        else:
            # 转置卷积实现上采样
            # 输出通道数减半，宽高增加一倍
            self.up = nn.ConvTranspose2d(in_channels,out_channels//2,kernel_size=2,stride=2)
            self.conv = doubleConv(in_channels,out_channels)

    def forward(self,x1,x2):
        # 上采样
        x1 = self.up(x1)
        # 拼接
        x = torch.cat([x1,x2],dim=1)
        # 经历双卷积
        x = self.conv(x)
        return x

# 双卷积层
def doubleConv(in_channels,out_channels,mid_channels=None):
    '''
    :param in_channels: 输入通道数 
    :param out_channels: 双卷积后输出的通道数
    :param mid_channels: 中间的通道数，这个主要针对的是最后一个下采样和上采样层
    :return: 
    '''
    if mid_channels is None:
        mid_channels = out_channels
    layer = []
    layer.append(nn.Conv2d(in_channels,mid_channels,kernel_size=3,padding=1,bias=False))
    layer.append(nn.BatchNorm2d(mid_channels))
    layer.append(nn.ReLU(inplace=True))
    layer.append(nn.Conv2d(mid_channels,out_channels,kernel_size=3,padding=1,bias=False))
    layer.append(nn.BatchNorm2d(out_channels))
    layer.append(nn.ReLU(inplace=True))
    return nn.Sequential(*layer)

# 下采样
def down(in_channels,out_channels):
    # 池化 + 双卷积
    layer = []
    layer.append(nn.MaxPool2d(2,stride=2))
    layer.append(doubleConv(in_channels,out_channels))
    return nn.Sequential(*layer)

# 整个网络架构
class U_net(nn.Module):
    def __init__(self,in_channels,out_channels,bilinear=True,base_channel=64):
        '''
        :param in_channels: 输入通道数，一般为3，即彩色图像
        :param out_channels: 输出通道数，即网络最后输出的通道数，一般为2，即进行2分类
        :param bilinear: 是否采用双线性插值来上采样，这里默认采取
        :param base_channel: 第一个卷积后的通道数，即64
        '''
        super(U_net, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.bilinear = bilinear

        # 输入
        self.in_conv = doubleConv(self.in_channels,base_channel)
        # 下采样
        self.down1 = down(base_channel,base_channel*2) # 64,128
        self.down2 = down(base_channel*2,base_channel*4) # 128,256
        self.down3 = down(base_channel*4,base_channel*8) # 256,512
        # 最后一个下采样，通道数不翻倍（因为双线性差值，不会改变通道数的，为了可以简单拼接，就不改变通道数）
        # 当然，是否采取双线新差值，还是由我们自己决定
        factor = 2  if self.bilinear else 1
        self.down4 = down(base_channel*8,base_channel*16 // factor) # 512,512
        # 上采样 + 拼接
        self.up1 = Up(base_channel*16 ,base_channel*8 // factor,self.bilinear) # 1024(双卷积的输入),256（双卷积的输出）
        self.up2 = Up(base_channel*8 ,base_channel*4 // factor,self.bilinear)
        self.up3 = Up(base_channel*4 ,base_channel*2 // factor,self.bilinear)
        self.up4 = Up(base_channel*2 ,base_channel,self.bilinear)
        # 输出
        self.out = nn.Conv2d(in_channels=base_channel,out_channels=self.out_channels,kernel_size=1)

    def forward(self,x):
        x1 = self.in_conv(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)
        out = self.out(x)

        return {'out':out}

6. 损失函数：

​ 其实，完成了数据加载、网络结构后，其实基本上已经完成了大部分内容了，因为图像分割就是像素级的分类，因此损失函数完全可以简简单单的用个交叉熵损失函数就行了。

​ 但是，源码上不仅使用交叉熵作为损失函数，还使用了dice作为损失函数（度量两个集合的相似性）。

​ 其中，dice_loss的计算公式为：

​ 具体的实现，我就不多说了，因为不是重点，感兴趣的可以看看注释，不感兴趣的可以直接拿过来用：

import torch
from torch import nn

# dice : 度量两个集合的相似性

# loss
def criterion(inputs, target, loss_weight=None, num_classes: int = 2, dice: bool = True, ignore_index: int = -100):
    losses = {}
    for name, x in inputs.items():
        # x = torch.Size([4, 2, 480, 480]) ，4为batch，2表示输出通道数，480*480表示和原图大小一样
        # target = torch.Size([4, 480, 480]) 真实mask值
        # ignore_index = 255，表示忽略那些不感兴趣的部分
        loss = nn.functional.cross_entropy(x, target, ignore_index=ignore_index, weight=loss_weight)
        if dice is True:
            # 计算每个类别的dice（背景、前景），然后求均值
            # 因此需要为每个类别构建一个target
            dice_target = build_target(target, num_classes, ignore_index)
            loss += dice_loss(x, dice_target, multiclass=True, ignore_index=ignore_index)
        losses[name] = loss

    if len(losses) == 1:
        return losses['out']

    return losses['out'] + 0.5 * losses['aux']


def build_target(target: torch.Tensor, num_classes: int = 2, ignore_index: int = -100):
    """build target for dice coefficient"""
    dice_target = target.clone()
    if ignore_index >= 0:
        # 将255的区域设置为0，因为不感兴趣的不需要计算dice，因此先设置为0
        ignore_mask = torch.eq(target, ignore_index)
        dice_target[ignore_mask] = 0
        # target转为one-hot编码形式 [1 0]表示一个类别 and [0 1]表示一个类别
        dice_target = nn.functional.one_hot(dice_target, num_classes).float()
        # 将255的值又填充回去
        dice_target[ignore_mask] = ignore_index
    else:
        dice_target = nn.functional.one_hot(dice_target, num_classes).float()

    # [N, H, W] -> [N, H, W, C]
    return dice_target.permute(0, 3, 1, 2)


def dice_coeff(x: torch.Tensor, target: torch.Tensor, ignore_index: int = -100, epsilon=1e-6):
    # 计算一个batch中所有图片某个类别的dice_coefficient
    d = 0.
    batch_size = x.shape[0]
    # 遍历每张图片
    for i in range(batch_size):
        x_i = x[i].reshape(-1)
        t_i = target[i].reshape(-1)
        if ignore_index >= 0:
            # 找出mask中不为ignore_index的区域
            roi_mask = torch.ne(t_i, ignore_index)
            x_i = x_i[roi_mask]
            t_i = t_i[roi_mask]
        inter = torch.dot(x_i, t_i)
        sets_sum = torch.sum(x_i) + torch.sum(t_i)
        if sets_sum == 0:
            sets_sum = 2 * inter

        # 计算dice
        d += (2 * inter + epsilon) / (sets_sum + epsilon)

    return d / batch_size


def multiclass_dice_coeff(x: torch.Tensor, target: torch.Tensor, ignore_index: int = -100, epsilon=1e-6):
    """Average of Dice coefficient for all classes"""
    dice = 0.
    # 遍历每个channel，计算每个类别的dice
    for channel in range(x.shape[1]):
        dice += dice_coeff(x[:, channel, ...], target[:, channel, ...], ignore_index, epsilon)

    # 求均值
    return dice / x.shape[1]


def dice_loss(x: torch.Tensor, target: torch.Tensor, multiclass: bool = False, ignore_index: int = -100):
    # 在channel方向做softmax
    x = nn.functional.softmax(x, dim=1)
    # 选择采用的方法
    fn = multiclass_dice_coeff if multiclass else dice_coeff
    return 1 - fn(x, target, ignore_index=ignore_index)

7. 预处理方法：

​ 这个文件，我是直接拷贝过来的，因为里面实现的方法功能和torchvision.transforms一致，这里之所以重新定义一次，是因为我们需要同时处理图像和mask两个值。

​ 感兴趣的可以自己研究一下：

import numpy as np
import random
import torch
from torchvision import transforms as T
from torchvision.transforms import functional as F

# 填充
def pad_if_smaller(img, size, fill=0):
    # 获取图像最短边
    min_size = min(img.size)
    # 如果最短边小于给定的尺寸
    if min_size < size:
        # 用fill值进行填充
        ow, oh = img.size
        padh = size - oh if oh < size else 0
        padw = size - ow if ow < size else 0
        img = F.pad(img, (0, 0, padw, padh), fill=fill)
    return img


# 将所有预处理方法集合在一起，然后运行
class Compose(object):
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, image, target):
        for t in self.transforms:
            image, target = t(image, target)
        return image, target

# 随机缩放
class RandomResize(object):
    def __init__(self, min_size, max_size=None):
        self.min_size = min_size
        if max_size is None:
            max_size = min_size
        self.max_size = max_size

    def __call__(self, image, target):
        size = random.randint(self.min_size, self.max_size)
        # 这里size传入的是int类型，所以是将图像的最小边长缩放到size大小
        image = F.resize(image, size)
        # 这里的interpolation注意下，在torchvision(0.9.0)以后才有InterpolationMode.NEAREST
        # 如果是之前的版本需要使用PIL.Image.NEAREST
        target = F.resize(target, size, interpolation=T.InterpolationMode.NEAREST)
        return image, target

# 水平翻转
class RandomHorizontalFlip(object):
    def __init__(self, flip_prob):
        self.flip_prob = flip_prob

    def __call__(self, image, target):
        if random.random() < self.flip_prob:
            image = F.hflip(image)
            target = F.hflip(target)
        return image, target

# 垂直翻转
class RandomVerticalFlip(object):
    def __init__(self, flip_prob):
        self.flip_prob = flip_prob

    def __call__(self, image, target):
        if random.random() < self.flip_prob:
            image = F.vflip(image)
            target = F.vflip(target)
        return image, target

# 随机裁剪
class RandomCrop(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, image, target):
        image = pad_if_smaller(image, self.size)
        target = pad_if_smaller(target, self.size, fill=255)
        crop_params = T.RandomCrop.get_params(image, (self.size, self.size))
        image = F.crop(image, *crop_params)
        target = F.crop(target, *crop_params)
        return image, target


# 中心裁剪
class CenterCrop(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, image, target):
        image = F.center_crop(image, self.size)
        target = F.center_crop(target, self.size)
        return image, target


# 转为tensor
class ToTensor(object):
    def __call__(self, image, target):
        image = F.to_tensor(image)
        target = torch.as_tensor(np.array(target), dtype=torch.int64)
        return image, target


# 归一化
class Normalize(object):
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, image, target):
        image = F.normalize(image, mean=self.mean, std=self.std)
        return image, target

8. 训练：

​ 这里，简单的实现了一下训练的方法，主要的实现思路如下：

​ 按照上述思路，可以实现代码：

from network_files.u_net import U_net
from utils import transforms as T
from My_Dataset import My_Dataset
from utils import Loss

import torch
from torch import nn
from torch.utils.data import DataLoader
from torch import optim

# 训练的预处理方法获取
class SegmentationPresetTrain:
    def __init__(self, base_size, crop_size, hflip_prob=0.5, vflip_prob=0.5,
                 mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
        # 初始化
        min_size = int(0.5 * base_size)
        max_size = int(1.2 * base_size)

        # 随机裁剪
        trans = [T.RandomResize(min_size, max_size)]
        if hflip_prob > 0:
            trans.append(T.RandomHorizontalFlip(hflip_prob))
        if vflip_prob > 0:
            trans.append(T.RandomVerticalFlip(vflip_prob))
        trans.extend([
            T.RandomCrop(crop_size),
            T.ToTensor(),
            T.Normalize(mean=mean, std=std),
        ])
        self.transforms = T.Compose(trans)

    def __call__(self, img, target):
        return self.transforms(img, target)

# 获取预处理方法
def get_transform(train, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
    base_size = 565
    crop_size = 480

    if train:
        # 训练模式获取的预处理方法
        return SegmentationPresetTrain(base_size, crop_size, mean=mean, std=std)


def main():
    # 设置基本参数信息
    # device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    # 如果GPU不够用，可以用cpu
    device = torch.device('cpu')
    batch_size= 4
    epoch = 20
    num_classes = 2 # 1(object)+1(background)
    if num_classes == 2:
        # 设置cross_entropy中背景和前景的loss权重(根据自己的数据集进行设置)
        loss_weight = torch.as_tensor([1.0, 2.0], device=device)
    else:
        loss_weight = None
    # 加载数据
    mean = (0.709, 0.381, 0.224)
    std = (0.127, 0.079, 0.043)
    train_dataset = My_Dataset('../data',train=True,transforms=get_transform(train=True, mean=mean, std=std))
    val_dataset = My_Dataset('../data',train=False,transforms=get_transform(train=False, mean=mean, std=std))
    train_loader = DataLoader(train_dataset,batch_size=batch_size,shuffle=True)
    val_loader = DataLoader(val_dataset,batch_size=batch_size,shuffle=True)
    # 创建模型
    model = U_net(3,num_classes) # 输入通道数，输出通道数
    model.to(device)
    # 定义优化器
    params = [p for p in model.parameters() if p.requires_grad] # 定义需要优化的参数
    sgd = optim.SGD(params,lr=0.01,momentum=0.9,weight_decay=1e-4)
    # 开始训练
    model.train()
    for e in range(epoch):
        loss_temp = 0
        for i,(image,mask) in enumerate(train_loader):
            image,mask = image.to(device),mask.to(device)
            output = model(image)
            loss = Loss.criterion(output, mask, loss_weight, num_classes=num_classes, ignore_index=255)
            loss_temp += loss.item()
            sgd.zero_grad()
            loss.backward()
            sgd.step()
        print(f'第{e+1}个epoch,平均损失loss={loss_temp/(i+1)}')
    # 保存权重
    name = 'save_weights/u_net.pth'
    torch.save(model.state_dict(),name)



if __name__ == '__main__':
    main()

​ 运行结果展示，这里就展示一下正常运行的结果，大家可以自己去运行：

9. 预测：

​ 上面训练完后，就保存了相关的参数值，我们可以利用它来进行预测。

​ 预测实现的思路如下：

​ 代码如下：

import os
import time

import torch
from utils import transforms as  T
import numpy as np
from PIL import Image

from network_files.u_net import U_net

def main():
    # 设置基本参数,需要自己改变路径参数
    classes = 1
    weights_path = "./save_weights/u_net.pth"
    img_path = "../data/DRIVE/test/images/03_test.tif" # ，另外可以改变预测的图像对象
    roi_mask_path = "../data/DRIVE/test/mask/03_test_mask.gif"
    # 这个是图像归一化的参数
    mean = (0.709, 0.381, 0.224)
    std = (0.127, 0.079, 0.043)
    # 获取设备
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    # 创建模型，输入通道数为3，输出为2
    model = U_net(3,classes+1)
    # 加载权重
    model.load_state_dict(torch.load(weights_path))
    model.to(device)
    # 打开相关路径
    roi_img = Image.open(roi_mask_path).convert('L')
    original_img = Image.open(img_path).convert('RGB')
    # 预处理
    data_transform = T.Compose([T.RandomCrop(480),
                                         T.ToTensor(),
                                         T.Normalize(mean=mean, std=std)])
    img,roi_img = data_transform(original_img,roi_img)
    # 将三维的数据转为四维，因为需要添加batch这个维度
    img = torch.unsqueeze(img, dim=0)
    # 将roi转为array，方便后期处理
    roi_img = np.array(roi_img)
    model.eval()  # 进入验证模式
    with torch.no_grad():
        # 初始化一个全黑的图像，后期将白色的添加进去，就可以得到预测的mask值
        img_height, img_width = img.shape[-2:]
        init_img = torch.zeros((1, 3, img_height, img_width), device=device)
        model(init_img)
        output = model(img.to(device))
        prediction = output['out'].argmax(1).squeeze(0)
        prediction = prediction.to("cpu").numpy().astype(np.uint8)
        # 将前景对应的像素值改成255(白色)
        prediction[prediction == 1] = 255
        # 将不敢兴趣的区域像素设置成0(黑色)
        prediction[roi_img == 0] = 0
        mask = Image.fromarray(prediction)
        mask.save("test_result.png")

if __name__ == '__main__':
    main()

​ 结果展示：

​ 说明：由于预测的时候，我们将图像裁剪为了480*480，因此和真实结果看起来大小有一定的差距。

​ 从图中，也可以看出，那些比较大的血管，基本都分割出来了，而较细微的血管，部分没有分割出来。综合来看，效果还是不错。

10. 总结：

​ 这里只是简单的实现了u-net网络，还有很多地方可以改进，比如：

• 学习率没有调整
• 没有添加验证集部分
• 预测仅仅只是针对一张图片预测，没有实现集成的预测函数
• …

​ 当然，代码写到这里，基本上也实现了我预期的内容了，再次感谢大佬们的开源工作，让我可以轻松copy一些代码过来。

全部代码和训练好的参数：

链接：https://pan.baidu.com/s/1Ro9MS9OnxJTUkiJLzQrcmA 
提取码：e3ol