refactor: reversed unet update

Author: anna-grim

src/aind_exaspim_image_compression/machine_learning/unet3d.py

@@ -1,3 +1,4 @@
+
 """
 Created on Fri Aug 14 15:00:00 2025
 
@@ -51,7 +52,7 @@ def __init__(self, width_multiplier=1, trilinear=True):
         super(UNet, self).__init__()
 
         # Initializations
-        _channels = (32, 64, 128, 256)
+        _channels = (32, 64, 128, 256, 512)
         factor = 2 if trilinear else 1
 
         # Instance attributes
@@ -62,12 +63,14 @@ def __init__(self, width_multiplier=1, trilinear=True):
         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] // factor)
+        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[3], self.channels[2] // factor, trilinear)
-        self.up2 = Up(self.channels[2], self.channels[1] // factor, trilinear)
-        self.up3 = Up(self.channels[1], self.channels[0], trilinear)
+        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.outc = OutConv(self.channels[0], 1)
 
     def forward(self, x):
@@ -90,11 +93,13 @@ def forward(self, x):
         x2 = self.down1(x1)
         x3 = self.down2(x2)
         x4 = self.down3(x3)
+        x5 = self.down4(x4)
 
         # Expanding layers
-        x = self.up1(x4, x3)
-        x = self.up2(x, x2)
-        x = self.up3(x, x1)
+        x = self.up1(x5, x4)
+        x = self.up2(x, x3)
+        x = self.up3(x, x2)
+        x = self.up4(x, x1)
         logits = self.outc(x)
         return logits
 
@@ -188,14 +193,7 @@ def __init__(self, in_channels, out_channels):
 
         # Instance attributes
         self.maxpool_conv = nn.Sequential(
-            nn.Conv3d(
-                in_channels,
-                out_channels,
-                kernel_size=3,
-                stride=2,
-                padding=1
-            ),
-            DoubleConv(out_channels, out_channels)
+            nn.MaxPool3d(2), DoubleConv(in_channels, out_channels)
         )
 
     def forward(self, x):
@@ -275,7 +273,7 @@ def forward(self, x1, x2):
 
         Returns
         -------
-        x : torch.Tensor
+        torch.Tensor
             Output tensor after upsampling, concatenation with the skip
             connection, and double convolution. The output shape is
             (B, out_channels, D, H2, W2).
@@ -289,8 +287,7 @@ def forward(self, x1, x2):
             [diffX // 2, diffX - diffX // 2, diffY // 2, diffY - diffY // 2],
         )
         x = torch.cat([x2, x1], dim=1)
-        x = self.conv(x)
-        return x
+        return self.conv(x)
 
 
 class OutConv(nn.Module):