一.图像语义分割问题概述
????????图像语义分割(Semantic Segmentation) 是图像处理和机器学习视觉技术中关于图像理解的重要一环,也是 AI 领域中一个重要的分支。图像语义分割问题就是对图像中的不同物体打上语义标签(用不同的颜色代表不同类别的物体),其本质即是对图像中每一个像素点进行分类,确定每个点的类别(如属于背景、人或车等),从而进行区域划分。目前,语义分割已经被广泛应用于自动驾驶、无人机落点判定等场景中。其实,图像语义分割问题相对于传统的神经网络分类就是将概率分类问题扩展到了二维空间的每个像素点上,因此,我们的神经网络结构和学习策略也应该随之发生改变。
二.全卷积神经网络 FCN 网络模型
?????????FCN(Fully Convolutional Networks)全称为全卷积神经网络。传统的CNN网络(比如VGG、Resnet等)使用的网络模型通常会在最后连接几层全连接层,它会将原来二维的矩阵(图片)压扁成一维的,从而丢失了空间信息,最后训练输出一个标量,这就是我们的分类概率。但是这种方式只能标识整个图片的类别,而不能标识每个像素点的类别,所以这种后接全连接的分类方法不适用于图像分割。而FCN的提出抛弃了全连接层,全部使用卷积层,这样最后就可以获得一张2维的feature map,从而对像素点进行分类。
1.全卷积化?
? ? ? ? FCN网络模型中全部为卷积层,不包含全连接层。这样使得网络可以适应任意尺寸的输入,并且输入输出均为同维度数据,成为图像语义分割实现的前提条件。
2. 上采样
(1)上采样概述
? ? ? ? 在卷积层中,我们一般包含卷积和池化两个操作,经过卷积层会使得我们的数据尺寸不断变小,这是一个信息聚合的过程(或者称为编码),我们将这个过程叫做下采样。但是在图像语义分割过程中,我们最终需要对原图尺寸上训练的像素点进行分类,为了得到和原图等大的分割图,我们需要将下采样的数据进行扩大恢复,这可以看作是下采样操作的逆过程(或者解码过程),我们将这个过程称为上采样。
????????FCN的思想就是使用反卷积层对最后一个卷积层的特征图(feature map)进行上采样,使它恢复到输入图像相同的尺寸,从而对每一个像素都产生一个预测,同时保留了原始输入图像中的空间信息,最后在上采样的特征图进行像素的分类。
?
? ? ? ? ?在深度学习中,上采样(upsampling)主要包括2种方式:
- Resize:Resize是一种简单的上采样方法,其类似于图像的缩放操作。其关键是所使用的采样算法,比如nearest算法就是复制邻近的值来扩张,linear是通过线性插值来扩张等,这种方式往往不需要进行参数的学习。FCN使用的就是双线性插值算法进行上采样操作。
- Deconvolution:反卷积/逆卷积,也叫做转置卷积。这种方式在效果上可以看作是卷积操作的一个逆过程(但本质上不是),其需要进行参数学习和训练,具体解释我们在U-Net中进行说明。
(2)Pytorch中对Resize上采样的支持?
????????上采样有很多算法,Pytorch提供了一种最基本和通用的上采样函数 torch.nn.UpSimple()
torch.nn.Upsample(size=None, scale_factor=None, mode='nearest', align_corners=None)
?
-
输入:一维数据 [N, C, W]、二维数据 [N, C, H, W]、三维数据 [N, C, D, H, W]。输入数据的格式为minibatch x channels x size -
size (int or Tuple**[int] or Tuple**[int, int] or Tuple**[int, int, int]**, optional) – 根据不同的输入类型制定的输出大小 -
scale_factor (float or Tuple**[float] or Tuple**[float, float] or Tuple**[float, float, float]**, optional) – 指定输出为输入的多少倍数。如果输入为tuple,其也要制定为tuple类型 -
mode (str, optional) – 可使用的上采样算法,有'nearest' , 'linear' , 'bilinear' , 'bicubic' and 'trilinear' . 默认使用``'nearest' -
align_corners (bool, optional) – 如果为True,输入的角像素将与输出张量对齐,因此将保存下来这些像素的值。仅当使用的算法为'linear' , 'bilinear' or 'trilinear'时可以使用。``默认设置为``False
3.跳跃结构(Skip)?
? ? ? ? 在FCN层层的下采样和上采样操作过程中(编解码过程中),会丢失原图像很多的语义结构信息,如果直接将最后上采样的结果输出,会使得学习效果很粗糙。为了丰富和保留语义信息,优化输出结果,FCN在网络模型中添加了跳跃结构,即将下层结果上采样后与上层结果逐点相加,进行多层特征融合,吸收上层的语义信息。经过实践证明,多层融合后的输出结果更加精确!
?三.U-Net 网络模型
????????U-Net网络是2015年提出来的一种全卷积神经网络,该网络最初提出是用来解决医学细胞图像分割问题,后来被广泛应用在图像的语义分割上。U-Net网络是一种U形的编码-解码结构,使用下采样和上采样操作来对特征进行压缩和提取。并通过跳连接将原始特征拼接到解码阶段,进一步丰富语义信息。相比于FCN网络模型,U-Net的特点如下:
- 采用U型的对称式编-解码结构
- 采用转置卷积/反卷积的上采样方式
- 采用跳跃拼接的特征融合方式
1.U型的编解码结构?
? ? ? ? U-Net网络模型结构中的每个处理单元包括两次3x3卷积核的卷积操作,输入数据经过U-net网络进行了四次下采样(Pooling)操作,然后又经过了四次上采样操作恢复到原始尺寸。最终结果通过一个sigmod输出为每个像素点的二分类概率,并通过BCE二元交叉熵作为损失来训练网络模型。
?2.上采样--转置卷积/反卷积
? ? ? ? 在FCN中上采样使用的是简单的双线性插值来扩大尺寸,而在U-Net中,我们使用转置卷积/反卷积的方式来进行上采样操作。关于什么是反卷积,首先举一个例子,将一个4x4的输入通过3x3的卷积核在进行普通卷积(无padding, stride=1),将得到一个2x2的输出。而转置卷积将一个2x2的输入通过同样3x3大小的卷积核将得到一个4x4的输出,看起来似乎是普通卷积的逆过程。但事实上两者在理论上并没有什么关系,操作的过程也不是可逆的。只是这个过程比较像一个正逆操作的效果,所以我们习惯于称他为反卷积/逆卷积。接下来我们看一下转置卷积的概念由来(来自于文章:https://blog.csdn.net/tsyccnh/article/details/87357447)?
(1)转置卷积推导过程?
?
?(2)Pytorch中的转置卷积/反卷积操作
????????转置卷积有时被称作反卷积,结果上它并不是卷积的逆,但操作上的确是卷积的反过程。它是一种上采样技术,可以理解成把输入尺寸放大的技术,被广泛的应用在图像分割,超分辨率等应用中。与传统的上采样技术,如线性插值,双线性插值等方法相比,转置卷积是一种需要训练,可学习的方法。Pytorch提供了转置卷积函数ConvTranspose(我们这里只讨论二维的nn.ConvTranspose2d)
CLASS torch.nn.ConvTranspose2d(in_channels: int,
out_channels: int,
kernel_size: Union[int, Tuple[int, int]],
stride: Union[int, Tuple[int, int]] = 1,
padding: Union[int, Tuple[int, int]] = 0,
output_padding: Union[int, Tuple[int, int]]=0,
groups: int = 1,
bias: bool = True,
dilation: int = 1,
padding_mode: str = 'zeros')
- ?stride: 当
stride=1 、padding=0 、kernel=1*1 并且权值为1 时,由图5方法可得输出应该和输入一模一样。代码验证后发现的确如此。维持其他参数不变,让stride=2 时,可以发现输出是在输入的每两个元素间都插入了一个0 的结果,此时两个相邻非零元素的间距是2 。维持其他参数不变,让stride=3 时,可以发现输出是在输入的每两个元素间都插入了两个0 的结果,此时两个相邻非零元素的间距是3 …。可见转置卷积层的stride ,就等价于在input 每两个元素之间插入(stride-1) 个0 ,即将输入膨胀。简略:在输入图像中元素之间添加(stride - 1)个0来扩张图像
-
padding: 与上面同样的输入和卷积核,让stride = 1 ,修改padding ,padding=1 时,输出是输入上下少一行,左右少一列的结果。padding=2 时,输出是输入上下少两行行,左右少两行的结果。可见转置卷积层的padding ,就等价于在input 上,边缘处上下各去掉padding 行,左右各去掉padding 列的结果。即将输入向中心缩小。简略:图像收缩的大小尺寸。p_new = k - p - 1(实际图像扩张的尺寸,这也是卷积能扩大尺寸的原因,填充了好几圈0) -
output_padding: 默认在输出的最右边拓展output_padding 列,在输出的最下边拓展output_padding 行。output_padding 必须比stride 或者dilation 小,否则会报错:RuntimeError: output padding must be smaller than either stride or dilation.简略:输出扩充大小 -
kernel_size:卷积核大小,转置卷积本质上还是用的卷积运算来实现转置卷积,只不过通过填充0来先将原图扩大再卷积缩小 ?
(3)转置卷积/反卷积计算公式与计算过程?
?3.特征融合--跳跃拼接
? ? ? ? FCN中是通过将输出结果上采样后与上层结果进行逐点相加的方式进行特征融合。而在U-Net中是通过将解码阶段输出与同层编码阶段的输出在通道上进行直接拼接进行特征融合,拼接后再通过两次3x3卷积进行训练学习。
?4.U-Net Pytorch实现
? ? ? ? 本次U-Net网络通过Pytorch框架实现,以DUTS数据集作为图像分割的训练和测试数据。输入为RGB格式三通道的256x256大小的彩色图片,输出为黑白两类像素点的图像分割结果。
(1)nn.Sequential()
????????nn.Sequential是一个PyTorch提供的有序容器,神经网络模块将按照在传入构造器的顺序依次被添加到计算图中执行,同时以神经网络模块为元素的有序字典也可以作为传入参数。该容器相当于一个大的组合计算模块,在执行forward计算时,因为每一个module都继承于nn.Module,都会实现__call__ 与forward 函数,所以forward函数中通过for循环依次调用添加到self._module 中的子模块,最后输出经过所有神经网络层的结果。
#最简单的模式(每一层默认名称为0、1、2、3...)
model = nn.Sequential(
nn.Conv2d(1,20,5),
nn.ReLU(),
nn.Conv2d(20,64,5),
nn.ReLU()
)
#字典传入(自定义命名每一层)
model = nn.Sequential(OrderedDict([
('conv1', nn.Conv2d(1,20,5)),
('relu1', nn.ReLU()),
('conv2', nn.Conv2d(20,64,5)),
('relu2', nn.ReLU())
]))
(2)utils.py 工具类
from PIL import Image
#图片标准化等比缩放尺寸 256*256
def UnifyImage(image,size=(256,256)):
max_len = max(image.size) # 取得最大边长
#print(image.mode)
if image.mode == "RGB":
img_back = Image.new("RGB",(max_len,max_len),(0,0,0)) #新建一个RGB三通道指定Size的全黑背景
img_back.paste(image,(0,0)) #把image原图贴到背景左上角(0,0)
return img_back.resize(size) #等比缩放到Size
elif image.mode == "1":
img_back = Image.new("1", (max_len, max_len), 0) # 新建一个二值指定Size的全黑背景
img_back.paste(image, (0, 0)) # 把image原图贴到背景左上角(0,0)
return img_back.resize(size) # 等比缩放到Size
(3)dataset.py 数据模型
import os
import torch
import numpy as np
from PIL import Image
from torch.utils.data import Dataset
import matplotlib.pyplot as plt
import torchvision.transforms as transforms
import utils
#构造自定义数据集
class ModelDataset(Dataset):
def __init__(self,image_dir,mask_dir):
super(ModelDataset, self).__init__()
self.transTensor = transforms.ToTensor()# 将PIL Image自动转化并归一化为tensor(c,h,w)
self.image_dir = image_dir
self.mask_dir = mask_dir
self.images = os.listdir(self.image_dir) #os.listdir()函数用于返回指定的文件夹包含的文件或文件夹的[名字]的列表。
def __len__(self):
return len(self.images)
def __getitem__(self, index):
#加载图片
img_path = os.path.join(self.image_dir,self.images[index]) #os.path.join函数用于将字符串按照系统盘符拼接为路径
mask_path = os.path.join(self.mask_dir,self.images[index].replace(".jpg",".png")) #mask数据集为png后缀
image = Image.open(img_path).convert("RGB") #图片为 "RGB"格式 真彩图像 三通道 (h,w,c) 范围[0,255]
mask = Image.open(mask_path).convert("1") #图片默认为 "L"格式 灰度图像 单通道 (h,w) 范围[0,255] 其中0为黑色,255为白色; => 转换为二值图像,范围{true,false}(转换为numpy),其中true为白色,false为黑色(转换为tensor后,0为黑色,1为白色)
#对图片做填充后,等比缩放为 256*256
segment_image = utils.UnifyImage(image)
segment_mask = utils.UnifyImage(mask)
#返回值 image : 256*256 标准大小的 (3,h,w) 的 [0,1] float32范围 的归一化输入
#返回值 mask : 256*256 标准大小的 (1,h,w) 的 {0,1} float32范围 的归一化二值输出
return self.transTensor(segment_image),self.transTensor(segment_mask) #转为Tensor归一化后返回
'''
1.PIL Image的一些操作
(1)Image.open(path):打开指定路径的图像,返回Image对象
(2)Image.convert(mode):图像格式转换 1:二值图像.1像素(非黑即白,true或false) L:灰度图像([0,255]) RGB:真彩图像
- RGB -> L灰度转换使用了 YUV模型中Y通道的转换公式 L = R * 299/1000 + G * 587/1000+ B * 114/1000
(3)显示图像:
- 使用matplotlib显示:PIL读取的图像为PIL.Image.Image对象,无法用matplotlib直接显示,需要先转为numpy.ndarray对象。
import matplotlib.pyplot as plt
image = np.array(image) #(h,w,c)格式
plt.imshow(img)
plt.show()
- 使用Image.show():调用本地的图片浏览器,不方便
(4)Image.save(path,format):保存图像 路径和格式(如png、jpg)
(5)Image.new(mode,size,color):创建新图像,(模式,二元组尺寸(w,h),颜色)
- 颜色字符串设置:支持“red”、“blue”等方式
- 颜色通道设置:(x,y,z) 三个数字分别对应"RGB"通道的数值, 比如(0,0,255)是蓝色
(6)Image.paste(box,point):图像黏贴(合并),将box沾到Image的point位置
img_back = Image.new("RGB",(400,400),(0,0,0))
img_back.paste(image,(0,0))
img_back.show()
(7)Image.resize:返回改变尺寸的图像的拷贝(通道模式不变)。变量size是所要求的尺寸,是一个二元组:(width, height)。Resize会将图片进行拉伸,非等比改变(有可能会改变图像信息)
(8)Image.size:返回图像的尺寸(w,h)
(9)Image转Tensor:torchvision.transforms.ToTensor
(10)Tensor转Image:torchvision.transforms.ToPILImage
2.torchvision库的用处:
(1)torchvision.datasets:提供现成的一些流行数据集,比如MNIST、CIFAR等
(2)torchvision.models:提供一些定义好的流行网络模型,比如AlexNet、VGG、ResNet
(3)torchvision.transforms:提供图片转换和处理工具
- torchvision.transforms.ToTensor:把一个取值范围是[0,255]的PIL.Image或者shape为(H,W,C)的numpy.ndarray,转换成形状为[C,H,W],取值范围是[0,1.0]的torch.FloadTensor(除以255,自动进行归一化)
- torchvision.transforms.Normalize(mean, std):给定均值:(R,G,B) 方差:(R,G,B),将会把Tensor正则化。即:Normalized_image=(image-mean)/std。
- torchvision.transforms.Compose(transforms):组合器,将多个transform组合起来使用。序列化依次进行处理
- torchvision.transforms.ToPILImage:将shape为(C,H,W)的Tensor或shape为(H,W,C)的numpy.ndarray转换成PIL.Image。(每个元素自动乘以255,恢复原格式)
- torchvision.utils.save_image(tensor, filename, nrow=8, padding=2):将tensor自动转换为image,如何保存到本地路径
'''
(4)net.py 网络模型结构
import torch
import torch.nn as nn
#卷积模块
class ConvBlock(nn.Module):
def __init__(self,in_channel,out_channel):
super(ConvBlock, self).__init__()
#构建 卷积块(进行两次卷积操作)
self.layer = nn.Sequential(
#第一次卷积操作
nn.Conv2d(in_channel,out_channel,kernel_size=3,stride=1,padding=1),#卷积操作 (batch,in_ch,h,w) -> (batch,out_ch,h,w) 不改变大小
nn.BatchNorm2d(out_channel),#批标准化 将数据标准化到正态分布
nn.ReLU(inplace=True),#激活函数 inplace=True表示覆盖输入数据(避免了临时变量频繁释放,提高效率)
#第二次卷积操作
nn.Conv2d(out_channel, out_channel, kernel_size=3, stride=1, padding=1),
# 卷积操作 (batch,in_ch,h,w) -> (batch,out_ch,h,w) 不改变大小
nn.BatchNorm2d(out_channel), # 批标准化 将数据标准化到正态分布
nn.ReLU(inplace=True), # 激活函数 inplace=True表示覆盖输入数据(避免了临时变量频繁释放,提高效率)
)
def forward(self,x):
return self.layer(x)
#上采样模块:反卷积+拼接
class UpBlock(nn.Module):
def __init__(self,in_channel,out_channel):
super(UpBlock, self).__init__()
#上采样:反卷积 (batch,in_ch,h,w) -> (batch,out_ch,2h,2w)
self.up = nn.ConvTranspose2d(in_channel,out_channel,kernel_size=2,stride=2)
def forward(self,input_1,input_2):
output_2 = self.up(input_2) #先上采样
merge = torch.cat([input_1,output_2],dim=1) #跳跃连接,合并输入
return merge #返回合并后结果
#U-Net网络:编码器+解码器
class Unet(nn.Module):
def __init__(self,in_channel,out_channel):
super(Unet, self).__init__()
self.pool = nn.MaxPool2d(2)
#编码器
self.encoderConv1 = ConvBlock(in_channel,64)
self.encoderConv2 = ConvBlock(64,128)
self.encoderConv3 = ConvBlock(128,256)
self.encoderConv4 = ConvBlock(256,512)
self.encoderConv5 = ConvBlock(512,1024)
#解码器
self.upSimple1 = UpBlock(1024,512)
self.decoderConv1 = ConvBlock(1024,512)
self.upSimple2 = UpBlock(512,256)
self.decoderConv2 = ConvBlock(512,256)
self.upSimple3 = UpBlock(256, 128)
self.decoderConv3 = ConvBlock(256, 128)
self.upSimple4 = UpBlock(128, 64)
self.decoderConv4 = ConvBlock(128, 64)
#输出
self.out = nn.Conv2d(64,out_channel,kernel_size=1)
self.cls = nn.Sigmoid() #输出每个像素点的分类概率 0黑色,1白色(转换为二分类问题)
def forward(self,x):
#编码,下采样过程
en_x1 = self.encoderConv1(x) #输出 (batch,64,256,256)
down_x1 = self.pool(en_x1) #输出 (batch,64,128,128)
en_x2 = self.encoderConv2(down_x1) #输出 (batch,128,128,128)
down_x2 = self.pool(en_x2) #输出 (batch,128,64,64)
en_x3 = self.encoderConv3(down_x2) #输出 (batch,256,64,64)
down_x3 = self.pool(en_x3) #输出(batch,256,32,32)
en_x4 = self.encoderConv4(down_x3) #输出(batch,512,32,32)
down_x4 = self.pool(en_x4) #输出(batch,512,16,16)
en_x5 = self.encoderConv5(down_x4) #输出(batch,1024,16,16)
#解码,上采样过程
up_x1 = self.upSimple1(en_x4,en_x5) #输出 (batch,1024,32,32)
de_x1 = self.decoderConv1(up_x1) #输出 (batch,512,32,32)
up_x2 = self.upSimple2(en_x3, de_x1) #输出 (batch,512,64,64)
de_x2 = self.decoderConv2(up_x2) #输出 (batch,256,64,64)
up_x3 = self.upSimple3(en_x2, de_x2)# 输出 (batch,256,128,128)
de_x3 = self.decoderConv3(up_x3)# 输出 (batch,128,128,128)
up_x4 = self.upSimple4(en_x1, de_x3)# 输出 (batch,128,256,256)
de_x4 = self.decoderConv4(up_x4)# 输出 (batch,64,256,256)
#输出
return self.cls(self.out(de_x4)) #输出(batch,1,256,256)
(5)train.py 训练模块
import torch
from torch.utils.data import DataLoader
from dataset import *
from net import *
#.cuda() 与 .to(device)在指定GPU上没有区别。
# 1..to(device) 可以指定CPU 或者GPU
# 2.cuda() 只能指定GPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')#调整设备,优先使用gpu
img_dir = r"D:\日常材料\作业报告\机器学习\DUTS数据集\train_data"
label_dir = r"D:\日常材料\作业报告\机器学习\DUTS数据集\train_label"
weight_path = "params/unet.pth"
epoch = 10
if __name__ == "__main__":
#1.加载自己训练数据集的数据加载器
data_loader = DataLoader(ModelDataset(img_dir,label_dir),batch_size=4,shuffle=True)
#2.将模型加载到设备上
unet = Unet(3,1).to(device)
#2.1加载预训练权重(如果有的话)
if os.path.exists(weight_path):
unet.load_state_dict(torch.load(weight_path))
print("successful load weight")
else:
print("weight is not exist")
#3.设置优化器和损失
optim = torch.optim.Adam(unet.parameters(),lr=0.001)
loss = torch.nn.BCELoss() #默认对一个batch里面的数据每个相应位置做二元交叉熵损失并且求平均。
#4.开始训练
for i in range(epoch):
loss_sum = 0.0
for index,(input_img,label_img) in enumerate(data_loader):
input_img,label_img = input_img.to(device),label_img.to(device) #将数据放到设备上
#计算损失
output_img = unet(input_img)
train_loss = loss(output_img,label_img)
loss_sum += train_loss.item()
#清空梯度
optim.zero_grad()
#反馈计算梯度
train_loss.backward()
#更新参数权重
optim.step()
torch.save(unet.state_dict(),weight_path)#每一轮都保存训练参数weight
print("[epoch %d]: loss is %.3f" % (i+1,loss_sum))
'''
一.模型的保存和加载有两种方式:
1.只保存和加载模型参数(推荐使用):自定义网络,并且参数名称与结构要与保存的模型中的一致
#保存模型参数
torch.save(the_model.state_dict(), PATH)
#加载模型参数(这种方式将会提取所有的参数, 然后再放到你的新建网络中)
the_model = TheModelClass(*args, **kwargs)
the_model.load_state_dict(torch.load(PATH))
2. 保存和加载整个模型:无需自定义网络
#保存网络结构+参数
torch.save(the_model, PATH)
#加载网络模型
the_model = torch.load(PATH)
3.pytorch保存文件后缀:.pkl 或 .pth (二者都是采用二进制序列化存储,只是编码方式不同)
二.eval()和train():如果Model里有BN或者 Dropout层
1.model.eval():切换为测试模式,固定BN层和Drop层,使用训练好的参数。
- 在eval模式下,dropout层会让所有的激活单元都通过,而batchnorm层会停止计算和更新mean和var,直接使用在训练阶段已经学出的mean和std值。
2.model.train():切换为训练模式,此时 dropout和batch normalization的操作在训练时起到防止网络过拟合的问题(默认)
- 在train模式下,dropout网络层会按照设定的参数p设置保留激活单元的概率(保留概率=p); batchnorm层会继续计算数据的mean和std等参数并更新。
3.注意:该模式不会影响各层的gradient计算行为,即gradient计算和存储与training模式一样
三.with torch.no_grad():
1.含义:在该模块下,所有计算得出的tensor的requires_grad都自动设置为False。数据不会梯度计算,即不会进行反相传播,不参与构建计算图;
2.作用:起到加速和节省显存的作用,具体行为就是停止gradient计算,从而节省了GPU算力和显存,但是并不会影响dropout和batchnorm层的行为。
3.使用:如果不在意显存大小和计算时间的话,仅仅使用model.eval()已足够得到正确的validation的结果;而with torch.zero_grad()则是更进一步加速和节省gpu空间。
4.推荐:with torch.no_grad()和model.eval()在测试时搭配使用
'''
(6)test.py 测试模块
import os
import torch
from net import *
from dataset import *
from torch.utils.data import DataLoader
from torchvision.utils import save_image
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
test_dir = r"D:\日常材料\作业报告\机器学习\DUTS数据集\test_data"
true_dir = r"D:\日常材料\作业报告\机器学习\DUTS数据集\test_label"
save_dir = r"D:\日常材料\作业报告\机器学习\DUTS数据集\result"
weight_path = "params/unet.pth"
unet = Unet(3,1).to(device)
if os.path.exists(weight_path):
unet.load_state_dict(torch.load(weight_path))
print("weights load successful")
else:
print("weights is not exist")
unet.eval() #切换训练模式
test_loader = DataLoader(ModelDataset(test_dir,true_dir),batch_size=1,shuffle=False) #(1,c,h,w) 网络必须四维输入
#不进行计算图构建
with torch.no_grad():
for index,(test_img,test_label) in enumerate(test_loader):
test_img, test_label = test_img.to(device), test_label.to(device) # 将数据放到设备上
# 输出 (1,1,h,w) 的概率矩阵
out = unet(test_img)
# 将概率转化为像素黑白(0黑,1白)
out[out <= 0.5] = 0.0
out[out > 0.5] = 1.0
# 将真实图像和预测图象拼接(拼接到batch上,构造雪碧网格图),也可以降维输出单个图像
# 1.torch.stack(tuple,dim):将tensor在dim维上进行堆叠(升维操作)
# 1.tuple:拼接tensor的元组,注意Size必须相同
# 2.dim:维度,默认为0
# 2.cat :原维度上拼接
img = torch.cat([test_label,out],dim=0)
#保存图像
print(img.shape)
save_path = os.path.join(save_dir,str(index)+".png")
#注意,save_image将图像保存为RGB三通道,如果是二值图像则三个通道数值都相同,即伪灰度图像
save_image(img,save_path) #将给定的Tensor保存成image文件。如果给定的是mini-batch tensor,那就用make-grid做成雪碧图,再保存。
# 张量的压缩的膨胀(升维和降维)
# 1.squeeze()函数 : 对张量的维度进行减少的操作(注意压缩或者扩充的维度为1)
# - torch.squeeze(tensorA,dim) :在第dim维上进行降维
# - tensorA.squeeze(dim):同上
# 注意:(1)如果指定的维度大于1,那么将操作无效
# (2)如果不指定维度dim,那么将删除所有维度为1的维度
# 2.unsqueeze函数: 对张量的维度进行膨胀的操作
# - torch.unsqueeze(tensorA,dim) :在第dim维上进行升维
# - tensorA.unsqueeze(dim):同上
(7)训练效果
????????左边为label,右边为网络输出。可以看出有一些效果,但还是比较粗糙,因该是训练轮数太少所致,加大训练轮数和训练集规模应该会得到改善。
?
|