医学图像语义分割unet++实现代码

以下是基于Pytorch实现的医学图像语义分割Unet++的代码：

import torch
import torch.nn as nn
from torchvision import models


class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True, bn=True):
        super(ConvBlock, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias)
        self.bn = nn.BatchNorm2d(out_channels) if bn else None
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        x = self.relu(x)
        return x


class UpConv(nn.Module):
    def __init__(self, in_channels, out_channels, upscale_factor, mode='transpose', align_corners=True):
        super(UpConv, self).__init__()
        self.upscale_factor = upscale_factor
        self.align_corners = align_corners
        if mode == 'transpose':
            self.conv = nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2*self.upscale_factor,
                                           stride=self.upscale_factor, padding=self.upscale_factor//2,
                                           output_padding=self.upscale_factor%2, bias=True)
        else:
            self.conv = nn.Sequential(
                nn.Upsample(scale_factor=self.upscale_factor, mode=mode, align_corners=self.align_corners),
                nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=True)
            )

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


class NestedUNet(nn.Module):
    def __init__(self, in_channels=1, out_channels=2, init_features=32):
        super(NestedUNet, self).__init__()

        self.down1 = nn.Sequential(
            ConvBlock(in_channels, init_features, bn=False),
            ConvBlock(init_features, init_features*2)
        )
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)

        self.down2 = nn.Sequential(
            ConvBlock(init_features*2, init_features*2*2),
            ConvBlock(init_features*2*2, init_features*2*2*2)
        )
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)

        self.down3 = nn.Sequential(
            ConvBlock(init_features*2*2*2, init_features*2*2*2*2),
            ConvBlock(init_features*2*2*2*2, init_features*2*2*2*2*2)
        )
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)

        self.down4 = nn.Sequential(
            ConvBlock(init_features*2*2*2*2, init_features*2*2*2*2*2),
            ConvBlock(init_features*2*2*2*2*2, init_features*2*2*2*2*2*2)
        )
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)

        self.bottom = nn.Sequential(
            ConvBlock(init_features*2*2*2*2*2, init_features*2*2*2*2*2*2),
            ConvBlock(init_features*2*2*2*2*2*2, init_features*2*2*2*2*2*2),
            UpConv(init_features*2*2*2*2*2*2, init_features*2*2*2*2*2, upscale_factor=2)
        )

        self.up4 = nn.Sequential(
            ConvBlock(init_features*2*2*2*2*2, init_features*2*2*2*2*2),
            ConvBlock(init_features*2*2*2*2*2, init_features*2*2*2*2),
            UpConv(init_features*2*2*2*2, init_features*2*2, upscale_factor=2)
        )
        self.up3 = nn.Sequential(
            ConvBlock(init_features*2*2*2*2, init_features*2*2),
            ConvBlock(init_features*2*2, init_features*2),
            UpConv(init_features*2, init_features, upscale_factor=2)
        )
        self.up2 = nn.Sequential(
            ConvBlock(init_features*2*2, init_features),
            ConvBlock(init_features, init_features),
            UpConv(init_features, init_features//2, upscale_factor=2)
        )
        self.up1 = nn.Sequential(
            ConvBlock(init_features, init_features//2),
            ConvBlock(init_features//2, out_channels)
        )

    def forward(self, x):
        x1 = self.down1(x)
        x2 = self.pool1(x1)
        x2 = self.down2(x2)
        x3 = self.pool2(x2)
        x3 = self.down3(x3)
        x4 = self.pool3(x3)
        x4 = self.down4(x4)
        btm = self.pool4(x4)

        btm = self.bottom(btm)

        x4 = torch.cat([btm, x4], dim=1)
        x4 = self.up4(x4)
        x3 = torch.cat([x4, x3], dim=1)
        x3 = self.up3(x3)
        x2 = torch.cat([x3, x2], dim=1)
        x2 = self.up2(x2)
        x1 = torch.cat([x2, x1], dim=1)
        x1 = self.up1(x1)

        return x1

这段代码实现了一个基于Pytorch的NestedUNet模型，包含四个下采样/池化层和四个上采样/卷积层。其中，第一个下采样/池化层的输出被送入第二个下采样/池化层，以此类推。最后的bottom层会将最后一层下采样/池化层的输出送入上采样/卷积层，以生成最终的分割结果。