Allow convolutional unit padding to be configurable in unext

src/virtual_stain_flow/models/blocks/blocks.py

@@ -31,7 +31,7 @@
         Activation layers, commonly used in UNet architectures.
 """
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Optional, Literal
 
 import timm
 import torch.nn as nn
@@ -44,15 +44,72 @@
     ActivationType
 )
 
+PaddingMode = Literal['zeros', 'reflect']
+
+
+def _validate_padding_mode(padding_mode: PaddingMode) -> None:
+    """
+    Validate the padding mode.
+
+    :param padding_mode: Padding mode to validate.
+    :raises ValueError: If the padding mode is not 'zeros' or 'reflect'.
+    """
+    if padding_mode not in ('zeros', 'reflect'):
+        raise ValueError(
+            f"Unsupported padding_mode: {padding_mode!r}. "
+            "Expected 'zeros' or 'reflect'."
+        )
+
+
+def _is_padded_conv(conv: nn.Conv2d) -> bool:
+    """
+    Check if a convolution layer has padding applied.
+
+    :param conv: Convolution layer to check.
+    :return: True if the convolution has padding, False otherwise.
+    """
+    if isinstance(conv.padding, str):
+        return conv.padding != 'valid'
+    return any(p > 0 for p in conv.padding)
+
+
+def _set_padding_mode_on_padded_convs(module: nn.Module, padding_mode: PaddingMode) -> None:
+    """
+    Set the padding mode on all padded convolution layers within the module.
+
+    :param module: Module containing convolution layers.
+    :param padding_mode: Padding mode to set ('zeros' or 'reflect').
+    """
+    if padding_mode == 'zeros':
+        return
+    for layer in module.modules():
+        if isinstance(layer, nn.Conv2d) and _is_padded_conv(layer):
+            layer.padding_mode = padding_mode
+
+
 class AbstractBlock(ABC, nn.Module):
+    """
+    Abstract base class for all computational blocks.
 
+    This class centralizes input type checking and provides common properties
+    for blocks, such as in_channels, out_channels, and num_units.
+    """
     def __init__(
         self,
         in_channels: int,
         out_channels: Optional[int] = None,
         num_units: int = 1,
         **kwargs: dict
     ):
+        """
+        Initialize the block with input and output channels, number of units, 
+            and additional keyword arguments.
+
+        :param in_channels: Number of input channels.
+        :param out_channels: Number of output channels. Defaults to in_channels if None.
+        :param num_units: Number of units in the block.
+        :param kwargs: Additional keyword arguments. Not used by the base class.
+        """
 
         super().__init__()
 
@@ -86,6 +143,9 @@ def __init__(
     def forward(self, x: Tensor) -> Tensor:
         """
         Forward pass of the block.
+
+        :param x: Input tensor.
+        :return: Output tensor after applying the block.
         """
         raise NotImplementedError
 
@@ -100,15 +160,33 @@ def out_channels(self) -> int:
     def num_units(self) -> int:
         return self._num_units
 
-    # 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.
+    """
+    Compute the output spatial dimensions of the block.
+
+    By default, the block is assumed to preserve spatial dimensions,
+    i.e., the output height and width are the same as the input.
+
+    Should be overridden by blocks that change spatial dimensions.
+    """
     def out_h(self, in_h: int) -> int:
+        """
+        Compute the output height given the input height.
+
+        :param in_h: Input height.
+        :return: Output height after applying the block.
+        """
         return in_h
 
     def out_w(self, in_w: int) -> int:
+        """
+        Compute the output width given the input width.
+
+        :param in_w: Input width.
+        :return: Output width after applying the block.
+        """
         return in_w    
 
+
 class Conv2DConvNeXtBlock(AbstractBlock):    
     """
     A ConvNeXt block that applies a sequence of ConvNeXt units 
@@ -123,7 +201,8 @@ def __init__(
         num_units: int = 1,
         norm_type: NormType = 'layer',
         conv_kernel_size: int = 1,
-        convnext_kernel_size: int = 7
+        convnext_kernel_size: int = 7,
+        padding_mode: PaddingMode = 'zeros',
     ):
         """
         Initializes a ConvNeXt block with the specified parameters.
@@ -142,7 +221,10 @@ def __init__(
         :param convnext_kernel_size: Kernel size for the depth-wise convolution
             layers inside the ConvNeXtBlock. Default is 7, 
             as recommended by Liu et al. (2022), for large receptive fields.
+        :param padding_mode: Padding mode for padded convolutional layers in
+            this block. Supported values are 'zeros' and 'reflect'.
         """
+        _validate_padding_mode(padding_mode)
 
         out_channels = out_channels or in_channels
 
@@ -165,7 +247,8 @@ def __init__(
                     out_channels=out_channels,
                     kernel_size=conv_kernel_size,
                     stride=1, # fixed
-                    padding='same' # ensure spatial dimensions are preserved
+                    padding='same', # ensure spatial dimensions are preserved
+                    padding_mode=padding_mode,
                 )
             )
             # timm.models.convnext.ConvNeXtStage adds a normalization
@@ -191,8 +274,8 @@ def __init__(
         # spatial and channel-wise feature extraction computations, 
         # contrasting with the standard Conv2D whose kernel does both
         # simultaneously.
-        layers = layers + [
-            timm.models.convnext.ConvNeXtBlock(
+        for _ in range(num_units):
+            convnext_block = timm.models.convnext.ConvNeXtBlock(
                 in_chs=out_channels, # same input/output channels
                 out_chs=out_channels, 
                 kernel_size=convnext_kernel_size,
@@ -205,8 +288,9 @@ def __init__(
                 conv_mlp=True, 
                 use_grn=True, # GlobalResponseNorm for mlp layers
                 norm_layer=timm.layers.LayerNorm2d, 
-            ) for _ in range(num_units)
-        ]
+            )
+            _set_padding_mode_on_padded_convs(convnext_block, padding_mode)
+            layers.append(convnext_block)
 
         self.network = nn.Sequential(*layers)
 
@@ -233,7 +317,8 @@ def __init__(
         out_channels: Optional[int] = None,
         num_units: int = 1,
         norm_type: NormType = 'batch',
-        act_type: ActivationType = 'relu'
+        act_type: ActivationType = 'relu',
+        padding_mode: PaddingMode = 'zeros',
     ):
         """
         Initializes a Conv2D block with the specified parameters.
@@ -245,7 +330,10 @@ def __init__(
         :param norm_type: Type of normalization to apply. Default is 'batch'
             (GroupNorm with 1 group).
         :param act_type: Type of activation function to apply. Default is 'relu'.
+        :param padding_mode: Padding mode for Conv2d layers. Supported values
+            are 'zeros' and 'reflect'.
         """
+        _validate_padding_mode(padding_mode)
 
         out_channels = out_channels or in_channels
 
@@ -271,7 +359,8 @@ def __init__(
                     out_channels=_out,
                     kernel_size=3, # fixed for now
                     stride=1, # fixed
-                    padding='same' # indicate spatial preserving unit
+                    padding='same', # indicate spatial preserving unit
+                    padding_mode=padding_mode,
                 )
             )
             layers.append(

src/virtual_stain_flow/models/unext.py

@@ -25,6 +25,7 @@
     PixelShuffle2DUpBlock,
     ActivationType
 )
+from .blocks.blocks import PaddingMode, _set_padding_mode_on_padded_convs
 
 class ConvNeXtUNet(BaseGeneratorModel):
     """
@@ -50,6 +51,9 @@ def __init__(
         pretrained_backbone: bool = False,
         _num_units: Union[List[int], int] = 2,
         _pixel_shuffle_preserve_channels: bool = False,
+        padding_mode: PaddingMode = 'zeros',
+        _encoder_padding_mode: PaddingMode | None = None,
+        _decoder_padding_mode: PaddingMode | None = None,
     ):
         """
         Initializes the ConvNeXtUNet model.
@@ -62,20 +66,49 @@ def __init__(
         :param decoder_compute_block: Type of computation block to use in the
             decoder. Can be 'convnext' for Conv2DConvNeXtBlock or 'conv2d' for
             Conv2DNormActBlock. Default is 'convnext'.
-        :param pretrained_backbone: Whether to use a pretrained ConvNeXtV2_tiny backbone.
-            Default is False.
         :param act_type: Type of activation function to use in the output layer.
             Default is 'sigmoid'.
+        :param pretrained_backbone: Whether to use a pretrained ConvNeXtV2_tiny backbone.
+            Default is False.
         :param _num_units: Number of computation units in each stage.
             Can be an integer for uniform number of units in all stages,
             or a list of integers specifying the number of units for each stage.
+        :param _pixel_shuffle_preserve_channels: Whether to preserve the number of channels
+            when using PixelShuffle2DUpBlock in the decoder. Default is False.
+        :param padding_mode: Shared padding mode for padded encoder convolutions
+            and decoder compute blocks (not upsampling blocks). Supports 'zeros'
+            (default) and 'reflect'. Reflection
+            padding requires feature-map dimensions larger than the padding
+            size (3 for the encoder's 7x7 depthwise convolutions).
+        :param _encoder_padding_mode: Padding mode used by padded convolution
+            layers in the ConvNeXtV2 encoder. Supports 'zeros' and 'reflect'.
+            Defaults to None, inheriting padding_mode; otherwise overrides it.
+        :param _decoder_padding_mode: Padding mode used by padded convolution
+            layers in decoder compute blocks. Supports 'zeros' and 'reflect'.
+            Defaults to None, inheriting padding_mode; otherwise overrides it.
         """
+        if _encoder_padding_mode is None:
+            _encoder_padding_mode = padding_mode
+        if _decoder_padding_mode is None:
+            _decoder_padding_mode = padding_mode
+
+        for mode_name, mode_value in (
+            ('padding_mode', padding_mode),
+            ('_encoder_padding_mode', _encoder_padding_mode),
+            ('_decoder_padding_mode', _decoder_padding_mode),
+        ):
+            if mode_value not in ('zeros', 'reflect'):
+                raise ValueError(
+                    f"Unsupported {mode_name}: {mode_value!r}. "
+                    "Expected 'zeros' or 'reflect'."
+                )
 
         super().__init__(
             in_channels=in_channels,
             out_channels=out_channels,
             out_activation=act_type,
         )
+        self._padding_mode = padding_mode
 
         # timm implementation of the depth 4 ConvNeXtV2 model
         # with only down-sampling path, originally optimized
@@ -101,6 +134,10 @@ def __init__(
 
         self.encoder = convnextv2_model
 
+        # Leave unpadded stem, downsampling, and pointwise convolutions unchanged.
+        _set_padding_mode_on_padded_convs(self.encoder, _encoder_padding_mode)
+        self._encoder_padding_mode = _encoder_padding_mode
+
         if decoder_up_block == 'pixelshuffle':
             in_block_handles = [PixelShuffle2DUpBlock] * (depth - 1)
             in_block_kwargs = [{'preserve_channels': _pixel_shuffle_preserve_channels}] * (depth - 1)
@@ -114,6 +151,7 @@ def __init__(
             )
         self._pixel_shuffle_preserve_channels = _pixel_shuffle_preserve_channels
         self._decoder_up_block = decoder_up_block
+        self._decoder_padding_mode = _decoder_padding_mode
 
         if decoder_compute_block == 'convnext':
             comp_block_handles = [Conv2DConvNeXtBlock] * (depth - 1)
@@ -127,14 +165,25 @@ def __init__(
         self._decoder_compute_block = decoder_compute_block
 
         if isinstance(_num_units, int):
-            comp_block_kwargs = [{'num_units': _num_units}] * (depth - 1)
+            comp_block_kwargs = [
+                {
+                    'num_units': _num_units,
+                    'padding_mode': _decoder_padding_mode,
+                }
+            ] * (depth - 1)
         elif isinstance(_num_units, list):
             if len(_num_units) != (depth - 1):
                 raise ValueError(
                     f"Expected _num_units to be a list of length {(depth - 1)}, "
                     f"got {len(_num_units)}"
                 )
-            comp_block_kwargs = [{'num_units': n} for n in _num_units]
+            comp_block_kwargs = [
+                {
+                    'num_units': n,
+                    'padding_mode': _decoder_padding_mode,
+                }
+                for n in _num_units
+            ]
         else:
             raise TypeError(
                 f"Expected _num_units to be int or list, "
@@ -202,8 +251,11 @@ def to_config(self) -> Dict[str, Any]:
                 "decoder_compute_block": self._decoder_compute_block,
                 "act_type": self._act_type,
                 "pretrained_backbone": self._pretrained_backbone,
+                "padding_mode": self._padding_mode,
                 "_num_units": self._num_units_cfg,
                 "_pixel_shuffle_preserve_channels": self._pixel_shuffle_preserve_channels,
+                "_decoder_padding_mode": self._decoder_padding_mode,
+                "_encoder_padding_mode": self._encoder_padding_mode,
             },
         }
 
@@ -214,12 +266,9 @@ def from_config(cls, config: Dict[str, Any]) -> "ConvNeXtUNet":
         Accepts either the full dict or just the "init" sub-dict.
         """
 
-        init_cfg = config.get("init", config)
-        if "_pixel_shuffle_preserve_channels" not in init_cfg:
-            # For backward compatibility with configs that don't have this key
-            init_cfg["_pixel_shuffle_preserve_channels"] = True
-        if "pretrained_backbone" not in init_cfg:
-            # For backward compatibility with configs that don't have this key
-            init_cfg["pretrained_backbone"] = False
+        init_cfg = dict(config.get("init", config))
+        # Preserve historical pixel-shuffle behavior for configs without this key.
+        init_cfg.setdefault("_pixel_shuffle_preserve_channels", True)
+        # All other missing options use constructor defaults, including padding.
 
         return cls(**init_cfg)

tests/conftest.py

@@ -194,9 +194,9 @@ def file_index(tmp_path):
     for index in range(3):
         for channel, start_value in (("input_ch1", 1), ("input_ch2", 2), ("target_ch1", 1)):
             value = start_value + index * (2 if channel.startswith("input") else 1)
-            Image.fromarray(np.full((10, 10), value, dtype=np.uint16), mode="I;16").save(
-                paths[channel][index]
-            )
+            Image.fromarray(
+                np.full((10, 10), value, dtype=np.uint16),
+            ).save(paths[channel][index])
 
     return pd.DataFrame(paths)
 

tests/datasets/conftest.py

@@ -51,21 +51,21 @@ def file_index(tmp_path):
     for idx, path in enumerate(paths["input_ch1"]):
         value = 1 + idx * 2  # 1, 3, 5
         img_array = np.full((10, 10), value, dtype=np.uint16)
-        img = Image.fromarray(img_array, mode='I;16')
+        img = Image.fromarray(img_array)
         img.save(path)
 
     # input_ch2: 2, 4, 6 for images 0, 1, 2
     for idx, path in enumerate(paths["input_ch2"]):
         value = 2 + idx * 2  # 2, 4, 6
         img_array = np.full((10, 10), value, dtype=np.uint16)
-        img = Image.fromarray(img_array, mode='I;16')
+        img = Image.fromarray(img_array)
         img.save(path)
 
     # target_ch1: 1, 2, 3 for images 0, 1, 2
     for idx, path in enumerate(paths["target_ch1"]):
         value = 1 + idx  # 1, 2, 3
         img_array = np.full((10, 10), value, dtype=np.uint16)
-        img = Image.fromarray(img_array, mode='I;16')
+        img = Image.fromarray(img_array)
         img.save(path)
 
     return pd.DataFrame(paths)