careamics 0.0.13__py3-none-any.whl → 0.0.15__py3-none-any.whl

careamics/models/unet.py

@@ -205,16 +205,23 @@ class UnetDecoder(nn.Module):
         decoder_blocks: list[nn.Module] = []
         for n in range(depth):
             decoder_blocks.append(upsampling)
-            in_channels = (num_channels_init * 2 ** (depth - n)) * groups
-            out_channels = in_channels // 2
+
+            in_channels = (num_channels_init * 2 ** (depth - n - 1)) * groups
+            # final decoder block has the same number in and out features
+            out_channels = in_channels // 2 if n != depth - 1 else in_channels
+            if not (n2v2 and (n == depth - 1)):
+                in_channels = in_channels * 2  # accounting for skip connection concat
+
             decoder_blocks.append(
                 Conv_Block(
                     conv_dim,
-                    in_channels=(
-                        in_channels + in_channels // 2 if n > 0 else in_channels
-                    ),
+                    in_channels=in_channels,
                     out_channels=out_channels,
-                    intermediate_channel_multiplier=2,
+                    # TODO: Tensorflow n2v implementation has intermediate channel
+                    #   multiplication for skip_skipone=True but not skip_skipone=False
+                    #   this needs to be benchmarked.
+                    # final decoder block doesn't multiply the intermediate features
+                    intermediate_channel_multiplier=2 if n != depth - 1 else 1,
                     dropout_perc=dropout,
                     activation="ReLU",
                     use_batch_norm=use_batch_norm,
@@ -241,6 +248,7 @@ class UnetDecoder(nn.Module):
         """
         x: torch.Tensor = features[0]
         skip_connections: tuple[torch.Tensor, ...] = features[-1:0:-1]
+        depth = len(skip_connections)
 
         x = self.bottleneck(x)
 
@@ -249,10 +257,8 @@ class UnetDecoder(nn.Module):
             if isinstance(module, nn.Upsample):
                 # divide index by 2 because of upsampling layers
                 skip_connection: torch.Tensor = skip_connections[i // 2]
-                if self.n2v2:
-                    if x.shape != skip_connections[-1].shape:
-                        x = self._interleave(x, skip_connection, self.groups)
-                else:
+                # top level skip connection not added for n2v2
+                if (not self.n2v2) or (self.n2v2 and (i // 2 < depth - 1)):
                     x = self._interleave(x, skip_connection, self.groups)
         return x
 

careamics/prediction_utils/lvae_prediction.py

@@ -1,6 +1,6 @@
 """Module containing pytorch implementations for obtaining predictions from an LVAE."""
 
-from typing import Any, Optional
+from typing import Any
 
 import torch
 
@@ -18,7 +18,7 @@ def lvae_predict_single_sample(
     model: LVAE,
     likelihood_obj: LikelihoodModule,
     input: torch.Tensor,
-) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+) -> tuple[torch.Tensor, torch.Tensor | None]:
     """
     Generate a single sample prediction from an LVAE model, for a given input.
 
@@ -57,7 +57,7 @@ def lvae_predict_tiled_batch(
     model: LVAE,
     likelihood_obj: LikelihoodModule,
     input: tuple[Any],
-) -> tuple[tuple[Any], Optional[tuple[Any]]]:
+) -> tuple[tuple[Any], tuple[Any] | None]:
     # TODO: fix docstring return types, ... too many output options
     """
     Generate a single sample prediction from an LVAE model, for a given input.
@@ -98,7 +98,7 @@ def lvae_predict_mmse_tiled_batch(
     likelihood_obj: LikelihoodModule,
     input: tuple[Any],
     mmse_count: int,
-) -> tuple[tuple[Any], tuple[Any], Optional[tuple[Any]]]:
+) -> tuple[tuple[Any], tuple[Any], tuple[Any] | None]:
     # TODO: fix docstring return types, ... hard to make readable
     """
     Generate the MMSE (minimum mean squared error) prediction, for a given input.
@@ -137,7 +137,7 @@ def lvae_predict_mmse_tiled_batch(
 
     input_shape = x.shape
     output_shape = (input_shape[0], model.target_ch, *input_shape[2:])
-    log_var: Optional[torch.Tensor] = None
+    log_var: torch.Tensor | None = None
     # pre-declare empty array to fill with individual sample predictions
     sample_predictions = torch.zeros(size=(mmse_count, *output_shape))
     for mmse_idx in range(mmse_count):

careamics/transforms/compose.py

@@ -1,6 +1,6 @@
 """A class chaining transforms together."""
 
-from typing import Optional, Union, cast
+from typing import Union, cast
 
 from numpy.typing import NDArray
 
@@ -64,8 +64,8 @@ class Compose:
         ]
 
     def _chain_transforms(
-        self, patch: NDArray, target: Optional[NDArray]
-    ) -> tuple[Optional[NDArray], ...]:
+        self, patch: NDArray, target: NDArray | None
+    ) -> tuple[NDArray | None, ...]:
         """Chain transforms on the input data.
 
         Parameters
@@ -80,7 +80,7 @@ class Compose:
         tuple[np.ndarray, Optional[np.ndarray]]
             The output of the transformations.
         """
-        params: Union[tuple[NDArray, Optional[NDArray]],] = (patch, target)
+        params: Union[tuple[NDArray, NDArray | None],] = (patch, target)
 
         for t in self.transforms:
             *params, _ = t(*params)  # ignore additional_arrays dict
@@ -92,9 +92,9 @@ class Compose:
     def _chain_transforms_additional_arrays(
         self,
         patch: NDArray,
-        target: Optional[NDArray],
+        target: NDArray | None,
         **additional_arrays: NDArray,
-    ) -> tuple[NDArray, Optional[NDArray], dict[str, NDArray]]:
+    ) -> tuple[NDArray, NDArray | None, dict[str, NDArray]]:
         """Chain transforms on the input data, with additional arrays.
 
         Parameters
@@ -121,7 +121,7 @@ class Compose:
         return patch, target, additional_arrays
 
     def __call__(
-        self, patch: NDArray, target: Optional[NDArray] = None
+        self, patch: NDArray, target: NDArray | None = None
     ) -> tuple[NDArray, ...]:
         """Apply the transforms to the input data.
 
@@ -143,9 +143,9 @@ class Compose:
     def transform_with_additional_arrays(
         self,
         patch: NDArray,
-        target: Optional[NDArray] = None,
+        target: NDArray | None = None,
         **additional_arrays: NDArray,
-    ) -> tuple[NDArray, Optional[NDArray], dict[str, NDArray]]:
+    ) -> tuple[NDArray, NDArray | None, dict[str, NDArray]]:
         """Apply the transforms to the input data, including additional arrays.
 
         Parameters

careamics/transforms/n2v_manipulate.py

@@ -1,6 +1,6 @@
 """N2V manipulation transform."""
 
-from typing import Any, Literal, Optional
+from typing import Any, Literal
 
 import numpy as np
 from numpy.typing import NDArray
@@ -61,7 +61,7 @@ class N2VManipulate(Transform):
         remove_center: bool = True,
         struct_mask_axis: Literal["horizontal", "vertical", "none"] = "none",
         struct_mask_span: int = 5,
-        seed: Optional[int] = None,
+        seed: int | None = None,
     ):
         """Constructor.
 
@@ -88,7 +88,7 @@ class N2VManipulate(Transform):
         self.remove_center = remove_center  # TODO is this ever used?
 
         if struct_mask_axis == SupportedStructAxis.NONE:
-            self.struct_mask: Optional[StructMaskParameters] = None
+            self.struct_mask: StructMaskParameters | None = None
         else:
             self.struct_mask = StructMaskParameters(
                 axis=0 if struct_mask_axis == SupportedStructAxis.HORIZONTAL else 1,

careamics/transforms/n2v_manipulate_torch.py

@@ -1,7 +1,7 @@
 """N2V manipulation transform for PyTorch."""
 
 import platform
-from typing import Any, Optional
+from typing import Any
 
 import torch
 
@@ -49,8 +49,8 @@ class N2VManipulateTorch:
     def __init__(
         self,
         n2v_manipulate_config: N2VManipulateModel,
-        seed: Optional[int] = None,
-        device: Optional[str] = None,
+        seed: int | None = None,
+        device: str | None = None,
     ):
         """Constructor.
 
@@ -69,7 +69,7 @@ class N2VManipulateTorch:
         self.remove_center = n2v_manipulate_config.remove_center
 
         if n2v_manipulate_config.struct_mask_axis == SupportedStructAxis.NONE:
-            self.struct_mask: Optional[StructMaskParameters] = None
+            self.struct_mask: StructMaskParameters | None = None
         else:
             self.struct_mask = StructMaskParameters(
                 axis=(

careamics/transforms/normalize.py

@@ -1,7 +1,5 @@
 """Normalization and denormalization transforms for image patches."""
 
-from typing import Optional
-
 import numpy as np
 from numpy.typing import NDArray
 
@@ -66,8 +64,8 @@ class Normalize(Transform):
         self,
         image_means: list[float],
         image_stds: list[float],
-        target_means: Optional[list[float]] = None,
-        target_stds: Optional[list[float]] = None,
+        target_means: list[float] | None = None,
+        target_stds: list[float] | None = None,
     ):
         """Constructor.
 
@@ -92,9 +90,9 @@ class Normalize(Transform):
     def __call__(
         self,
         patch: np.ndarray,
-        target: Optional[NDArray] = None,
+        target: NDArray | None = None,
         **additional_arrays: NDArray,
-    ) -> tuple[NDArray, Optional[NDArray], dict[str, NDArray]]:
+    ) -> tuple[NDArray, NDArray | None, dict[str, NDArray]]:
         """Apply the transform to the source patch and the target (optional).
 
         Parameters

careamics/transforms/pixel_manipulation.py

@@ -5,8 +5,6 @@ Pixel manipulation is used in N2V and similar algorithm to replace the value of
 masked pixels.
 """
 
-from typing import Optional
-
 import numpy as np
 
 from .struct_mask_parameters import StructMaskParameters
@@ -16,7 +14,7 @@ def _apply_struct_mask(
     patch: np.ndarray,
     coords: np.ndarray,
     struct_params: StructMaskParameters,
-    rng: Optional[np.random.Generator] = None,
+    rng: np.random.Generator | None = None,
 ) -> np.ndarray:
     """Apply structN2V masks to patch.
 
@@ -108,7 +106,7 @@ def _odd_jitter_func(step: float, rng: np.random.Generator) -> np.ndarray:
 def _get_stratified_coords(
     mask_pixel_perc: float,
     shape: tuple[int, ...],
-    rng: Optional[np.random.Generator] = None,
+    rng: np.random.Generator | None = None,
 ) -> np.ndarray:
     """
     Generate coordinates of the pixels to mask.
@@ -241,8 +239,8 @@ def uniform_manipulate(
     mask_pixel_percentage: float,
     subpatch_size: int = 11,
     remove_center: bool = True,
-    struct_params: Optional[StructMaskParameters] = None,
-    rng: Optional[np.random.Generator] = None,
+    struct_params: StructMaskParameters | None = None,
+    rng: np.random.Generator | None = None,
 ) -> tuple[np.ndarray, np.ndarray]:
     """
     Manipulate pixels by replacing them with a neighbor values.
@@ -321,8 +319,8 @@ def median_manipulate(
     patch: np.ndarray,
     mask_pixel_percentage: float,
     subpatch_size: int = 11,
-    struct_params: Optional[StructMaskParameters] = None,
-    rng: Optional[np.random.Generator] = None,
+    struct_params: StructMaskParameters | None = None,
+    rng: np.random.Generator | None = None,
 ) -> tuple[np.ndarray, np.ndarray]:
     """
     Manipulate pixels by replacing them with the median of their surrounding subpatch.

careamics/transforms/pixel_manipulation_torch.py

@@ -1,7 +1,5 @@
 """N2V manipulation functions for PyTorch."""
 
-from typing import Optional
-
 import torch
 
 from .struct_mask_parameters import StructMaskParameters
@@ -11,7 +9,7 @@ def _apply_struct_mask_torch(
     patch: torch.Tensor,
     coords: torch.Tensor,
     struct_params: StructMaskParameters,
-    rng: Optional[torch.Generator] = None,
+    rng: torch.Generator | None = None,
 ) -> torch.Tensor:
     """Apply structN2V masks to patch.
 
@@ -154,8 +152,8 @@ def uniform_manipulate_torch(
     mask_pixel_percentage: float,
     subpatch_size: int = 11,
     remove_center: bool = True,
-    struct_params: Optional[StructMaskParameters] = None,
-    rng: Optional[torch.Generator] = None,
+    struct_params: StructMaskParameters | None = None,
+    rng: torch.Generator | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Manipulate pixels by replacing them with a neighbor values.
@@ -256,8 +254,8 @@ def median_manipulate_torch(
     batch: torch.Tensor,
     mask_pixel_percentage: float,
     subpatch_size: int = 11,
-    struct_params: Optional[StructMaskParameters] = None,
-    rng: Optional[torch.Generator] = None,
+    struct_params: StructMaskParameters | None = None,
+    rng: torch.Generator | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """
     Manipulate pixels by replacing them with the median of their surrounding subpatch.

careamics/transforms/xy_flip.py

@@ -1,7 +1,5 @@
 """XY flip transform."""
 
-from typing import Optional
-
 import numpy as np
 from numpy.typing import NDArray
 
@@ -43,7 +41,7 @@ class XYFlip(Transform):
         flip_x: bool = True,
         flip_y: bool = True,
         p: float = 0.5,
-        seed: Optional[int] = None,
+        seed: int | None = None,
     ) -> None:
         """Constructor.
 
@@ -81,9 +79,9 @@ class XYFlip(Transform):
     def __call__(
         self,
         patch: NDArray,
-        target: Optional[NDArray] = None,
+        target: NDArray | None = None,
         **additional_arrays: NDArray,
-    ) -> tuple[NDArray, Optional[NDArray], dict[str, NDArray]]:
+    ) -> tuple[NDArray, NDArray | None, dict[str, NDArray]]:
         """Apply the transform to the source patch and the target (optional).
 
         Parameters

careamics/transforms/xy_random_rotate90.py

@@ -1,7 +1,5 @@
 """Patch transform applying XY random 90 degrees rotations."""
 
-from typing import Optional
-
 import numpy as np
 from numpy.typing import NDArray
 
@@ -30,7 +28,7 @@ class XYRandomRotate90(Transform):
         Random seed, by default None.
     """
 
-    def __init__(self, p: float = 0.5, seed: Optional[int] = None):
+    def __init__(self, p: float = 0.5, seed: int | None = None):
         """Constructor.
 
         Parameters
@@ -52,9 +50,9 @@ class XYRandomRotate90(Transform):
     def __call__(
         self,
         patch: NDArray,
-        target: Optional[NDArray] = None,
+        target: NDArray | None = None,
         **additional_arrays: NDArray,
-    ) -> tuple[NDArray, Optional[NDArray], dict[str, NDArray]]:
+    ) -> tuple[NDArray, NDArray | None, dict[str, NDArray]]:
         """Apply the transform to the source patch and the target (optional).
 
         Parameters

careamics/utils/logging.py

@@ -9,7 +9,7 @@ import sys
 import time
 from collections.abc import Generator
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any, Union
 
 LOGGERS: dict = {}
 
@@ -17,7 +17,7 @@ LOGGERS: dict = {}
 def get_logger(
     name: str,
     log_level: int = logging.INFO,
-    log_path: Optional[Union[str, Path]] = None,
+    log_path: Union[str, Path] | None = None,
 ) -> logging.Logger:
     """
     Create a python logger instance with configured handlers.
@@ -97,10 +97,10 @@ class ProgressBar:
 
     def __init__(
         self,
-        max_value: Optional[int] = None,
-        epoch: Optional[int] = None,
-        num_epochs: Optional[int] = None,
-        stateful_metrics: Optional[list] = None,
+        max_value: int | None = None,
+        epoch: int | None = None,
+        num_epochs: int | None = None,
+        stateful_metrics: list | None = None,
         always_stateful: bool = False,
         mode: str = "train",
     ) -> None:
@@ -159,7 +159,7 @@ class ProgressBar:
             self.message = "Denoising"
 
     def update(
-        self, current_step: int, batch_size: int = 1, values: Optional[list] = None
+        self, current_step: int, batch_size: int = 1, values: list | None = None
     ) -> None:
         """
         Update the progress bar.
@@ -264,7 +264,7 @@ class ProgressBar:
 
         self._last_update = now
 
-    def add(self, n: int, values: Optional[list] = None) -> None:
+    def add(self, n: int, values: list | None = None) -> None:
         """
         Update the progress bar by n steps.
 

careamics/utils/metrics.py

@@ -5,7 +5,7 @@ This module contains various metrics and a metrics tracking class.
 """
 
 from collections.abc import Callable
-from typing import Optional, Union
+from typing import Union
 
 import numpy as np
 import torch
@@ -210,7 +210,7 @@ class RunningPSNR:
         self.mse_sum += torch.nansum(elementwise_mse)
         self.N += len(elementwise_mse) - torch.sum(torch.isnan(elementwise_mse))
 
-    def get(self) -> Optional[torch.Tensor]:
+    def get(self) -> torch.Tensor | None:
         """Get the actual PSNR value given the running statistics.
 
         Returns

careamics/utils/plotting.py

@@ -1,7 +1,5 @@
 """Plotting utilities."""
 
-from typing import Optional
-
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
@@ -14,7 +12,7 @@ def plot_noise_model_probability_distribution(
     noise_model: GaussianMixtureNoiseModel,
     signalBinIndex: int,
     histogram: NDArray,
-    channel: Optional[str] = None,
+    channel: str | None = None,
     number_of_bins: int = 100,
 ) -> None:
     """Plot probability distribution P(x|s) for a certain ground truth signal.

{careamics-0.0.13.dist-info → careamics-0.0.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: careamics
-Version: 0.0.13
+Version: 0.0.15
 Summary: Toolbox for running N2V and friends.
 Project-URL: homepage, https://careamics.github.io/
 Project-URL: repository, https://github.com/CAREamics/careamics
@@ -15,7 +15,7 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Typing :: Typed
 Requires-Python: >=3.10
-Requires-Dist: bioimageio-core==0.7
+Requires-Dist: bioimageio-core==0.9.0
 Requires-Dist: matplotlib<=3.10.3
 Requires-Dist: numpy<2.0.0
 Requires-Dist: pillow<=11.2.1
@@ -28,7 +28,6 @@ Requires-Dist: tifffile<=2025.5.10
 Requires-Dist: torch<=2.7.1,>=2.0
 Requires-Dist: torchvision<=0.22.1
 Requires-Dist: typer<=0.16.0,>=0.12.3
-Requires-Dist: xarray<2025.3.0
 Requires-Dist: zarr<3.0.0
 Provides-Extra: czi
 Requires-Dist: pylibczirw<6.0.0,>=4.1.2; extra == 'czi'