计算机视觉——基于深度学习UNet实现的复杂背景文档二值化算法实现与模型训练

1. 引言

阈值分割可以被视为一个分类问题,通常涉及两个类别,这也是为什么阈值分割也被称为二值化。对于文档图像,我们期望阈值算法能够正确地将墨水分类为黑色,将纸张分类为白色,从而得到二值化图像。对于数字灰度图像,最简单的实现方法是选择一个阈值值,比如图像二值化,并将高于这个值的灰度级别分配为白色,将剩余的级别分配为黑色。问题在于正确找到这个值,以便能够完美匹配前景和背景元素。
在这里插入图片描述
对于文档图像,我们知道预期的结果,尽管存在一些问题使得这个领域变得非常具有挑战性,比如老化退化。老化伪影包括斑点(出现在纸面上的棕褐色斑点)、前后墨水干扰、皱褶纸、胶带标记、折叠痕迹等。
在这里插入图片描述

老化过程引起的许多问题的例子,如(右上和右下)老化,(左下和右中)背对前方的干扰和(左上)人为操纵作为胶带和(右下)皱纸造成的问题

一些二值化技术,如 Otsu[2]、Sauvola[3] 和 Niblack[4],但它们在处理如图 1.2 所示的文档时效果并不理想:
在这里插入图片描述
在这里将探讨如何通过使用基于卷积神经网络(CNN)的U-Net架构训练的模型进行分类,来实现具有不同类型问题的文档二值化。CNN的典型用途在于分类任务,其中对图像的输出是一个单一的类别标签。然而,在许多视觉任务中,期望的结果不仅包括图像中物体是否存在,还包括其定位,即每个像素都应该被分配到一个类别标签。

训练代码地址:https://download.csdn.net/download/matt45m/89112642

2.开源数据集

数据集由总共 5,027 张图像及其相应的真实情况(二值参考图像)组成。所使用的图像来自以下数据集:

  • 文档图像二值化 (DIB) - Nabuco 数据集:15 张图像 [5]
  • DIBCO 和 H-DIBCO (年份: 2009, 2010, 2011, 2012, 2013, 2014, 2016, 2017):116 张图像 [6]
  • ICFHR 2016 棕榈叶手稿图像二值化挑战赛:99 张图像 [7]
  • ICDAR2017 历史文档作家识别竞赛:4,782 张图像 [8]
  • PHIBD 2012 波斯遗产图像二值化数据集:15 张图像 [9]

为了增加样本数量,将对原始图像和二值参考图像应用数据增强。由于模型只接受 256×256 尺寸的图像,将对图像进行分割而不是调整大小;这样,即不丢失信息,并且增加训练样本的数量。

2.1 数据增强

数据增强过程从对原始图像和其真实情况进行转换开始。选择了仅翻转(垂直和水平)和旋转(90°, 180° 和 270°)增强转换。其他转换也是可能的,例如模糊、添加噪声、改变亮度和对比度等。只是要记住,对于这些类型的转换,相应的真实情况不应接受它们。

应用转换后,结果图像和原始图像经过切割过程,产生 256×256 像素的图像。在生成增强的代码中,有可能增加结果图像的数量,使切割步骤小于 256,生成重叠的切割。

下图通过动画展示了切割过程。灰色线条显示了切割器将如何分割图像。左侧图像的步长为 256,而右侧为 128。请注意,在两种情况下裁剪图像的最终大小都是相同的(256×256),如白色方框中所示。
在这里插入图片描述
在这里插入图片描述

使用 5,027 张图像的数据集,步长为 256 的切割产生 27,630 张图像,而步长为 128 产生 97,146 张图像。请注意,这个数字甚至不到四倍,原因是接近图像边缘时的原因,不需要处理。
在这里插入图片描述

3. U-Net 模型架构

U-Net 是从传统的卷积神经网络演变而来的,最初由 Olaf Ronneberger 等人在 2015 年为生物医学图像分割设计和应用。一般的卷积神经网络专注于图像分类,其中输入是一张图像,输出是一个或多个模式标签。然而,在许多视觉任务中,期望的结果不仅包括图像中物体是否存在,还包括其定位,即每个像素都应该被分配到一个类别标签。

3.1 卷积

更正式地说,卷积是一个积分,它表示一个函数 g 在另一个函数 f 上移动时的重叠量,但在数字图像处理和深度学习中,卷积是将两个图像结合起来形成第三个图像的数学方法。通常,两个结合的图像中有一个不是图像,而是一个滤波器(或核心),一个值矩阵,其大小和值决定了卷积过程的效果类型。主要思想是将核心放置在每个像素上(横跨整个图像),并将其值与目标像素及其局部邻居相乘并求和。
在这里插入图片描述
卷积在数字图像处理中最常见的用途是边缘检测、模糊和去噪。虽然在 CNN 中卷积的中间层的效果众所周知,但可以通过展示最终应用来使其更加明显。
在这里插入图片描述

卷积在数字图像处理中最常见的用途是边缘检测、模糊和去噪。虽然在 CNN 中卷积的中间层的效果众所周知,但可以通过展示最终应用来使其更加明显。

上图 显示了在左侧图像上应用的不同核心的卷积,首先是模糊,然后是 Sobel 边缘检测。我使用了这两个卷积在原始图像的灰度版本上。使用 CNN 时,卷积分别应用于每个 RGB 通道,这在图像处理中并不常见,因为它会产生奇怪的结果。

填充:如上例所示,结果图像比原始图像小,差距与核心大小有关;核心越大,中心离图像边缘越远。为了产生与输入相同大小的输出,在边缘添加额外的像素进行填充。这样,当滑动时,核心可以让原始边缘像素位于其中心,同时延伸到边缘之外的额外像素。

在这里插入图片描述
如上图显示了使用 OpenCV 的 copyMakeBorder 函数的一些填充方法。原始图像来自核心滤波器,每个角落有四个颜色点以帮助展示每种方法的差异。
在这里插入图片描述

步幅:步幅是核心窗口在输入矩阵上每次移动的像素数。步幅为一意味着每次滑动间隔一个像素,因此每个单独的滑动都作为一个标准卷积。步幅为二意味着每次滑动间隔两个像素,在此过程中跳过每隔一个滑动,大致缩小了两倍。步幅为三意味着跳过每两个滑动,大致缩小了三倍,以此类推。

3.2 U-Net 的工作原理

该架构包含两条路径,如下图所示。第一条路径是收缩路径(也称为编码器),用于捕捉图像中的上下文。编码器只是传统的卷积和最大池化层堆栈。第二条路径是对称的扩展路径(也称为解码器),用于使用转置卷积进行精确定位,以对输入特征图进行上采样。它是一个端到端的全卷积网络,没有任何密集层。
在这里插入图片描述

U-Net 架构(以最低分辨率的 32x32 像素为例)。每个蓝色框对应一个多通道特征图。通道数在框的顶部标注。x-y 大小在框的左下角提供。白色框表示复制的特征图。箭头表示不同的操作

每个编码器块接收输入,与相应裁剪的特征图进行连接,应用两个 3×3 卷积(无填充),每个卷积后跟一个修正线性单元(ReLu),然后是一个 2×2 最大池化操作,步幅为 2 进行下采样,如图 3.7 详细说明。
在这里插入图片描述
每个解码器块由两个卷积层组成,在本项目中,输入图像的形状从 256×256×1 变为 256×256×64,在第一个卷积过程中,特征通道的深度增加了 64 倍。请注意,在上图所示的代码中,使用了 padding=’same’(用零填充额外的边界,但可以在使用 padding=’valid’ 的 Keras 中进行非填充卷积);这样,卷积过程不会减小图像尺寸。Ronneberger 等人使用了非填充卷积,原因如下:‘由于每次卷积都会丢失边界像素,所以裁剪是必要的’ 。在 U-Net 的原始实现中,输入是一张 128×128×1 图像,编码器输出一个 8×8×256 形状。
在这里插入图片描述
解码器由扩展块组成,每个块包括使用 2×2 上采样层(转置卷积)对特征图进行上采样,该层将特征通道的数量减半,与编码器路径中相应裁剪的特征图进行连接,以及两个 3×3 卷积,每个卷积后跟一个 ReLU,如上图像所示。由于每次卷积都会丢失边界像素,裁剪是必要的。由于我在编码器中采用了填充卷积(padding=’same’),因此裁剪过程是不必要的。
在 U-Net 的原始实现中,解码器将形状从 8×8×256 增加到 128×128×1。该过程以一个 1×1 卷积结束,将每个 64 个组件的特征向量映射到所需的类别数。

4. 改进算法

在算法实现的过程中,发现很多文档图像都是整张A4左右的大小,而Unet的模型的输入最大为256,为了改进输入的局限,这里模型结构借签了Enet,代码使用的框架是Pytorch。
Unet网络部分:

import torch.nn as nn
import torchclass Inception(nn.Module):def __init__(self, in_ch, out_ch):super(Inception, self).__init__()hide_ch = out_ch // 2self.inception = nn.Sequential(nn.Conv2d(in_ch, hide_ch, 1),nn.BatchNorm2d(hide_ch),nn.ReLU(inplace=True),nn.Conv2d(hide_ch, hide_ch, 3, padding=1, groups=hide_ch),nn.BatchNorm2d(hide_ch),nn.ReLU(inplace=True),nn.Conv2d(hide_ch, out_ch, 1))def forward(self, x):return self.inception(x)class DoubleConv(nn.Module):def __init__(self, in_ch, out_ch):super(DoubleConv, self).__init__()self.doubleConv = nn.Sequential(Inception(in_ch, out_ch),nn.BatchNorm2d(out_ch),nn.ReLU(inplace=True),Inception(out_ch, out_ch),nn.BatchNorm2d(out_ch),nn.ReLU(inplace=True))def forward(self, x):return self.doubleConv(x)class UNet(nn.Module):def __init__(self, in_ch, out_ch):super(UNet, self).__init__()# downself.conv1 = DoubleConv(in_ch, 64)self.pool1 = nn.Conv2d(64, 64, 2, 2, groups=64)self.conv2 = DoubleConv(64, 128)self.pool2 = nn.Conv2d(128, 128, 2, 2, groups=128)self.bottom = DoubleConv(128, 256)# upself.up3 = nn.ConvTranspose2d(256, 128, 2, 2)self.conv3 = DoubleConv(128 * 2, 128)self.up4 = nn.ConvTranspose2d(128, 64, 2, 2)self.conv4 = DoubleConv(64 * 2, 64)self.out = nn.Conv2d(64, out_ch, 1)def forward(self, x):# downconv1 = self.conv1(x)pool1 = self.pool1(conv1)conv2 = self.conv2(pool1)pool2 = self.pool2(conv2)bottom = self.bottom(pool2)# upup3 = self.up3(bottom)merge3 = torch.cat([up3, conv2], dim=1)conv3 = self.conv3(merge3)up4 = self.up4(conv3)merge4 = torch.cat([up4, conv1], dim=1)conv4 = self.conv4(merge4)out = self.out(conv4)return nn.Sigmoid()(out)if __name__ == '__main__':net = UNet(1, 2)inputs = torch.zeros((1, 1, 512, 512), dtype=torch.float32)output = net(inputs)print(output.size())

使用ENet实现:

import torch.nn as nn
import torch
class InitialBlock(nn.Module):def __init__(self,in_channels,out_channels,bias=False,relu=True):super().__init__()if relu:activation = nn.ReLUelse:activation = nn.PReLU# Main branch - As stated above the number of output channels for this# branch is the total minus 3, since the remaining channels come from# the extension branchself.main_branch = nn.Conv2d(in_channels,out_channels - 1,kernel_size=3,stride=2,padding=1,bias=bias)# Extension branchself.ext_branch = nn.MaxPool2d(3, stride=2, padding=1)# Initialize batch normalization to be used after concatenationself.batch_norm = nn.BatchNorm2d(out_channels)# PReLU layer to apply after concatenating the branchesself.out_activation = activation()def forward(self, x):main = self.main_branch(x)ext = self.ext_branch(x)# Concatenate branchesout = torch.cat((main, ext), 1)# Apply batch normalizationout = self.batch_norm(out)return self.out_activation(out)class RegularBottleneck(nn.Module):def __init__(self,channels,internal_ratio=4,kernel_size=3,padding=0,dilation=1,asymmetric=False,dropout_prob=0,bias=False,relu=True):super().__init__()# Check in the internal_scale parameter is within the expected range# [1, channels]if internal_ratio <= 1 or internal_ratio > channels:raise RuntimeError("Value out of range. Expected value in the ""interval [1, {0}], got internal_scale={1}.".format(channels, internal_ratio))internal_channels = channels // internal_ratioif relu:activation = nn.ReLUelse:activation = nn.PReLU# Main branch - shortcut connection# Extension branch - 1x1 convolution, followed by a regular, dilated or# asymmetric convolution, followed by another 1x1 convolution, and,# finally, a regularizer (spatial dropout). Number of channels is constant.# 1x1 projection convolutionself.ext_conv1 = nn.Sequential(nn.Conv2d(channels,internal_channels,kernel_size=1,stride=1,bias=bias), nn.BatchNorm2d(internal_channels), activation())# If the convolution is asymmetric we split the main convolution in# two. Eg. for a 5x5 asymmetric convolution we have two convolution:# the first is 5x1 and the second is 1x5.if asymmetric:self.ext_conv2 = nn.Sequential(nn.Conv2d(internal_channels,internal_channels,kernel_size=(kernel_size, 1),stride=1,padding=(padding, 0),dilation=dilation,bias=bias), nn.BatchNorm2d(internal_channels), activation(),nn.Conv2d(internal_channels,internal_channels,kernel_size=(1, kernel_size),stride=1,padding=(0, padding),dilation=dilation,bias=bias), nn.BatchNorm2d(internal_channels), activation())else:self.ext_conv2 = nn.Sequential(nn.Conv2d(internal_channels,internal_channels,kernel_size=kernel_size,stride=1,padding=padding,dilation=dilation,bias=bias), nn.BatchNorm2d(internal_channels), activation())# 1x1 expansion convolutionself.ext_conv3 = nn.Sequential(nn.Conv2d(internal_channels,channels,kernel_size=1,stride=1,bias=bias), nn.BatchNorm2d(channels), activation())self.ext_regul = nn.Dropout2d(p=dropout_prob)# PReLU layer to apply after adding the branchesself.out_activation = activation()def forward(self, x):# Main branch shortcutmain = x# Extension branchext = self.ext_conv1(x)ext = self.ext_conv2(ext)ext = self.ext_conv3(ext)ext = self.ext_regul(ext)# Add main and extension branchesout = main + extreturn self.out_activation(out)class DownsamplingBottleneck(nn.Module):def __init__(self,in_channels,out_channels,internal_ratio=4,return_indices=False,dropout_prob=0,bias=False,relu=True):super().__init__()# Store parameters that are needed laterself.return_indices = return_indices# Check in the internal_scale parameter is within the expected range# [1, channels]if internal_ratio <= 1 or internal_ratio > in_channels:raise RuntimeError("Value out of range. Expected value in the ""interval [1, {0}], got internal_scale={1}. ".format(in_channels, internal_ratio))internal_channels = in_channels // internal_ratioif relu:activation = nn.ReLUelse:activation = nn.PReLU# Main branch - max pooling followed by feature map (channels) paddingself.main_max1 = nn.MaxPool2d(2,stride=2,return_indices=return_indices)# Extension branch - 2x2 convolution, followed by a regular, dilated or# asymmetric convolution, followed by another 1x1 convolution. Number# of channels is doubled.# 2x2 projection convolution with stride 2self.ext_conv1 = nn.Sequential(nn.Conv2d(in_channels,internal_channels,kernel_size=2,stride=2,bias=bias), nn.BatchNorm2d(internal_channels), activation())# Convolutionself.ext_conv2 = nn.Sequential(nn.Conv2d(internal_channels,internal_channels,kernel_size=3,stride=1,padding=1,bias=bias), nn.BatchNorm2d(internal_channels), activation())# 1x1 expansion convolutionself.ext_conv3 = nn.Sequential(nn.Conv2d(internal_channels,out_channels,kernel_size=1,stride=1,bias=bias), nn.BatchNorm2d(out_channels), activation())self.ext_regul = nn.Dropout2d(p=dropout_prob)# PReLU layer to apply after concatenating the branchesself.out_activation = activation()def forward(self, x):# Main branch shortcutif self.return_indices:main, max_indices = self.main_max1(x)else:main = self.main_max1(x)# Extension branchext = self.ext_conv1(x)ext = self.ext_conv2(ext)ext = self.ext_conv3(ext)ext = self.ext_regul(ext)# Main branch channel paddingn, ch_ext, h, w = ext.size()ch_main = main.size()[1]padding = torch.zeros(n, ch_ext - ch_main, h, w)# Before concatenating, check if main is on the CPU or GPU and# convert padding accordinglyif main.is_cuda:padding = padding.cuda()# Concatenateif int(ch_ext - ch_main) == 48:main = torch.cat((main, main), 1)main = torch.cat((main, main), 1)else:main = torch.cat((main, main), 1)# main = torch.cat((main, padding), 1)# Add main and extension branchesout = main + extreturn self.out_activation(out), max_indicesclass UpsamplingBottleneck(nn.Module):def __init__(self,in_channels,out_channels,internal_ratio=4,dropout_prob=0,bias=False,relu=True):super().__init__()# Check in the internal_scale parameter is within the expected range# [1, channels]if internal_ratio <= 1 or internal_ratio > in_channels:raise RuntimeError("Value out of range. Expected value in the ""interval [1, {0}], got internal_scale={1}. ".format(in_channels, internal_ratio))internal_channels = in_channels // internal_ratioif relu:activation = nn.ReLUelse:activation = nn.PReLU# Main branch - max pooling followed by feature map (channels) paddingself.main_conv1 = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=bias),nn.BatchNorm2d(out_channels))# Remember that the stride is the same as the kernel_size, just like# the max pooling layersself.main_unpool1 = nn.MaxUnpool2d(kernel_size=2)# Extension branch - 1x1 convolution, followed by a regular, dilated or# asymmetric convolution, followed by another 1x1 convolution. Number# of channels is doubled.# self.up1 = nn.ConvTranspose2d(out_channels, out_channels, 2, 2)# 1x1 projection convolution with stride 1self.ext_conv1 = nn.Sequential(nn.Conv2d(in_channels, internal_channels, kernel_size=1, bias=bias),nn.BatchNorm2d(internal_channels), activation())# Transposed convolutionself.ext_tconv1 = nn.ConvTranspose2d(internal_channels,internal_channels,kernel_size=2,stride=2,bias=bias)self.ext_tconv1_bnorm = nn.BatchNorm2d(internal_channels)self.ext_tconv1_activation = activation()# 1x1 expansion convolutionself.ext_conv2 = nn.Sequential(nn.Conv2d(internal_channels, out_channels, kernel_size=1, bias=bias),nn.BatchNorm2d(out_channels), activation())self.ext_regul = nn.Dropout2d(p=dropout_prob)# PReLU layer to apply after concatenating the branchesself.out_activation = activation()def forward(self, x, max_indices, output_size):# Main branch shortcutmain = self.main_conv1(x)# main = self.up1(main)main = nn.Upsample(size=output_size[2:], mode='bilinear')(main)# main = self.main_unpool1(#     main, max_indices, output_size=output_size)# Extension branchext = self.ext_conv1(x)ext = self.ext_tconv1(ext, output_size=output_size)ext = self.ext_tconv1_bnorm(ext)ext = self.ext_tconv1_activation(ext)ext = self.ext_conv2(ext)ext = self.ext_regul(ext)# Add main and extension branchesout = main + extreturn self.out_activation(out)class ENet(nn.Module):def __init__(self, num_classes, encoder_relu=True, decoder_relu=True):super().__init__()self.initial_block = InitialBlock(1, 16, relu=encoder_relu)# Stage 1 - Encoderself.downsample1_0 = DownsamplingBottleneck(16,64,return_indices=True,dropout_prob=0.01,relu=encoder_relu)self.regular1_1 = RegularBottleneck(64, padding=1, dropout_prob=0.01, relu=encoder_relu)self.regular1_2 = RegularBottleneck(64, padding=1, dropout_prob=0.01, relu=encoder_relu)self.regular1_3 = RegularBottleneck(64, padding=1, dropout_prob=0.01, relu=encoder_relu)self.regular1_4 = RegularBottleneck(64, padding=1, dropout_prob=0.01, relu=encoder_relu)# Stage 2 - Encoderself.downsample2_0 = DownsamplingBottleneck(64,128,return_indices=True,dropout_prob=0.1,relu=encoder_relu)self.regular2_1 = RegularBottleneck(128, padding=1, dropout_prob=0.1, relu=encoder_relu)self.dilated2_2 = RegularBottleneck(128, dilation=2, padding=2, dropout_prob=0.1, relu=encoder_relu)self.asymmetric2_3 = RegularBottleneck(128,kernel_size=5,padding=2,asymmetric=True,dropout_prob=0.1,relu=encoder_relu)self.dilated2_4 = RegularBottleneck(128, dilation=4, padding=4, dropout_prob=0.1, relu=encoder_relu)self.regular2_5 = RegularBottleneck(128, padding=1, dropout_prob=0.1, relu=encoder_relu)self.dilated2_6 = RegularBottleneck(128, dilation=8, padding=8, dropout_prob=0.1, relu=encoder_relu)self.asymmetric2_7 = RegularBottleneck(128,kernel_size=5,asymmetric=True,padding=2,dropout_prob=0.1,relu=encoder_relu)self.dilated2_8 = RegularBottleneck(128, dilation=16, padding=16, dropout_prob=0.1, relu=encoder_relu)# Stage 3 - Encoderself.regular3_0 = RegularBottleneck(128, padding=1, dropout_prob=0.1, relu=encoder_relu)self.dilated3_1 = RegularBottleneck(128, dilation=2, padding=2, dropout_prob=0.1, relu=encoder_relu)self.asymmetric3_2 = RegularBottleneck(128,kernel_size=5,padding=2,asymmetric=True,dropout_prob=0.1,relu=encoder_relu)self.dilated3_3 = RegularBottleneck(128, dilation=4, padding=4, dropout_prob=0.1, relu=encoder_relu)self.regular3_4 = RegularBottleneck(128, padding=1, dropout_prob=0.1, relu=encoder_relu)self.dilated3_5 = RegularBottleneck(128, dilation=8, padding=8, dropout_prob=0.1, relu=encoder_relu)self.asymmetric3_6 = RegularBottleneck(128,kernel_size=5,asymmetric=True,padding=2,dropout_prob=0.1,relu=encoder_relu)self.dilated3_7 = RegularBottleneck(128, dilation=16, padding=16, dropout_prob=0.1, relu=encoder_relu)# Stage 4 - Decoderself.upsample4_0 = UpsamplingBottleneck(128, 64, dropout_prob=0.1, relu=decoder_relu)self.regular4_1 = RegularBottleneck(128, padding=1, dropout_prob=0.1, relu=decoder_relu)self.regular4_2 = RegularBottleneck(128, padding=1, dropout_prob=0.1, relu=decoder_relu)# Stage 5 - Decoderself.upsample5_0 = UpsamplingBottleneck(128, 16, dropout_prob=0.1, relu=decoder_relu)self.regular5_1 = RegularBottleneck(32, padding=1, dropout_prob=0.1, relu=decoder_relu)self.regular5_2 = RegularBottleneck(32, padding=1, dropout_prob=0.1, relu=decoder_relu)self.transposed_conv = nn.ConvTranspose2d(32,num_classes,kernel_size=3,stride=2,padding=1,bias=False)def forward(self, x):# Initial blockinput_size = x.size()x = self.initial_block(x)# Stage 1 - Encoderstage1_input_size = x.size()stage1_input = xx, max_indices1_0 = self.downsample1_0(x)x = self.regular1_1(x)x = self.regular1_2(x)x = self.regular1_3(x)x = self.regular1_4(x)# Stage 2 - Encoderstage2_input_size = x.size()stage2_input = xx, max_indices2_0 = self.downsample2_0(x)x = self.regular2_1(x)x = self.dilated2_2(x)x = self.asymmetric2_3(x)x = self.dilated2_4(x)x = self.regular2_5(x)x = self.dilated2_6(x)x = self.asymmetric2_7(x)x = self.dilated2_8(x)# Stage 3 - Encoderx = self.regular3_0(x)x = self.dilated3_1(x)x = self.asymmetric3_2(x)x = self.dilated3_3(x)x = self.regular3_4(x)x = self.dilated3_5(x)x = self.asymmetric3_6(x)x = self.dilated3_7(x)# Stage 4 - Decoderx = self.upsample4_0(x, max_indices2_0, output_size=stage2_input_size)x = torch.cat([stage2_input, x], dim=1)x = self.regular4_1(x)x = self.regular4_2(x)# Stage 5 - Decoderx = self.upsample5_0(x, max_indices1_0, output_size=stage1_input_size)x = torch.cat([stage1_input, x], dim=1)x = self.regular5_1(x)x = self.regular5_2(x)x = self.transposed_conv(x, output_size=input_size)return nn.Sigmoid()(x)
if __name__ == '__main__':net = ENet(2)inputs = torch.zeros((1, 1, 512, 512), dtype=torch.float32)output = net(inputs)print(output.size())

6.数据集

6.1 数据集准备

数据集是使用PS或者类似的工具做去掉背景处理出来的标签文件,如下图:
在这里插入图片描述
在这里插入图片描述
数据集标注好之后,把清除背景的图像放一个目录,原图放一个目录,图像的名称要相同,层级结构如下图:
在这里插入图片描述
data/test存放测试图像
data/train存放训练原始图像
data/train_cleaned存放去背景后的图像
在这里插入图片描述
在这里插入图片描述

7.环境安装与训练

7.1 环境安装

conda create --name unet python=3.7
conda activate unet
conda install pytorch torchvision cudatoolkit=10.2 -c pytorch
pip install cython matplotlib tqdm opencv-python scipy pillow

7.2 训练

运行以下命令:

python train.py --data_path data --batch_size 64 --epoch 100 --checkpoint_path checkpoints --lr 0.001

参数说明:
•–data_path:训练数据目录,在该目录下需要有train和train_cleaned两个文件夹
•–batch_size:训练batch size
•–epoch:训练epoch数
•–checkpoint_path:模型保存目录
•–lr:学习率

7.3 测试

python predict.py --image_path data/test --weight checkpoints/last.pth --output_path results

参数说明:
•–weight: 训练得到的模型
•–image_path:图像文件夹路径
•–output_path:输出文件夹路径

7.4 转成onnx模型

为了之后推理方便,这里把模型转成onnx模型,onnx模型可以使用onnxruntime或者opencv dnn进行推理,也可以转成ncnn或者别的移动端可部署的模型进行推理,详细步骤可参考我之前的博客:使用深度学习解决拍照文档复杂背景二值化问题
运行以下命令:

python convert_to_onnx.py --input checkpoints/last.pth --output checkpoints/last.onnx

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.mzph.cn/news/806166.shtml

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

【产品】ANET智能通信管理机 物联网网关 电力监控/能耗监测/能源管理系统

产品概述 本系列智能通信管理机是一款采用嵌入式硬件计算机平台&#xff0c;具有多个下行通信接口及一个或者多个上行网络接口&#xff0c;用于将一个目标区域内所有的智能监控/保护装置的通信数据整理汇总后&#xff0c;实时上传主站系统&#xff0c;完成遥信、遥测等能源数据…

【实战解析】YOLOv9全流程训练至优化终极指南

【实战解析】YOLOv9全流程训练至优化终极指南 0.引言1.环境准备2.数据预处理&#xff08;1&#xff09;数据准备&#xff08;2&#xff09;按比例划分数据集&#xff08;3&#xff09;xml转txt脚本&#xff08;4&#xff09;配置文件 3.模型训练&#xff08;1&#xff09;单GPU…

打开游戏缺少dll文件怎么办,dll文件一键修复方法

在我们日常操作电脑&#xff0c;经常会遇到各种各样的问题。比如想玩一会游戏的时候&#xff0c;电脑屏幕上却赫然弹出一则令人颇为扫兴的提示&#xff1a;“打开游戏缺少dll文件”。这个问题可能会让我们感到困惑和沮丧&#xff0c;但是幸运的是&#xff0c;有很多方法可以帮助…

电压继电器SRMUVS-220VAC-2H2D 导轨安装 JOSEF约瑟

系列型号&#xff1a; SRMUVS-58VAC-2H欠电压监视继电器&#xff1b;SRMUVS-100VAC-2H欠电压监视继电器&#xff1b; SRMUVS-110VAC-2H欠电压监视继电器&#xff1b;SRMUVS-220VAC-2H欠电压监视继电器&#xff1b; SRMUVS-58VAC-2H2D欠电压监视继电器&#xff1b;SRMUVS-100…

数据仓库与数据挖掘(第三版)陈文伟思维导图1-5章作业

第一章 概述 8.基于数据仓库的决策支持系统与传统决策支持系统有哪些区别&#xff1f; 决策支持系统经历了4个阶段。 1.基本决策支持系统 是在运筹学单模型辅助决策的基础上发展起来的&#xff0c;以模型库系统为核心&#xff0c;以多模型和数据库的组合形成方案辅助决策。 它…

如何监控容器或K8s中的OpenSearch

概述 当前 OpenSearch 使用的越来越多, 但是 OpenSearch 生态还不尽完善. 针对如下情况: 监控容器化或运行在 K8s 中的 OpenSearch 我查了下, 官方还没有提供完备的方案. 这里如何监控 K8s 中的 OpenSearch, 包括安装 exporter 插件、采集、展示全环节。 OpenSearch 简介…

2017NOIP普及组真题 4. 跳房子

线上OJ&#xff1a; 一本通&#xff1a;http://ybt.ssoier.cn:8088/problem_show.php?pid1417\ 核心思想 首先、本题中提到 “ 至少 要花多少金币改造机器人&#xff0c;能获得 至少 k分 ”。看到这样的话语&#xff0c;基本可以考虑要使用 二分答案。 那么&#xff0c;本题中…

用vue3写一个AI聊天室

效果图如下&#xff1a; 1、页面布局&#xff1a; <template><div class"body" style"background-color: rgb(244, 245, 248); height: 730px"><div class"container"><div class"right"><div class"…

如何用electron(vue)搜索电脑本地wifi

对于搜索本地 WiFi 网络&#xff0c;可以使用 Electron 结合 Node.js 来编写一个简单的应用程序。 以下是一个基本的示例&#xff0c;它使用 Node.js 的 wifi 模块来搜索并列出附近的 WiFi 网络&#xff1a; 首先&#xff0c;确保你已经安装了 Node.js 和 Electron。 然后&am…

数据结构——线性表(链式存储结构)

语言&#xff1a;C语言软件&#xff1a;Visual Studio 2022笔记书籍&#xff1a;数据结构——用C语言描述如有错误&#xff0c;感谢指正。若有侵权请联系博主 一、线性表的逻辑结构 线性表是n个类型相同的数据元素的有限序列&#xff0c;对n>0&#xff0c;除第一元素无直接…

蓝桥杯刷题 二分-[2145]求阶乘(C++)

问题描述 满足 N! 的末尾恰好有 K 个 0 的最小的 N 是多少? 如果这样的 N 不存在输出 −1。 输入格式 一个整数 K。 输出格式 一个整数代表答案。 样例输入 2 样例输出 10 评测用例规模与约定 对于 30% 的数据&#xff0c;1 ≤ K ≤ 10的6次方 对于 100% 的数据&…

结合 tensorflow.js 、opencv.js 与 Ant Design 创建美观且高性能的人脸动捕组件并发布到InsCode

系列文章目录 如何在前端项目中使用opencv.js | opencv.js入门如何使用tensorflow.js实现面部特征点检测tensorflow.js 如何从 public 路径加载人脸特征点检测模型tensorflow.js 如何使用opencv.js通过面部特征点估算脸部姿态并绘制示意图tensorflow.js 使用 opencv.js 将人脸…

uniapp:聊天消息列表(好友列表+私人单聊)支持App、H5、小程序

&#x1f3ac; 江城开朗的豌豆&#xff1a;个人主页 &#x1f525; 个人专栏 :《 VUE 》 《 javaScript 》 &#x1f4dd; 个人网站 :《 江城开朗的豌豆&#x1fadb; 》 ⛺️ 生活的理想&#xff0c;就是为了理想的生活 ! 目录 ⭐ 文章简介&#xff08;效果图展示&#xff…

2024-04-10 Linux gzip 和 gunzip 命令,gzip 压缩的文件通常比原始文件小得多。

一、gzip 是 Linux 系统中用于压缩文件的命令&#xff0c;它通常用于将单个文件压缩成 .gz 格式的文件。gzip 压缩的文件通常比原始文件小得多&#xff0c;因此它在节省磁盘空间和减少文件传输时间方面非常有用。 gzip 命令的基本语法如下&#xff1a; gzip [选项] [文件]复制…

Vue3学习01 Vue3核心语法

Vue3学习 1. Vue3新的特性 2. 创建Vue3工程2.1 基于 vue-cli 创建项目文件说明 2.2 基于 vite 创建具体操作项目文件说明 2.3 简单案例(vite) 3. Vue3核心语法3.1 OptionsAPI 与 CompositionAPIOptions API 弊端Composition API 优势 ⭐3.2 setup小案例setup返回值setup 与 Opt…

ssm038汽车养护管理系统+jsp

汽车养护管理系统设计与实现 摘 要 现代经济快节奏发展以及不断完善升级的信息化技术&#xff0c;让传统数据信息的管理升级为软件存储&#xff0c;归纳&#xff0c;集中处理数据信息的管理方式。本汽车养护管理系统就是在这样的大环境下诞生&#xff0c;其可以帮助管理者在短…

保姆级教程带你实现HarmonyOS手语猜一猜元服务(二)

由于文章篇幅较长&#xff0c;共分为了三篇发布&#xff1a; 保姆级教程带你实现HarmonyOS手语猜一猜元服务&#xff08;一&#xff09; 保姆级教程带你实现HarmonyOS手语猜一猜元服务&#xff08;二&#xff09; 保姆级教程带你实现HarmonyOS手语猜一猜元服务&#xff08;三&…

微信小程序页面交互综合练习 (重点:解决“setData of undefined”报错问题)

一、写一个注册表单&#xff0c;点击“注册”按钮将用户输入的数据带到服务器&#xff0c;并且能在控制台显示参数。 &#xff08;1&#xff09;首先&#xff0c;我需要在vscode里面创建一个简易的node.js服务器 //第一步:引入http模块 var http require(http); //第二步:创建…

自动驾驶定位算法-粒子滤波(Particle Filter)

自动驾驶定位算法-粒子滤波(Particle Filter) 自动驾驶对定位的精度的要求在厘米级的&#xff0c;如何实现厘米级的高精度定位呢&#xff1f;一种众所周知的定位方法是利用全球定位系统(GPS)&#xff0c;利用多颗卫星的测量结果&#xff0c;通过三角测量(Triangulation)机制确…

spring-cloud微服务openfeign

Spring Cloud openfeign对Feign进行了增强&#xff0c;使其支持Spring MVC注解&#xff0c;另外还整合了Ribbon和Nacos&#xff0c;从而使得Feign的使用更加方便 优势&#xff0c;openfeign可以做到使用HTTP请求远程服务时就像洞用本地方法一样的体验&#xff0c;开发者完全感…