net.py

# -*- coding:utf-8 -*-
# @Time : 2022/4/7
# @Author : 
# @Note :
import torch
from torch import nn
import torch.nn.functional as F


class Conv_Block(nn.Module):
    def __init__(self, in_channel, out_channel):
        super(Conv_Block, self).__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(in_channel, out_channel, (3,3), 1, 1, padding_mode='reflect', bias=False),      # padding_mode默认的话，周围填充的是0，但是0是没有特征的，所以用reflect模式；padding_mode='reflect'是一种填充方式，在周围填充对称的值
            nn.BatchNorm2d(out_channel),
            nn.Dropout2d(0.3),
            nn.LeakyReLU(),
            nn.Conv2d(out_channel, out_channel, (3, 3), 1, 1, padding_mode='reflect', bias=False),
            nn.BatchNorm2d(out_channel),
            nn.Dropout2d(0.3),
            nn.LeakyReLU()
        )

    def forward(self, x):
        return self.layer(x)


class DownSample(nn.Module):
    def __init__(self, channel):
        super(DownSample, self).__init__()
        self.layer = nn.Sequential(
            nn.Conv2d(channel, channel, 3, 2, 1, padding_mode='reflect',bias='False'),
            nn.BatchNorm2d(channel),
            nn.LeakyReLU()
        )

    def forward(self, x):
        return self.layer(x)


class UpSample(nn.Module):
    def __init__(self, channel):
        super(UpSample, self).__init__()
        self.layer = nn.Sequential(                 # 卷积来降通道
            nn.Conv2d(channel, channel//2, 1, 1)    # 1*1的卷积不会影响feature map的大小
        )

    def forward(self, x, feature_map):     # 之前的特征图
        # 采用临近插值法
        up = F.interpolate(x, scale_factor=2, mode='nearest')   # 变为原来的2倍
        out = self.layer(up)
        return torch.cat((out, feature_map), dim=1)     # 卷积喂入数据格式是(N, C, H, W),cat在通道维度C上进行，所以是第一维度


class UNet(nn.Module):
    def __init__(self):
        super(UNet, self).__init__()
        self.c1 = Conv_Block(3, 64)
        self.d1 = DownSample(64)
        self.c2 = Conv_Block(64, 128)
        self.d2 = DownSample(128)
        self.c3 = Conv_Block(128, 256)
        self.d3 = DownSample(256)
        self.c4 = Conv_Block(256, 512)
        self.d4 = DownSample(512)
        self.c5 = Conv_Block(512, 1024)

        self.u1 = UpSample(1024)
        self.c6 = Conv_Block(1024, 512)
        self.u2 = UpSample(512)
        self.c7 = Conv_Block(512, 256)
        self.u3 = UpSample(256)
        self.c8 = Conv_Block(256, 128)
        self.u4 = UpSample(128)
        self.c9 = Conv_Block(128, 64)

        self.out = nn.Conv2d(64, 3, (3,3), 1, 1)
        self.Th = nn.Sigmoid()
        # note:使用sigmoid的原因， 对像素点进行二分类就行

    def forward(self, x):
        R1 = self.c1(x)             # 主干特征1
        R2 = self.c2(self.d1(R1))   # 主干特征2
        R3 = self.c3(self.d2(R2))   # 主干特征3
        R4 = self.c4(self.d3(R3))   # 主干特征4
        R5 = self.c5(self.d4(R4))

        O1 = self.c6(self.u1(R5, R4))   # 主干特征4
        O2 = self.c7(self.u2(O1, R3))   # 主干特征3
        O3 = self.c8(self.u3(O2, R2))   # 主干特征2
        O4 = self.c9(self.u4(O3, R1))   # 主干特征1

        return self.Th(self.out(O4))


unet = UNet()


if __name__ == '__main__':
    x = torch.randn(2, 3, 128, 128)     # 喂入随机数据，测试输出形状
    print(unet(x).shape)