当前位置：首页 > article >正文

基于Python实现的结合U - Net与Transformer的神经网络用于视网膜血管分割的示例代码

article 2025/3/18 23:28:21

以下是一个基于Python实现的结合U - Net与Transformer的神经网络用于视网膜血管分割的示例代码。我们将使用PyTorch深度学习框架。

1. 安装依赖库

确保你已经安装了以下库：

pip install torch torchvision numpy

2. 代码实现

import torch
import torch.nn as nn
import torch.nn.functional as F


# Transformer模块
class TransformerBlock(nn.Module):
    def __init__(self, dim, num_heads, mlp_dim, dropout=0.1):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)
        self.attn = nn.MultiheadAttention(dim, num_heads, dropout=dropout)
        self.norm2 = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, mlp_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(mlp_dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        x1 = self.norm1(x)
        attn_output, _ = self.attn(x1, x1, x1)
        x = x + attn_output
        x2 = self.norm2(x)
        x = x + self.mlp(x2)
        return x


# 编码器块
class EncoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        skip = x
        x = self.pool(x)
        return x, skip


# 解码器块
class DecoderBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.up = nn.ConvTranspose2d(in_channels, in_channels // 2, kernel_size=2, stride=2)
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)

    def forward(self, x, skip):
        x = self.up(x)
        diffY = skip.size()[2] - x.size()[2]
        diffX = skip.size()[3] - x.size()[3]
        x = F.pad(x, [diffX // 2, diffX - diffX // 2,
                      diffY // 2, diffY - diffY // 2])
        x = torch.cat([skip, x], dim=1)
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        return x


# 结合U - Net与Transformer的模型
class UNetTransformer(nn.Module):
    def __init__(self, in_channels=3, out_channels=1, features=[64, 128, 256, 512]):
        super().__init__()
        self.encoders = nn.ModuleList(
            [EncoderBlock(in_channels, features[0])] +
            [EncoderBlock(features[i], features[i + 1]) for i in range(len(features) - 1)]
        )
        self.transformer = TransformerBlock(features[-1], num_heads=8, mlp_dim=2048)
        self.decoders = nn.ModuleList(
            [DecoderBlock(features[i], features[i - 1]) for i in range(len(features) - 1, 0, -1)]
        )
        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)

    def forward(self, x):
        skips = []
        for encoder in self.encoders:
            x, skip = encoder(x)
            skips.append(skip)
        b, c, h, w = x.shape
        x = x.flatten(2).transpose(1, 2)
        x = self.transformer(x)
        x = x.transpose(1, 2).view(b, c, h, w)
        skips = skips[::-1]
        for i, decoder in enumerate(self.decoders):
            skip = skips[i]
            x = decoder(x, skip)
        x = self.final_conv(x)
        return x


# 测试代码
if __name__ == "__main__":
    model = UNetTransformer(in_channels=3, out_channels=1)
    x = torch.randn(1, 3, 256, 256)
    output = model(x)
    print(output.shape)

3. 代码解释

TransformerBlock：实现了一个标准的Transformer块，包含多头自注意力机制和前馈神经网络。
EncoderBlock：U - Net的编码器块，包含两个卷积层和一个最大池化层。
DecoderBlock：U - Net的解码器块，包含一个反卷积层和两个卷积层，用于上采样和特征融合。
UNetTransformer：结合了U - Net和Transformer的模型。编码器部分使用U - Net的编码器块，中间使用Transformer块进行特征提取，解码器部分使用U - Net的解码器块。