careamics 0.0.19__py3-none-any.whl

careamics/models/unet.py

@@ -0,0 +1,449 @@
+"""
+UNet model.
+
+A UNet encoder, decoder and complete model.
+"""
+
+from typing import Any, Union
+
+import torch
+import torch.nn as nn
+
+from ..config.support import SupportedActivation
+from .activation import get_activation
+from .layers import Conv_Block, MaxBlurPool
+
+
+class UnetEncoder(nn.Module):
+    """
+    Unet encoder pathway.
+
+    Parameters
+    ----------
+    conv_dim : int
+        Number of dimension of the convolution layers, 2 for 2D or 3 for 3D.
+    in_channels : int, optional
+        Number of input channels, by default 1.
+    depth : int, optional
+        Number of encoder blocks, by default 3.
+    num_channels_init : int, optional
+        Number of channels in the first encoder block, by default 64.
+    use_batch_norm : bool, optional
+        Whether to use batch normalization, by default True.
+    dropout : float, optional
+        Dropout probability, by default 0.0.
+    pool_kernel : int, optional
+        Kernel size for the max pooling layers, by default 2.
+    n2v2 : bool, optional
+        Whether to use N2V2 architecture, by default False.
+    groups : int, optional
+        Number of blocked connections from input channels to output
+        channels, by default 1.
+    """
+
+    def __init__(
+        self,
+        conv_dim: int,
+        in_channels: int = 1,
+        depth: int = 3,
+        num_channels_init: int = 64,
+        use_batch_norm: bool = True,
+        dropout: float = 0.0,
+        pool_kernel: int = 2,
+        n2v2: bool = False,
+        groups: int = 1,
+    ) -> None:
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        conv_dim : int
+            Number of dimension of the convolution layers, 2 for 2D or 3 for 3D.
+        in_channels : int, optional
+            Number of input channels, by default 1.
+        depth : int, optional
+            Number of encoder blocks, by default 3.
+        num_channels_init : int, optional
+            Number of channels in the first encoder block, by default 64.
+        use_batch_norm : bool, optional
+            Whether to use batch normalization, by default True.
+        dropout : float, optional
+            Dropout probability, by default 0.0.
+        pool_kernel : int, optional
+            Kernel size for the max pooling layers, by default 2.
+        n2v2 : bool, optional
+            Whether to use N2V2 architecture, by default False.
+        groups : int, optional
+            Number of blocked connections from input channels to output
+            channels, by default 1.
+        """
+        super().__init__()
+
+        self.pooling = (
+            getattr(nn, f"MaxPool{conv_dim}d")(kernel_size=pool_kernel)
+            if not n2v2
+            else MaxBlurPool(dim=conv_dim, kernel_size=3, max_pool_size=pool_kernel)
+        )
+
+        encoder_blocks = []
+
+        for n in range(depth):
+            out_channels = num_channels_init * (2**n) * groups
+            in_channels = in_channels if n == 0 else out_channels // 2
+            encoder_blocks.append(
+                Conv_Block(
+                    conv_dim,
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    dropout_perc=dropout,
+                    use_batch_norm=use_batch_norm,
+                    groups=groups,
+                )
+            )
+            encoder_blocks.append(self.pooling)
+        self.encoder_blocks = nn.ModuleList(encoder_blocks)
+
+    def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
+        """
+        Forward pass.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor.
+
+        Returns
+        -------
+        list[torch.Tensor]
+            Output of each encoder block (skip connections) and final output of the
+            encoder.
+        """
+        encoder_features = []
+        for module in self.encoder_blocks:
+            x = module(x)
+            if isinstance(module, Conv_Block):
+                encoder_features.append(x)
+        features = [x, *encoder_features]
+        return features
+
+
+class UnetDecoder(nn.Module):
+    """
+    Unet decoder pathway.
+
+    Parameters
+    ----------
+    conv_dim : int
+        Number of dimension of the convolution layers, 2 for 2D or 3 for 3D.
+    depth : int, optional
+        Number of decoder blocks, by default 3.
+    num_channels_init : int, optional
+        Number of channels in the first encoder block, by default 64.
+    use_batch_norm : bool, optional
+        Whether to use batch normalization, by default True.
+    dropout : float, optional
+        Dropout probability, by default 0.0.
+    n2v2 : bool, optional
+        Whether to use N2V2 architecture, by default False.
+    groups : int, optional
+        Number of blocked connections from input channels to output
+        channels, by default 1.
+    """
+
+    def __init__(
+        self,
+        conv_dim: int,
+        depth: int = 3,
+        num_channels_init: int = 64,
+        use_batch_norm: bool = True,
+        dropout: float = 0.0,
+        n2v2: bool = False,
+        groups: int = 1,
+    ) -> None:
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        conv_dim : int
+            Number of dimension of the convolution layers, 2 for 2D or 3 for 3D.
+        depth : int, optional
+            Number of decoder blocks, by default 3.
+        num_channels_init : int, optional
+            Number of channels in the first encoder block, by default 64.
+        use_batch_norm : bool, optional
+            Whether to use batch normalization, by default True.
+        dropout : float, optional
+            Dropout probability, by default 0.0.
+        n2v2 : bool, optional
+            Whether to use N2V2 architecture, by default False.
+        groups : int, optional
+            Number of blocked connections from input channels to output
+            channels, by default 1.
+        """
+        super().__init__()
+
+        upsampling = nn.Upsample(
+            scale_factor=2, mode="bilinear" if conv_dim == 2 else "trilinear"
+        )
+        in_channels = out_channels = num_channels_init * groups * (2 ** (depth - 1))
+
+        self.n2v2 = n2v2
+        self.groups = groups
+
+        self.bottleneck = Conv_Block(
+            conv_dim,
+            in_channels=in_channels,
+            out_channels=out_channels,
+            intermediate_channel_multiplier=2,
+            use_batch_norm=use_batch_norm,
+            dropout_perc=dropout,
+            groups=self.groups,
+        )
+
+        decoder_blocks: list[nn.Module] = []
+        for n in range(depth):
+            decoder_blocks.append(upsampling)
+
+            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,
+                    out_channels=out_channels,
+                    # 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,
+                    groups=groups,
+                )
+            )
+
+        self.decoder_blocks = nn.ModuleList(decoder_blocks)
+
+    def forward(self, *features: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Parameters
+        ----------
+        *features :  list[torch.Tensor]
+            List containing the output of each encoder block(skip connections) and final
+            output of the encoder.
+
+        Returns
+        -------
+        torch.Tensor
+            Output of the decoder.
+        """
+        x: torch.Tensor = features[0]
+        skip_connections: tuple[torch.Tensor, ...] = features[-1:0:-1]
+        depth = len(skip_connections)
+
+        x = self.bottleneck(x)
+
+        for i, module in enumerate(self.decoder_blocks):
+            x = module(x)
+            if isinstance(module, nn.Upsample):
+                # divide index by 2 because of upsampling layers
+                skip_connection: torch.Tensor = skip_connections[i // 2]
+                # 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
+
+    @staticmethod
+    def _interleave(A: torch.Tensor, B: torch.Tensor, groups: int) -> torch.Tensor:
+        """Interleave two tensors.
+
+        Splits the tensors `A` and `B` into equally sized groups along the channel
+        axis (axis=1); then concatenates the groups in alternating order along the
+        channel axis, starting with the first group from tensor A.
+
+        Parameters
+        ----------
+        A : torch.Tensor
+            First tensor.
+        B : torch.Tensor
+            Second tensor.
+        groups : int
+            The number of groups.
+
+        Returns
+        -------
+        torch.Tensor
+            Interleaved tensor.
+
+        Raises
+        ------
+        ValueError:
+            If either of `A` or `B`'s channel axis is not divisible by `groups`.
+        """
+        if (A.shape[1] % groups != 0) or (B.shape[1] % groups != 0):
+            raise ValueError(f"Number of channels not divisible by {groups} groups.")
+
+        m = A.shape[1] // groups
+        n = B.shape[1] // groups
+
+        A_groups: list[torch.Tensor] = [
+            A[:, i * m : (i + 1) * m] for i in range(groups)
+        ]
+        B_groups: list[torch.Tensor] = [
+            B[:, i * n : (i + 1) * n] for i in range(groups)
+        ]
+
+        interleaved = torch.cat(
+            [
+                tensor_list[i]
+                for i in range(groups)
+                for tensor_list in [A_groups, B_groups]
+            ],
+            dim=1,
+        )
+
+        return interleaved
+
+
+class UNet(nn.Module):
+    """
+    UNet model.
+
+    Adapted for PyTorch from:
+    https://github.com/juglab/n2v/blob/main/n2v/nets/unet_blocks.py.
+
+    Parameters
+    ----------
+    conv_dims : int
+        Number of dimensions of the convolution layers (2 or 3).
+    num_classes : int, optional
+        Number of classes to predict, by default 1.
+    in_channels : int, optional
+        Number of input channels, by default 1.
+    depth : int, optional
+        Number of downsamplings, by default 3.
+    num_channels_init : int, optional
+        Number of filters in the first convolution layer, by default 64.
+    use_batch_norm : bool, optional
+        Whether to use batch normalization, by default True.
+    dropout : float, optional
+        Dropout probability, by default 0.0.
+    pool_kernel : int, optional
+        Kernel size of the pooling layers, by default 2.
+    final_activation : Optional[Callable], optional
+        Activation function to use for the last layer, by default None.
+    n2v2 : bool, optional
+        Whether to use N2V2 architecture, by default False.
+    independent_channels : bool
+        Whether to train the channels independently, by default True.
+    **kwargs : Any
+        Additional keyword arguments, unused.
+    """
+
+    def __init__(
+        self,
+        conv_dims: int,
+        num_classes: int = 1,
+        in_channels: int = 1,
+        depth: int = 3,
+        num_channels_init: int = 64,
+        use_batch_norm: bool = True,
+        dropout: float = 0.0,
+        pool_kernel: int = 2,
+        final_activation: Union[SupportedActivation, str] = SupportedActivation.NONE,
+        n2v2: bool = False,
+        independent_channels: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        conv_dims : int
+            Number of dimensions of the convolution layers (2 or 3).
+        num_classes : int, optional
+            Number of classes to predict, by default 1.
+        in_channels : int, optional
+            Number of input channels, by default 1.
+        depth : int, optional
+            Number of downsamplings, by default 3.
+        num_channels_init : int, optional
+            Number of filters in the first convolution layer, by default 64.
+        use_batch_norm : bool, optional
+            Whether to use batch normalization, by default True.
+        dropout : float, optional
+            Dropout probability, by default 0.0.
+        pool_kernel : int, optional
+            Kernel size of the pooling layers, by default 2.
+        final_activation : Optional[Callable], optional
+            Activation function to use for the last layer, by default None.
+        n2v2 : bool, optional
+            Whether to use N2V2 architecture, by default False.
+        independent_channels : bool
+            Whether to train parallel independent networks for each channel, by
+            default True.
+        **kwargs : Any
+            Additional keyword arguments, unused.
+        """
+        super().__init__()
+
+        groups = in_channels if independent_channels else 1
+
+        self.encoder = UnetEncoder(
+            conv_dims,
+            in_channels=in_channels,
+            depth=depth,
+            num_channels_init=num_channels_init,
+            use_batch_norm=use_batch_norm,
+            dropout=dropout,
+            pool_kernel=pool_kernel,
+            n2v2=n2v2,
+            groups=groups,
+        )
+
+        self.decoder = UnetDecoder(
+            conv_dims,
+            depth=depth,
+            num_channels_init=num_channels_init,
+            use_batch_norm=use_batch_norm,
+            dropout=dropout,
+            n2v2=n2v2,
+            groups=groups,
+        )
+        self.final_conv = getattr(nn, f"Conv{conv_dims}d")(
+            in_channels=num_channels_init * groups,
+            out_channels=num_classes,
+            kernel_size=1,
+            groups=groups,
+        )
+        self.final_activation = get_activation(final_activation)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Parameters
+        ----------
+        x :  torch.Tensor
+            Input tensor.
+
+        Returns
+        -------
+        torch.Tensor
+            Output of the model.
+        """
+        encoder_features = self.encoder(x)
+        x = self.decoder(*encoder_features)
+        x = self.final_conv(x)
+        x = self.final_activation(x)
+        return x

careamics/nm_training_placeholder.py

@@ -0,0 +1,203 @@
+"""Placeholder code snippets for noise model training integration.
+
+This module contains template/placeholder code that demonstrates how noise model
+training could be integrated into CAREamist. These are reference implementations
+and should not be imported or used directly.
+"""
+
+import logging
+from pathlib import Path
+from typing import Union
+
+from numpy.typing import NDArray
+from pytorch_lightning.callbacks import Callback
+
+from careamics.config.configuration import Configuration
+from careamics.models.lvae.noise_models import (
+    GaussianMixtureNoiseModel,
+)
+
+logger = logging.getLogger(__name__)
+
+
+# In src/careamics/careamist.py (newly added section only)
+def __init__(
+    self,
+    source: Union[Path, str, Configuration],
+    work_dir: Union[Path, str] | None = None,
+    callbacks: list[Callback] | None = None,
+    enable_progress_bar: bool = True,
+) -> None:
+    """Placeholder __init__ method showing noise model initialization.
+
+    Parameters
+    ----------
+    self : object
+        CAREamist instance.
+    source : Union[Path, str, Configuration]
+        Configuration source.
+    work_dir : Union[Path, str] | None, optional
+        Working directory, by default None.
+    callbacks : list[Callback] | None, optional
+        List of callbacks, by default None.
+    enable_progress_bar : bool, optional
+        Whether to show progress bar, by default True.
+    """
+    # ... existing initialization code ...
+
+    # Initialize untrained noise models if needed
+    self.untrained_noise_models = None
+    if (
+        hasattr(self.cfg.algorithm_config, "train_noise_model")
+        and self.cfg.algorithm_config.train_noise_model_from_data
+    ):
+        self._initialize_noise_models_for_training()
+
+
+# In src/careamics/careamist.py
+def train_noise_model(
+    self,
+    clean_data: Union[Path, str, NDArray],
+    noisy_data: Union[Path, str, NDArray],
+    learning_rate: float = 1e-1,
+    batch_size: int = 250000,
+    n_epochs: int = 2000,
+    lower_clip: float = 0.0,
+    upper_clip: float = 100.0,
+    save_noise_models: bool = True,
+) -> None:
+    """Train noise models from clean/noisy data pairs.
+
+    Parameters
+    ----------
+    self : object
+        CAREamist instance.
+    clean_data : Union[Path, str, NDArray]
+        Clean (signal) data for training noise models.
+    noisy_data : Union[Path, str, NDArray]
+        Noisy (observation) data for training noise models.
+    learning_rate : float, default=1e-1
+        Learning rate for noise model training.
+    batch_size : int, default=250000
+        Batch size for noise model training.
+    n_epochs : int, default=2000
+        Number of epochs for noise model training.
+    lower_clip : float, default=0.0
+        Lower percentile for clipping training data.
+    upper_clip : float, default=100.0
+        Upper percentile for clipping training data.
+    save_noise_models : bool, default=True
+        Whether to save trained noise models to disk.
+
+    Raises
+    ------
+    ValueError
+        If noise models are not initialized for training.
+    ValueError
+        If data shapes don't match expectations.
+    """
+    # Check if noise model is initialized (config should have MultiChannelNMConfig)
+    if self.cfg.algorithm_config.noise_model is None:
+        raise ValueError(
+            "No untrained noise models found. Set `train_noise_model=True` "
+            "in configuration."
+        )
+
+    # Load data if paths provided (currently NM expects only numpy)
+    if isinstance(clean_data, (str, Path)):
+        clean_data = self._load_data(clean_data)
+    if isinstance(noisy_data, (str, Path)):
+        noisy_data = self._load_data(noisy_data)
+
+    # Type narrowing for mypy
+    assert not isinstance(clean_data, (str, Path))
+    assert not isinstance(noisy_data, (str, Path))
+
+    # Validate data shapes
+    if clean_data.shape != noisy_data.shape:
+        raise ValueError(
+            f"Clean and noisy data shapes must match. "
+            f"Got clean: {clean_data.shape}, noisy: {noisy_data.shape}"
+        )
+    # TODO other data shape checks
+
+    # parameter controlling the number of channels to split for MS, for HDN it's 1
+    output_channels = self.cfg.algorithm_config.model.output_channels
+
+    # Train noise model for each channel
+    trained_noise_models = []
+    for channel_idx in range(output_channels):
+        logger.info(
+            f"Training noise model for channel {channel_idx + 1}/{output_channels}"
+        )
+
+        # Extract single channel data
+        clean_channel = clean_data[:, channel_idx]  # (N, H, W)
+        noisy_channel = noisy_data[:, channel_idx]  # (N, H, W)
+
+        # Train noise model for this channel
+        noise_model = self.untrained_noise_models[channel_idx]
+        noise_model.fit(
+            signal=clean_channel,
+            observation=noisy_channel,
+            learning_rate=learning_rate,
+            batch_size=batch_size,
+            n_epochs=n_epochs,
+            lower_clip=lower_clip,
+            upper_clip=upper_clip,
+        )
+
+        trained_noise_models.append(noise_model)
+
+        # Save individual noise model if requested
+        if save_noise_models:
+            save_path = self.work_dir / "noise_models"
+            noise_model.save(str(save_path), f"noise_model_ch{channel_idx}.npz")
+            logger.info(f"Saved noise model for channel {channel_idx} to {save_path}")
+
+    # Update the algorithm configuration with trained noise models
+    self._update_config_with_trained_noise_models(trained_noise_models)
+
+    logger.info("Noise model training completed successfully")
+
+
+def _update_config_with_trained_noise_models(
+    self, trained_models: list[GaussianMixtureNoiseModel]
+) -> None:
+    """Update algorithm config with trained noise models.
+
+    Parameters
+    ----------
+    self : object
+        CAREamist instance.
+    trained_models : list[GaussianMixtureNoiseModel]
+        List of trained noise models, one per channel.
+    """
+    # Currently the model is initialized in the __init__ of CAREamist
+    # multichannel_noise_model_factory inside VAEModule expects paths to noise models
+    # Ideally, we change that and call multichannel_noise_model_factory here after the
+    # model init and update the parameters of noise models right in the
+    # MultiChannelNoiseModel
+
+
+def _load_data(self, data_path: Union[Path, str]) -> NDArray:
+    """Load data from file path.
+
+    Parameters
+    ----------
+    self : object
+        CAREamist instance.
+    data_path : Union[Path, str]
+        Path to data file.
+
+    Returns
+    -------
+    NDArray
+        Loaded data array.
+
+    Raises
+    ------
+    NotImplementedError
+        This is a placeholder method.
+    """
+    raise NotImplementedError("Data loading not yet implemented")

careamics/prediction_utils/__init__.py

@@ -0,0 +1,21 @@
+"""Package to house various prediction utilies."""
+
+__all__ = [
+    "convert_outputs",
+    "convert_outputs_microsplit",
+    "convert_outputs_pn2v",
+    "stitch_prediction",
+    "stitch_prediction_single",
+    "stitch_prediction_vae",
+]
+
+from .prediction_outputs import (
+    convert_outputs,
+    convert_outputs_microsplit,
+    convert_outputs_pn2v,
+)
+from .stitch_prediction import (
+    stitch_prediction,
+    stitch_prediction_single,
+    stitch_prediction_vae,
+)