Refactor AbstractBlock and Stage classes to improve output channel handling and streamline output dimension calculations

Author: wli51

src/virtual_stain_flow/models/blocks.py

@@ -49,7 +49,7 @@ class AbstractBlock(ABC, nn.Module):
     def __init__(
         self,
         in_channels: int,
-        out_channels: int,
+        out_channels: Optional[int] = None,
         num_units: int = 1,
         **kwargs: dict
     ):
@@ -63,6 +63,8 @@ def __init__(
         if in_channels <= 0:
             raise ValueError("Expected in_channels to be positive, "
                              f"got {in_channels}")
+        if out_channels is None:
+            out_channels = in_channels
         if not isinstance(out_channels, int):
             raise TypeError("Expected out_channels to be int, "
                             f"got {type(out_channels).__name__}")
@@ -101,10 +103,9 @@ def num_units(self) -> int:
     # These 2 below should be overriden to reflect the actual spatial dimension
     # changes the block applies. By default they indicate spatial preserving
     # blocks, i.e. the height and width of the input tensor remain unchanged.
-    @property
     def out_h(self, in_h: int) -> int:
         return in_h
-    @property
+    
     def out_w(self, in_w: int) -> int:
         return in_w    
 

src/virtual_stain_flow/models/stages.py

@@ -173,18 +173,16 @@ def skip_channels(self) -> int:
     def out_channels(self) -> int:
         return self._out_channels
 
-    @property
     def out_h(self, in_h: int) -> int:
         _out_h = in_h
-        for block in self.blocks:
+        for block in [self.in_block, self.comp_block]:
             if isinstance(block, Conv2DDownBlock):
                 _out_h = block.out_h(_out_h)
         return _out_h
 
-    @property
     def out_w(self, in_w: int) -> int:
         _out_w = in_w
-        for block in self.blocks:
+        for block in [self.in_block, self.comp_block]:
             if isinstance(block, Conv2DDownBlock):
                 _out_w = block.out_w(_out_w)
         return _out_w