【Pytorch项目实战】之语义分割:U-Net、UNet++、U2Net

文章目录

博主精品专栏导航

一、前言

1.1、什么是图像分割?

对图像中属于特定类别的像素进行分类的过程,即逐像素分类

  • 图像分类:识别图像中存在的内容。
  • 目标检测:识别图像中的内容和位置(通过边界框)。
  • 语义分割:识别图像中存在的内容以及位置(通过查找属于它的所有像素)。

(1)传统的图像分割算法:灰度分割,条件随机场等。
(2)深度学习的图像分割算法:利用卷积神经网络,来理解图像中的每个像素所代表的真实世界物体。在这里插入图片描述

1.2、语义分割与实例分割的区别

基于深度学习的图像分割技术主要分为两类:语义分割及实例分割。

语义分割(Semantic Segmentation):对图像中的每个像素点都进行分类预测,得到像素化的密集分类。然后提取具有感兴趣区域Mask。

  • 特点语义分割只能判断类别,无法区分个体。(只能将属于人的像素位置分割出来,但是无法分辨出图中有多少个人)
    在这里插入图片描述

实例分割(Instance Segmentation):不需要对每个像素点进行标记,只需要找到感兴趣物体的边缘轮廓即可

  • 详细过程:即同时利用目标检测和语义分割的结果,通过目标检测提供的目标最高置信度类别的索引,将语义分割中目标对应的Mask抽取出来。
  • 区别:目标检测输出目标的边界框和类别,实例分割输出的是目标的Mask和类别。
  • 特点可以区分个体。 (可以区分图像中有多少个人,不同人的轮廓都是不同颜色)
    在这里插入图片描述

1.3、语义分割的上下文信息

  • 上下文:指的是图像中的每一个像素点不可能是孤立的,一个像素一定和周围像素是有一定的关系的,大量像素的互相联系才产生了图像中的各种物体。
  • 上下文特征:指像素以及周边像素的某种联系。 即在判断某一个位置上的像素属于哪种类别的时候,不仅考察到该像素的灰度值,还充分考虑和它临近的像素。

1.4、语义分割的网络架构

一个通用的语义分割网络结构可以被广泛认为是一个:编码器 – 解码器(Encoder-Decoder)

  • (1)编码器:负责特征提取,通常是一个预训练的分类网络(如:VGG、ResNet)。
  • (2)解码器:将编码器学习到的可判别特征(低分辨率)从语义上投影到像素空间(高分辨率),以获得密集分类。

二、网络 + 数据集

在这里插入图片描述

2.1、经典网络的发展史(模型详解)

论文下载:史上最全语义分割综述(FCN、UNet、SegNet、Deeplab、ASPP…)
参考链接:经典网络 + 评价指标 + Loss损失(超详细介绍)

在这里插入图片描述

2.2、分割数据集下载

下载链接:【语义分割】FCN、UNet、SegNet、DeepLab

数据集简介
CamVid32个类别:367张训练图,101张验证图,233张测试图。
PascalVOC 2012(1)支持 5 类任务:分类、分割、检测、姿势识别、人体。(2)对于分割任务,共支持 21 个类别,训练和验证各 1464 和 1449 张图
NYUDv240个类别:795张训练图,645张测试图。
Cityscapes(1)50个不同城市的街景数据集,train/val/test的城市都不同。(2)包含:5k 精细标注数据,20k 粗糙标注数据。标注了 30 个类别。(3)5000张精细标注:2975张训练图,500张验证图,1525张测试图。(4)图像大小:1024×2048
Sun-RGBD37个类别:10355张训练图,2860张测试图。
MS COCO91个类别,328k 图像,2.5 million 带 label 的实例。
ADE20K150个类别,20k张训练图,2k张验证图。

三、算法详解

3.1、U-Net

论文地址:U-Net:Convolutional Networks for Biomedical Image Segmentation

论文源码:论文源码已开源,可惜是基于MATLAB的Caffe版本。 U-Net的实验是一个比较简单的ISBI cell tracking数据集,由于本身的任务比较简单,U-Net紧紧通过30张图片并辅以数据扩充策略便达到非常低的错误率,拿了当届比赛的冠军。

Unet 发表于 2015 年,属于 FCN 的一种变体,是一个经典的全卷积神经网络(即没有全连接层)。采用编码器 – 解码器(下采样 – 上采样)的对称U形结构和跳跃连接结构

  • 全卷积神经网络(FCN)是图像分割的开天辟地之作。
    • 为什么引入FCN:CNN浅层网络得到图像的纹理特征,深层得到轮廓特征等,但无法做到更精细的分割(像素级)。为了弥补这一缺陷,引入FCN。
    • FCN与CNN的不同点:FCN将CNN最后的全连接层替换为卷积层,故FCN可以输入任意尺寸的图像。
  • 而U-Net的初衷是为了解决生物医学图像问题。由于效果好,也被广泛的应用在卫星图像分割,工业瑕疵检测等。目前已有许多新的卷积神经网络设计方法,但仍延续了U-Net的核心思想。
    在这里插入图片描述

3.1.1、网络框架(U形结构+跳跃连接结构)

在这里插入图片描述
具体过程:

  • 输入图像大小为572 x 572。FCN可以输入任意尺寸的图像,且输出也是图像。
  • (1)压缩路径(Contracting path):由4个block组成,每个block使用2个(conv 3×3,ReLU)和1个MaxPooling 2×2。
    • 每次降采样之后的Feature Map的尺寸减半、数量翻倍。经过四次后,最终得到32×32的Feature Map。
  • (2)扩展路径(Expansive path):由4个block组成,每个block使用2个(conv 3×3,ReLU)和1个反卷积(up-conv 2×2)。
    • 11、每次上采样之后的Feature Map的尺寸翻倍、数量减半
    • 22、跳跃连接结构(skip connections):将左侧对称的压缩路径的Feature Map进行拼接(copy and crop)。由于左右两侧的Feature Map尺寸不同,将压缩路径的Feature Map裁剪到和扩展路径的Feature Map相同尺寸(左:虚线裁剪。右:白色块拼接)。
    • 33、逐层上采样 :经过四次后,得到392X392的Feature Map。
    • 44、卷积分类:再经过两次(conv 3×3,ReLU),一次(conv 1×1)。由于该任务是一个二分类任务,最后得到两张Feature Map(388x388x2)。

3.1.2、镜像扩大(保留边缘信息)

在不断的卷积过程中,图像会越来越小。为了避免数据丢失,在模型训练前,每一小块的四个边需要进行镜像扩大(不是直接补0扩大),以保留更多边缘信息。

在这里插入图片描述
由于当时计算机的内存较小,无法直接对整张图片进行处理(医学图像通常都很大),会采取把大图进行分块输入的训练方式,最后将结果一块块拼起来。

3.1.3、数据增强(变形)

医学影像数据普遍特点,就是样本量较少。当只有很少的训练样本可用时,数据增强对于教会网络所需的不变性和鲁棒性财产至关重要。

  • 对于显微图像,主要需要平移和旋转不变性,以及对变形和灰度值变化的鲁棒性。特别是训练样本的随机弹性变形,是训练具有很少注释图像的关键。
  • 在生物医学分割中,变形是组织中最常见的变化,并且可以有效地模拟真实的变形。在这里插入图片描述
    论文中的具体操作:使用粗糙的3乘3网格上的随机位移向量生成平滑变形。位移从具有10像素标准偏差的高斯分布中采样。然后使用双三次插值计算每个像素的位移。收缩路径末端的丢弃层执行进一步的隐式数据扩充。

3.1.4、损失函数(交叉熵)

论文的相关配置:Caffe框架,SGD优化器,每个batch一张图片,动量=0.99,交叉熵损失函数。

在这里插入图片描述

3.1.5、性能表现

在这里插入图片描述
用DIC(微分干涉对比)显微镜记录玻璃上的HeLa细胞。
(a) 原始图像。
(b) 覆盖地面真实分割。不同的颜色表示HeLa细胞的不同实例。
(c) 生成的分割掩码(白色:前景,黑色:背景)。
(d) 使用像素级损失权重映射,以迫使网络学习边界像素。

3.2、UNet++

论文地址:UNet++:A Nested U-Net Architecture for Medical Image Segmentation

UNet++ 发表于 2018 年,基于U-Net,采用一系列嵌套的密集的跳跃连接结构,并通过深度监督进行剪枝

  • UNet++的初衷是为了解决 " U-Net对病变或异常的医学图像缺乏更高的精确性 " 问题。

3.2.1、网络框架(U型结构+密集跳跃连接结构)

黑、红、绿、蓝色的组件将UNet++与U-Net区分开来。【语义分割】UNet++

  • 黑色:U-Net网络
  • 红色:深度监督(deep supervision)。可以进行模型剪枝 (model pruning)
  • 绿色:在跳跃连接(skip connections)设置卷积层,在 Encoder 和 Decoder 网络之间架起语义鸿沟。
  • 蓝色:一系列嵌套的密集的跳跃连接,改善了梯度流动。

在这里插入图片描述

3.2.2、改进的跳跃连接结构(融合+拼接)

Encoder 网络通过下采样提取低级特征;Decoder 网络通过上采样提取高级特征

  • U-Net 网络:(作者认为会产生语义鸿沟)
    • 特点:跳跃连接,又叫长连接或直接跳跃连接。将左右两边对称的特征图通过裁剪的方式进行拼接,有助于还原降采样所带来的信息损失(与残差块非常类似)。
    • 缺点裁剪将导致图像的深层细节丢失(如:人的毛发、小瘤附近的微刺等),影响细胞的微小特征(如:小瘤附近的微刺,可能预示着恶性瘤)。
  • UNet++网络:
    • 特点:一系列嵌套的,密集的跳跃连接。包括L1、L2、L3、L4四个U-Net网络,分别抓取浅层到深层特征。将左右两边对称的特征图先融合,再拼接,进而可以获取不同层次的特征。
      【备注】不同大小的感受野,对不同大小的目标,其敏感度也不同,获取图像的特征也不同。浅层(小感受野)对小目标更敏感;深层(大感受野)对大目标更敏感。

3.2.3、深度监督Deep supervision(剪枝)

此概念在对 U-Net 改进的多篇论文中都有使用,并不是该论文首先提出。

在结构 在这里插入图片描述后加上1×1卷积,相当于去监督每个分支的 U-Net 输出。在深度监督中,因为每个子网络的输出都是图像分割结果,所以通过剪枝使得网络有两种模式。

  • (1)精确模式:对所有分割分支的输出求平均值
  • (2)快速模式:从所有分割分支中选择一个分割图。剪枝越多参数越少,在不影响准确率的前提下,剪枝可以降低计算时间。
    在这里插入图片描述

(1)为什么可以剪枝?

  • 测试阶段:输入图像只有前向传播,剪掉部分对前面的输出完全没有影响;
  • 训练阶段:输入图像既有前向,又有反向传播,剪掉部分对剩余部分有影响 (绿色方框为剪掉部分) ,会帮助其他部分做权重更新。

(2)为什么要在测试时剪枝,而不是直接拿剪完的L1、L2、L3训练?

  • 剪掉的那部分对训练时的反向传播时时有贡献的,如果直接拿L1、L2、L3训练,就相当于只训练不同深度的U-NET,最后的结果会很差。

(3)如何进行剪枝?

  • 将数据分为训练集、验证集和测试集。
    训练集是需要训练的,测试集是不能碰的,所以根据选择的子网络在验证集的结果来决定剪多少。

3.2.4、损失函数

在这里插入图片描述

3.2.5、性能表现

如图显示:U-Net、宽U-Net和UNet++结果之间的定性比较。
在这里插入图片描述

如图显示:U-Net、宽U-Net和UNet++(在肺结节分割、结肠息肉分割、肝脏分割和细胞核分割任务中)的数量参数和分割精度。
在这里插入图片描述

  • 结论:
    (1)宽U-Net始终优于U-Net,除了两种架构表现相当的肝脏分割。这一改进归因于宽U-Net中的参数数量更大。
    (2)在没有深度监督的情况下,UNet++比UNet和宽U-Net都取得了显著的性能提升,IoU平均提高了2.8和3.3个点。
    (3)与没有深度监督的UNet++相比,具有深度监督的UNet++平均提高0.6分。

如图显示:在不同级别处修剪的UNet++分割性能。使用 UNet++ Li 表示在级别 i 处修剪的UNet++。

  • 结论:UNet++ L3平均减少了32.2%的推断时间,同时仅将IoU降低了0.6个点。更积极的修剪进一步减少了推断时间,但代价是显著的精度降低。
    在这里插入图片描述

3.3、U2-Net

论文地址:U2-Net:Going Deeper with Nested U-Structure for Salient Object Detection
代码下载:U2-Net-master

U2-Net 于 2020 年在CVPR上发表 ,主要针对显著性目标检测任务提出(Salient Object Detetion,SOD)。

显著性目标检测任务与语义分割任务非常相似,其是二分类任务,将图像中最吸引人的目标或区域分割出来,故只有前景和背景两类
在这里插入图片描述
第一列为原始图像,第二列为GT,第三列为U2-net结果、第四列为轻量级U2-net结果,其他列为其他比较主流的显著性目标检测网络模型。

  • 结论:无论是U2-net,还是轻量级U2-net,结果都比其他模型更出色。

U2-Net 基于 U-Net 提出了一种残余U形块(ReSidual U-blocks,RSU)结构。每个RSU就是一个缩版的 U-net,最后通过FPN的跳跃连接构建完整模型。

  • U2-Net 中的每一个block里面也是 U-Net,故称为 U2-Net 结构
  • 经过测试,对于分割物体前背景取得了惊人的效果。同样具有较好的实时性,经过测试在P100上前向时间仅为18ms(56fps)。

3.3.1、网络框架(RSU结构+U型结构+跳跃连接结构)

U2-Net包括6个编码器+5个解码器。除编码器En-6,其余的模型都是对称结构。通过跳跃连接结构进行特征拼接,并得到7个基于深度监督的损失值(Sup6-Sup0)。(6个block输出结果、1个特征融合后的结果)
在这里插入图片描述

3.3.2、残余U形块RSU

残余U形块RSU与现有卷积块的对比图:
(a)普通卷积块:PLN
(b)残余块:RES
(c)密集块:DSE
(d)初始块:INC
(e)残余U形块:RSU
在这里插入图片描述

  • RSU:每通过一个block后,Eecoder都会通过最大池化层下采样2倍,Decoder都会采用双线性插值进行上采样。

残余U形块RSU与残差模块的对比图:
(1)残差模块的权重层替换为U形模块;
(2)原始特征替换为本地特征;
在这里插入图片描述

3.3.3、损失函数(交叉熵)

由于U2net分成了多个block,故每个block都将输出一个loss值。7个loss相加(6个block输出结果、1个特征融合后的结果)

  • 公式(1):叠加损失值loss。l表示二值交叉熵损失函数,w表示每个损失的权重。
  • 公式(2):采用二值交叉熵损失函数。

在训练过程中,使用类似于HED的深度监督[45]。其有效性已在HED和DSS中得到验证。U2-net网络详解在这里插入图片描述

3.3.4、性能表现

U2-Net与其他最先进SOD模型的模型大小和性能比较。

四、项目实战

实战一:U-Net(不训练版)

由于模型未训练,故每次运行得到的结果都不同。原因:每次运行的初始化卷积核不同。代码剖析在这里插入图片描述

import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import torchvision.transforms as transforms
from PIL import Image
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'		# "OMP: Error #15: Initializing libiomp5md.dll"


class Encoder(nn.Module):
	def __init__(self, in_channels, out_channels):
		super(Encoder, self).__init__()
		self.block1 = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=0), nn.ReLU(inplace=True))
		self.block2 = nn.Sequential(nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=0), nn.ReLU(inplace=True))
		self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

	def forward(self, x):
		x = self.block1(x)
		x = self.block2(x)
		x_pooled = self.pool(x)
		return x, x_pooled


class Decoder(nn.Module):
	def __init__(self, in_channels, out_channels):
		super(Decoder, self).__init__()
		self.up_sample = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2)
		self.block1 = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=0), nn.ReLU(inplace=True))
		self.block2 = nn.Sequential(nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=0), nn.ReLU(inplace=True))

	def forward(self, x_prev, x):
		x = self.up_sample(x)						# 上采样
		x_shape = x.shape[2:]
		x_prev_shape = x.shape[2:]
		h_diff = x_prev_shape[0] - x_shape[0]
		w_diff = x_prev_shape[1] - x_shape[1]
		x_tmp = torch.zeros(x_prev.shape).to(x.device)
		x_tmp[:, :, h_diff//2: h_diff+x_shape[0], w_diff//2: x_shape[1]] = x
		x = torch.cat([x_prev, x_tmp], dim=1)		# 拼接
		x = self.block1(x)							# 卷积+ReLU
		x = self.block2(x)							# 卷积+ReLU
		return x


class UNet(nn.Module):
	def __init__(self, num_classes=2):
		super(UNet, self).__init__()
		"""
		padding=1。		输出图像大小=((572-3 + 2*1) / 1) + 1 = 572		# 卷积前后图像大小不变
		padding=0。		输出图像大小=((572-3) / 1) + 1 = 570			# 原论文每次卷积后,图像长宽各减2
		"""		

		"""编码器(4) —— 通道变化[3, 64, 128, 256, 512]"""
		self.down_sample1 = Encoder(in_channels=3, out_channels=64)
		self.down_sample2 = Encoder(in_channels=64, out_channels=128)
 		self.down_sample3 = Encoder(in_channels=128, out_channels=256)
		self.down_sample4 = Encoder(in_channels=256, out_channels=512)

		"""中间过渡层 —— 通道变化512, 1024]"""
		self.mid1 = nn.Sequential(nn.Conv2d(512, 1024, 3, bias=False), nn.ReLU(inplace=True))
		self.mid2 = nn.Sequential(nn.Conv2d(1024, 1024, 3, bias=False), nn.ReLU(inplace=True))

		"""解码器(4) —— 通道变化[1024, 512, 256, 128, 64]"""
		self.up_sample1 = Decoder(in_channels=1024, out_channels=512)
		self.up_sample2 = Decoder(in_channels=512, out_channels=256)
		self.up_sample3 = Decoder(in_channels=256, out_channels=128)
		self.up_sample4 = Decoder(in_channels=128, out_channels=64)

		"""分类器    —— 通道变化[64, 类别数]"""
		self.classifier = nn.Conv2d(64, num_classes, 1)

	def forward(self, x):
		x1, x = self.down_sample1(x)
		x2, x = self.down_sample2(x)
		x3, x = self.down_sample3(x)
		x4, x = self.down_sample4(x)

		x = self.mid1(x)
		x = self.mid2(x)

		x = self.up_sample1(x4, x)
		x = self.up_sample2(x3, x)
		x = self.up_sample3(x2, x)
		x = self.up_sample4(x1, x)

		x = self.classifier(x)
		return x


def image_loader(image_path):
	"""模型训练前的格式转换:[3, 384, 384] -> [1, 3, 384, 384]"""
	image = Image.open(image_path)			# 打开图像(numpy格式)
	loader = transforms.ToTensor()			# 数据预处理(Tensor格式)
	image = loader(image).unsqueeze(0)		# tensor.unsqueeze():增加一个维度,其值为1。
	return image.to(device, torch.float)


def image_trans(tensor):
	"""绘制图像前的格式转换:[1, 3, 384, 384] -> [3, 384, 384]"""
	image = tensor.clone()					# clone():复制
	image = torch.squeeze(image, 0)			# tensor.squeeze():减少一个维度,其值为1。
	unloader = transforms.ToPILImage()		# 数据预处理(PILImage格式)
	image = unloader(image)					# 图像转换
	return image
	
	
if __name__ == '__main__':
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")		# 可用设备
	raw_image = image_loader(r"大黄蜂.jpg")										# 导入图像
	
	model = UNet(4)																# 模型实例化
	new_image = model(raw_image)												# 前向传播
	print("输入图像维度: ", raw_image.shape)
	print("输出图像维度: ", new_image.shape)

	raw_image = image_trans(raw_image)
	new_image = image_trans(new_image)
	# 由于模型未训练,故每次运行得到的结果都不同。原因:每次运行的初始化卷积核不同。
	plt.subplot(121), plt.imshow(raw_image, 'gray'), plt.title('raw_image')
	plt.subplot(122), plt.imshow(new_image, 'gray'), plt.title('new_image')
	plt.show()

实战二:U2-Net(不训练版)

由于模型未训练,故每次运行得到的结果都不同。原因:每次运行的初始化卷积核不同。图像分割之U-Net、U2-Net及其Pytorch代码构建
在这里插入图片描述

import torch.nn.functional as F
import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import torchvision.transforms as transforms
from PIL import Image
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'  # "OMP: Error #15: Initializing libiomp5md.dll"


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, 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)


def image_loader(image_path):
    """模型训练前的格式转换:[3, 384, 384] -> [1, 3, 384, 384]"""
    image = Image.open(image_path)              # 打开图像(numpy格式)
    loader = transforms.ToTensor()              # 数据预处理(Tensor格式)
    image = loader(image).unsqueeze(0)          # tensor.unsqueeze():增加一个维度,其值为1。
    return image.to(device, torch.float)


def image_trans(tensor):
    """绘制图像前的格式转换:[1, 3, 384, 384] -> [3, 384, 384]"""
    image = tensor.clone()                      # clone():复制
    image = torch.squeeze(image, 0)             # tensor.squeeze():减少一个维度,其值为1。
    unloader = transforms.ToPILImage()          # 数据预处理(PILImage格式)
    image = unloader(image)                     # 图像转换
    return image


if __name__ == '__main__':
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")       # 可用设备
    raw_image = image_loader(r"大黄蜂.jpg")                                      # 导入图像

    model = U2Net(3, 1)                                                         # 模型实例化
    new_image = model(raw_image)                                                # 前向传播
    print("输入图像维度: ", raw_image.shape)
    print("输出图像维度: ", new_image.shape)

    raw_image = image_trans(raw_image)
    new_image = image_trans(new_image)
    # 由于模型未训练,故每次运行得到的结果都不同。原因:每次运行的初始化卷积核不同。
    plt.subplot(121), plt.imshow(raw_image, 'gray'), plt.title('raw_image')
    plt.subplot(122), plt.imshow(new_image, 'gray'), plt.title('new_image')
    plt.show()

实战三:基于U-Net实现目标检测(数据集:PASCAL VOC)

在GitCode上,基于Pascal VOC数据集的U-Net、PSP-Net、deeplabv3+三个网络模型的开源代码。
代码链接:基于Pytorch的目标分割:中文详细教程 + Pascal VOC数据集 + 完整代码

在这里插入图片描述

PASCAL VOC是由欧盟组织的世界级计算机视觉挑战赛。2005年举办第一场挑战赛,2012年停止举办。每年的内容都有所不同,从目标分类,到检测,分割,人体布局,动作识别等等,数据集的容量以及种类也在不断的增加和改善。

  • PASCAL全称:Pattern Analysis,Statical Modeling and Computational Learning(模式分析,静态建模和计算学习)。
  • VOC全称:Visual Object Classes(可视化对象类)。
  • 近年来,目标检测或分割模型更倾向于使用MS COCO数据集Computer Vision Datasets。但PASCAL VOC数据集对于目标检测或分割类型具有先驱者地位PASCAL VOC Datasets
  • 最重要两个年份的数据集:PASCAL VOC 2007 与 PASCAL VOC 2012。PASCAL VOC Datasets的详细介绍
  • 有兴趣的小伙伴还可以尝试自己制作训练集。语义分割:VOC数据集的制作教程

实战四:基于U2-Net的服装裤子分割(数据集:pants_data)

网盘链接:https://pan.baidu.com/s/1p32LsehWk8RmgvMOKxWsrw?pwd=2aem
提取码:2aem

U2-Net网络实现目标边缘检测(pants_data数据集)

  1. 训练图像(服装裤子) —— 训练标签(服装裤子的轮廓图)
  2. 构建模型:将数据集与U2-Net官方开源代码进行整合,并对u2net_train.py以及u2net_test.py进行了详细的整理与备注。

超参数设置:epoch=10000,batch_size=10,iter=5000的演示图。
由于服装裤子的轮廓图相对简单,验证发现:iter = 200可以得到最优模型,而 iter = 500 生成的图1和图2裤脚有灰痕在这里插入图片描述

❤️ u2net_train.py

import torch
import torchvision
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import torch.optim as optim
import torchvision.transforms as standard_transforms

import numpy as np
import glob

from data_loader import Rescale
from data_loader import RescaleT
from data_loader import RandomCrop
from data_loader import ToTensor
from data_loader import ToTensorLab
from data_loader import SalObjDataset

from model import U2NET
from model import U2NETP

import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'         # "OMP: Error #15: Initializing libiomp5md.dll"
########################################################################################################


def muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v):
    """损失函数"""
    bce_loss = nn.BCELoss(size_average=True)
    loss0 = bce_loss(d0, labels_v)
    loss1 = bce_loss(d1, labels_v)
    loss2 = bce_loss(d2, labels_v)
    loss3 = bce_loss(d3, labels_v)
    loss4 = bce_loss(d4, labels_v)
    loss5 = bce_loss(d5, labels_v)
    loss6 = bce_loss(d6, labels_v)
    loss = loss0 + loss1 + loss2 + loss3 + loss4 + loss5 + loss6
    print("l0: %3f, l1: %3f, l2: %3f, l3: %3f, l4: %3f, l5: %3f, l6: %3f\n"
          %(loss0.data.item(), loss1.data.item(), loss2.data.item(), loss3.data.item(),
            loss4.data.item(), loss5.data.item(), loss6.data.item()))
    return loss0, loss


if __name__ == '__main__':
    ########################################################################################################
    # (1)导入训练集
    data_dir = os.path.join(os.getcwd(), 'train_data' + os.sep)                         # 数据路径(train_data:存放图像+标签的文件夹)
    tra_image_dir = os.path.join('train_img' + os.sep)                                  # 训练图像(train_img:存放图像的文件夹)
    tra_label_dir = os.path.join('train_label' + os.sep)                                # 训练标签(train_label:存放标签的文件夹)
    model_name = 'u2net'                    # 定义了两种模型:u2net、轻量级u2netp
    model_dir = os.path.join(os.getcwd(), 'saved_models' + os.sep)                      # 预训练模型(saved_models:存放预训练模型的文件夹)。os.sep不可删除

    image_ext = '.jpg'						# 注意:图像与标签的后缀(tif、gif、jpg、png)
    label_ext = '.png'						# 注意:图像与标签的文件名需一一对应
    tra_img_name_list = glob.glob(data_dir + tra_image_dir + '*' + image_ext)           # 获取图像
    tra_lbl_name_list = []
    for img_path in tra_img_name_list:
        img_name = img_path.split(os.sep)[-1]
        aaa = img_name.split(".")
        bbb = aaa[0:-1]
        img_idx = bbb[0]
        for i in range(1, len(bbb)):
            img_idx = img_idx + "." + bbb[i]
        tra_lbl_name_list.append(data_dir + tra_label_dir + img_idx + label_ext)        # 获取图像对应的标签
    print("train images: ", len(tra_img_name_list))
    print("train labels: ", len(tra_lbl_name_list))
    ########################################################################################################
    # (2)超参数设置 ———— 图像增强 + 数据分配器
    epoch_num = 10
    batch_size = 10
    salobj_dataset = SalObjDataset(img_name_list=tra_img_name_list, lbl_name_list=tra_lbl_name_list,
                                   transform=transforms.Compose([RescaleT(320), RandomCrop(288), ToTensorLab(flag=0)]))
    salobj_dataloader = DataLoader(salobj_dataset, batch_size=batch_size, shuffle=True, num_workers=1)
    ########################################################################################################
    # (3)模型选择
    if model_name == 'u2net':
        net = U2NET(3, 1)
    elif model_name == 'u2netp':
        net = U2NETP(3, 1)
    if torch.cuda.is_available():
        net.cuda()

    optimizer = optim.Adam(net.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)
    ########################################################################################################
    # (4)开始训练
    print("start training", "..."*25)
    train_num = len(tra_img_name_list)      # 训练图像的总数
    ite_num = 0                             # 迭代次数
    ite_num4val = 0
    running_loss = 0.0                      # 训练损失(总)
    running_tar_loss = 0.0                  # 训练损失(loss0)
    save_frq = 100                          # 每100次迭代训练,保存预训练模型

    for epoch in range(0, epoch_num):
        net.train()     # 模型训练

        for i, data in enumerate(salobj_dataloader):
            ite_num = ite_num + 1
            ite_num4val = ite_num4val + 1
            inputs, labels = data['image'], data['label']
            inputs = inputs.type(torch.FloatTensor)
            labels = labels.type(torch.FloatTensor)
            if torch.cuda.is_available():
                inputs_v, labels_v = Variable(inputs.cuda(), requires_grad=False), Variable(labels.cuda(), requires_grad=False)
            else:
                inputs_v, labels_v = Variable(inputs, requires_grad=False), Variable(labels, requires_grad=False)

            optimizer.zero_grad()                                                           # 梯度清零
            d0, d1, d2, d3, d4, d5, d6 = net(inputs_v)                                      # 前向传播
            loss2, loss = muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v)        # 损失函数
            loss.backward()                                                                 # 反向传播
            optimizer.step()                                                                # 梯度更新

            running_loss += loss.data.item()                # 累加损失值(总)
            running_tar_loss += loss2.data.item()           # 累加损失值(loss0)
            del d0, d1, d2, d3, d4, d5, d6, loss2, loss     # 删除临时变量

            print("[epoch: %3d/%3d, batch: %5d/%5d, ite: %d] train loss: %3f, tar: %3f "
                  % (epoch + 1, epoch_num, (i + 1) * batch_size, train_num, ite_num,
                     running_loss / ite_num4val, running_tar_loss / ite_num4val))

            if ite_num % save_frq == 0:
                torch.save(net.state_dict(), model_dir + model_name + "_itr_%d_train_%3f_tar_%3f.pth"
                           % (ite_num, running_loss / ite_num4val, running_tar_loss / ite_num4val))
                running_loss = 0.0
                running_tar_loss = 0.0
                net.train()         # 继续训练
                ite_num4val = 0

❤️ u2net_test.py

import os
from skimage import io, transform
import torch
import torchvision
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
# import torch.optim as optim

import numpy as np
from PIL import Image
import glob

from data_loader import RescaleT
from data_loader import ToTensor
from data_loader import ToTensorLab
from data_loader import SalObjDataset

from model import U2NET         # full size version 173.6 MB
from model import U2NETP        # small version u2net 4.7 MB


def normPRED(d):
    # normalize the predicted SOD probability map
    ma = torch.max(d)
    mi = torch.min(d)
    dn = (d-mi)/(ma-mi)
    return dn


def save_output(image_name, pred, d_dir):
    predict = pred
    predict = predict.squeeze()
    predict_np = predict.cpu().data.numpy()

    im = Image.fromarray(predict_np*255).convert('RGB')
    img_name = image_name.split(os.sep)[-1]
    image = io.imread(image_name)
    imo = im.resize((image.shape[1], image.shape[0]), resample=Image.BILINEAR)

    pb_np = np.array(imo)

    aaa = img_name.split(".")
    bbb = aaa[0:-1]
    img_idx = bbb[0]
    for i in range(1, len(bbb)):
        img_idx = img_idx + "." + bbb[i]

    imo.save(d_dir + img_idx + '.png')


def main():
    ########################################################################################################
    # (1)导入测试集
    model_name = 'u2net'                                                    # 定义了两种模型:u2net、轻量级u2netp
    pre_model_name = 'u2net_itr_4_train_6.046402_tar_0.528644.pth'          # 预训练模型
    data_dir = 'test_images'                                                # 存放测试图像的文件夹

    image_dir = os.path.join(os.getcwd(), 'test_data', data_dir)                                # 测试图像地址(test_data存放测试图像的上一级文件夹)
    prediction_dir = os.path.join(os.getcwd(), 'test_data', data_dir + '_results' + os.sep)     # 结果存放地址(若无,则自动新建文件夹)
    model_dir = os.path.join(os.getcwd(), 'saved_models', pre_model_name)                       # 预训练模型地址(saved_models存放预训练模型的文件夹)
    img_name_list = glob.glob(image_dir + os.sep + '*')                                         # 获取图像
    print(img_name_list)
    ########################################################################################################
    # (2)超参数设置 ———— 图像增强 + 数据分配器
    test_salobj_dataset = SalObjDataset(img_name_list=img_name_list, lbl_name_list=[],
                                        transform=transforms.Compose([RescaleT(320), ToTensorLab(flag=0)]))
    test_salobj_dataloader = DataLoader(test_salobj_dataset, batch_size=1, shuffle=False, num_workers=1)
    ########################################################################################################
    # (3)模型选择
    if model_name == 'u2net':
        print("load U2NET = 173.6 MB")
        net = U2NET(3, 1)
    elif model_name == 'u2netp':
        print("load U2NEP = 4.7 MB")
        net = U2NETP(3, 1)
    if torch.cuda.is_available():
        net.load_state_dict(torch.load(model_dir))
        net.cuda()
    else:
        net.load_state_dict(torch.load(model_dir, map_location='cpu'))
    ########################################################################################################
    # (4)开始训练
    print("start testing", "..."*25)
    net.eval()      # 测试模型

    for i_test, data_test in enumerate(test_salobj_dataloader):
        print("Extracting image:", img_name_list[i_test].split(os.sep)[-1])      # 提取图像(逐张)
        inputs_test = data_test['image']
        inputs_test = inputs_test.type(torch.FloatTensor)

        # 判断可用设备类型,并进行图像格式转换
        if torch.cuda.is_available():
            inputs_test = Variable(inputs_test.cuda())
        else:
            inputs_test = Variable(inputs_test)

        d1, d2, d3, d4, d5, d6, d7 = net(inputs_test)                   # 前向传播
        pred = d1[:, 0, :, :]
        pred = normPRED(pred)                                           # 归一化

        # 判断文件夹是否存在,若不存在则新建
        if not os.path.exists(prediction_dir):
            os.makedirs(prediction_dir, exist_ok=True)
        save_output(img_name_list[i_test], pred, prediction_dir)        # 保存预测图像

        del d1, d2, d3, d4, d5, d6, d7


if __name__ == "__main__":
    main()

实战五:基于U2-Net的视网膜血管分割(数据集:DRIVE_data)

网盘链接:https://pan.baidu.com/s/1q-vbgDFsDnabhOXQyqNYtw?pwd=znry
提取码:znry

DRIVE(Digital Retinal Images for Vessel Extraction)数据集来自于荷兰的糖尿病视网膜病变筛查计划,用于视网膜血管分割,进而研究病变原理。数据集于 2004 年由图像科学研究所发布,筛查人群为25-90岁的糖尿病受试者。共包括40张图像(训练集20、测试机20),33张未显示任何糖尿病视网膜病变迹象,7张显示轻度早期糖尿病视网膜病变迹象。
在这里插入图片描述

深度学习框架Keras:基于U-Net的眼底图像血管分割实例(DRIVE数据集)

  1. 构建模型:博主将数据集与Pytorch下的U2-Net官方开源代码进行了整合,将u2net_train.py以及u2net_test.py进行了详细的整理与备注。
  2. 可以将眼部图像分别与眼部轮廓图像、手工标注血管图像进行训练,得到两个预训练模型,然后进行图像测试。

模型一:眼部图像作为训练集(Images)、手工标注血管图像作为训练掩膜(mask)

超参数设置:epoch=10000,batch_size=10,iter=5000的演示图。
由于手工标注血管图像相对简单,验证发现,iter = 100 可以得到最优模型
在这里插入图片描述

模型二:眼部图像作为训练集(Images)、眼部轮廓图像作为训练掩膜(manual)

超参数设置:epoch=10000,batch_size=10,iter=5000的演示图。
由于眼部轮廓图像相对手工标注血管图像比较复杂,验证发现:iter = 1000可以得到最优模型,而 iter = 100 生成的结果会有点模糊
在这里插入图片描述

❤️ u2net_train.py

import torch
import torchvision
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
import torch.optim as optim
import torchvision.transforms as standard_transforms

import numpy as np
import glob

from data_loader import Rescale
from data_loader import RescaleT
from data_loader import RandomCrop
from data_loader import ToTensor
from data_loader import ToTensorLab
from data_loader import SalObjDataset

from model import U2NET
from model import U2NETP

import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'         # "OMP: Error #15: Initializing libiomp5md.dll"
########################################################################################################


def muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v):
    """损失函数"""
    bce_loss = nn.BCELoss(size_average=True)
    loss0 = bce_loss(d0, labels_v)
    loss1 = bce_loss(d1, labels_v)
    loss2 = bce_loss(d2, labels_v)
    loss3 = bce_loss(d3, labels_v)
    loss4 = bce_loss(d4, labels_v)
    loss5 = bce_loss(d5, labels_v)
    loss6 = bce_loss(d6, labels_v)
    loss = loss0 + loss1 + loss2 + loss3 + loss4 + loss5 + loss6
    print("l0: %3f, l1: %3f, l2: %3f, l3: %3f, l4: %3f, l5: %3f, l6: %3f\n"
          % (loss0.data.item(), loss1.data.item(), loss2.data.item(), loss3.data.item(),
             loss4.data.item(), loss5.data.item(), loss6.data.item()))
    return loss0, loss


if __name__ == '__main__':
    ########################################################################################################
    # (1)导入训练集
    data_dir = os.path.join(os.getcwd(), 'train_data' + os.sep)                         # 数据路径(train_data:存放图像+标签的文件夹)
    tra_image_dir = os.path.join('images' + os.sep)                                     # 训练图像(train_img:存放图像的文件夹)
    tra_label_dir = os.path.join('mask' + os.sep)                                       # 训练标签(train_label:存放标签的文件夹)
    model_name = 'u2net'                    # 定义了两种模型:u2net、轻量级u2netp
    model_dir = os.path.join(os.getcwd(), 'saved_models' + os.sep)                      # 预训练模型(saved_models:存放预训练模型的文件夹)。os.sep不可删除

    image_ext = '.tif'						# 注意:图像与标签的后缀(tif、gif、jpg、png)
    label_ext = '_mask.gif'					# 注意:图像与标签的文件名需一一对应
    tra_img_name_list = glob.glob(data_dir + tra_image_dir + '*' + image_ext)           # 获取图像
    tra_lbl_name_list = []
    for img_path in tra_img_name_list:
        img_name = img_path.split(os.sep)[-1]
        aaa = img_name.split(".")
        bbb = aaa[0:-1]
        img_idx = bbb[0]
        for i in range(1, len(bbb)):
            img_idx = img_idx + "." + bbb[i]
        tra_lbl_name_list.append(data_dir + tra_label_dir + img_idx + label_ext)        # 获取图像对应的标签
    print("train images: ", len(tra_img_name_list))
    print("train labels: ", len(tra_lbl_name_list))
    ########################################################################################################
    # (2)超参数设置 ———— 图像增强 + 数据分配器
    epoch_num = 10000
    batch_size = 10
    salobj_dataset = SalObjDataset(img_name_list=tra_img_name_list, lbl_name_list=tra_lbl_name_list,
                                   transform=transforms.Compose([RescaleT(320), RandomCrop(288), ToTensorLab(flag=0)]))
    salobj_dataloader = DataLoader(salobj_dataset, batch_size=batch_size, shuffle=True, num_workers=1)
    ########################################################################################################
    # (3)模型选择
    if model_name == 'u2net':
        net = U2NET(3, 1)
    elif model_name == 'u2netp':
        net = U2NETP(3, 1)
    if torch.cuda.is_available():
        net.cuda()

    optimizer = optim.Adam(net.parameters(), lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0)
    ########################################################################################################
    # (4)开始训练
    print("start training", "..."*25)
    train_num = len(tra_img_name_list)      # 训练图像的总数
    ite_num = 0                             # 迭代次数
    ite_num4val = 0
    running_loss = 0.0                      # 训练损失(总)
    running_tar_loss = 0.0                  # 训练损失(loss0)
    save_frq = 100                          # 每100次迭代训练,保存预训练模型

    for epoch in range(0, epoch_num):
        net.train()     # 模型训练

        for i, data in enumerate(salobj_dataloader):
            ite_num = ite_num + 1
            ite_num4val = ite_num4val + 1
            inputs, labels = data['image'], data['label']
            inputs = inputs.type(torch.FloatTensor)
            labels = labels.type(torch.FloatTensor)
            if torch.cuda.is_available():
                inputs_v, labels_v = Variable(inputs.cuda(), requires_grad=False), Variable(labels.cuda(), requires_grad=False)
            else:
                inputs_v, labels_v = Variable(inputs, requires_grad=False), Variable(labels, requires_grad=False)

            optimizer.zero_grad()                                                           # 梯度清零
            d0, d1, d2, d3, d4, d5, d6 = net(inputs_v)                                      # 前向传播
            loss2, loss = muti_bce_loss_fusion(d0, d1, d2, d3, d4, d5, d6, labels_v)        # 损失函数
            loss.backward()                                                                 # 反向传播
            optimizer.step()                                                                # 梯度更新

            running_loss += loss.data.item()                # 累加损失值(总)
            running_tar_loss += loss2.data.item()           # 累加损失值(loss0)
            del d0, d1, d2, d3, d4, d5, d6, loss2, loss     # 删除临时变量

            print("[epoch: %3d/%3d, batch: %5d/%5d, ite: %d] train loss: %3f, tar: %3f "
                  % (epoch + 1, epoch_num, (i + 1) * batch_size, train_num, ite_num,
                     running_loss / ite_num4val, running_tar_loss / ite_num4val))

            if ite_num % save_frq == 0:
                torch.save(net.state_dict(), model_dir + model_name + "_itr_%d_train_%3f_tar_%3f.pth"
                           % (ite_num, running_loss / ite_num4val, running_tar_loss / ite_num4val))
                running_loss = 0.0
                running_tar_loss = 0.0
                net.train()         # 继续训练
                ite_num4val = 0

❤️ u2net_test.py

import os
from skimage import io, transform
import torch
import torchvision
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
# import torch.optim as optim

import numpy as np
from PIL import Image
import glob

from data_loader import RescaleT
from data_loader import ToTensor
from data_loader import ToTensorLab
from data_loader import SalObjDataset

from model import U2NET         # full size version 173.6 MB
from model import U2NETP        # small version u2net 4.7 MB


def normPRED(d):
    # normalize the predicted SOD probability map
    ma = torch.max(d)
    mi = torch.min(d)
    dn = (d-mi)/(ma-mi)
    return dn


def save_output(image_name, pred, d_dir):
    predict = pred
    predict = predict.squeeze()
    predict_np = predict.cpu().data.numpy()

    im = Image.fromarray(predict_np*255).convert('RGB')
    img_name = image_name.split(os.sep)[-1]
    image = io.imread(image_name)
    imo = im.resize((image.shape[1], image.shape[0]), resample=Image.BILINEAR)

    pb_np = np.array(imo)

    aaa = img_name.split(".")
    bbb = aaa[0:-1]
    img_idx = bbb[0]
    for i in range(1, len(bbb)):
        img_idx = img_idx + "." + bbb[i]

    imo.save(d_dir + img_idx + '.png')


def main():
    ########################################################################################################
    # (1)导入测试集
    model_name = 'u2net'                                                    # 定义了两种模型:u2net、轻量级u2netp
    pre_model_name = 'u2net_itr_10_train_0.494240_tar_0.077563.pth'         # 预训练模型
    data_dir = 'images'                                                     # 存放测试图像的文件夹

    image_dir = os.path.join(os.getcwd(), 'test_data', data_dir)                                # 测试图像地址(test_data存放测试图像的上一级文件夹)
    prediction_dir = os.path.join(os.getcwd(), 'test_data', data_dir + '_results' + os.sep)     # 结果存放地址(若无,则自动新建文件夹)
    model_dir = os.path.join(os.getcwd(), 'saved_models', pre_model_name)                       # 预训练模型地址(saved_models存放预训练模型的文件夹)
    img_name_list = glob.glob(image_dir + os.sep + '*')                                         # 获取图像
    print(img_name_list)
    ########################################################################################################
    # (2)超参数设置 ———— 图像增强 + 数据分配器
    test_salobj_dataset = SalObjDataset(img_name_list=img_name_list, lbl_name_list=[],
                                        transform=transforms.Compose([RescaleT(320), ToTensorLab(flag=0)]))
    test_salobj_dataloader = DataLoader(test_salobj_dataset, batch_size=1, shuffle=False, num_workers=1)
    ########################################################################################################
    # (3)模型选择
    if model_name == 'u2net':
        print("load U2NET = 173.6 MB")
        net = U2NET(3, 1)
    elif model_name == 'u2netp':
        print("load U2NEP = 4.7 MB")
        net = U2NETP(3, 1)
    if torch.cuda.is_available():
        net.load_state_dict(torch.load(model_dir))
        net.cuda()
    else:
        net.load_state_dict(torch.load(model_dir, map_location='cpu'))
    ########################################################################################################
    # (4)开始训练
    print("start testing", "..."*25)
    net.eval()      # 测试模型

    for i_test, data_test in enumerate(test_salobj_dataloader):
        print("Extracting image:", img_name_list[i_test].split(os.sep)[-1])      # 提取图像(逐张)
        inputs_test = data_test['image']
        inputs_test = inputs_test.type(torch.FloatTensor)

        # 判断可用设备类型,并进行图像格式转换
        if torch.cuda.is_available():
            inputs_test = Variable(inputs_test.cuda())
        else:
            inputs_test = Variable(inputs_test)

        d1, d2, d3, d4, d5, d6, d7 = net(inputs_test)                   # 前向传播
        pred = d1[:, 0, :, :]
        pred = normPRED(pred)                                           # 归一化

        # 判断文件夹是否存在,若不存在则新建
        if not os.path.exists(prediction_dir):
            os.makedirs(prediction_dir, exist_ok=True)
        save_output(img_name_list[i_test], pred, prediction_dir)        # 保存预测图像

        del d1, d2, d3, d4, d5, d6, d7


if __name__ == "__main__":
    main()

文章出处登录后可见!

已经登录?立即刷新

共计人评分,平均

到目前为止还没有投票!成为第一位评论此文章。

(0)
乘风的头像乘风管理团队
上一篇 2023年2月22日 下午10:03
下一篇 2023年2月23日 上午8:40

相关推荐