careamics 0.0.1__py3-none-any.whl → 0.1.0rc2__py3-none-any.whl

careamics/losses/__init__.py

@@ -0,0 +1,4 @@
+"""Losses module."""
+
+
+from .loss_factory import create_loss_function as create_loss_function

careamics/losses/loss_factory.py

@@ -0,0 +1,38 @@
+"""
+Loss factory module.
+
+This module contains a factory function for creating loss functions.
+"""
+from typing import Callable
+
+from careamics.config import Configuration
+from careamics.config.algorithm import Loss
+
+from .losses import n2v_loss
+
+
+def create_loss_function(config: Configuration) -> Callable:
+    """
+    Create loss function based on Configuration.
+
+    Parameters
+    ----------
+    config : Configuration
+        Configuration.
+
+    Returns
+    -------
+    Callable
+        Loss function.
+
+    Raises
+    ------
+    NotImplementedError
+        If the loss is unknown.
+    """
+    loss_type = config.algorithm.loss
+
+    if loss_type == Loss.N2V:
+        return n2v_loss
+    else:
+        raise NotImplementedError(f"Loss {loss_type} is not yet supported.")

careamics/losses/losses.py

@@ -0,0 +1,34 @@
+"""
+Loss submodule.
+
+This submodule contains the various losses used in CAREamics.
+"""
+import torch
+
+
+def n2v_loss(
+    samples: torch.Tensor, labels: torch.Tensor, masks: torch.Tensor, device: str
+) -> torch.Tensor:
+    """
+    N2V Loss function (see Eq.7 in Krull et al).
+
+    Parameters
+    ----------
+    samples : torch.Tensor
+        Patches with manipulated pixels.
+    labels : torch.Tensor
+        Noisy patches.
+    masks : torch.Tensor
+        Array containing masked pixel locations.
+    device : str
+        Device to use.
+
+    Returns
+    -------
+    torch.Tensor
+        Loss value.
+    """
+    errors = (labels - samples) ** 2
+    # Average over pixels and batch
+    loss = torch.sum(errors * masks) / torch.sum(masks)
+    return loss

careamics/manipulation/__init__.py

@@ -0,0 +1,4 @@
+"""Pixel manipulation functions for N2V."""
+
+
+from .pixel_manipulation import default_manipulate as default_manipulate

careamics/manipulation/pixel_manipulation.py

@@ -0,0 +1,158 @@
+"""
+Pixel manipulation methods.
+
+Pixel manipulation is used in N2V and similar algorithm to replace the value of
+masked pixels.
+"""
+from typing import Callable, Optional, Tuple
+
+import numpy as np
+
+
+def _odd_jitter_func(step: float, rng: np.random.Generator) -> np.ndarray:
+    """
+    Randomly sample a jitter to be applied to the masking grid.
+
+    This is done to account for cases where the step size is not an integer.
+
+    Parameters
+    ----------
+    step : float
+        Step size of the grid, output of np.linspace.
+    rng : np.random.Generator
+        Random number generator.
+
+    Returns
+    -------
+    np.ndarray
+        Array of random jitter to be added to the grid.
+    """
+    # Define the random jitter to be added to the grid
+    odd_jitter = np.where(np.floor(step) == step, 0, rng.integers(0, 2))
+
+    # Round the step size to the nearest integer depending on the jitter
+    return np.floor(step) if odd_jitter == 0 else np.ceil(step)
+
+
+def get_stratified_coords(
+    mask_pixel_perc: float,
+    shape: Tuple[int, ...],
+) -> np.ndarray:
+    """
+    Generate coordinates of the pixels to mask.
+
+    Randomly selects the coordinates of the pixels to mask in a stratified way, i.e.
+    the distance between masked pixels is approximately the same.
+
+    Parameters
+    ----------
+    mask_pixel_perc : float
+        Actual (quasi) percentage of masked pixels across the whole image. Used in
+        calculating the distance between masked pixels across each axis.
+    shape : Tuple[int, ...]
+        Shape of the input patch.
+
+    Returns
+    -------
+    np.ndarray
+        Array of coordinates of the masked pixels.
+    """
+    rng = np.random.default_rng()
+
+    # Define the approximate distance between masked pixels
+    mask_pixel_distance = np.round((100 / mask_pixel_perc) ** (1 / len(shape))).astype(
+        np.int32
+    )
+
+    # Define a grid of coordinates for each axis in the input patch and the step size
+    pixel_coords = []
+    for axis_size in shape:
+        # make sure axis size is evenly divisible by box size
+        num_pixels = int(np.ceil(axis_size / mask_pixel_distance))
+        axis_pixel_coords, step = np.linspace(
+            0, axis_size, num_pixels, dtype=np.int32, endpoint=False, retstep=True
+        )
+        # explain
+        pixel_coords.append(axis_pixel_coords.T)
+
+    # Create a meshgrid of coordinates for each axis in the input patch
+    coordinate_grid_list = np.meshgrid(*pixel_coords)
+    coordinate_grid = np.array(coordinate_grid_list).reshape(len(shape), -1).T
+
+    grid_random_increment = rng.integers(
+        _odd_jitter_func(float(step), rng)
+        * np.ones_like(coordinate_grid).astype(np.int32)
+        - 1,
+        size=coordinate_grid.shape,
+        endpoint=True,
+    )
+    coordinate_grid += grid_random_increment
+    coordinate_grid = np.clip(coordinate_grid, 0, np.array(shape) - 1)
+    return coordinate_grid
+
+
+def default_manipulate(
+    patch: np.ndarray,
+    mask_pixel_percentage: float,
+    roi_size: int = 11,
+    augmentations: Optional[Callable] = None,
+) -> Tuple[np.ndarray, ...]:
+    """
+    Manipulate pixel in a patch, i.e. replace the masked value.
+
+    Parameters
+    ----------
+    patch : np.ndarray
+        Image patch, 2D or 3D, shape (y, x) or (z, y, x).
+    mask_pixel_percentage : floar
+        Approximate percentage of pixels to be masked.
+    roi_size : int
+        Size of the ROI the new pixel value is sampled from, by default 11.
+    augmentations : Callable, optional
+        Augmentations to apply, by default None.
+
+    Returns
+    -------
+    Tuple[np.ndarray]
+        Tuple containing the manipulated patch, the original patch and the mask.
+    """
+    original_patch = patch.copy()
+
+    # Get the coordinates of the pixels to be replaced
+    roi_centers = get_stratified_coords(mask_pixel_percentage, patch.shape)
+    rng = np.random.default_rng()
+
+    # Generate coordinate grid for ROI
+    roi_span_full = np.arange(-np.floor(roi_size / 2), np.ceil(roi_size / 2)).astype(
+        np.int32
+    )
+    # Remove the center pixel from the grid
+    roi_span_wo_center = roi_span_full[roi_span_full != 0]
+
+    # Randomly select coordinates from the grid
+    random_increment = rng.choice(roi_span_wo_center, size=roi_centers.shape)
+
+    # Clip the coordinates to the patch size
+    replacement_coords = np.clip(
+        roi_centers + random_increment,
+        0,
+        [patch.shape[i] - 1 for i in range(len(patch.shape))],
+    )
+    # Get the replacement pixels from all rois
+    replacement_pixels = patch[tuple(replacement_coords.T.tolist())]
+
+    # Replace the original pixels with the replacement pixels
+    patch[tuple(roi_centers.T.tolist())] = replacement_pixels
+    mask = np.where(patch != original_patch, 1, 0).astype(np.uint8)
+
+    patch, original_patch, mask = (
+        (patch, original_patch, mask)
+        if augmentations is None
+        else augmentations(patch, original_patch, mask)
+    )
+
+    return (
+        np.expand_dims(patch, 0),
+        np.expand_dims(original_patch, 0),
+        np.expand_dims(mask, 0),
+    )

careamics/models/__init__.py

@@ -0,0 +1,4 @@
+"""Models package."""
+
+from .model_factory import create_model as create_model
+from .unet import UNet as UNet

careamics/models/layers.py

@@ -0,0 +1,152 @@
+"""
+Layer module.
+
+This submodule contains layers used in the CAREamics models.
+"""
+import torch
+import torch.nn as nn
+
+
+class Conv_Block(nn.Module):
+    """
+    Convolution block used in UNets.
+
+    Convolution block consist of two convolution layers with optional batch norm,
+    dropout and with a final activation function.
+
+    The parameters are directly mapped to PyTorch Conv2D and Conv3d parameters, see
+    PyTorch torch.nn.Conv2d and torch.nn.Conv3d for more information.
+
+    Parameters
+    ----------
+    conv_dim : int
+        Number of dimension of the convolutions, 2 or 3.
+    in_channels : int
+        Number of input channels.
+    out_channels : int
+        Number of output channels.
+    intermediate_channel_multiplier : int, optional
+        Multiplied for the number of output channels, by default 1.
+    stride : int, optional
+        Stride of the convolutions, by default 1.
+    padding : int, optional
+        Padding of the convolutions, by default 1.
+    bias : bool, optional
+        Bias of the convolutions, by default True.
+    groups : int, optional
+        Controls the connections between inputs and outputs, by default 1.
+    activation : str, optional
+        Activation function, by default "ReLU".
+    dropout_perc : float, optional
+        Dropout percentage, by default 0.
+    use_batch_norm : bool, optional
+        Use batch norm, by default False.
+    """
+
+    def __init__(
+        self,
+        conv_dim: int,
+        in_channels: int,
+        out_channels: int,
+        intermediate_channel_multiplier: int = 1,
+        stride: int = 1,
+        padding: int = 1,
+        bias: bool = True,
+        groups: int = 1,
+        activation: str = "ReLU",
+        dropout_perc: float = 0,
+        use_batch_norm: bool = False,
+    ) -> None:
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        conv_dim : int
+            Number of dimension of the convolutions, 2 or 3.
+        in_channels : int
+            Number of input channels.
+        out_channels : int
+            Number of output channels.
+        intermediate_channel_multiplier : int, optional
+            Multiplied for the number of output channels, by default 1.
+        stride : int, optional
+            Stride of the convolutions, by default 1.
+        padding : int, optional
+            Padding of the convolutions, by default 1.
+        bias : bool, optional
+            Bias of the convolutions, by default True.
+        groups : int, optional
+            Controls the connections between inputs and outputs, by default 1.
+        activation : str, optional
+            Activation function, by default "ReLU".
+        dropout_perc : float, optional
+            Dropout percentage, by default 0.
+        use_batch_norm : bool, optional
+            Use batch norm, by default False.
+        """
+        super().__init__()
+        self.use_batch_norm = use_batch_norm
+        self.conv1 = getattr(nn, f"Conv{conv_dim}d")(
+            in_channels,
+            out_channels * intermediate_channel_multiplier,
+            kernel_size=3,
+            stride=stride,
+            padding=padding,
+            bias=bias,
+            groups=groups,
+        )
+
+        self.conv2 = getattr(nn, f"Conv{conv_dim}d")(
+            out_channels * intermediate_channel_multiplier,
+            out_channels,
+            kernel_size=3,
+            stride=stride,
+            padding=padding,
+            bias=bias,
+            groups=groups,
+        )
+
+        self.batch_norm1 = getattr(nn, f"BatchNorm{conv_dim}d")(
+            out_channels * intermediate_channel_multiplier
+        )
+        self.batch_norm2 = getattr(nn, f"BatchNorm{conv_dim}d")(out_channels)
+
+        self.dropout = (
+            getattr(nn, f"Dropout{conv_dim}d")(dropout_perc)
+            if dropout_perc > 0
+            else None
+        )
+        self.activation = (
+            getattr(nn, f"{activation}")() if activation is not None else nn.Identity()
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor.
+        """
+        if self.use_batch_norm:
+            x = self.conv1(x)
+            x = self.batch_norm1(x)
+            x = self.activation(x)
+            x = self.conv2(x)
+            x = self.batch_norm2(x)
+            x = self.activation(x)
+        else:
+            x = self.conv1(x)
+            x = self.activation(x)
+            x = self.conv2(x)
+            x = self.activation(x)
+        if self.dropout is not None:
+            x = self.dropout(x)
+        return x

careamics/models/model_factory.py

@@ -0,0 +1,251 @@
+"""
+Model factory.
+
+Model creation factory functions.
+"""
+from pathlib import Path
+from typing import Dict, Optional, Tuple, Union
+
+import torch
+
+from careamics.bioimage import import_bioimage_model
+from careamics.config import Configuration
+from careamics.config.algorithm import Models
+from careamics.utils.logging import get_logger
+
+from .unet import UNet
+
+logger = get_logger(__name__)
+
+
+def model_registry(model_name: str) -> torch.nn.Module:
+    """
+    Model factory.
+
+    Supported models are defined in config.algorithm.Models.
+
+    Parameters
+    ----------
+    model_name : str
+        Name of the model.
+
+    Returns
+    -------
+    torch.nn.Module
+        Model class.
+
+    Raises
+    ------
+    NotImplementedError
+        If the requested model is not implemented.
+    """
+    if model_name == Models.UNET:
+        return UNet
+    else:
+        raise NotImplementedError(f"Model {model_name} is not implemented")
+
+
+def create_model(
+    *,
+    model_path: Optional[Union[str, Path]] = None,
+    config: Optional[Configuration] = None,
+    device: Optional[torch.device] = None,
+) -> Tuple[
+    torch.nn.Module,
+    torch.optim.Optimizer,
+    Union[
+        torch.optim.lr_scheduler.LRScheduler,
+        torch.optim.lr_scheduler.ReduceLROnPlateau,  # not a subclass of LRScheduler
+    ],
+    torch.cuda.amp.GradScaler,
+    Configuration,
+]:
+    """
+    Instantiate a model from a model path or configuration.
+
+    If both path and configuration are provided, the model path is used. The model
+    path should point to either a checkpoint (created during training) or a model
+    exported to the bioimage.io format.
+
+    Parameters
+    ----------
+    model_path : Optional[Union[str, Path]], optional
+        Path to a checkpoint or bioimage.io archive, by default None.
+    config : Optional[Configuration], optional
+        Configuration, by default None.
+    device : Optional[torch.device], optional
+        Torch device, by default None.
+
+    Returns
+    -------
+    torch.nn.Module
+        Instantiated model.
+
+    Raises
+    ------
+    ValueError
+        If the checkpoint path is invalid.
+    ValueError
+        If the checkpoint is invalid.
+    ValueError
+        If neither checkpoint nor configuration are provided.
+    """
+    if model_path is not None:
+        # Create model from checkpoint
+        model_path = Path(model_path)
+        if not model_path.exists() or model_path.suffix not in [".pth", ".zip"]:
+            raise ValueError(
+                f"Invalid model path: {model_path}. Current working dir: \
+                              {Path.cwd()!s}"
+            )
+
+        if model_path.suffix == ".zip":
+            model_path = import_bioimage_model(model_path)
+
+        # Load checkpoint
+        checkpoint = torch.load(model_path, map_location=device)
+
+        # Load the configuration
+        if "config" in checkpoint:
+            config = Configuration(**checkpoint["config"])
+            algo_config = config.algorithm
+            model_config = algo_config.model_parameters
+            model_name = algo_config.model
+        else:
+            raise ValueError("Invalid checkpoint format, no configuration found.")
+
+        # Create model
+        model: torch.nn.Module = model_registry(model_name)(
+            depth=model_config.depth,
+            conv_dim=algo_config.get_conv_dim(),
+            num_channels_init=model_config.num_channels_init,
+        )
+        model.to(device)
+
+        # Load the model state dict
+        if "model_state_dict" in checkpoint:
+            model.load_state_dict(checkpoint["model_state_dict"])
+            logger.info("Loaded model state dict")
+        else:
+            raise ValueError("Invalid checkpoint format")
+
+        # Load the optimizer and scheduler
+        optimizer, scheduler = get_optimizer_and_scheduler(
+            config, model, state_dict=checkpoint
+        )
+        scaler = get_grad_scaler(config, state_dict=checkpoint)
+
+    elif config is not None:
+        # Create model from configuration
+        algo_config = config.algorithm
+        model_config = algo_config.model_parameters
+        model_name = algo_config.model
+
+        # Create model
+        model = model_registry(model_name)(
+            depth=model_config.depth,
+            conv_dim=algo_config.get_conv_dim(),
+            num_channels_init=model_config.num_channels_init,
+        )
+        model.to(device)
+        optimizer, scheduler = get_optimizer_and_scheduler(config, model)
+        scaler = get_grad_scaler(config)
+        logger.info("Engine initialized from configuration")
+
+    else:
+        raise ValueError("Either config or model_path must be provided")
+
+    return model, optimizer, scheduler, scaler, config
+
+
+def get_optimizer_and_scheduler(
+    config: Configuration, model: torch.nn.Module, state_dict: Optional[Dict] = None
+) -> Tuple[
+    torch.optim.Optimizer,
+    Union[
+        torch.optim.lr_scheduler.LRScheduler,
+        torch.optim.lr_scheduler.ReduceLROnPlateau,  # not a subclass of LRScheduler
+    ],
+]:
+    """
+    Create optimizer and learning rate schedulers.
+
+    If a checkpoint state dictionary is provided, the optimizer and scheduler are
+    instantiated to the same state as the checkpoint's optimizer and scheduler.
+
+    Parameters
+    ----------
+    config : Configuration
+        Configuration.
+    model : torch.nn.Module
+        Model.
+    state_dict : Optional[Dict], optional
+        Checkpoint state dictionary, by default None.
+
+    Returns
+    -------
+    Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LRScheduler]
+        Optimizer and scheduler.
+    """
+    # retrieve optimizer name and parameters from config
+    optimizer_name = config.training.optimizer.name
+    optimizer_params = config.training.optimizer.parameters
+
+    # then instantiate it
+    optimizer_func = getattr(torch.optim, optimizer_name)
+    optimizer = optimizer_func(model.parameters(), **optimizer_params)
+
+    # same for learning rate scheduler
+    scheduler_name = config.training.lr_scheduler.name
+    scheduler_params = config.training.lr_scheduler.parameters
+    scheduler_func = getattr(torch.optim.lr_scheduler, scheduler_name)
+    scheduler = scheduler_func(optimizer, **scheduler_params)
+
+    # load state from ther checkpoint if available
+    if state_dict is not None:
+        if "optimizer_state_dict" in state_dict:
+            optimizer.load_state_dict(state_dict["optimizer_state_dict"])
+            logger.info("Loaded optimizer state dict")
+        else:
+            logger.warning(
+                "No optimizer state dict found in checkpoint. Optimizer not loaded."
+            )
+        if "scheduler_state_dict" in state_dict:
+            scheduler.load_state_dict(state_dict["scheduler_state_dict"])
+            logger.info("Loaded LR scheduler state dict")
+        else:
+            logger.warning(
+                "No LR scheduler state dict found in checkpoint. "
+                "LR scheduler not loaded."
+            )
+    return optimizer, scheduler
+
+
+def get_grad_scaler(
+    config: Configuration, state_dict: Optional[Dict] = None
+) -> torch.cuda.amp.GradScaler:
+    """
+    Instantiate gradscaler.
+
+    If a checkpoint state dictionary is provided, the scaler is instantiated to the
+    same state as the checkpoint's scaler.
+
+    Parameters
+    ----------
+    config : Configuration
+        Configuration.
+    state_dict : Optional[Dict], optional
+        Checkpoint state dictionary, by default None.
+
+    Returns
+    -------
+    torch.cuda.amp.GradScaler
+        Instantiated gradscaler.
+    """
+    use = config.training.amp.use
+    scaling = config.training.amp.init_scale
+    scaler = torch.cuda.amp.GradScaler(init_scale=scaling, enabled=use)
+    if state_dict is not None and "scaler_state_dict" in state_dict:
+        scaler.load_state_dict(state_dict["scaler_state_dict"])
+        logger.info("Loaded GradScaler state dict")
+    return scaler