careamics 0.1.0rc2__py3-none-any.whl → 0.1.0rc4__py3-none-any.whl

careamics/models/activation.py

@@ -0,0 +1,35 @@
+from typing import Callable, Union
+
+import torch.nn as nn
+
+from ..config.support import SupportedActivation
+
+
+def get_activation(activation: Union[SupportedActivation, str]) -> Callable:
+    """
+    Get activation function.
+
+    Parameters
+    ----------
+    activation : str
+        Activation function name.
+
+    Returns
+    -------
+    Callable
+        Activation function.
+    """
+    if activation == SupportedActivation.RELU:
+        return nn.ReLU()
+    elif activation == SupportedActivation.LEAKYRELU:
+        return nn.LeakyReLU()
+    elif activation == SupportedActivation.TANH:
+        return nn.Tanh()
+    elif activation == SupportedActivation.SIGMOID:
+        return nn.Sigmoid()
+    elif activation == SupportedActivation.SOFTMAX:
+        return nn.Softmax(dim=1)
+    elif activation == SupportedActivation.NONE:
+        return nn.Identity()
+    else:
+        raise ValueError(f"Activation {activation} not supported.")

careamics/models/layers.py

@@ -3,8 +3,11 @@ Layer module.
 
 This submodule contains layers used in the CAREamics models.
 """
+from typing import List, Optional, Tuple, Union
+
 import torch
 import torch.nn as nn
+from torch.nn import functional as F
 
 
 class Conv_Block(nn.Module):
@@ -150,3 +153,244 @@ class Conv_Block(nn.Module):
         if self.dropout is not None:
             x = self.dropout(x)
         return x
+
+
+def _unpack_kernel_size(
+    kernel_size: Union[Tuple[int, ...], int], dim: int
+) -> Tuple[int, ...]:
+    """Unpack kernel_size to a tuple of ints.
+
+    Inspired by Kornia implementation. TODO: link
+    """
+    if isinstance(kernel_size, int):
+        kernel_dims = tuple([kernel_size for _ in range(dim)])
+    else:
+        kernel_dims = kernel_size
+    return kernel_dims
+
+
+def _compute_zero_padding(
+    kernel_size: Union[Tuple[int, ...], int], dim: int
+) -> Tuple[int, ...]:
+    """Utility function that computes zero padding tuple."""
+    kernel_dims = _unpack_kernel_size(kernel_size, dim)
+    return tuple([(kd - 1) // 2 for kd in kernel_dims])
+
+
+def get_pascal_kernel_1d(
+    kernel_size: int,
+    norm: bool = False,
+    *,
+    device: Optional[torch.device] = None,
+    dtype: Optional[torch.dtype] = None,
+) -> torch.Tensor:
+    """Generate Yang Hui triangle (Pascal's triangle) for a given number.
+
+    Inspired by Kornia implementation. TODO link
+
+    Parameters
+    ----------
+    kernel_size: height and width of the kernel.
+    norm: if to normalize the kernel or not. Default: False.
+    device: tensor device
+    dtype: tensor dtype
+
+    Returns
+    -------
+    kernel shaped as :math:`(kernel_size,)`
+
+    Examples
+    --------
+    >>> get_pascal_kernel_1d(1)
+    tensor([1.])
+    >>> get_pascal_kernel_1d(2)
+    tensor([1., 1.])
+    >>> get_pascal_kernel_1d(3)
+    tensor([1., 2., 1.])
+    >>> get_pascal_kernel_1d(4)
+    tensor([1., 3., 3., 1.])
+    >>> get_pascal_kernel_1d(5)
+    tensor([1., 4., 6., 4., 1.])
+    >>> get_pascal_kernel_1d(6)
+    tensor([ 1.,  5., 10., 10.,  5.,  1.])
+    """
+    pre: List[float] = []
+    cur: List[float] = []
+    for i in range(kernel_size):
+        cur = [1.0] * (i + 1)
+
+        for j in range(1, i // 2 + 1):
+            value = pre[j - 1] + pre[j]
+            cur[j] = value
+            if i != 2 * j:
+                cur[-j - 1] = value
+        pre = cur
+
+    out = torch.tensor(cur, device=device, dtype=dtype)
+
+    if norm:
+        out = out / out.sum()
+
+    return out
+
+
+def _get_pascal_kernel_nd(
+    kernel_size: Union[Tuple[int, int], int],
+    norm: bool = True,
+    dim: int = 2,
+    *,
+    device: Optional[torch.device] = None,
+    dtype: Optional[torch.dtype] = None,
+) -> torch.Tensor:
+    """Generate pascal filter kernel by kernel size.
+
+    Inspired by Kornia implementation.
+
+    Parameters
+    ----------
+    kernel_size: height and width of the kernel.
+    norm: if to normalize the kernel or not. Default: True.
+    device: tensor device
+    dtype: tensor dtype
+
+    Returns
+    -------
+    if kernel_size is an integer the kernel will be shaped as (kernel_size, kernel_size)
+    otherwise the kernel will be shaped as kernel_size
+
+    Examples
+    --------
+    >>> _get_pascal_kernel_nd(1)
+    tensor([[1.]])
+    >>> _get_pascal_kernel_nd(4)
+    tensor([[0.0156, 0.0469, 0.0469, 0.0156],
+            [0.0469, 0.1406, 0.1406, 0.0469],
+            [0.0469, 0.1406, 0.1406, 0.0469],
+            [0.0156, 0.0469, 0.0469, 0.0156]])
+    >>> _get_pascal_kernel_nd(4, norm=False)
+    tensor([[1., 3., 3., 1.],
+            [3., 9., 9., 3.],
+            [3., 9., 9., 3.],
+            [1., 3., 3., 1.]])
+    """
+    kernel_dims = _unpack_kernel_size(kernel_size, dim)
+
+    kernel = [
+        get_pascal_kernel_1d(kd, device=device, dtype=dtype) for kd in kernel_dims
+    ]
+
+    if dim == 2:
+        kernel = kernel[0][:, None] * kernel[1][None, :]
+    elif dim == 3:
+        kernel = (
+            kernel[0][:, None, None]
+            * kernel[1][None, :, None]
+            * kernel[2][None, None, :]
+        )
+    if norm:
+        kernel = kernel / torch.sum(kernel)
+    return kernel
+
+
+def _max_blur_pool_by_kernel2d(
+    x: torch.Tensor,
+    kernel: torch.Tensor,
+    stride: int,
+    max_pool_size: int,
+    ceil_mode: bool,
+) -> torch.Tensor:
+    """Compute max_blur_pool by a given :math:`CxC_(out, None)xNxN` kernel.
+
+    Inspired by Kornia implementation.
+    """
+    # compute local maxima
+    x = F.max_pool2d(
+        x, kernel_size=max_pool_size, padding=0, stride=1, ceil_mode=ceil_mode
+    )
+    # blur and downsample
+    padding = _compute_zero_padding((kernel.shape[-2], kernel.shape[-1]), dim=2)
+    return F.conv2d(x, kernel, padding=padding, stride=stride, groups=x.size(1))
+
+
+def _max_blur_pool_by_kernel3d(
+    x: torch.Tensor,
+    kernel: torch.Tensor,
+    stride: int,
+    max_pool_size: int,
+    ceil_mode: bool,
+) -> torch.Tensor:
+    """Compute max_blur_pool by a given :math:`CxC_(out, None)xNxNxN` kernel.
+
+    Inspired by Kornia implementation.
+    """
+    # compute local maxima
+    x = F.max_pool3d(
+        x, kernel_size=max_pool_size, padding=0, stride=1, ceil_mode=ceil_mode
+    )
+    # blur and downsample
+    padding = _compute_zero_padding(
+        (kernel.shape[-3], kernel.shape[-2], kernel.shape[-1]), dim=3
+    )
+    return F.conv3d(x, kernel, padding=padding, stride=stride, groups=x.size(1))
+
+
+class MaxBlurPool(nn.Module):
+    """Compute pools and blurs and downsample a given feature map.
+
+    Inspired by Kornia MaxBlurPool implementation. Equivalent to
+    ```nn.Sequential(nn.MaxPool2d(...), BlurPool2D(...))```
+
+    Parameters
+    ----------
+    dim: int
+        Toggles between 2D and 3D
+    kernel_size: Union[Tuple[int, int], int]
+        Kernel size for max pooling.
+    stride: int
+        Stride for pooling.
+    max_pool_size: int
+        Max kernel size for max pooling.
+    ceil_mode: bool
+        Should be true to match output size of conv2d with same kernel size.
+
+    Returns
+    -------
+    torch.Tensor
+        The pooled and blurred tensor.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        kernel_size: Union[Tuple[int, int], int],
+        stride: int = 2,
+        max_pool_size: int = 2,
+        ceil_mode: bool = False,
+    ) -> None:
+        super().__init__()
+        self.dim = dim
+        self.kernel_size = kernel_size
+        self.stride = stride
+        self.max_pool_size = max_pool_size
+        self.ceil_mode = ceil_mode
+        self.kernel = _get_pascal_kernel_nd(kernel_size, norm=True, dim=self.dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass of the function."""
+        self.kernel = torch.as_tensor(self.kernel, device=x.device, dtype=x.dtype)
+        if self.dim == 2:
+            return _max_blur_pool_by_kernel2d(
+                x,
+                self.kernel.repeat((x.size(1), 1, 1, 1)),
+                self.stride,
+                self.max_pool_size,
+                self.ceil_mode,
+            )
+        else:
+            return _max_blur_pool_by_kernel3d(
+                x,
+                self.kernel.repeat((x.size(1), 1, 1, 1, 1)),
+                self.stride,
+                self.max_pool_size,
+                self.ceil_mode,
+            )

careamics/models/model_factory.py

@@ -3,31 +3,30 @@ Model factory.
 
 Model creation factory functions.
 """
-from pathlib import Path
-from typing import Dict, Optional, Tuple, Union
+from typing import 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 ..config.architectures import CustomModel, UNetModel, VAEModel, get_custom_model
+from ..config.support import SupportedArchitecture
+from ..utils import get_logger
 from .unet import UNet
 
 logger = get_logger(__name__)
 
 
-def model_registry(model_name: str) -> torch.nn.Module:
+def model_factory(
+    model_configuration: Union[UNetModel, VAEModel, CustomModel]
+) -> torch.nn.Module:
     """
-    Model factory.
+    Deep learning model factory.
 
-    Supported models are defined in config.algorithm.Models.
+    Supported models are defined in careamics.config.SupportedArchitecture.
 
     Parameters
     ----------
-    model_name : str
-        Name of the model.
+    model_configuration : Union[UNetModel, VAEModel]
+        Model configuration
 
     Returns
     -------
@@ -37,215 +36,16 @@ def model_registry(model_name: str) -> torch.nn.Module:
     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 the requested architecture is not implemented.
     """
-    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")
+    if model_configuration.architecture == SupportedArchitecture.UNET:
+        return UNet(**model_configuration.model_dump())
+    elif model_configuration.architecture == SupportedArchitecture.CUSTOM:
+        assert isinstance(model_configuration, CustomModel)
+        model = get_custom_model(model_configuration.name)
 
+        return model(**model_configuration.model_dump())
     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
+        raise NotImplementedError(
+            f"Model {model_configuration.architecture} is not implemented or unknown."
+        )