位置: IT常识 - 正文

图像分割之U-Net、U2-Net及其Pytorch代码构建(图像分割 unet)

编辑:rootadmin
图像分割之U-Net、U2-Net及其Pytorch代码构建 图像分割之U-Net、U2 -Net及其Pytorch代码构建1、图像分割

推荐整理分享图像分割之U-Net、U2-Net及其Pytorch代码构建(图像分割 unet),希望有所帮助,仅作参考,欢迎阅读内容。

文章相关热门搜索词:图像分割的经典算法fcn,图像分割的经典算法fcn,图像分割ncut,图像分割csdn,图像分割otsu,图像分割otsu,图像分割otsu,图像分割 unet,内容如对您有帮助,希望把文章链接给更多的朋友!

图像分割就是把图像分成若干个特定的、具有独特性质的区域并提出感兴趣目标的技术和过程。

做法便是对图片中的每一个像素进行分类。

在自动驾驶、自动抠图、医疗影像等领域有着比较广泛的应用。

图像分割大致可分为以下三类:

普通分割:将不同分属不同物体的像素区域分开。比如前景和背景分割开,狗的区域和猫的区域与背景分割开。语义分割:在普通分割的基础上,分类出每一块区域的语义(即这块区域是什么物体)。如把画面中的所有物体都指出他们各自的类别。实例分割:在语义分割的基础上,给每一个物体编号。如这个是该画面中的狗A,那个是画面中的狗B。普通分割语义分割实例分割

可以看出,图像分割是由一张图片到另一张图片。因此,神经网络的输入是图片,输出也是同样的图片,Encoder-Decoder的结构是合适的。U-Net、U2 -Net可作为语义分割使用,可以按照生成图像的方式,生成分割图。也可以按通道划分类,每一个通道就是一个类别,使用sigmoid激活。

2、U-Net

U-Net即使用Encoder-Decoder的结构,首先下采样,然后上采样,中间每一级由残差组成。

则可构建网络的代码如下:

首先是卷积层,可以看出,网络在每一级,均有两层卷积组成。因此构建卷积层如下:

from torch import nnimport torchclass 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 torchimport randomimport cv2from torch.utils.data import Datasetclass 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 nnimport torchfrom torch.utils.data import DataLoaderfrom torchvision.utils import save_imagefrom u_net import UNetfrom dataset import DriveDatasetimport osclass 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()

训练过程见下图。左边为原图,中间为标签,右边为网络预测值

epochimages012图像分割之U-Net、U2-Net及其Pytorch代码构建(图像分割 unet)

完整代码: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 torchimport torch.nn as nnimport torch.nn.functional as Fclass 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 torchimport torch.nn as nnimport torch.nn.functional as Fclass 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 srcclass 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 + x0class 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 + x0class 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 + x0class 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 + x0class 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 + x0class 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)
本文链接地址:https://www.jiuchutong.com/zhishi/287058.html 转载请保留说明!

上一篇:前后端RSA互相加解密、加签验签、密钥对生成(Java)(rsa前端解密)

下一篇:fetch的基本用法、请求参数及响应结果(fetchall的用法)

  • 抖音互关好友有浏览记录吗(抖音互关好友有浏览记录会提醒别人吗)

    抖音互关好友有浏览记录吗(抖音互关好友有浏览记录会提醒别人吗)

  • 手机qq如何显示qq音乐(手机QQ如何显示群头衔)

    手机qq如何显示qq音乐(手机QQ如何显示群头衔)

  • 苹果手机镜像模式是什么意思(苹果手机镜像模式)

    苹果手机镜像模式是什么意思(苹果手机镜像模式)

  • smg9350是什么型号

    smg9350是什么型号

  • sn码和序列号是一个意思吗(苹果序列号怎么查询)

    sn码和序列号是一个意思吗(苹果序列号怎么查询)

  • ios13怎么下载不用密码(苹果13怎么下载不了app)

    ios13怎么下载不用密码(苹果13怎么下载不了app)

  • 电脑开机速度慢(电脑开机速度慢了一秒)

    电脑开机速度慢(电脑开机速度慢了一秒)

  • 怎么让微信一直运行不关闭(怎么让微信一直通话中)

    怎么让微信一直运行不关闭(怎么让微信一直通话中)

  • 抖音切换账号草稿箱会丢吗(抖音切换账号草稿箱不见了怎么恢复)

    抖音切换账号草稿箱会丢吗(抖音切换账号草稿箱不见了怎么恢复)

  • 情侣黄钻一个人能开吗(情侣黄钻一个人过期了另一个人还会显示嘛)

    情侣黄钻一个人能开吗(情侣黄钻一个人过期了另一个人还会显示嘛)

  • 微信注销多久生效(微信注销多久生效 手机号)

    微信注销多久生效(微信注销多久生效 手机号)

  • hwitl00是华为什么型号手机(hwltl00华为什么型号)

    hwitl00是华为什么型号手机(hwltl00华为什么型号)

  • 华为荣耀7怎么关机(华为荣耀7怎么恢复出厂设置)

    华为荣耀7怎么关机(华为荣耀7怎么恢复出厂设置)

  • 支付宝集五福可以集几次(支付宝集五福可以多次合成吗)

    支付宝集五福可以集几次(支付宝集五福可以多次合成吗)

  • 爱奇艺会员可以几个人同时看(爱奇艺会员可以几个人一起登录)

    爱奇艺会员可以几个人同时看(爱奇艺会员可以几个人一起登录)

  • 微信朋友圈怎么看到完整信息(微信朋友圈怎么发长视频)

    微信朋友圈怎么看到完整信息(微信朋友圈怎么发长视频)

  • 华为手机怎么卸载下载软件(华为手机怎么卸载不需要的软件)

    华为手机怎么卸载下载软件(华为手机怎么卸载不需要的软件)

  • 小米双4g是什么意思啊(小米双4g有什么用)

    小米双4g是什么意思啊(小米双4g有什么用)

  • 开关字母o代表什么(开关o表示什么)

    开关字母o代表什么(开关o表示什么)

  • 三星s8手机什么时候上市的(三星s8手机什么处理器)

    三星s8手机什么时候上市的(三星s8手机什么处理器)

  • 图虫怎么删除自己的视频(图虫聊天怎么删除)

    图虫怎么删除自己的视频(图虫聊天怎么删除)

  • 鸟瞰高耸入云的洛根山,加拿大克鲁瓦尼国家公园 (© plainpicture/Design Pics/Robert Postma)(鸟瞰高耸入云的成语)

    鸟瞰高耸入云的洛根山,加拿大克鲁瓦尼国家公园 (© plainpicture/Design Pics/Robert Postma)(鸟瞰高耸入云的成语)

  • 小规模纳税人发票可以抵扣吗
  • 餐具报损表格样板
  • 固定资产税前一次性扣除
  • 总公司和子公司可以有业务往来吗
  • 有限责任公司需要几个股东
  • 代开专票收入未超30万税务怎么处理
  • 土地摊销计入在建工程吗
  • 企业所得税年度申报
  • 怎么算应纳税收多少
  • 红字发票票据怎么处理
  • 固定资产的原价包括
  • 固定资产正常报废如何处理
  • 资本化支出加计摊销形成的差异
  • 非正常损失进项转出额如何计算
  • 股东分配红利交什么税
  • 收到社保局生育津贴
  • 资产负债表如何调平
  • 外购货物用于职工福利企业所得税怎么调整
  • 个人独资企业租车缴纳税金
  • 货物销售价比采购价高
  • 单位汽车按揭贷款怎么贷
  • 到银行开公司账户需要多少手续费
  • 预缴地实现的月销售额未超过
  • 企业停产没有收入,费用可以计入长期待摊费用吗
  • 管理费用与税金的比例
  • dir852迅雷路由器
  • 商标转让过程中公司转让
  • php封装composer包
  • win10电脑时间改了又跳回去
  • gradle视频教程
  • 鸟瞰高耸入云的成语
  • 怎么核算购进商品
  • 埃莫尔斯
  • php+mysql+jquery实现日历签到功能
  • php怎么取数组
  • 会计科目结构什么意思
  • php psr
  • 销售产品是销项税额还是进项税额
  • springboot jni
  • 支付宝服务窗支付是啥
  • php访问sqlserver
  • 发票说明格式
  • 农副产品普通发票可以抵扣增值税吗
  • 含税金额和税额的关系
  • 简易征收应纳税额
  • 企业以现金形式发工资
  • 建筑行业未取得收入如何结转成本
  • 购销合同印花税按70%
  • 逾期贷款利息收入增值税和企业所得税纳税义务时间
  • 增值税普通发票有什么用
  • 员工异地工作
  • 担保公司预计负债
  • 记账凭证的审核与编制可以是一个人吗
  • 检查sql脚本语法
  • java调用jni
  • windows update更新卡住不动了
  • ubuntu常用的50个命令
  • centos如何挂载镜像
  • 虚拟光驱安装
  • ie10怎么设置ie8兼容模式
  • mac怎么使用islide
  • mac自带计算器
  • shnlog.exe - shnlog是什么进程 有什么用
  • kazaalite.exe是什么进程 kazaalite进程有什么用
  • perfom.exe
  • [置顶]公主大人接下来是拷问时间31
  • linux 指定动态库连接位置
  • 日历插件vue
  • python 邻接表
  • JavaScript运算符都有哪些
  • django分层
  • shell脚本显示进度条
  • nodejs开发博客
  • 用jquery实现动态添加
  • python xml 解析
  • unity3d游戏开发笔记本推荐
  • jquery基本知识
  • 税务局上班吗今天
  • 补税务登记证报税要钱吗
  • 天津定额发票查询真伪查询
  • 免责声明:网站部分图片文字素材来源于网络,如有侵权,请及时告知,我们会第一时间删除,谢谢! 邮箱:opceo@qq.com

    鄂ICP备2023003026号

    网站地图: 企业信息 工商信息 财税知识 网络常识 编程技术

    友情链接: 武汉网站建设