From fc50d7918b0686cff59f11a40d70cbd857d2c708 Mon Sep 17 00:00:00 2001 From: Weishan Li Date: Wed, 19 Aug 2026 16:22:34 -0600 Subject: [PATCH 01/11] Fix output channel calculation in PixelShuffle2DUpBlock to account for spatial downsample factor. Erroenous behavior from previous refactor due to lack of caution. With pixel shuffle upsampling, the channel should reduce proportional to the combined 2D sptial expansion as opposed to remaining identical. This also fixes the problem of older (up to v0.7) model loading. --- src/virtual_stain_flow/models/blocks/up_down_blocks.py | 2 +- tests/models/test_up_down_blocks.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/virtual_stain_flow/models/blocks/up_down_blocks.py b/src/virtual_stain_flow/models/blocks/up_down_blocks.py index b4078a9..206fa9c 100644 --- a/src/virtual_stain_flow/models/blocks/up_down_blocks.py +++ b/src/virtual_stain_flow/models/blocks/up_down_blocks.py @@ -303,7 +303,7 @@ def __init__( # out_channel is determined by the number of input channels # as the pixel shuffle operation merely rearranges the channels # to the spatial dimensions - out_channels = in_channels + out_channels = in_channels // (scale_factor ** spatial_dims) super().__init__( in_channels=in_channels, diff --git a/tests/models/test_up_down_blocks.py b/tests/models/test_up_down_blocks.py index e9051a8..f9ed360 100644 --- a/tests/models/test_up_down_blocks.py +++ b/tests/models/test_up_down_blocks.py @@ -23,8 +23,8 @@ 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), (Bilinear2DUpsampleBlock, {}, 3, 3, 2), (Bilinear2DUpsampleBlock, {"out_channels": 8}, 3, 3, 2), ], From 44d1c1f3752d5b7afa6e64dd553d7a89f79caeff Mon Sep 17 00:00:00 2001 From: Weishan Li Date: Wed, 19 Aug 2026 16:42:30 -0600 Subject: [PATCH 02/11] Add preserve_channels option to PixelShuffle2DUpBlock to support both channel perserving and unpreserving behavior for maximized backward compatibility. The default behavior for unext initialization is channel non-preserving which is the more reasonable yet lower capacity version. --- src/virtual_stain_flow/models/blocks/up_down_blocks.py | 5 ++++- src/virtual_stain_flow/models/unext.py | 7 +++++-- tests/models/test_up_down_blocks.py | 2 ++ 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/src/virtual_stain_flow/models/blocks/up_down_blocks.py b/src/virtual_stain_flow/models/blocks/up_down_blocks.py index 206fa9c..9587e7c 100644 --- a/src/virtual_stain_flow/models/blocks/up_down_blocks.py +++ b/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 ): """ @@ -303,7 +304,9 @@ def __init__( # out_channel is determined by the number of input channels # as the pixel shuffle operation merely rearranges the channels # to the spatial dimensions - out_channels = in_channels // (scale_factor ** spatial_dims) + out_channels = in_channels + if not preserve_channels: + out_channels = out_channels // (scale_factor ** spatial_dims) super().__init__( in_channels=in_channels, diff --git a/src/virtual_stain_flow/models/unext.py b/src/virtual_stain_flow/models/unext.py index 572a248..99d5b1a 100644 --- a/src/virtual_stain_flow/models/unext.py +++ b/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,8 +99,10 @@ 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}. " @@ -138,7 +141,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, ) diff --git a/tests/models/test_up_down_blocks.py b/tests/models/test_up_down_blocks.py index f9ed360..a566cce 100644 --- a/tests/models/test_up_down_blocks.py +++ b/tests/models/test_up_down_blocks.py @@ -25,6 +25,8 @@ class TestUpDownBlocks: (ConvTrans2DUpBlock, {"out_channels": 3}, 4, 3, 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), ], From 62e01aa53062984c9ab1918373c022db5a1fcb97 Mon Sep 17 00:00:00 2001 From: Weishan Li Date: Fri, 21 Aug 2026 11:05:25 -0600 Subject: [PATCH 03/11] Add _pixel_shuffle_preserve_channels attribute to ConvNeXtUNet for improved configuration handling and backward compatibility --- src/virtual_stain_flow/models/unext.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/virtual_stain_flow/models/unext.py b/src/virtual_stain_flow/models/unext.py index 99d5b1a..da82d94 100644 --- a/src/virtual_stain_flow/models/unext.py +++ b/src/virtual_stain_flow/models/unext.py @@ -108,6 +108,7 @@ def __init__( 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': @@ -197,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, }, } @@ -208,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"] = False return cls(**init_cfg) From a7c0dd6d1d100d27b711c433721aae8509d4920c Mon Sep 17 00:00:00 2001 From: Weishan Li Date: Fri, 21 Aug 2026 14:58:10 -0600 Subject: [PATCH 04/11] Fix potential bug with out_h and out_w methods in Stage class where out_h and out_w only works for a specific directional of sampling against a very specific block type. --- src/virtual_stain_flow/models/stages.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/virtual_stain_flow/models/stages.py b/src/virtual_stain_flow/models/stages.py index 212118d..4ee41af 100644 --- a/src/virtual_stain_flow/models/stages.py +++ b/src/virtual_stain_flow/models/stages.py @@ -176,15 +176,13 @@ def out_channels(self) -> int: def out_h(self, in_h: int) -> int: _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: _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 """ From 24f534d727a3eaa028c6f35f711de014583946ce Mon Sep 17 00:00:00 2001 From: Weishan Li Date: Fri, 11 Sep 2026 14:56:50 -0600 Subject: [PATCH 05/11] Set default value of _pixel_shuffle_preserve_channels to True for backward compatibility --- src/virtual_stain_flow/models/unext.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/virtual_stain_flow/models/unext.py b/src/virtual_stain_flow/models/unext.py index da82d94..d36f78b 100644 --- a/src/virtual_stain_flow/models/unext.py +++ b/src/virtual_stain_flow/models/unext.py @@ -212,6 +212,6 @@ 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"] = False + init_cfg["_pixel_shuffle_preserve_channels"] = True return cls(**init_cfg) From efd6cd67c9d91f17e5d7b175b0b738d767606fc2 Mon Sep 17 00:00:00 2001 From: Weishan Li Date: Fri, 11 Sep 2026 14:59:24 -0600 Subject: [PATCH 06/11] Refactor out channel calculation in PixelShuffle2DUpBlock to use in_channels instead of overwriting out channels for extra clarity --- src/virtual_stain_flow/models/blocks/up_down_blocks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/virtual_stain_flow/models/blocks/up_down_blocks.py b/src/virtual_stain_flow/models/blocks/up_down_blocks.py index 9587e7c..51172f6 100644 --- a/src/virtual_stain_flow/models/blocks/up_down_blocks.py +++ b/src/virtual_stain_flow/models/blocks/up_down_blocks.py @@ -306,7 +306,7 @@ def __init__( # to the spatial dimensions out_channels = in_channels if not preserve_channels: - out_channels = out_channels // (scale_factor ** spatial_dims) + out_channels = in_channels // (scale_factor ** spatial_dims) super().__init__( in_channels=in_channels, From fa44ca8a137866b1c51b711cfdd9f99a1db3b73d Mon Sep 17 00:00:00 2001 From: Weishan Li Date: Fri, 11 Sep 2026 15:29:30 -0600 Subject: [PATCH 07/11] Add docstrings to out_h and out_w methods in Stage class for clarity --- src/virtual_stain_flow/models/stages.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/virtual_stain_flow/models/stages.py b/src/virtual_stain_flow/models/stages.py index 4ee41af..a769f3e 100644 --- a/src/virtual_stain_flow/models/stages.py +++ b/src/virtual_stain_flow/models/stages.py @@ -174,12 +174,14 @@ 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]: _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]: _out_w = block.out_w(_out_w) From 0b8311a74d1978f6bb2eb53786b9b72e4ec10671 Mon Sep 17 00:00:00 2001 From: Weishan Li Date: Fri, 11 Sep 2026 15:46:01 -0600 Subject: [PATCH 08/11] Add channel count validation to PixelShuffle2DUpBlock and corresponding test to ensure lower than expected input channel count throws an error with explanation. --- src/virtual_stain_flow/models/blocks/up_down_blocks.py | 10 +++++++++- tests/models/test_up_down_blocks.py | 7 +++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/src/virtual_stain_flow/models/blocks/up_down_blocks.py b/src/virtual_stain_flow/models/blocks/up_down_blocks.py index 51172f6..408d76a 100644 --- a/src/virtual_stain_flow/models/blocks/up_down_blocks.py +++ b/src/virtual_stain_flow/models/blocks/up_down_blocks.py @@ -306,7 +306,15 @@ def __init__( # to the spatial dimensions out_channels = in_channels if not preserve_channels: - out_channels = in_channels // (scale_factor ** spatial_dims) + 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, diff --git a/tests/models/test_up_down_blocks.py b/tests/models/test_up_down_blocks.py index a566cce..34db12d 100644 --- a/tests/models/test_up_down_blocks.py +++ b/tests/models/test_up_down_blocks.py @@ -50,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) From f0f3b0a6d751b6feda004ae0fe57d2522a7d6954 Mon Sep 17 00:00:00 2001 From: Weishan Li Date: Fri, 11 Sep 2026 15:56:50 -0600 Subject: [PATCH 09/11] Update CHANGELOG.md to document changes for version 0.4.10, including backward compatibility for ConvNeXtUNet, fixes to PixelShuffle2DUpBlock output channel calculation, and updates to Stage spatial-shape propagation and tests. --- CHANGELOG.md | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1de026b..8263770 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,10 +1,37 @@ # 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`. + +--- + ## [0.4.9] - 2026-09-11 ### Added From 9c15c29faeab2ad2bc5803c1a24d9c793ffc4777 Mon Sep 17 00:00:00 2001 From: Weishan Li Date: Fri, 11 Sep 2026 16:06:15 -0600 Subject: [PATCH 10/11] Add clipping to MaxScaleNormalize apply method to ensure output is within [0, 1] range --- .../transforms/normalizations.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/src/virtual_stain_flow/transforms/normalizations.py b/src/virtual_stain_flow/transforms/normalizations.py index f333d14..23d16e6 100644 --- a/src/virtual_stain_flow/transforms/normalizations.py +++ b/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) From 77149a6aade976fcd8e832e6721c8dbe28b5b076 Mon Sep 17 00:00:00 2001 From: Weishan Li Date: Fri, 11 Sep 2026 16:09:48 -0600 Subject: [PATCH 11/11] Update CHANGELOG.md to include clipping in max scale normalization and bump version to 0.4.10 --- CHANGELOG.md | 3 +++ pyproject.toml | 2 +- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8263770..0101f68 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -30,6 +30,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - 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 diff --git a/pyproject.toml b/pyproject.toml index a9497a2..628269a 100644 --- a/pyproject.toml +++ b/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 = [