| import torch |
| import torch.nn as nn |
| from einops import rearrange |
|
|
| class SEBlock(nn.Module): |
| """Squeeze-and-Excitation Block.""" |
| def __init__(self, channels, reduction=16): |
| super(SEBlock, self).__init__() |
| self.fc1 = nn.Linear(channels, channels // reduction, bias=False) |
| self.fc2 = nn.Linear(channels // reduction, channels, bias=False) |
| self.relu = nn.ReLU(inplace=True) |
| self.sigmoid = nn.Sigmoid() |
|
|
| def forward(self, x): |
| b, c, d, h, w = x.shape |
| y = x.mean(dim=[2, 3, 4]) |
| y = self.fc1(y) |
| y = self.relu(y) |
| y = self.fc2(y) |
| y = self.sigmoid(y).view(b, c, 1, 1, 1) |
| return x * y |
|
|
|
|
| class ResNetBlock(nn.Module): |
| """Residual Block with convolutional layers and an SE module.""" |
| def __init__(self, in_channels, out_channels, dropout_rate=0.1): |
| super(ResNetBlock, self).__init__() |
| self.conv1 = nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1) |
| self.in1 = nn.InstanceNorm3d(out_channels, affine=True) |
| self.prelu = nn.PReLU() |
| self.dropout = nn.Dropout3d(dropout_rate) |
|
|
| self.conv2 = nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1) |
| self.in2 = nn.InstanceNorm3d(out_channels, affine=True) |
|
|
| self.se = SEBlock(out_channels) |
| self.residual = ( |
| nn.Conv3d(in_channels, out_channels, kernel_size=1) |
| if in_channels != out_channels else nn.Identity() |
| ) |
|
|
| def forward(self, x): |
| residual = self.residual(x) |
| x = self.conv1(x) |
| x = self.in1(x) |
| x = self.prelu(x) |
| x = self.dropout(x) |
| x = self.conv2(x) |
| x = self.in2(x) |
| x = self.se(x) |
| x += residual |
| return self.prelu(x) |
|
|
|
|
| class Down(nn.Module): |
| """Downsampling block with strided convolution, then a ResNetBlock.""" |
| def __init__(self, in_channels, out_channels, stride=(2,2,2), dropout_rate=0.1): |
| super().__init__() |
| self.stride_conv = nn.Conv3d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1) |
| self.res_block = ResNetBlock(out_channels, out_channels, dropout_rate) |
|
|
| def forward(self, x): |
| x = self.stride_conv(x) |
| x = self.res_block(x) |
| return x |
|
|
|
|
| class Up(nn.Module): |
| """Upsampling block: skip-connection concat, convolutions, then transposed conv.""" |
| def __init__(self, in_channels, out_channels, upsample_channels, |
| stride=(2,2,2), dropout_rate=0.1): |
| super().__init__() |
| self.double_conv = nn.Sequential( |
| nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1), |
| nn.InstanceNorm3d(out_channels, affine=True), |
| nn.PReLU(), |
| nn.Dropout3d(dropout_rate), |
| nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1), |
| nn.InstanceNorm3d(out_channels, affine=True), |
| nn.PReLU(), |
| nn.Dropout3d(dropout_rate), |
| ) |
| self.up = nn.ConvTranspose3d(out_channels, upsample_channels, |
| kernel_size=stride, stride=stride) |
|
|
| def forward(self, decoder_x, skip_x): |
| x = torch.cat([skip_x, decoder_x], dim=1) |
| x = self.double_conv(x) |
| x = self.up(x) |
| return x |
|
|
|
|
| class Bottleneck(nn.Module): |
| """Bottleneck block with strided conv, ResNetBlock, and transposed conv.""" |
| def __init__(self, in_channels, out_channels, stride=(2,2,2), dropout_rate=0.1): |
| super().__init__() |
| self.stride_conv = nn.Conv3d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1) |
| self.res_block = ResNetBlock(out_channels, out_channels, dropout_rate) |
| self.up = nn.ConvTranspose3d(out_channels, out_channels, kernel_size=stride, stride=stride) |
|
|
| def forward(self, x): |
| x = self.stride_conv(x) |
| x = self.res_block(x) |
| x = self.up(x) |
| return x |
|
|
|
|
| class Final(nn.Module): |
| """Final block that merges the top-level up with the inc output, then outputs num_classes.""" |
| def __init__(self, num_classes): |
| super().__init__() |
| self.final_conv = nn.Sequential( |
| nn.Conv3d(64, 32, kernel_size=3, padding=1), |
| nn.InstanceNorm3d(32, affine=True), |
| nn.PReLU(), |
| nn.Conv3d(32, 32, kernel_size=3, padding=1), |
| nn.InstanceNorm3d(32, affine=True), |
| nn.PReLU(), |
| nn.Conv3d(32, num_classes, kernel_size=1), |
| ) |
|
|
| def forward(self, top_x, inc_x): |
| x = torch.cat([inc_x, top_x], dim=1) |
| return self.final_conv(x) |
|
|
|
|
| class NNUNet3DEncoder(nn.Module): |
| """ |
| 3D UNet-like architecture with dynamic depth. |
| If depth=4, it replicates the original structure: |
| inc -> down1..down4 -> bottleneck -> up4..up1 -> final |
| plus 4 auxiliary outputs. |
| """ |
| def __init__( |
| self, |
| dropout_rate=0.1, |
| depth=5, |
| stride_list=[[1, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2], [2, 2, 2]], |
| channels=[32, 64, 128, 256, 320, 512] |
| ): |
| super().__init__() |
| self.depth = depth |
| self.channels = channels |
|
|
| if depth >= len(channels): |
| raise ValueError(f"depth={depth} too large for {channels}") |
|
|
| self.inc = ResNetBlock(1, channels[0], dropout_rate) |
|
|
| self.down_blocks = nn.ModuleList() |
| for i in range(depth): |
| self.down_blocks.append( |
| Down(channels[i], channels[i+1], |
| stride=stride_list[i], dropout_rate=dropout_rate) |
| ) |
|
|
| def forward(self, x, return_encoded_tokens=False): |
| """ |
| Args: |
| x: (batch_size, channels, depth, height, width) |
| return_encoded_tokens: (bool) if True, return the encoded tokens |
| Returns: |
| encoded_tokens: (batch_size, depth, height, width, channels) |
| """ |
| x0 = self.inc(x) |
|
|
| downs = x0 |
| for i in range(self.depth): |
| downs = self.down_blocks[i](downs) |
|
|
| downs = rearrange(downs, 'b c d h w -> b d h w c') |
| if return_encoded_tokens: |
| return downs |
|
|
| return downs |
|
|
|
|
