Merge pull request #32 from AllenNeuralDynamics/feat-n2v2unet-subclass

feat: n2v2unet subclass

Author: anna-grim

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()

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")