图像分割之U-Net、U2 -Net及其Pytorch代码构建
1、图像分割
图像分割就是把图像分成若干个特定的、具有独特性质的区域并提出感兴趣目标的技术和过程。
做法便是对图片中的每一个像素进行分类。
在自动驾驶、自动抠图、医疗影像等领域有着比较广泛的应用。
图像分割大致可分为以下三类：
- 普通分割：将不同分属不同物体的像素区域分开。比如前景和背景分割开，狗的区域和猫的区域与背景分割开。
- 语义分割：在普通分割的基础上，分类出每一块区域的语义（即这块区域是什么物体）。如把画面中的所有物体都指出他们各自的类别。
- 实例分割：在语义分割的基础上，给每一个物体编号。如这个是该画面中的狗A，那个是画面中的狗B。
| 普通分割 |  |
|---|---|
| 语义分割 |  |
| 实例分割 |  |
可以看出，图像分割是由一张图片到另一张图片。因此，神经网络的输入是图片，输出也是同样的图片，Encoder-Decoder的结构是合适的。U-Net、U2 -Net可作为语义分割使用，可以按照生成图像的方式，生成分割图。也可以按通道划分类，每一个通道就是一个类别，使用sigmoid激活。
2、U-Net
U-Net即使用Encoder-Decoder的结构，首先下采样，然后上采样，中间每一级由残差组成。
则可构建网络的代码如下：
首先是卷积层，可以看出，网络在每一级，均有两层卷积组成。因此构建卷积层如下：
from torch import nn
import torch
class ConvolutionLayer(nn.Module):
def __init__(self, in_channels, out_channels):
"""
卷积层
:param in_channels: 输入通道
:param out_channels: 输出通道
"""
super(ConvolutionLayer, self).__init__()
self.layer = nn.Sequential(
# 卷积层
nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), stride=(1, 1), padding=1, bias=False),
nn.BatchNorm2d(out_channels), # BN层
nn.ReLU(), # 激活
nn.Conv2d(out_channels, out_channels, kernel_size=(3, 3), stride=(1, 1), padding=1, bias=False),
nn.BatchNorm2d(out_channels),
nn.ReLU(),
)
def forward(self, x):
return self.layer(x)
同时与图示不同的地方在于，使用了Padding，以免图片在卷积中的尺寸缩小。这样，横向的灰色箭头可以直接使用cat进行两个特征图的拼接。
模型图中，红色箭头的max pool 2×2，使用的是池化窗口为2×2的最大值池化。这里的目的是进行下采样，因此可以定义一个下采样如下：
class DownSample(nn.Module):
def __init__(self,):
"""
最大池化层构成的下采样，池化窗口为2×2
"""
super(DownSample, self).__init__()
self.layer = nn.MaxPool2d(kernel_size=2, stride=2)
def forward(self, x):
return self.layer(x)
模型图中，绿色箭头的up-conv 2×2，使用的是反卷积。这里的目的是进行上采样，因此可以定义一个上采样如下：
class UpSample(nn.Module):
def __init__(self, in_channels):
"""
反卷积，上采样，通道数将会减半，
:param in_channels: 输入通道数
"""
super(UpSample, self).__init__()
self.layer = nn.Sequential(
nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=(2, 2), stride=(2, 2)),
nn.LeakyReLU(),
)
def forward(self, x):
return self.layer(x)
首先定义各个网络层：
class UNet(nn.Module):
def __init__(self, in_channels, out_channels):
super(UNet, self).__init__()
self.conv1 = ConvolutionLayer(in_channels, 64) # 三通道拓展至64通道
self.down1 = DownSample() # 下采样至1/2
self.conv2 = ConvolutionLayer(64, 128) # 64通道==>128通道
self.down2 = DownSample() # 下采样至1/4
self.conv3 = ConvolutionLayer(128, 256) # 128通道==>256通道
self.down3 = DownSample() # 下采样至1/8
self.conv4 = ConvolutionLayer(256, 512) # 256通道==>512通道
self.down4 = DownSample() # 下采样至1/16
self.conv5 = ConvolutionLayer(512, 1024) # 512通道==>1024通道
self.up1 = UpSample(1024) # 上采样至1/8
self.conv6 = ConvolutionLayer(1024, 512) # 1024通道==>512通道
self.up2 = UpSample(512) # 上采样至1/4
self.conv7 = ConvolutionLayer(512, 256) # 512通道==>256通道
self.up3 = UpSample(256) # 上采样至1/2
self.conv8 = ConvolutionLayer(256, 128) # 256通道==>128通道
self.up4 = UpSample(128) # 上采样至1/1
self.conv9 = ConvolutionLayer(128, 64) # 128通道==>64通道
self.predict = nn.Sequential( # 输出层，由sigmoid函数激活
nn.Conv2d(64, out_channels, kernel_size=(3,3), stride=(1,1), padding=1),
nn.Sigmoid()
)
def forward(self, image_tensor):
pass
对应于模型图如下：
class UNet(nn.Module):
def __init__(self, in_channels, out_channels):
super(UNet, self).__init__()
"""
......
"""
def forward(self, x):
"""下采样"""
x1 = self.conv1(x) # ===> 1/1 64
d1 = self.down1(x1) # ===> 1/2 64
x2 = self.conv2(d1) # ===> 1/2 128
d2 = self.down2(x2) # ===> 1/4 128
x3 = self.conv3(d2) # ===> 1/4 256
d3 = self.down3(x3) # ===> 1/8 256
x4 = self.conv4(d3) # ===> 1/8 512
d4 = self.down4(x4) # ===> 1/16 512
x5 = self.conv5(d4) # ===> 1/16 1024
"""上采样"""
up1 = self.up1(x5) # ===> 1/8 512
x6 = self.conv6(torch.cat((x4, up1), dim=1)) # ===> 1/8 512
up2 = self.up2(x6 # ===> 1/4 256
x7 = self.conv7(torch.cat((x3, up2), dim=1)) # ===> 1/4 256
up3 = self.up3(x7) # ===> 1/2 128
x8 = self.conv8(torch.cat((x2, up3), dim=1)) # ===> 1/2 128
up4 = self.up4(x8) # ===> 1/1 64
x9 = self.conv9(torch.cat((x1, up4), dim=1)) # ===> 1/1 64
mask = self.predict(x9) # ===> 1/1 out_channels
return mask
以一张512×512的3通道图片为例，其张量的形状为（1,3,512,512），经过conv1得到x1 （1, 64, 512, 512），下采样至（1, 64, 256, 256）；经过conv2得到x2 （1, 128, 256, 256），下采样至（1, 128, 128, 128）；经过conv3得到x3 （1, 256, 128, 128），下采样至（1, 256, 64, 64）；经过conv4得到x4 （1, 512, 64, 64），下采样至（1, 512, 32, 32）；经过conv5得到x5 （1, 1024, 32, 32）。下采样过程完成，开始上采样还原至原始图片大小。
x5经过up1得到up1 （1, 512, 64, 64）,同x4 拼接（cat）在一起 组成（1, 1024, 64, 64）的张量，经过conv6得到x6（1, 512, 64, 64）;
x6经过up2得到up2 （1, 256, 128, 128）,同x3 拼接在一起 组成（1, 512, 128, 128）的张量，经过conv7得到x7（1, 256, 128, 128）;
x7经过up3得到up3 （1, 128, 256, 256）,同x2 拼接在一起 组成（1, 256, 256, 256）的张量，经过conv8得到x8（1, 128, 256, 256）;
x8经过up4得到up4 （1, 64, 512, 512）,同x1 拼接在一起 组成（1, 128, 512, 512）的张量，经过conv6得到x9（1, 64, 512, 512）;
最后，x9经过预测层predict输出，得到分割图mask。
以drive数据集为例训练网络，数据示例如下。
标签如下：
输入数据为3通道的图片，而输出数据为1通道的二值图。一张图片的原始尺寸是565×584
可以在原始图像中随机裁剪256×256大小的图片，进行训练，而在使用时，图像尺寸只要是16的倍数即可。
定义数据加载函数如下：
import torch
import random
import cv2
from torch.utils.data import Dataset
class DriveDataset(Dataset):
def __init__(self,root='data/training'):
super(DriveDataset, self).__init__()
self.dataset = []
start = 20
for i in range(1, 21): # 按照一一对应的原则，加载图像和标签的路径
image_path = f'{root}/images/{i+start}_training.tif'
label_path = f'{root}/1st_manual/{i + start}_manual1.gif'
self.dataset.append((image_path, label_path))
def __len__(self):
return len(self.dataset)
def __getitem__(self, item):
image_path, label_path = self.dataset[item] # 获取图像路径
image = cv2.imread(image_path) # 图片
video = cv2.VideoCapture(label_path)
_, mask_label = video.read() # 读取标签掩码图
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
mask_label = cv2.cvtColor(mask_label, cv2.COLOR_BGR2GRAY) # 转换至单通道图
"""随即裁剪256×256的图幅，图片和标签裁剪相同的位置"""
h, w = mask_label.shape
w = random.randint(0, w-256)
h = random.randint(0, h-256)
image = image[h:h+256, w:w+256]
mask_label = mask_label[h:h + 256, w:w + 256]
"""转换至tensor"""
image = torch.from_numpy(image).float().permute(2, 0, 1)/255
mask_label = torch.from_numpy(mask_label).unsqueeze(0).float()/255
return image, mask_label
读取相对应的图片和标签，转换为张量，供网络学习。其中，标签的读取使用了OpenCV的视频捕获（VideoCapture）读取首帧完成标签的数据加载。
定义训练器如下：
from torch import nn
import torch
from torch.utils.data import DataLoader
from torchvision.utils import save_image
from u_net import UNet
from dataset import DriveDataset
import os
class Trainer:
def __init__(self):
self.device = torch.device('cuda:0' if torch.cuda.is_available() else "cpu") # 设置设备
self.net = UNet(3, 1).to(self.device) # 实例U-Net
if os.path.exists('unet.pth'): # 加载权重，如果存在的话
self.net.load_state_dict(torch.load('unet.pth', map_location='cpu'))
self.dataset = DriveDataset() # 实例数据集
self.data_loader = DataLoader(self.dataset, 3, True, drop_last=True) # 实例数据加载器
self.loss_func = nn.BCELoss() # 实例二值交叉熵
self.optimizer = torch.optim.Adam(self.net.parameters()) # 实例adam优化器
def train(self): # 训练
for epoch in range(100000): # 迭代epoch
for i, (image, target) in enumerate(self.data_loader):
image = image.to(self.device)
target = target.to(self.device)
out = self.net(image) # 预测
loss = self.loss_func(out, target) # 计算损失
self.optimizer.zero_grad() # 清空梯度
loss.backward() # 反向传播
self.optimizer.step() # 优化
print(epoch, loss.item())
if epoch % 5 == 0:
torch.save(self.net.state_dict(),'unet.pth')
save_image([image[0], target[0].expand(3, 256, 256), out[0].expand(3, 256, 256)], f'{epoch}.jpg',normalize=True,range=(0,1))
二值交叉熵做损失，adam优化器优化网络。
class Trainer:
"""
......
"""
if __name__ == '__main__':
trainer = Trainer()
trainer.train()
训练过程见下图。左边为原图，中间为标签，右边为网络预测值
| epoch | images |
|---|---|
| 0 |  |
| 1 |  |
| 2 |  |
完整代码：https://github.com/HibikiJie/UNetAndU2Net
3、U2-Net
而U2-Net，就是U-Net的堆叠，类似于，将U-Net中的conv块，替换成完整的U-Net网络。
其网络图如下：
其中EN_1与De_1一致，EN_2与De_2一致，EN_3与De_3一致，EN_4与De_4一致，EN_5、En6和De_5一致。
先分别定义，EN_1、EN_2、EN_3、EN_4、EN_5为UNet1、UNet2、UNet3、UNet4、UNet5.
首先定义UNet1：
注意到，图中的白色的方块示意的，卷积使用到了dilation参数，因此，定义ConvolutionLayer为：
import torch
import torch.nn as nn
import torch.nn.functional as F
class ConvolutionLayer(nn.Module):
def __init__(self, in_channels, out_channels, dilation=1):
super(ConvolutionLayer, self).__init__()
self.layer = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), padding=1 * dilation,
dilation=(1 * dilation, 1 * dilation)), # 卷积
nn.BatchNorm2d(out_channels), # BN
nn.ReLU(inplace=True) # 激活函数
)
def forward(self, x):
return self.layer(x)
卷积层由Conv、BN、ReLU构成。
上采样使用机器学习算法，由双线性插值法完成上采样：
def upsample_like(src, tar):
src = F.upsample(src, size=tar.shape[2:], mode='bilinear')
return src
该方法，将使src上采样至tar相同的尺寸大小。
而下采样同样使用最大池化完成，这里可以使用与U-Net相同的代码。
因此，UNet1：
class UNet1(nn.Module):
def __init__(self, in_channels, mid_channels, out_channels):
super(UNet1, self).__init__()
self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1)
self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1)
self.down1 = DownSample()
self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down2 = DownSample()
self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down3 = DownSample()
self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down4 = DownSample()
self.conv5 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down5 = DownSample()
self.conv6 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.conv7 = ConvolutionLayer(mid_channels, mid_channels, dilation=2)
self.conv8 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv9 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv10 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv11 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv12 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv13 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1)
def forward(self, x):
"""下采样，编码encode的过程"""
x0 = self.conv0(x)
x1 = self.conv1(x0)
d1 = self.down1(x1)
x2 = self.conv2(d1)
d2 = self.down2(x2)
x3 = self.conv3(d2)
d3 = self.down3(x3)
x4 = self.conv4(d3)
d4 = self.down4(x4)
x5 = self.conv5(d4)
d5 = self.down5(x5)
x6 = self.conv6(d5)
x7 = self.conv7(x6)
"""上采样，解码decode的过程"""
x8 = self.conv8(torch.cat((x7, x6), dim=1))
up1 = upsample_like(x8, x5)
x9 = self.conv9(torch.cat((up1, x5), dim=1))
up2 = upsample_like(x9, x4)
x10 = self.conv10(torch.cat((up2, x4), dim=1))
up3 = upsample_like(x10, x3)
x11 = self.conv11(torch.cat((up3, x3), dim=1))
up4 = upsample_like(x11, x2)
x12 = self.conv12(torch.cat((up4, x2), dim=1))
up5 = upsample_like(x12, x1)
x13 = self.conv13(torch.cat((up5, x1), dim=1))
return x13 + x0
按照上图所示的方式编码，可见，与写UNet的代码是非常类似的。可以对比着看。可见，U2-Net是U-Net的堆叠。
于是类似的，UNet2的代码为:
class UNet2(nn.Module):
def __init__(self, in_channels, mid_channels, out_channels):
super(UNet2, self).__init__()
self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1)
self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1)
self.down1 = DownSample()
self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down2 = DownSample()
self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down3 = DownSample()
self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down4 = DownSample()
self.conv5 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.conv6 = ConvolutionLayer(mid_channels, mid_channels, dilation=2)
self.conv7 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv8 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv9 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv10 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv11 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1)
def forward(self, x):
"""encode"""
x0 = self.conv0(x)
x1 = self.conv1(x0)
d1 = self.down1(x1)
x2 = self.conv2(d1)
d2 = self.down2(x2)
x3 = self.conv3(d2)
d3 = self.down3(x3)
x4 = self.conv4(d3)
d4 = self.down4(x4)
x5 = self.conv5(d4)
x6 = self.conv6(x5)
"""decode"""
x7 = self.conv7(torch.cat((x6, x5), dim=1))
up1 = upsample_like(x7, x4)
x8 = self.conv8(torch.cat((up1, x4), dim=1))
up2 = upsample_like(x8, x3)
x9 = self.conv9(torch.cat((up2, x3), dim=1))
up3 = upsample_like(x9, x2)
x10 = self.conv10(torch.cat((up3, x2), dim=1))
up4 = upsample_like(x10, x1)
x11 = self.conv11(torch.cat((up4, x1), dim=1))
return x11 + x0
UNet3为：
class UNet3(nn.Module):
def __init__(self, in_channels, mid_channels, out_channels):
super(UNet3, self).__init__()
self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1)
self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1)
self.down1 = DownSample()
self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down2 = DownSample()
self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down3 = DownSample()
self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.conv5 = ConvolutionLayer(mid_channels, mid_channels, dilation=2)
self.conv6 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv7 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv8 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv9 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1)
def forward(self, x):
"""encode"""
x0 = self.conv0(x)
x1 = self.conv1(x0)
d1 = self.down1(x1)
x2 = self.conv2(d1)
d2 = self.down2(x2)
x3 = self.conv3(d2)
d3 = self.down3(x3)
x4 = self.conv4(d3)
x5 = self.conv5(x4)
"""decode"""
x6 = self.conv6(torch.cat((x5, x4), dim=1))
up1 = upsample_like(x6, x3)
x7 = self.conv7(torch.cat((up1, x3), dim=1))
up2 = upsample_like(x7, x2)
x8 = self.conv8(torch.cat((up2, x2), dim=1))
up3 = upsample_like(x8, x1)
x9 = self.conv9(torch.cat((up3, x1), dim=1))
return x9 + x0
UNet4为:
class UNet4(nn.Module):
def __init__(self, in_channels, mid_channels, out_channels):
super(UNet4, self).__init__()
self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1)
self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1)
self.down1 = DownSample()
self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down2 = DownSample()
self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=2)
self.conv5 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv6 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv7 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1)
def forward(self, x):
"""encode"""
x0 = self.conv0(x)
x1 = self.conv1(x0)
d1 = self.down1(x1)
x2 = self.conv2(d1)
d2 = self.down2(x2)
x3 = self.conv3(d2)
x4 = self.conv4(x3)
"""decode"""
x5 = self.conv5(torch.cat((x4, x3), dim=1))
up1 = upsample_like(x5, x2)
x6 = self.conv6(torch.cat((up1, x2), dim=1))
up2 = upsample_like(x6, x1)
x7 = self.conv7(torch.cat((up2, x1), dim=1))
return x7 + x0
UNet5为:
class UNet5(nn.Module):
def __init__(self, in_channels, mid_channels, out_channels):
super(UNet5, self).__init__()
self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1)
self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1)
self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=2)
self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=4)
self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=8)
self.conv5 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=4)
self.conv6 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=2)
self.conv7 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1)
def forward(self, x):
x0 = self.conv0(x)
x1 = self.conv1(x0)
x2 = self.conv2(x1)
x3 = self.conv3(x2)
x4 = self.conv4(x3)
x5 = self.conv5(torch.cat((x4, x3), dim=1))
x6 = self.conv6(torch.cat((x5, x2), dim=1))
x7 = self.conv7(torch.cat((x6, x1), dim=1))
return x7 + x0
于是将
UNet1、UNet2、UNet3、UNet4、UNet5.组装成为U2-Net
再看一下网络结构图：
其中EN_1与De_1一致，EN_2与De_2一致，EN_3与De_3一致，EN_4与De_4一致，EN_5、En6和De_5一致。
先分别定义，EN_1、EN_2、EN_3、EN_4、EN_5为UNet1、UNet2、UNet3、UNet4、UNet5.
于是EN_1与De_1使用UNet1；
EN_2与De_2使用UNet2；
EN_3与De_3使用UNet3；
EN_4与De_4使用UNet4；
EN_5、EN_6、De_5使用UNet1。
故，构建网络U2-Net：
class U2Net(nn.Module):
def __init__(self, in_channels=3, out_channels=1):
super(U2Net, self).__init__()
self.en_1 = UNet1(in_channels, 32, 64)
self.down1 = DownSample()
self.en_2 = UNet2(64, 32, 128)
self.down2 = DownSample()
self.en_3 = UNet3(128, 64, 256)
self.down3 = DownSample()
self.en_4 = UNet4(256, 128, 512)
self.down4 = DownSample()
self.en_5 = UNet5(512, 256, 512)
self.down5 = DownSample()
self.en_6 = UNet5(512, 256, 512)
# decoder
self.de_5 = UNet5(1024, 256, 512)
self.de_4 = UNet4(1024, 128, 256)
self.de_3 = UNet3(512, 64, 128)
self.de_2 = UNet2(256, 32, 64)
self.de_1 = UNet1(128, 16, 64)
self.side1 = nn.Conv2d(64, out_channels, kernel_size=(3, 3), padding=1)
self.side2 = nn.Conv2d(64, out_channels, kernel_size=(3, 3), padding=1)
self.side3 = nn.Conv2d(128, out_channels, kernel_size=(3, 3), padding=1)
self.side4 = nn.Conv2d(256, out_channels, kernel_size=(3, 3), padding=1)
self.side5 = nn.Conv2d(512, out_channels, kernel_size=(3, 3), padding=1)
self.side6 = nn.Conv2d(512, out_channels, kernel_size=(3, 3), padding=1)
self.out_conv = nn.Conv2d(6, out_channels, kernel_size=(1, 1))
def forward(self, x):
# ------encode ------
x1 = self.en_1(x)
d1 = self.down1(x1)
x2 = self.en_2(d1)
d2 = self.down2(x2)
x3 = self.en_3(d2)
d3 = self.down3(x3)
x4 = self.en_4(d3)
d4 = self.down4(x4)
x5 = self.en_5(d4)
d5 = self.down5(x5)
x6 = self.en_6(d5)
up1 = upsample_like(x6, x5)
# ------decode ------
x7 = self.de_5(torch.cat((up1, x5), dim=1))
up2 = upsample_like(x7, x4)
x8 = self.de_4(torch.cat((up2, x4), dim=1))
up3 = upsample_like(x8, x3)
x9 = self.de_3(torch.cat((up3, x3), dim=1))
up4 = upsample_like(x9, x2)
x10 = self.de_2(torch.cat((up4, x2), dim=1))
up5 = upsample_like(x10, x1)
x11 = self.de_1(torch.cat((up5, x1), dim=1))
# side output
sup1 = self.side1(x11)
sup2 = self.side2(x10)
sup2 = upsample_like(sup2, sup1)
sup3 = self.side3(x9)
sup3 = upsample_like(sup3, sup1)
sup4 = self.side4(x8)
sup4 = upsample_like(sup4, sup1)
sup5 = self.side5(x7)
sup5 = upsample_like(sup5, sup1)
sup6 = self.side6(x6)
sup6 = upsample_like(sup6, sup1)
sup0 = self.out_conv(torch.cat((sup1, sup2, sup3, sup4, sup5, sup6), dim=1))
return torch.sigmoid(sup0)
U2-Net完整代码：
import torch
import torch.nn as nn
import torch.nn.functional as F
class ConvolutionLayer(nn.Module):
def __init__(self, in_channels, out_channels, dilation=1):
super(ConvolutionLayer, self).__init__()
self.layer = nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), padding=1 * dilation,
dilation=(1 * dilation, 1 * dilation)),
nn.BatchNorm2d(out_channels),
nn.ReLU(inplace=True)
)
self.conv_s1 = nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), padding=1 * dilation,
dilation=(1 * dilation, 1 * dilation))
self.bn_s1 = nn.BatchNorm2d(out_channels)
self.relu_s1 = nn.ReLU(inplace=True)
def forward(self, x):
return self.layer(x)
def upsample_like(src, tar):
src = F.interpolate(src, size=tar.shape[2:], mode='bilinear')
return src
class DownSample(nn.Module):
def __init__(self, ):
super(DownSample, self).__init__()
self.layer = nn.MaxPool2d(kernel_size=2, stride=2)
def forward(self, x):
return self.layer(x)
class UNet1(nn.Module):
def __init__(self, in_channels, mid_channels, out_channels):
super(UNet1, self).__init__()
self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1)
self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1)
self.down1 = DownSample()
self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down2 = DownSample()
self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down3 = DownSample()
self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down4 = DownSample()
self.conv5 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down5 = DownSample()
self.conv6 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.conv7 = ConvolutionLayer(mid_channels, mid_channels, dilation=2)
self.conv8 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv9 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv10 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv11 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv12 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv13 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1)
def forward(self, x):
x0 = self.conv0(x)
x1 = self.conv1(x0)
d1 = self.down1(x1)
x2 = self.conv2(d1)
d2 = self.down2(x2)
x3 = self.conv3(d2)
d3 = self.down3(x3)
x4 = self.conv4(d3)
d4 = self.down4(x4)
x5 = self.conv5(d4)
d5 = self.down5(x5)
x6 = self.conv6(d5)
x7 = self.conv7(x6)
x8 = self.conv8(torch.cat((x7, x6), 1))
up1 = upsample_like(x8, x5)
x9 = self.conv9(torch.cat((up1, x5), 1))
up2 = upsample_like(x9, x4)
x10 = self.conv10(torch.cat((up2, x4), 1))
up3 = upsample_like(x10, x3)
x11 = self.conv11(torch.cat((up3, x3), 1))
up4 = upsample_like(x11, x2)
x12 = self.conv12(torch.cat((up4, x2), 1))
up5 = upsample_like(x12, x1)
x13 = self.conv13(torch.cat((up5, x1), 1))
return x13 + x0
class UNet2(nn.Module):
def __init__(self, in_channels, mid_channels, out_channels):
super(UNet2, self).__init__()
self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1)
self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1)
self.down1 = DownSample()
self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down2 = DownSample()
self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down3 = DownSample()
self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down4 = DownSample()
self.conv5 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.conv6 = ConvolutionLayer(mid_channels, mid_channels, dilation=2)
self.conv7 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv8 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv9 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv10 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv11 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1)
def forward(self, x):
x0 = self.conv0(x)
x1 = self.conv1(x0)
d1 = self.down1(x1)
x2 = self.conv2(d1)
d2 = self.down2(x2)
x3 = self.conv3(d2)
d3 = self.down3(x3)
x4 = self.conv4(d3)
d4 = self.down4(x4)
x5 = self.conv5(d4)
x6 = self.conv6(x5)
x7 = self.conv7(torch.cat((x6, x5), dim=1))
up1 = upsample_like(x7, x4)
x8 = self.conv8(torch.cat((up1, x4), dim=1))
up2 = upsample_like(x8, x3)
x9 = self.conv9(torch.cat((up2, x3), dim=1))
up3 = upsample_like(x9, x2)
x10 = self.conv10(torch.cat((up3, x2), dim=1))
up4 = upsample_like(x10, x1)
x11 = self.conv11(torch.cat((up4, x1), dim=1))
return x11 + x0
class UNet3(nn.Module):
def __init__(self, in_channels, mid_channels, out_channels):
super(UNet3, self).__init__()
self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1)
self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1)
self.down1 = DownSample()
self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down2 = DownSample()
self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down3 = DownSample()
self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.conv5 = ConvolutionLayer(mid_channels, mid_channels, dilation=2)
self.conv6 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv7 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv8 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv9 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1)
def forward(self, x):
x0 = self.conv0(x)
x1 = self.conv1(x0)
d1 = self.down1(x1)
x2 = self.conv2(d1)
d2 = self.down2(x2)
x3 = self.conv3(d2)
d3 = self.down3(x3)
x4 = self.conv4(d3)
x5 = self.conv5(x4)
x6 = self.conv6(torch.cat((x5, x4), 1))
up1 = upsample_like(x6, x3)
x7 = self.conv7(torch.cat((up1, x3), 1))
up2 = upsample_like(x7, x2)
x8 = self.conv8(torch.cat((up2, x2), 1))
up3 = upsample_like(x8, x1)
x9 = self.conv9(torch.cat((up3, x1), 1))
return x9 + x0
class UNet4(nn.Module):
def __init__(self, in_channels, mid_channels=12, out_channels):
super(UNet4, self).__init__()
self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1)
self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1)
self.down1 = DownSample()
self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.down2 = DownSample()
self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=1)
self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=2)
self.conv5 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv6 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=1)
self.conv7 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1)
def forward(self, x):
"""encode"""
x0 = self.conv0(x)
x1 = self.conv1(x0)
d1 = self.down1(x1)
x2 = self.conv2(d1)
d2 = self.down2(x2)
x3 = self.conv3(d2)
x4 = self.conv4(x3)
"""decode"""
x5 = self.conv5(torch.cat((x4, x3), 1))
up1 = upsample_like(x5, x2)
x6 = self.conv6(torch.cat((up1, x2), 1))
up2 = upsample_like(x6, x1)
x7 = self.conv7(torch.cat((up2, x1), 1))
return x7 + x0
class UNet5(nn.Module):
def __init__(self, in_channels, mid_channels, out_channels):
super(UNet5, self).__init__()
self.conv0 = ConvolutionLayer(in_channels, out_channels, dilation=1)
self.conv1 = ConvolutionLayer(out_channels, mid_channels, dilation=1)
self.conv2 = ConvolutionLayer(mid_channels, mid_channels, dilation=2)
self.conv3 = ConvolutionLayer(mid_channels, mid_channels, dilation=4)
self.conv4 = ConvolutionLayer(mid_channels, mid_channels, dilation=8)
self.conv5 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=4)
self.conv6 = ConvolutionLayer(mid_channels * 2, mid_channels, dilation=2)
self.conv7 = ConvolutionLayer(mid_channels * 2, out_channels, dilation=1)
def forward(self, x):
x0 = self.conv0(x)
x1 = self.conv1(x0)
x2 = self.conv2(x1)
x3 = self.conv3(x2)
x4 = self.conv4(x3)
x5 = self.conv5(torch.cat((x4, x3), 1))
x6 = self.conv6(torch.cat((x5, x2), 1))
x7 = self.conv7(torch.cat((x6, x1), 1))
return x7 + x0
class U2Net(nn.Module):
def __init__(self, in_channels=3, out_channels=1):
super(U2Net, self).__init__()
self.en_1 = UNet1(in_channels, 32, 64)
self.down1 = DownSample()
self.en_2 = UNet2(64, 32, 128)
self.down2 = DownSample()
self.en_3 = UNet3(128, 64, 256)
self.down3 = DownSample()
self.en_4 = UNet4(256, 128, 512)
self.down4 = DownSample()
self.en_5 = UNet5(512, 256, 512)
self.down5 = DownSample()
self.en_6 = UNet5(512, 256, 512)
# decoder
self.de_5 = UNet5(1024, 256, 512)
self.de_4 = UNet4(1024, 128, 256)
self.de_3 = UNet3(512, 64, 128)
self.de_2 = UNet2(256, 32, 64)
self.de_1 = UNet1(128, 16, 64)
self.side1 = nn.Conv2d(64, out_channels, kernel_size=(3, 3), padding=1)
self.side2 = nn.Conv2d(64, out_channels, kernel_size=(3, 3), padding=1)
self.side3 = nn.Conv2d(128, out_channels, kernel_size=(3, 3), padding=1)
self.side4 = nn.Conv2d(256, out_channels, kernel_size=(3, 3), padding=1)
self.side5 = nn.Conv2d(512, out_channels, kernel_size=(3, 3), padding=1)
self.side6 = nn.Conv2d(512, out_channels, kernel_size=(3, 3), padding=1)
self.out_conv = nn.Conv2d(6, out_channels, kernel_size=(1, 1))
def forward(self, x):
# ------encode ------
x1 = self.en_1(x)
d1 = self.down1(x1)
x2 = self.en_2(d1)
d2 = self.down2(x2)
x3 = self.en_3(d2)
d3 = self.down3(x3)
x4 = self.en_4(d3)
d4 = self.down4(x4)
x5 = self.en_5(d4)
d5 = self.down5(x5)
x6 = self.en_6(d5)
up1 = upsample_like(x6, x5)
# ------decode ------
x7 = self.de_5(torch.cat((up1, x5), dim=1))
up2 = upsample_like(x7, x4)
x8 = self.de_4(torch.cat((up2, x4), dim=1))
up3 = upsample_like(x8, x3)
x9 = self.de_3(torch.cat((up3, x3), dim=1))
up4 = upsample_like(x9, x2)
x10 = self.de_2(torch.cat((up4, x2), dim=1))
up5 = upsample_like(x10, x1)
x11 = self.de_1(torch.cat((up5, x1), dim=1))
# side output
sup1 = self.side1(x11)
sup2 = self.side2(x10)
sup2 = upsample_like(sup2, sup1)
sup3 = self.side3(x9)
sup3 = upsample_like(sup3, sup1)
sup4 = self.side4(x8)
sup4 = upsample_like(sup4, sup1)
sup5 = self.side5(x7)
sup5 = upsample_like(sup5, sup1)
sup6 = self.side6(x6)
sup6 = upsample_like(sup6, sup1)
sup0 = self.out_conv(torch.cat((sup1, sup2, sup3, sup4, sup5, sup6), 1))
return torch.sigmoid(sup0)
if __name__ == '__main__':
u2net = U2Net(3, 1)
x = torch.randn(1,3, 512, 512)
print(u2net(x).shape)
x5 = self.en_5(d4)
d5 = self.down5(x5)
x6 = self.en_6(d5)
up1 = upsample_like(x6, x5)
# ------decode ------
x7 = self.de_5(torch.cat((up1, x5), dim=1))
up2 = upsample_like(x7, x4)
x8 = self.de_4(torch.cat((up2, x4), dim=1))
up3 = upsample_like(x8, x3)
x9 = self.de_3(torch.cat((up3, x3), dim=1))
up4 = upsample_like(x9, x2)
x10 = self.de_2(torch.cat((up4, x2), dim=1))
up5 = upsample_like(x10, x1)
x11 = self.de_1(torch.cat((up5, x1), dim=1))
# side output
sup1 = self.side1(x11)
sup2 = self.side2(x10)
sup2 = upsample_like(sup2, sup1)
sup3 = self.side3(x9)
sup3 = upsample_like(sup3, sup1)
sup4 = self.side4(x8)
sup4 = upsample_like(sup4, sup1)
sup5 = self.side5(x7)
sup5 = upsample_like(sup5, sup1)
sup6 = self.side6(x6)
sup6 = upsample_like(sup6, sup1)
sup0 = self.out_conv(torch.cat((sup1, sup2, sup3, sup4, sup5, sup6), 1))
return torch.sigmoid(sup0)
if __name__ == '__main__':
u2net = U2Net(3, 1)
x = torch.randn(1,3, 512, 512)
print(u2net(x).shape)