Fix output channel calculation in PixelShuffle2DUpBlock

CHANGELOG.md

@@ -1,10 +1,40 @@
 # Changelog
 
-All notable chagnes to this project will be documented in this file.
+All notable changes to this project will be documented in this file.
 
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),  
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.4.10] - 2026-09-11
+
+### Added
+
+#### Backward-compatibility config flag for ConvNeXtUNet (`virtual_stain_flow/models/unext.py`)
+- Added `_pixel_shuffle_preserve_channels` to ConvNeXtUNet init/config serialization.
+- `from_config` now defaults `_pixel_shuffle_preserve_channels=True` when loading older configs that do not include this key, preserving legacy behavior for previously trained models.
+
+### Fixed
+
+#### Output channel calculation in PixelShuffle2DUpBlock (`virtual_stain_flow/models/blocks/up_down_blocks.py`)
+- Corrected default output-channel inference for pixel-shuffle upsampling.
+- Output channels now reduce with the spatial expansion factor (for 2D: `in_channels // scale_factor^2`) instead of being preserved by default.
+- Added a validation error when channel reduction would produce fewer than 1 output channel.
+- Added `preserve_channels` to keep the previous behavior when needed for compatibility.
+
+#### Stage spatial-shape propagation for upsampling blocks (`virtual_stain_flow/models/stages.py`)
+- Fixed `Stage.out_h` and `Stage.out_w` to apply shape transforms from both stage blocks, not only `Conv2DDownBlock`.
+- This ensures correct output-shape reporting for stages that use upsampling blocks such as pixel shuffle and transposed convolution.
+
+#### Tests for PixelShuffle2DUpBlock behavior (`tests/models/test_up_down_blocks.py`)
+- Updated expected output-channel assertions to match corrected pixel-shuffle defaults.
+- Added tests for compatibility mode (`preserve_channels=True`).
+- Added a regression test asserting that insufficient `in_channels` raises a clear `ValueError`.
+
+#### Added clipping to [0,1] range in max scale normalization (`virtual_stain_flow/transforms/normalizations.py`)
+- Good to have for the sake of ensuring post normalization values fall within expected ranges even if normalization factor is misspecified.
+
+---
+
 ## [0.4.9] - 2026-09-11
 
 ### Added

pyproject.toml

@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "virtual_stain_flow"
-version = "0.4.9"
+version = "0.4.10"
 description = "For developing virtual staining models"
 requires-python = ">=3.9"
 dependencies = [

src/virtual_stain_flow/models/blocks/up_down_blocks.py

@@ -287,6 +287,7 @@ def __init__(
         self,
         in_channels: int,
         out_channels: Optional[int] = None,
+        preserve_channels: bool = False,
         **kwargs
     ):
         """
@@ -304,6 +305,16 @@ def __init__(
         # as the pixel shuffle operation merely rearranges the channels
         # to the spatial dimensions
         out_channels = in_channels
+        if not preserve_channels:
+            channel_factor = scale_factor ** spatial_dims
+            out_channels = in_channels // channel_factor
+            if out_channels < 1:
+                raise ValueError(
+                    "PixelShuffle2DUpBlock requires at least "
+                    f"{channel_factor} input channels for {scale_factor}x"
+                    f"{scale_factor} upsampling. Received in_channels={in_channels}, "
+                    "which would reduce the channel count below 1."
+                )
 
         super().__init__(
             in_channels=in_channels,

src/virtual_stain_flow/models/stages.py

@@ -174,17 +174,17 @@ def out_channels(self) -> int:
         return self._out_channels
 
     def out_h(self, in_h: int) -> int:
+        """Computes the output height after passing through the stage."""
         _out_h = in_h
         for block in [self.in_block, self.comp_block]:
-            if isinstance(block, Conv2DDownBlock):
-                _out_h = block.out_h(_out_h)
+            _out_h = block.out_h(_out_h)
         return _out_h
-    
+
     def out_w(self, in_w: int) -> int:
+        """Computes the output width after passing through the stage."""
         _out_w = in_w
         for block in [self.in_block, self.comp_block]:
-            if isinstance(block, Conv2DDownBlock):
-                _out_w = block.out_w(_out_w)
+            _out_w = block.out_w(_out_w)
         return _out_w
 
 """

src/virtual_stain_flow/models/unext.py

@@ -47,7 +47,8 @@ def __init__(
         decoder_up_block: Literal['pixelshuffle', 'convt'] = 'pixelshuffle',
         decoder_compute_block: Literal['convnext', 'conv2d'] = 'convnext',
         act_type: ActivationType = 'sigmoid',
-        _num_units: Union[List[int], int] = 2
+        _num_units: Union[List[int], int] = 2,
+        _pixel_shuffle_preserve_channels: bool = False,
     ):
         """
         Initializes the ConvNeXtUNet model.
@@ -98,13 +99,16 @@ def __init__(
 
         if decoder_up_block == 'pixelshuffle':
             in_block_handles = [PixelShuffle2DUpBlock] * (depth - 1)
+            in_block_kwargs = [{'preserve_channels': _pixel_shuffle_preserve_channels}] * (depth - 1)
         elif decoder_up_block == 'convt':
             in_block_handles = [ConvTrans2DUpBlock] * (depth - 1)
+            in_block_kwargs = [{'norm_type': 'layer'}] * (depth - 1)
         else:
             raise ValueError(
                 f"Unsupported decoder_up_block: {decoder_up_block!r}. "
                 "Expected 'pixelshuffle' or 'convt'."
             )
+        self._pixel_shuffle_preserve_channels = _pixel_shuffle_preserve_channels
         self._decoder_up_block = decoder_up_block
 
         if decoder_compute_block == 'convnext':
@@ -138,7 +142,7 @@ def __init__(
             encoder_feature_map_channels=convnextv2_model.feature_info.channels(),
             # use convolutional up-sampling blocks
             in_block_handles=in_block_handles,
-            in_block_kwargs=[{'norm_type': 'layer'}] * (depth - 1),
+            in_block_kwargs=in_block_kwargs,
             comp_block_handles=comp_block_handles,
             comp_block_kwargs=comp_block_kwargs,
         )
@@ -194,6 +198,7 @@ def to_config(self) -> Dict[str, Any]:
                 "decoder_compute_block": self._decoder_compute_block,
                 "act_type": self._act_type,
                 "_num_units": self._num_units_cfg,
+                "_pixel_shuffle_preserve_channels": self._pixel_shuffle_preserve_channels,
             },
         }
 
@@ -205,5 +210,8 @@ def from_config(cls, config: Dict[str, Any]) -> "ConvNeXtUNet":
         """
 
         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
 
         return cls(**init_cfg)

src/virtual_stain_flow/transforms/normalizations.py

@@ -90,10 +90,16 @@ def apply(
         img: np.ndarray, 
         **params
     ) -> np.ndarray:
+        """
+        Apply the normalization to the input image.
 
+        :param img: Input image as a NumPy array.
+        :param params: Additional parameters (not used here).
+        :return: Normalized image as a NumPy array.
+        """
         if isinstance(img, np.ndarray):
             # Normalize the image using the normalization factor
-            return img / self._normalization_factor
+            return np.clip(img / self._normalization_factor, 0.0, 1.0)
         else:
             raise TypeError(
                 "Expected input image to be a NumPy array, "
@@ -183,7 +189,13 @@ def __repr__(self) -> str:
         )
 
     def apply(self, img, **params):        
-
+        """
+        Apply Z-Score normalization to the input image.
+
+        :param img: Input image as a NumPy array.
+        :param params: Additional parameters (not used here).
+        :return: Normalized image as a NumPy array.
+        """
         if isinstance(img, np.ndarray):
 
             mean = self._mean or img.mean(axis=(1, 2), keepdims=True)

tests/models/test_up_down_blocks.py

@@ -23,8 +23,10 @@ class TestUpDownBlocks:
 			(MaxPool2DDownBlock, {"out_channels": 8}, 3, 3, 0.5),
 			(ConvTrans2DUpBlock, {}, 4, 2, 2),
 			(ConvTrans2DUpBlock, {"out_channels": 3}, 4, 3, 2),
-			(PixelShuffle2DUpBlock, {}, 4, 4, 2),
-			(PixelShuffle2DUpBlock, {"out_channels": 8}, 4, 4, 2),
+			(PixelShuffle2DUpBlock, {}, 4, 1, 2),
+			(PixelShuffle2DUpBlock, {"out_channels": 8}, 4, 1, 2),
+			(PixelShuffle2DUpBlock, {"preserve_channels": True}, 4, 4, 2),
+			(PixelShuffle2DUpBlock, {"preserve_channels": True, "out_channels": 8}, 4, 4, 2),
 			(Bilinear2DUpsampleBlock, {}, 3, 3, 2),
 			(Bilinear2DUpsampleBlock, {"out_channels": 8}, 3, 3, 2),
 		],
@@ -48,3 +50,10 @@ def test_output_channels_and_spatial_dimensions(
 		)
 		assert block.out_h(input_tensor.shape[2]) == output.shape[2]
 		assert block.out_w(input_tensor.shape[3]) == output.shape[3]
+
+	def test_pixel_shuffle_requires_enough_channels(self):
+		with pytest.raises(
+			ValueError,
+			match=r"requires at least 4 input channels[\s\S]*in_channels=1",
+		):
+			PixelShuffle2DUpBlock(in_channels=1)