| """Standalone PyTorch 3D model definitions and loader for Vesuvius Surface Segmentation.""" |
|
|
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
|
|
| class ResidualBlock3D(nn.Module): |
| def __init__(self, in_channels: int, out_channels: int): |
| super().__init__() |
| self.conv1 = nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1, bias=False) |
| self.norm1 = nn.InstanceNorm3d(out_channels, affine=True) |
| self.act1 = nn.LeakyReLU(negative_slope=0.01, inplace=True) |
| self.conv2 = nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1, bias=False) |
| self.norm2 = nn.InstanceNorm3d(out_channels, affine=True) |
| self.act2 = nn.LeakyReLU(negative_slope=0.01, inplace=True) |
| if in_channels != out_channels: |
| self.shortcut = nn.Sequential( |
| nn.Conv3d(in_channels, out_channels, kernel_size=1, bias=False), |
| nn.InstanceNorm3d(out_channels, affine=True), |
| ) |
| else: |
| self.shortcut = nn.Identity() |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| res = self.shortcut(x) |
| x = self.act1(self.norm1(self.conv1(x))) |
| x = self.norm2(self.conv2(x)) |
| return self.act2(x + res) |
|
|
|
|
| class DecoderBlock3D(nn.Module): |
| def __init__(self, in_channels: int, skip_channels: int, out_channels: int): |
| super().__init__() |
| self.upsample = nn.ConvTranspose3d(in_channels, in_channels // 2, kernel_size=2, stride=2, bias=False) |
| self.norm_up = nn.InstanceNorm3d(in_channels // 2, affine=True) |
| self.act_up = nn.LeakyReLU(negative_slope=0.01, inplace=True) |
| self.conv = ResidualBlock3D((in_channels // 2) + skip_channels, out_channels) |
|
|
| def forward(self, x: torch.Tensor, skip: torch.Tensor) -> torch.Tensor: |
| x = self.act_up(self.norm_up(self.upsample(x))) |
| if x.shape[2:] != skip.shape[2:]: |
| x = F.interpolate(x, size=skip.shape[2:], mode="trilinear", align_corners=False) |
| return self.conv(torch.cat([x, skip], dim=1)) |
|
|
|
|
| class ResEncUNet3D(nn.Module): |
| """5-stage 3D U-Net with Residual Encoder blocks (~15.91M parameters).""" |
|
|
| def __init__(self, in_channels: int = 1, out_channels: int = 1, channels=(32, 64, 128, 256, 320)): |
| super().__init__() |
| c0, c1, c2, c3, c4 = channels |
| self.enc0 = ResidualBlock3D(in_channels, c0) |
| self.down0 = nn.MaxPool3d(2, 2) |
| self.enc1 = ResidualBlock3D(c0, c1) |
| self.down1 = nn.MaxPool3d(2, 2) |
| self.enc2 = ResidualBlock3D(c1, c2) |
| self.down2 = nn.MaxPool3d(2, 2) |
| self.enc3 = ResidualBlock3D(c2, c3) |
| self.down3 = nn.MaxPool3d(2, 2) |
| self.bottleneck = ResidualBlock3D(c3, c4) |
| self.dec3 = DecoderBlock3D(c4, c3, c3) |
| self.dec2 = DecoderBlock3D(c3, c2, c2) |
| self.dec1 = DecoderBlock3D(c2, c1, c1) |
| self.dec0 = DecoderBlock3D(c1, c0, c0) |
| self.out_conv = nn.Conv3d(c0, out_channels, kernel_size=1) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| e0 = self.enc0(x) |
| e1 = self.enc1(self.down0(e0)) |
| e2 = self.enc2(self.down1(e1)) |
| e3 = self.enc3(self.down2(e2)) |
| b = self.bottleneck(self.down3(e3)) |
| u3 = self.dec3(b, e3) |
| u2 = self.dec2(u3, e2) |
| u1 = self.dec1(u2, e1) |
| u0 = self.dec0(u1, e0) |
| return self.out_conv(u0) |
|
|
|
|
| class _SegConvWrap(nn.Module): |
| def __init__(self, in_c: int, out_c: int, k: int = 3, s: int = 1, p: int = 1, bias: bool = False): |
| super().__init__() |
| self.conv = nn.Conv3d(in_c, out_c, kernel_size=k, stride=s, padding=p, bias=bias) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.conv(x) |
|
|
|
|
| class _SegResBlock(nn.Module): |
| def __init__(self, c: int): |
| super().__init__() |
| self.norm1 = nn.GroupNorm(8, c) |
| self.norm2 = nn.GroupNorm(8, c) |
| self.act = nn.ReLU(inplace=True) |
| self.conv1 = _SegConvWrap(c, c, 3, 1, 1, False) |
| self.conv2 = _SegConvWrap(c, c, 3, 1, 1, False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return x + self.conv2(self.act(self.norm2(self.conv1(self.act(self.norm1(x)))))) |
|
|
|
|
| class _SegUpSampleWrap(nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.upsample_non_trainable = nn.Upsample(scale_factor=(2.0, 2.0, 2.0), mode="trilinear", align_corners=False) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| return self.upsample_non_trainable(x) |
|
|
|
|
| class PureSegResNet3D(nn.Module): |
| """Pure PyTorch SegResNet 3D (~18.80M parameters, bit-for-bit MONAI SegResNet state_dict compatible).""" |
|
|
| def __init__(self, in_channels: int = 1, out_channels: int = 1, init_filters: int = 32, blocks_down=(1, 2, 2, 4), blocks_up=(1, 1, 1)): |
| super().__init__() |
| self.act_mod = nn.ReLU(inplace=True) |
| self.convInit = _SegConvWrap(in_channels, init_filters, 3, 1, 1, False) |
| self.down_layers = nn.ModuleList() |
| c = init_filters |
| for i, n_blk in enumerate(blocks_down): |
| layers = [nn.Identity() if i == 0 else _SegConvWrap(c // 2, c, 3, 2, 1, False)] |
| layers.extend(_SegResBlock(c) for _ in range(n_blk)) |
| self.down_layers.append(nn.Sequential(*layers)) |
| if i < len(blocks_down) - 1: |
| c *= 2 |
| self.up_layers = nn.ModuleList() |
| self.up_samples = nn.ModuleList() |
| for _, n_blk in enumerate(blocks_up): |
| half_c = c // 2 |
| self.up_samples.append(nn.Sequential(_SegConvWrap(c, half_c, 1, 1, 0, False), _SegUpSampleWrap())) |
| self.up_layers.append(nn.Sequential(*[_SegResBlock(half_c) for _ in range(n_blk)])) |
| c = half_c |
| self.conv_final = nn.Sequential(nn.GroupNorm(8, c), nn.ReLU(inplace=True), _SegConvWrap(c, out_channels, 1, 1, 0, True)) |
|
|
| def forward(self, x: torch.Tensor) -> torch.Tensor: |
| x = self.convInit(x) |
| skips = [] |
| for down in self.down_layers: |
| x = down(x) |
| skips.append(x) |
| skips.pop() |
| for up_s, up_l in zip(self.up_samples, self.up_layers): |
| x = up_s(x) + skips.pop() |
| x = up_l(x) |
| return self.conv_final(x) |
|
|
|
|
| def load_model(architecture: str, checkpoint_path: str, device: str = "cpu") -> nn.Module: |
| """Instantiate architecture ('PureSegResNet3D' or 'ResEncUNet3D') and load checkpoint weights.""" |
| if architecture == "PureSegResNet3D": |
| model = PureSegResNet3D() |
| elif architecture == "ResEncUNet3D": |
| model = ResEncUNet3D() |
| else: |
| raise ValueError(f"Unsupported architecture: {architecture}") |
|
|
| state = torch.load(checkpoint_path, map_location=device, weights_only=True) |
| if isinstance(state, dict) and ("model_state_dict" in state or "state_dict" in state): |
| state = state.get("model_state_dict", state.get("state_dict")) |
| model.load_state_dict(state) |
| model.to(device).eval() |
| return model |
|
|
