From fc50d7918b0686cff59f11a40d70cbd857d2c708 Mon Sep 17 00:00:00 2001 From: Weishan Li Date: Wed, 19 Aug 2026 16:22:34 -0600 Subject: [PATCH 1/4] 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 2/4] 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 3/4] 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 4/4] 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 """