From 20121187795b8b39c8091850770d96e72d2b252f Mon Sep 17 00:00:00 2001 From: anna-grim Date: Thu, 16 Jul 2026 22:49:49 +0000 Subject: [PATCH 1/3] feat: unet architecture from n2v2 --- .../machine_learning/train.py | 1 - .../machine_learning/unet3d.py | 353 +++++++++--------- 2 files changed, 186 insertions(+), 168 deletions(-) diff --git a/src/aind_exaspim_image_compression/machine_learning/train.py b/src/aind_exaspim_image_compression/machine_learning/train.py index d3dce49..ff97bc5 100644 --- a/src/aind_exaspim_image_compression/machine_learning/train.py +++ b/src/aind_exaspim_image_compression/machine_learning/train.py @@ -16,7 +16,6 @@ import numpy as np import os import torch -import torch.nn as nn import torch.optim as optim from skimage import io diff --git a/src/aind_exaspim_image_compression/machine_learning/unet3d.py b/src/aind_exaspim_image_compression/machine_learning/unet3d.py index cff18e8..1112a19 100644 --- a/src/aind_exaspim_image_compression/machine_learning/unet3d.py +++ b/src/aind_exaspim_image_compression/machine_learning/unet3d.py @@ -19,118 +19,135 @@ class UNet(nn.Module): """ - 3D U-Net architecture for 3D image data, suitable for tasks such as - denoising or segmentation. - - Attributes - ---------- - channels : List[int] - Number of channels in each layer after applying "width_multiplier". - trilinear : bool - Flag indicating whether trilinear upsampling is used. - inc : DoubleConv - Initial convolution block. - down1, down2, down3, down4 : Down - Downsampling blocks in the encoder path. - up1, up2, up3, up4 : Up - Upsampling blocks in the decoder path. - outc : OutConv - Final 1x1x1 convolution mapping features to the output channel. + 3D U-Net with optional MaxBlurPool and optional removal of the + highest-resolution skip connection. """ - def __init__(self, width_multiplier=1, trilinear=True, residual=True): - """ - Instantiates a UNet object. - - Parameters - ---------- - width_multiplier : int, optional - Positive integer factor that scales the number of channels in each - layer. Default is 1. - trilinear : bool, optional - If True, use trilinear interpolation for upsampling in decoder - blocks; otherwise, use transposed convolutions. Default is True. - residual : bool, optional - If True, the network predicts a residual added to the input, so it - learns to "remove noise" rather than reconstruct the full signal. - Default is True. - """ - # Call parent class - super(UNet, self).__init__() + def __init__( + self, + width_multiplier=1, + trilinear=True, + residual=True, + maxblurpool=False, + remove_top_skip=False, + ): + super().__init__() if ( - isinstance(width_multiplier, bool) - or not isinstance(width_multiplier, Real) + isinstance(width_multiplier, Real) or width_multiplier < 1 - or not float(width_multiplier).is_integer() + or int(width_multiplier) != width_multiplier ): - raise ValueError("width_multiplier must be a positive integer") + raise ValueError( + "width_multiplier must be a positive integer" + ) - # Initializations - _channels = (32, 64, 128, 256, 512) + base_channels = (32, 64, 128, 256, 512) factor = 2 if trilinear else 1 - # Instance attributes self.width_multiplier = int(width_multiplier) - self.channels = [c * self.width_multiplier for c in _channels] + self.channels = [ + c * self.width_multiplier + for c in base_channels + ] + self.trilinear = trilinear self.residual = residual + self.maxblurpool = maxblurpool + self.remove_top_skip = remove_top_skip - # Contracting layers + # Encoder self.inc = DoubleConv(1, self.channels[0]) - self.down1 = Down(self.channels[0], self.channels[1]) - self.down2 = Down(self.channels[1], self.channels[2]) - self.down3 = Down(self.channels[2], self.channels[3]) - self.down4 = Down(self.channels[3], self.channels[4] // factor) - - # Expanding layers - self.up1 = Up(self.channels[4], self.channels[3] // factor, trilinear) - self.up2 = Up(self.channels[3], self.channels[2] // factor, trilinear) - self.up3 = Up(self.channels[2], self.channels[1] // factor, trilinear) - self.up4 = Up(self.channels[1], self.channels[0], trilinear) + + self.down1 = Down( + self.channels[0], + self.channels[1], + maxblurpool=maxblurpool, + ) + + self.down2 = Down( + self.channels[1], + self.channels[2], + maxblurpool=maxblurpool, + ) + + self.down3 = Down( + self.channels[2], + self.channels[3], + maxblurpool=maxblurpool, + ) + + self.down4 = Down( + self.channels[3], + self.channels[4] // factor, + maxblurpool=maxblurpool, + ) + + # Decoder + self.up1 = Up( + self.channels[4], + self.channels[3] // factor, + trilinear=trilinear, + use_skip=True, + ) + + self.up2 = Up( + self.channels[3], + self.channels[2] // factor, + trilinear=trilinear, + use_skip=True, + ) + + self.up3 = Up( + self.channels[2], + self.channels[1] // factor, + trilinear=trilinear, + use_skip=True, + ) + + # N2V2 modification: remove ONLY the highest-resolution skip + self.up4 = Up( + self.channels[1], + self.channels[0], + trilinear=trilinear, + use_skip=not remove_top_skip, + ) + self.outc = OutConv(self.channels[0], 1) @property def config(self): - """Constructor arguments needed to recreate this model.""" return { "width_multiplier": self.width_multiplier, "trilinear": self.trilinear, "residual": self.residual, + "maxblurpool": self.maxblurpool, + "remove_top_skip": self.remove_top_skip, } def forward(self, x): - """ - Forward pass of the 3D U-Net. - - Parameters - ---------- - x : torch.Tensor - Input tensor with shape (B, 1, D, H, W). - - Returns - ------- - torch.Tensor - Output tensor with shape (B, 1, D, H, W), representing the - denoised image. - """ - # Contracting layers + # Encoder x1 = self.inc(x) x2 = self.down1(x1) x3 = self.down2(x2) x4 = self.down3(x3) x5 = self.down4(x4) - # Expanding layers + # Decoder d = self.up1(x5, x4) d = self.up2(d, x3) d = self.up3(d, x2) - d = self.up4(d, x1) + + if self.remove_top_skip: + d = self.up4(d) + else: + d = self.up4(d, x1) + logits = self.outc(d) - # Residual denoising: predict the correction added to the input if self.residual: return x + logits + return logits @@ -210,128 +227,130 @@ def forward(self, x): class Down(nn.Module): """ - A downsampling module for a 3D U-Net. - - Attributes - ---------- - maxpool_conv : nn.Sequential - Sequential module containing a MaxPool3d layer followed by a - DoubleConv block. + Downscaling followed by DoubleConv. """ - def __init__(self, in_channels, out_channels): - """ - Instantiates a Down object. - - Parameters - ---------- - in_channels : int - Number of input channels to this module. - out_channels : int - Number of output channels produced by this module. - """ - # Call parent class + def __init__( + self, + in_channels, + out_channels, + maxblurpool=False, + ): super().__init__() - # Instance attributes - self.maxpool_conv = nn.Sequential( - nn.MaxPool3d(2), DoubleConv(in_channels, out_channels) + if maxblurpool: + downsample = MaxBlurPool3D(in_channels) + else: + downsample = nn.MaxPool3d(2) + + self.block = nn.Sequential( + downsample, + DoubleConv(in_channels, out_channels), ) def forward(self, x): - """ - Forward pass of the downsampling block. - - Parameters - ---------- - x : torch.Tensor - Input tensor with shape (B, C, D, H, W). - - Returns - ------- - torch.Tensor - Output tensor after max pooling and double convolution. - """ - return self.maxpool_conv(x) + return self.block(x) class Up(nn.Module): """ - An upsampling block for a 3D U-Net that performs spatial upscaling - followed by a double convolution. - - Attributes - ---------- - up : nn.Module - Upsampling layer (either nn.Upsample or nn.ConvTranspose3d). - conv : DoubleConv - Double convolution block applied after concatenating the skip - connection. + Upscaling followed by DoubleConv. """ - def __init__(self, in_channels, out_channels, trilinear=True): - """ - Instantiates an Up object. - - Parameters - ---------- - in_channels : int - Number of input channels to this module. - out_channels : int - Number of output channels produced by this module. - trilinear : bool, optional - Indication of whether to use nn.Upsample or nn.ConvTranspose3d. - Default is True, meaning that nn.Upsample is used. - """ - # Call parent class + def __init__( + self, + in_channels, + out_channels, + trilinear=True, + use_skip=True, + ): super().__init__() - # Instance attributes + self.use_skip = use_skip + 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 + scale_factor=2, + mode="trilinear", + align_corners=True, ) else: self.up = nn.ConvTranspose3d( - in_channels, in_channels // 2, kernel_size=2, stride=2 + in_channels, + in_channels // 2, + kernel_size=2, + stride=2, ) - self.conv = DoubleConv(in_channels, out_channels) - def forward(self, x1, x2): - """ - Forward pass of the upsampling block in a 3D U-Net. + conv_in_channels = ( + in_channels if use_skip else in_channels // 2 + ) - Parameters - ---------- - x1 : torch.Tensor - Input tensor from the previous decoder layer with shape - (B, C1, D, H1, W1). - x2 : torch.Tensor - Skip connection tensor from the encoder path with shape - (B, C2, D, H2, W2). + self.conv = DoubleConv( + conv_in_channels, + out_channels, + ) - Returns - ------- - torch.Tensor - Output tensor after upsampling, concatenation with the skip - connection, and double convolution. The output shape is - (B, out_channels, D, H2, W2). - """ + def forward(self, x1, x2=None): x1 = self.up(x1) - diffY = x2.size()[2] - x1.size()[2] - diffX = x2.size()[3] - x1.size()[3] - x1 = F.pad( - x1, - [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2], - ) - x = torch.cat([x2, x1], dim=1) + if self.use_skip: + diff_z = x2.size(2) - x1.size(2) + diff_y = x2.size(3) - x1.size(3) + diff_x = x2.size(4) - x1.size(4) + + x1 = F.pad( + x1, + [ + diff_x // 2, + diff_x - diff_x // 2, + diff_y // 2, + diff_y - diff_y // 2, + diff_z // 2, + diff_z - diff_z // 2, + ], + ) + + x = torch.cat([x2, x1], dim=1) + + else: + x = x1 + return self.conv(x) +class MaxBlurPool3D(nn.Module): + def __init__(self, channels): + super().__init__() + + self.pool = nn.MaxPool3d(2, stride=1) + + kernel = torch.tensor([1., 2., 1.]) + kernel = ( + kernel[:, None, None] + * kernel[None, :, None] + * kernel[None, None, :] + ) + kernel /= kernel.sum() + + self.register_buffer( + "kernel", + kernel[None, None].repeat(channels, 1, 1, 1, 1), + ) + self.channels = channels + + def forward(self, x): + x = self.pool(x) + x = F.conv3d( + x, + self.kernel, + stride=2, + padding=1, + groups=self.channels, + ) + return x + + class OutConv(nn.Module): """ Final output convolution layer for a 3D U-Net. From 488ef5bd244f3547ab1afd8c1cf736da5290208e Mon Sep 17 00:00:00 2001 From: anna-grim Date: Thu, 16 Jul 2026 23:02:57 +0000 Subject: [PATCH 2/3] reverted changes --- .../machine_learning/unet3d.py | 142 ++++++++++++++++-- 1 file changed, 131 insertions(+), 11 deletions(-) diff --git a/src/aind_exaspim_image_compression/machine_learning/unet3d.py b/src/aind_exaspim_image_compression/machine_learning/unet3d.py index 1112a19..18a028a 100644 --- a/src/aind_exaspim_image_compression/machine_learning/unet3d.py +++ b/src/aind_exaspim_image_compression/machine_learning/unet3d.py @@ -18,9 +18,24 @@ class UNet(nn.Module): - """ - 3D U-Net with optional MaxBlurPool and optional removal of the - highest-resolution skip connection. +""" + 3D U-Net architecture for 3D image data, suitable for tasks such as + denoising or segmentation. + + Attributes + ---------- + channels : List[int] + Number of channels in each layer after applying "width_multiplier". + trilinear : bool + Flag indicating whether trilinear upsampling is used. + inc : DoubleConv + Initial convolution block. + down1, down2, down3, down4 : Down + Downsampling blocks in the encoder path. + up1, up2, up3, up4 : Up + Upsampling blocks in the decoder path. + outc : OutConv + Final 1x1x1 convolution mapping features to the output channel. """ def __init__( @@ -31,20 +46,37 @@ def __init__( maxblurpool=False, remove_top_skip=False, ): + """ + Instantiates a UNet object. + + Parameters + ---------- + width_multiplier : int, optional + Positive integer factor that scales the number of channels in each + layer. Default is 1. + trilinear : bool, optional + If True, use trilinear interpolation for upsampling in decoder + blocks; otherwise, use transposed convolutions. Default is True. + residual : bool, optional + If True, the network predicts a residual added to the input, so it + learns to "remove noise" rather than reconstruct the full signal. + Default is True. + """ + # Call parent class super().__init__() if ( isinstance(width_multiplier, Real) or width_multiplier < 1 - or int(width_multiplier) != width_multiplier + or not float(width_multiplier).is_integer() ): - raise ValueError( - "width_multiplier must be a positive integer" - ) + raise ValueError("width_multiplier must be a positive integer") + # Initializations base_channels = (32, 64, 128, 256, 512) factor = 2 if trilinear else 1 + # Instance attributes self.width_multiplier = int(width_multiplier) self.channels = [ c * self.width_multiplier @@ -117,6 +149,7 @@ def __init__( @property def config(self): + """Constructor arguments needed to recreate this model.""" return { "width_multiplier": self.width_multiplier, "trilinear": self.trilinear, @@ -126,6 +159,20 @@ def config(self): } def forward(self, x): + """ + Forward pass of the 3D U-Net. + + Parameters + ---------- + x : torch.Tensor + Input tensor with shape (B, 1, D, H, W). + + Returns + ------- + torch.Tensor + Output tensor with shape (B, 1, D, H, W), representing the + denoised image. + """ # Encoder x1 = self.inc(x) x2 = self.down1(x1) @@ -145,6 +192,7 @@ def forward(self, x): logits = self.outc(d) + # Residual denoising: predict the correction added to the input if self.residual: return x + logits @@ -227,7 +275,7 @@ def forward(self, x): class Down(nn.Module): """ - Downscaling followed by DoubleConv. + A downsampling module for a 3D U-Net. """ def __init__( @@ -236,25 +284,63 @@ def __init__( out_channels, maxblurpool=False, ): + """ + Instantiates a Down object. + + Parameters + ---------- + in_channels : int + Number of input channels to this module. + out_channels : int + Number of output channels produced by this module. + maxblurpool : bool, optional + True if max-blur pooling should be used to downsample. Default is + False. + """ + # Call parent class super().__init__() + # Initializations if maxblurpool: downsample = MaxBlurPool3D(in_channels) else: downsample = nn.MaxPool3d(2) + # Instance attributes self.block = nn.Sequential( downsample, DoubleConv(in_channels, out_channels), ) def forward(self, x): - return self.block(x) + """ + Forward pass of the downsampling block. + + Parameters + ---------- + x : torch.Tensor + Input tensor with shape (B, C, D, H, W). + + Returns + ------- + torch.Tensor + Output tensor after max pooling and double convolution. + """ + return self.maxpool_conv(x) class Up(nn.Module): - """ - Upscaling followed by DoubleConv. +""" + An upsampling block for a 3D U-Net that performs spatial upscaling + followed by a double convolution. + + Attributes + ---------- + up : nn.Module + Upsampling layer (either nn.Upsample or nn.ConvTranspose3d). + conv : DoubleConv + Double convolution block applied after concatenating the skip + connection. """ def __init__( @@ -264,8 +350,23 @@ def __init__( trilinear=True, use_skip=True, ): + """ + Instantiates an Up object. + + Parameters + ---------- + in_channels : int + Number of input channels to this module. + out_channels : int + Number of output channels produced by this module. + trilinear : bool, optional + Indication of whether to use nn.Upsample or nn.ConvTranspose3d. + Default is True, meaning that nn.Upsample is used. + """ + # Call parent class super().__init__() + # Instance attributes self.use_skip = use_skip if trilinear: @@ -292,6 +393,25 @@ def __init__( ) def forward(self, x1, x2=None): + """ + Forward pass of the upsampling block in a 3D U-Net. + + Parameters + ---------- + x1 : torch.Tensor + Input tensor from the previous decoder layer with shape + (B, C1, D, H1, W1). + x2 : torch.Tensor, optional + Skip connection tensor from the encoder path with shape + (B, C2, D, H2, W2). + + Returns + ------- + torch.Tensor + Output tensor after upsampling, concatenation with the skip + connection, and double convolution. The output shape is + (B, out_channels, D, H2, W2). + """ x1 = self.up(x1) if self.use_skip: From 254ed3678cb6b34b0f3b04d1b866b7fd76deb437 Mon Sep 17 00:00:00 2001 From: anna-grim Date: Thu, 16 Jul 2026 23:09:10 +0000 Subject: [PATCH 3/3] bug: comments, names --- .../machine_learning/unet3d.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/aind_exaspim_image_compression/machine_learning/unet3d.py b/src/aind_exaspim_image_compression/machine_learning/unet3d.py index 18a028a..b438225 100644 --- a/src/aind_exaspim_image_compression/machine_learning/unet3d.py +++ b/src/aind_exaspim_image_compression/machine_learning/unet3d.py @@ -18,7 +18,7 @@ class UNet(nn.Module): -""" + """ 3D U-Net architecture for 3D image data, suitable for tasks such as denoising or segmentation. @@ -137,7 +137,6 @@ def __init__( use_skip=True, ) - # N2V2 modification: remove ONLY the highest-resolution skip self.up4 = Up( self.channels[1], self.channels[0], @@ -307,7 +306,7 @@ def __init__( downsample = nn.MaxPool3d(2) # Instance attributes - self.block = nn.Sequential( + self.maxpool_conv = nn.Sequential( downsample, DoubleConv(in_channels, out_channels), ) @@ -330,7 +329,7 @@ def forward(self, x): class Up(nn.Module): -""" + """ An upsampling block for a 3D U-Net that performs spatial upscaling followed by a double convolution. @@ -440,6 +439,7 @@ def forward(self, x1, x2=None): class MaxBlurPool3D(nn.Module): + def __init__(self, channels): super().__init__()