图像分割1-UNet实现

目录

一、模型结构

1.卷积模块

2.上采样模块

二、模型实现

---

 

一、模型结构

模型的整体结构如上图，输入之后，通过多个卷积模块后进行maxpool，到底部之后，再通过多个卷积模块后进行上采样，同时对于每一层上采样之前，都把前面maxpool过程中的结果进行裁剪，之后和输入进模块的特征图进行合并。最终得到输出。

1.卷积模块

包含2个卷积层。这里没有把maxpool层直接添加进模块，而是写在了外面，也就是forward里面，这样的话，前五个就都可以使用同一个模块了，在中间添加maxpool层即可。

2.上采样模块

这个模块。实现的过程中，首先把输入进模块的特征图进行上采样，然后再拼接前面裁剪过后的特征图，这里要注意拼接特征图这部分的通道数的变化。

二、模型实现

import torch
import torch.nn as nn


def contracting_block(in_channels, out_channels):
    layer = nn.Sequential(
        nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=0),
        nn.ReLU(),
        nn.BatchNorm2d(out_channels),
        nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=0),
        nn.ReLU(),
        nn.BatchNorm2d(out_channels)
    )
    return layer


class expansive_block(nn.Module):
    def __init__(self, in_channels, mid_channels, out_channels):
        super(expansive_block, self).__init__()
        self.upsample = nn.ConvTranspose2d(in_channels, int(in_channels / 2), 3, stride=2, padding=1, output_padding=1)
        self.conv = nn.Sequential(
            nn.Conv2d(in_channels, mid_channels, kernel_size=3),
            nn.ReLU(),
            nn.BatchNorm2d(mid_channels),
            nn.Conv2d(mid_channels, out_channels, kernel_size=3),
            nn.ReLU(),
            nn.BatchNorm2d(out_channels)
        )

    def forward(self, e, d):
        out = self.upsample(e)
        dif = d.size(2) - out.size(2)
        crop1 = int(dif / 2)
        crop2 = int(d.size(2) - dif / 2)
        d_cropped = d[:, :, crop1:crop2, crop1:crop2]
        out = torch.cat([out, d_cropped], dim=1)
        out = self.conv(out)
        return out


def final_block(in_channels, out_channels):
    layer = nn.Sequential(
        nn.Conv2d(in_channels, out_channels, kernel_size=1),
        nn.ReLU(),
        nn.BatchNorm2d(out_channels)
    )
    return layer


class UNet(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(UNet, self).__init__()
        self.maxpool = nn.MaxPool2d(2, stride=2)
        self.conv1 = contracting_block(in_channels, 64)
        self.conv2 = contracting_block(64,128)
        self.conv3 = contracting_block(128,256)
        self.conv4 = contracting_block(256,512)
        self.bottleneck = contracting_block(512, 1024)

        self.upconv1 = expansive_block(1024,512,512)
        self.upconv2 = expansive_block(512, 256, 256)
        self.upconv3 = expansive_block(256, 128, 128)
        self.upconv4 = expansive_block(128, 64, 64)

        self.final_block = final_block(64,out_channels)

    def forward(self, x):
        out1_ = self.conv1(x)
        out1 = self.maxpool(out1_)
        out2_ = self.conv2(out1)
        out2 = self.maxpool(out2_)
        out3_ = self.conv3(out2)
        out3 = self.maxpool(out3_)
        out4_ = self.conv4(out3)
        out4 = self.maxpool(out4_)

        bottleneck = self.bottleneck(out4)

        out6 = self.upconv1(bottleneck,out4_)
        out7 = self.upconv2(out6, out3_)
        out8 = self.upconv3(out7, out2_)
        out9 = self.upconv4(out8, out1_)
        final_out = self.final_block(out9)
        return  final_out