diff --git a/src/aind_exaspim_image_compression/inference.py b/src/aind_exaspim_image_compression/inference.py index 1824603..6416019 100644 --- a/src/aind_exaspim_image_compression/inference.py +++ b/src/aind_exaspim_image_compression/inference.py @@ -274,6 +274,7 @@ def load_model(path, device="cuda"): transform : IntensityTransform The intensity transform the model was trained with. """ + # Read ckpt and create model config ckpt = torch.load(path, map_location=device) if isinstance(ckpt, dict) and "model" in ckpt: state_dict = ckpt["model"] @@ -284,7 +285,14 @@ def load_model(path, device="cuda"): transform_cfg = {"kind": "asinh"} model_cfg = {} - model = UNet(**model_cfg) + # Determine correct model class + model_type = model_cfg.pop("model", "UNet") + if model_type == "N2V2UNet": + model = N2V2UNet(**model_cfg) + else: + model = UNet(**model_cfg) + + # Load ckpt into model model.load_state_dict(state_dict) model.to(device) model.eval() diff --git a/src/aind_exaspim_image_compression/machine_learning/unet3d.py b/src/aind_exaspim_image_compression/machine_learning/unet3d.py index cff18e8..215ac88 100644 --- a/src/aind_exaspim_image_compression/machine_learning/unet3d.py +++ b/src/aind_exaspim_image_compression/machine_learning/unet3d.py @@ -321,14 +321,24 @@ def forward(self, x1, x2): (B, out_channels, D, H2, W2). """ x1 = self.up(x1) - diffY = x2.size()[2] - x1.size()[2] - diffX = x2.size()[3] - x1.size()[3] + + diffD = x2.size()[2] - x1.size()[2] + diffH = x2.size()[3] - x1.size()[3] + diffW = x2.size()[4] - x1.size()[4] x1 = F.pad( x1, - [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2], + [ + diffW // 2, + diffW - diffW // 2, + diffH // 2, + diffH - diffH // 2, + diffD // 2, + diffD - diffD // 2, + ], ) x = torch.cat([x2, x1], dim=1) + return self.conv(x) @@ -377,3 +387,204 @@ def forward(self, x): (B, 1, D, H, W). """ return self.conv(x) + + +class N2V2UNet(UNet): + """ + Noise2Void2 variant of the 3D U-Net. + + This implementation makes two architectural modifications relative to + the standard U-Net: + + 1. Replaces max pooling with anti-aliased MaxBlurPool3D. + 2. Removes the highest-resolution skip connection. + + Residual learning can be enabled to predict a residual correction + that is added to the input image. This is commonly used for denoising + tasks. + """ + + def __init__( + self, + width_multiplier=1, + trilinear=True, + residual=True, + ): + super().__init__( + width_multiplier=width_multiplier, + trilinear=trilinear, + residual=residual, + ) + + factor = 2 if trilinear else 1 + + self.down1 = DownBlur(self.channels[0], self.channels[1]) + self.down2 = DownBlur(self.channels[1], self.channels[2]) + self.down3 = DownBlur(self.channels[2], self.channels[3]) + self.down4 = DownBlur(self.channels[3], self.channels[4] // factor) + + # remove highest-resolution skip + self.up4 = UpNoSkip3D( + self.channels[1] // factor, + self.channels[0], + trilinear, + ) + + def forward(self, x): + input_size = x.shape[2:] + + x1 = self.inc(x) + x2 = self.down1(x1) + x3 = self.down2(x2) + x4 = self.down3(x3) + x5 = self.down4(x4) + + d = self.up1(x5, x4) + d = self.up2(d, x3) + d = self.up3(d, x2) + d = self.up4(d) + + logits = self.outc(d) + + # Restore original spatial size + diffD = input_size[0] - logits.size(2) + diffH = input_size[1] - logits.size(3) + diffW = input_size[2] - logits.size(4) + + logits = F.pad( + logits, + [ + diffW // 2, + diffW - diffW // 2, + diffH // 2, + diffH - diffH // 2, + diffD // 2, + diffD - diffD // 2, + ], + ) + + if self.residual: + return x + logits + + return logits + + @property + def config(self): + config = super().config + config["model"] = "N2V2UNet" + return config + + +class DownBlur(nn.Module): + """Downsampling block using MaxBlurPool3D.""" + + def __init__(self, in_channels, out_channels): + super().__init__() + + self.maxpool_conv = nn.Sequential( + MaxBlurPool3D(in_channels), + DoubleConv(in_channels, out_channels), + ) + + def forward(self, x): + return self.maxpool_conv(x) + + +class MaxBlurPool3D(nn.Module): + """ + Anti-aliased 3D max pooling (BlurPool). + """ + + def __init__(self, channels): + super().__init__() + + # Create separable binomial blur kernel + kernel = torch.tensor([1.0, 2.0, 1.0]) + kernel = ( + kernel[:, None, None] + * kernel[None, :, None] + * kernel[None, None, :] + ) + kernel /= kernel.sum() + + # Expand kernel for depthwise convolution + # Shape: (channels, 1, 3, 3, 3) + kernel = kernel[None, None].repeat(channels, 1, 1, 1, 1) + + self.register_buffer( + "kernel", + kernel, + persistent=True, + ) + + self.channels = channels + self.pool = nn.MaxPool3d(2, stride=1) + + def forward(self, x): + x = self.pool(x) + x = F.pad( + x, + [1, 1, 1, 1, 1, 1], + mode="replicate", + ) + return F.conv3d( + x, + self.kernel, + stride=2, + groups=self.channels, + ) + + +class UpNoSkip3D(nn.Module): + """ + Upsampling block without a skip connection. + """ + + def __init__(self, in_channels, out_channels, trilinear=True): + super().__init__() + + if trilinear: + self.up = nn.Upsample( + scale_factor=2, + mode="trilinear", + align_corners=True, + ) + self.conv = DoubleConv( + in_channels, + out_channels, + mid_channels=in_channels // 2, + ) + else: + self.up = nn.ConvTranspose3d( + in_channels, + in_channels // 2, + kernel_size=2, + stride=2, + ) + self.conv = DoubleConv( + in_channels // 2, + out_channels, + ) + + def forward(self, x): + x = self.up(x) + return self.conv(x) + + +def test_unets(): + for cls in [UNet, N2V2UNet]: + model = cls() + model.eval() + + for size in [32, 33, 64, 65, 128]: + x = torch.randn(1, 1, size, size, size) + + with torch.no_grad(): + y = model(x) + + assert y.shape == x.shape, ( + f"{cls.__name__}: " + f"{x.shape} -> {y.shape}" + ) + + print(f"{cls.__name__}: passed")