基于UNet的肝脏肿瘤分割
文章地址:arxiv.org/pdf/1505.04597.pdf
1.UNet概述
UNet是一种用于图像分割的深度学习网络架构,最初由Olaf Ronneberger、Philipp Fischer和Thomas Brox在2015年提出。它在许多图像分割任务中取得了显著的成功,尤其在医学影像分割领域得到广泛应用。UNet的设计灵感来自于全卷积网络(Fully Convolutional Network)。它采用了编码器-解码器结构,其中编码器用于逐步缩小输入图像的尺寸并提取高层次的特征表示,而解码器则用于逐步恢复分辨率并生成像素级别的预测结果。
UNet的架构由两部分组成:下采样路径(U型结构的左半部分)和上采样路径(U型结构的右半部分)。下采样路径通过堆叠卷积层和池化层来逐步减小输入图像的尺寸,并在每个步骤中提取更高级别的特征。这些特征映射会保存并用于上采样路径。
上采样路径是通过堆叠卷积层和上采样操作来逐步恢复分辨率。在每个步骤中,上采样操作将特征图扩展到较大的尺寸,然后与相应的下采样路径的特征图进行连接。这种设计允许高级别的特征与低级别的细节信息相结合,从而提高分割结果的准确性。
为了进一步改进分割性能,UNet在上采样路径中引入了跳跃连接(skip connections)。跳跃连接将下采样路径中的特征图与对应的上采样路径特征图进行连接,使得网络可以直接利用不同分辨率的特征信息。这样的设计有助于提取不同尺度的特征并改善边界细节的分割结果。
在训练过程中,UNet通常使用交叉熵损失函数来比较网络输出和真实标签之间的差异。通过反向传播和优化算法(如随机梯度下降),网络的参数可以被调整以最小化损失函数,从而实现准确的图像分割。
2.技术环境
- python3.7
- pytorch1.8.1
- cuda10.2
- 英伟达显卡
3.主要代码实现
unet.py
import torch
import torch.nn as nn
import torch.nn.functional as F
class DoubleConv(nn.Module):
def __init__(self, in_channels, out_channels, mid_channels=None):
super().__init__()
if not mid_channels:
mid_channels = out_channels
self.double_conv = nn.Sequential(
nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),
nn.BatchNorm2d(mid_channels),
nn.ReLU(inplace=True),
nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),
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__()