""" Full assembly of the parts to form the complete network """ import torch import torch.nn as nn import torchvision.transforms as transforms from .unet_parts import DoubleConv, Down, Up, OutConv from .swin_transformer import SwinTransformerBlock, Conv class UNet(nn.Module): def __init__(self, n_channels, n_classes, bilinear=False): super(UNet, self).__init__() self.n_channels = n_channels self.n_classes = n_classes self.bilinear = bilinear self.inc = (DoubleConv(n_channels, 64)) self.down1 = (Down(64, 128)) self.down2 = (Down(128, 256)) self.down3 = (Down(256, 512)) factor = 2 if bilinear else 1 self.down4 = (Down(512, 1024 // factor)) self.up1 = (Up(1024, 512 // factor, bilinear)) self.up2 = (Up(512, 256 // factor, bilinear)) self.up3 = (Up(256, 128 // factor, bilinear)) self.up4 = (Up(128, 64, bilinear)) self.outc = (OutConv(64, n_classes)) def forward(self, x): x1 = self.inc(x) x2 = self.down1(x1) x3 = self.down2(x2) x4 = self.down3(x3) x5 = self.down4(x4) x = self.up1(x5, x4) x = self.up2(x, x3) x = self.up3(x, x2) x = self.up4(x, x1) logits = self.outc(x) return logits def use_checkpointing(self): self.inc = torch.utils.checkpoint(self.inc) self.down1 = torch.utils.checkpoint(self.down1) self.down2 = torch.utils.checkpoint(self.down2) self.down3 = torch.utils.checkpoint(self.down3) self.down4 = torch.utils.checkpoint(self.down4) self.up1 = torch.utils.checkpoint(self.up1) self.up2 = torch.utils.checkpoint(self.up2) self.up3 = torch.utils.checkpoint(self.up3) self.up4 = torch.utils.checkpoint(self.up4) self.outc = torch.utils.checkpoint(self.outc) class SEUNet(nn.Module): def __init__(self, n_channels, n_classes, bilinear=False): super(SEUNet, self).__init__() self.n_channels = n_channels self.n_classes = n_classes self.bilinear = bilinear self.inc = (DoubleConv(n_channels, 64)) self.down1 = (Down(64, 128)) self.se1 = SEBlock(128) # 添加SE Block self.down2 = (Down(128, 256)) self.se2 = SEBlock(256) # 添加SE Block self.down3 = (Down(256, 512)) self.se3 = SEBlock(512) # 添加SE Block factor = 2 if bilinear else 1 self.down4 = (Down(512, 1024 // factor)) self.se4 = SEBlock(1024 // factor) # 添加SE Block self.up1 = (Up(1024, 512 // factor, bilinear)) self.se_up1 = SEBlock(512 // factor) # 添加SE Block self.up2 = (Up(512, 256 // factor, bilinear)) self.se_up2 = SEBlock(256 // factor) # 添加SE Block self.up3 = (Up(256, 128 // factor, bilinear)) self.se_up3 = SEBlock(128 // factor) # 添加SE Block self.up4 = (Up(128, 64, bilinear)) self.outc = (OutConv(64, n_classes)) def forward(self, x): x1 = self.inc(x) x2 = self.down1(x1) x2 = self.se1(x2) # 添加SE Block x3 = self.down2(x2) x3 = self.se2(x3) # 添加SE Block x4 = self.down3(x3) x4 = self.se3(x4) # 添加SE Block x5 = self.down4(x4) x5 = self.se4(x5) # 添加SE Block x = self.up1(x5, x4) x = self.se_up1(x) # 添加SE Block x = self.up2(x, x3) x = self.se_up2(x) # 添加SE Block x = self.up3(x, x2) x = self.se_up3(x) # 添加SE Block x = self.up4(x, x1) logits = self.outc(x) return logits def use_checkpointing(self): self.inc = torch.utils.checkpoint(self.inc) self.down1 = torch.utils.checkpoint(self.down1) self.down2 = torch.utils.checkpoint(self.down2) self.down3 = torch.utils.checkpoint(self.down3) self.down4 = torch.utils.checkpoint(self.down4) self.up1 = torch.utils.checkpoint(self.up1) self.up2 = torch.utils.checkpoint(self.up2) self.up3 = torch.utils.checkpoint(self.up3) self.up4 = torch.utils.checkpoint(self.up4) self.outc = torch.utils.checkpoint(self.outc) # SE注意力机制 class SEBlock(nn.Module): def __init__(self, in_channels, reduction=16): super(SEBlock, self).__init__() self.avg_pool = nn.AdaptiveAvgPool2d(1) self.fc = nn.Sequential( nn.Linear(in_channels, in_channels // reduction, bias=False), nn.ReLU(inplace=True), nn.Linear(in_channels // reduction, in_channels, bias=False), nn.Sigmoid() ) def forward(self, x): b, c, _, _ = x.size() # print(f"Input size to SEBlock: {x.size()}") y = self.avg_pool(x).view(b, c) # print(f"Size after avg_pool: {y.size()}") y = self.fc(y).view(b, c, 1, 1) # print(f"Size after fc: {y.size()}") return x * y.expand_as(x) class BasicBlock(nn.Module): expansion = 1 # 通道扩充比例 def __init__(self, in_channels, out_channels, stride=1): super().__init__() self.residual_function = nn.Sequential( nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True), nn.Conv2d(out_channels, out_channels * BasicBlock.expansion, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(out_channels * BasicBlock.expansion) ) self.shortcut = nn.Sequential() if stride != 1 or in_channels != BasicBlock.expansion * out_channels: self.shortcut = nn.Sequential( nn.Conv2d(in_channels, out_channels * BasicBlock.expansion, kernel_size=1, stride=stride, bias=False), nn.BatchNorm2d(out_channels * BasicBlock.expansion) ) def forward(self, x): return nn.ReLU(inplace=True)(self.residual_function(x) + self.shortcut(x)) class BottleNeck(nn.Module): expansion = 4 ''' expansion 是通道扩充的比例 注意实际输出channel = middle_channels * BottleNeck.expansion ''' def __init__(self, in_channels, middle_channels, stride=1): super().__init__() self.residual_function = nn.Sequential( nn.Conv2d(in_channels, middle_channels, kernel_size=1, bias=False), nn.BatchNorm2d(middle_channels), nn.ReLU(inplace=True), nn.Conv2d(middle_channels, middle_channels, stride=stride, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(middle_channels), nn.ReLU(inplace=True), nn.Conv2d(middle_channels, middle_channels * BottleNeck.expansion, kernel_size=1, bias=False), nn.BatchNorm2d(middle_channels * BottleNeck.expansion), ) self.shortcut = nn.Sequential() if stride != 1 or in_channels != middle_channels * BottleNeck.expansion: self.shortcut = nn.Sequential( nn.Conv2d(in_channels, middle_channels * BottleNeck.expansion, stride=stride, kernel_size=1, bias=False), nn.BatchNorm2d(middle_channels * BottleNeck.expansion) ) def forward(self, x): return nn.ReLU(inplace=True)(self.residual_function(x) + self.shortcut(x)) class Bottleneck(nn.Module): """Standard bottleneck.""" def __init__(self, c1, c2, shortcut=True, g=1, k=(3, 3), e=0.5): """Initializes a bottleneck module with given input/output channels, shortcut option, group, kernels, and expansion. """ super().__init__() c_ = int(c2 * e) # hidden channels self.cv1 = Conv(c1, c_, k[0], 1) self.cv2 = Conv(c_, c2, k[1], 1, g=g) self.add = shortcut and c1 == c2 def forward(self, x): """'forward()' applies the YOLO FPN to input data.""" return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) class C2f(nn.Module): """Faster Implementation of CSP Bottleneck with 2 convolutions.""" def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5): """Initialize CSP bottleneck layer with two convolutions with arguments ch_in, ch_out, number, shortcut, groups, expansion. """ super().__init__() self.c = int(c2 * e) # hidden channels self.cv1 = Conv(c1, 2 * self.c, 1, 1) self.cv2 = Conv((2 + n) * self.c, c2, 1) # optional act=FReLU(c2) self.m = nn.ModuleList(Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0) for _ in range(n)) def forward(self, x): """Forward pass through C2f layer.""" y = list(self.cv1(x).chunk(2, 1)) y.extend(m(y[-1]) for m in self.m) return self.cv2(torch.cat(y, 1)) def forward_split(self, x): """Forward pass using split() instead of chunk().""" y = list(self.cv1(x).split((self.c, self.c), 1)) y.extend(m(y[-1]) for m in self.m) return self.cv2(torch.cat(y, 1)) class C2fST(C2f): """C2f module with Swin TransformerBlock().""" def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5): super().__init__(c1, c2, n, shortcut, g, e) num_heads = self.c // 32 self.m = nn.ModuleList(SwinTransformerBlock(self.c, self.c, num_heads, n) for _ in range(n)) class UResnet(nn.Module): def __init__(self, block, layers, n_channels, n_classes, bilinear): super().__init__() self.n_channels = n_channels self.n_classes = n_classes self.bilinear = bilinear nb_filter = [64, 128, 256, 512, 1024] self.in_channel = nb_filter[0] self.pool = nn.MaxPool2d(2, 2) self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) self.conv0_0 = DoubleConv(n_channels, nb_filter[0], nb_filter[0]) self.conv1_0 = self._make_layer(block, nb_filter[1], layers[0], 1) self.trans1 = SwinTransformerBlock(nb_filter[1], nb_filter[1], num_heads = 4, num_layers = 2) self.conv2_0 = self._make_layer(block, nb_filter[2], layers[1], 1) self.trans2 = SwinTransformerBlock(nb_filter[2], nb_filter[2], num_heads = 8, num_layers = 2) self.conv3_0 = self._make_layer(block, nb_filter[3], layers[2], 1) self.trans3 = SwinTransformerBlock(nb_filter[3], nb_filter[3], num_heads = 16, num_layers = 2) self.conv4_0 = self._make_layer(block, nb_filter[4], layers[3], 1) self.trans4 = SwinTransformerBlock(nb_filter[4], nb_filter[4], num_heads = 32, num_layers = 2) self.conv3_1 = DoubleConv((nb_filter[3] + nb_filter[4]) * block.expansion, nb_filter[3] * block.expansion, nb_filter[3]) self.conv2_2 = DoubleConv((nb_filter[2] + nb_filter[3]) * block.expansion, nb_filter[2] * block.expansion, nb_filter[2]) self.conv1_3 = DoubleConv((nb_filter[1] + nb_filter[2]) * block.expansion, nb_filter[1] * block.expansion, nb_filter[1]) self.conv0_4 = DoubleConv(nb_filter[0] + nb_filter[1] * block.expansion, nb_filter[0], nb_filter[0]) self.final = nn.Conv2d(nb_filter[0], n_classes, kernel_size=1) def _make_layer(self, block, middle_channel, num_blocks, stride): """ middle_channels中间维度，实际输出channels = middle_channels * block.expansion num_blocks，一个Layer包含block的个数 """ strides = [stride] + [1] * (num_blocks - 1) layers = [] # for stride in strides: # layers.append(block(self.in_channel, middle_channel, stride)) # self.in_channel = middle_channel * block.expansion for stride in strides: layers.append(block(self.in_channel, middle_channel, stride)) self.in_channel = middle_channel * block.expansion # 更新输入通道数为当前层的输出通道数 return nn.Sequential(*layers) def forward(self, input): x0_0 = self.conv0_0(input) x1_0 = self.conv1_0(self.pool(x0_0)) x1_0 = self.trans1(x1_0) # print("conv1_0:", x1_0.shape) x2_0 = self.conv2_0(self.pool(x1_0)) x2_0 = self.trans2(x2_0) # print("conv2_0:", x2_0.shape) x3_0 = self.conv3_0(self.pool(x2_0)) x3_0 = self.trans3(x3_0) # print("conv3_0:", x3_0.shape) x4_0 = self.conv4_0(self.pool(x3_0)) x4_0 = self.trans4(x4_0) # print("conv4_0:", x4_0.shape) # x4_0 = self.trans1(self.pool(x3_0)) x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1)) x2_2 = self.conv2_2(torch.cat([x2_0, self.up(x3_1)], 1)) x1_3 = self.conv1_3(torch.cat([x1_0, self.up(x2_2)], 1)) x0_4 = self.conv0_4(torch.cat([x0_0, self.up(x1_3)], 1)) output = self.final(x0_4) return output class UResnet34(UResnet): def __init__(self, n_channels, n_classes=2, bilinear=False): super(UResnet34, self).__init__(block=BasicBlock,layers=[3,4,6,3], n_channels=n_channels, n_classes=n_classes, bilinear=bilinear) class UResnet50(UResnet): def __init__(self, n_channels, n_classes=2, bilinear=False): super(UResnet50, self).__init__(block=BottleNeck,layers=[3,4,6,3], n_channels=n_channels, n_classes=n_classes, bilinear=bilinear) class UResnet101(UResnet): def __init__(self, n_channels, n_classes=2, bilinear=False): super(UResnet101, self).__init__(block=BottleNeck,layers=[3,4,23,3], n_channels=n_channels, n_classes=n_classes, bilinear=bilinear) class UResnet152(UResnet): def __init__(self, n_channels, n_classes=2, bilinear=False): super(UResnet152, self).__init__(block=BottleNeck,layers=[3,8,36,3], n_channels=n_channels, n_classes=n_classes, bilinear=bilinear)