diff --git a/CHANGELOG.md b/CHANGELOG.md index 1de026b..0101f68 100644 --- a/CHANGELOG.md +++ b/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 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 = [ 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..408d76a 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 ): """ @@ -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, diff --git a/src/virtual_stain_flow/models/stages.py b/src/virtual_stain_flow/models/stages.py index 212118d..a769f3e 100644 --- a/src/virtual_stain_flow/models/stages.py +++ b/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 """ diff --git a/src/virtual_stain_flow/models/unext.py b/src/virtual_stain_flow/models/unext.py index 572a248..d36f78b 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,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) 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) diff --git a/tests/models/test_up_down_blocks.py b/tests/models/test_up_down_blocks.py index e9051a8..34db12d 100644 --- a/tests/models/test_up_down_blocks.py +++ b/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)