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

careamics/models/unet.py

@@ -0,0 +1,322 @@
+"""
+UNet model.
+
+A UNet encoder, decoder and complete model.
+"""
+from typing import Callable, List, Optional
+
+import torch
+import torch.nn as nn
+
+from .layers import Conv_Block
+
+
+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.
+    """
+
+    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,
+    ) -> 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.
+        """
+        super().__init__()
+
+        self.pooling = getattr(nn, f"MaxPool{conv_dim}d")(kernel_size=pool_kernel)
+
+        encoder_blocks = []
+
+        for n in range(depth):
+            out_channels = num_channels_init * (2**n)
+            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,
+                )
+            )
+            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.
+    """
+
+    def __init__(
+        self,
+        conv_dim: int,
+        depth: int = 3,
+        num_channels_init: int = 64,
+        use_batch_norm: bool = True,
+        dropout: float = 0.0,
+    ) -> 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.
+        """
+        super().__init__()
+
+        upsampling = nn.Upsample(
+            scale_factor=2, mode="bilinear" if conv_dim == 2 else "trilinear"
+        )
+        in_channels = out_channels = num_channels_init * 2 ** (depth - 1)
+        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,
+        )
+
+        decoder_blocks = []
+        for n in range(depth):
+            decoder_blocks.append(upsampling)
+            in_channels = num_channels_init * 2 ** (depth - n)
+            out_channels = num_channels_init
+            decoder_blocks.append(
+                Conv_Block(
+                    conv_dim,
+                    in_channels=in_channels,
+                    out_channels=out_channels,
+                    intermediate_channel_multiplier=2,
+                    dropout_perc=dropout,
+                    activation="ReLU",
+                    use_batch_norm=use_batch_norm,
+                )
+            )
+
+        self.decoder_blocks = nn.ModuleList(decoder_blocks)
+
+    def forward(self, *features: List[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 = features[0]
+        skip_connections = features[1:][::-1]
+        x = self.bottleneck(x)
+        for i, module in enumerate(self.decoder_blocks):
+            x = module(x)
+            if isinstance(module, nn.Upsample):
+                x = torch.cat([x, skip_connections[i // 2]], axis=1)
+        return x
+
+
+class UNet(nn.Module):
+    """
+    UNet model.
+
+    Adapted for PyTorch from
+    https://github.com/juglab/n2v/blob/main/n2v/nets/unet_blocks.py.
+
+    Parameters
+    ----------
+    conv_dim : 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.
+    last_activation : Optional[Callable], optional
+        Activation function to use for the last layer, by default None.
+    """
+
+    def __init__(
+        self,
+        conv_dim: 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,
+        last_activation: Optional[Callable] = None,
+    ) -> None:
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        conv_dim : 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.
+        last_activation : Optional[Callable], optional
+            Activation function to use for the last layer, by default None.
+        """
+        super().__init__()
+
+        self.encoder = UnetEncoder(
+            conv_dim,
+            in_channels=in_channels,
+            depth=depth,
+            num_channels_init=num_channels_init,
+            use_batch_norm=use_batch_norm,
+            dropout=dropout,
+            pool_kernel=pool_kernel,
+        )
+
+        self.decoder = UnetDecoder(
+            conv_dim,
+            depth=depth,
+            num_channels_init=num_channels_init,
+            use_batch_norm=use_batch_norm,
+            dropout=dropout,
+        )
+        self.final_conv = getattr(nn, f"Conv{conv_dim}d")(
+            in_channels=num_channels_init,
+            out_channels=num_classes,
+            kernel_size=1,
+        )
+        self.last_activation = last_activation if last_activation else nn.Identity()
+
+    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.last_activation(x)
+        return x

careamics/prediction/__init__.py

@@ -0,0 +1,9 @@
+"""Prediction functions."""
+
+__all__ = [
+    "stitch_prediction",
+    "tta_backward",
+    "tta_forward",
+]
+
+from .prediction_utils import stitch_prediction, tta_backward, tta_forward

careamics/prediction/prediction_utils.py

@@ -0,0 +1,106 @@
+"""
+Prediction convenience functions.
+
+These functions are used during prediction.
+"""
+from typing import List
+
+import numpy as np
+import torch
+
+
+def stitch_prediction(
+    tiles: List[np.ndarray],
+    stitching_data: List,
+) -> np.ndarray:
+    """
+    Stitch tiles back together to form a full image.
+
+    Parameters
+    ----------
+    tiles : List[Tuple[np.ndarray, List[int]]]
+        Cropped tiles and their respective stitching coordinates.
+    stitching_data : List
+        List of coordinates obtained from
+        dataset.tiling.compute_crop_and_stitch_coords_1d.
+
+    Returns
+    -------
+    np.ndarray
+        Full image.
+    """
+    # Get whole sample shape
+    input_shape = stitching_data[0][0]
+    predicted_image = np.zeros(input_shape, dtype=np.float32)
+    for tile, (_, overlap_crop_coords, stitch_coords) in zip(tiles, stitching_data):
+        # Compute coordinates for cropping predicted tile
+        slices = tuple([slice(c[0], c[1]) for c in overlap_crop_coords])
+
+        # Crop predited tile according to overlap coordinates
+        cropped_tile = tile.squeeze()[slices]
+
+        # Insert cropped tile into predicted image using stitch coordinates
+        predicted_image[
+            (..., *[slice(c[0], c[1]) for c in stitch_coords])
+        ] = cropped_tile
+    return predicted_image
+
+
+def tta_forward(x: torch.Tensor) -> List[torch.Tensor]:
+    """
+    Augment 8-fold an array.
+
+    The augmentation is performed using all 90 deg rotations and their flipped version,
+    as well as the original image flipped.
+
+    Tensors should be of shape SC(Z)YX, with S and C potentially singleton dimensions.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Data to augment.
+
+    Returns
+    -------
+    List
+        Stack of augmented images.
+    """
+    x_aug = [
+        x,
+        torch.rot90(x, 1, dims=(2, 3)),
+        torch.rot90(x, 2, dims=(2, 3)),
+        torch.rot90(x, 3, dims=(2, 3)),
+    ]
+    x_aug_flip = x_aug.copy()
+    for x_ in x_aug:
+        x_aug_flip.append(torch.flip(x_, dims=(1, 3)))
+    return x_aug_flip
+
+
+def tta_backward(x_aug: List[torch.Tensor]) -> np.ndarray:
+    """
+    Invert `tta_forward` and average the 8 images.
+
+    The function takes a list of torch tensors and returns a numpy array.
+
+    Parameters
+    ----------
+    x_aug : List[torch.Tensor]
+        Stack of 8-fold augmented images.
+
+    Returns
+    -------
+    np.ndarray
+        Average of de-augmented x_aug.
+    """
+    x_deaug = [
+        x_aug[0].numpy(),
+        np.rot90(x_aug[1], -1, axes=(2, 3)),
+        np.rot90(x_aug[2], -2, axes=(2, 3)),
+        np.rot90(x_aug[3], -3, axes=(2, 3)),
+        np.flip(x_aug[4].numpy(), axis=(1, 3)),
+        np.rot90(np.flip(x_aug[5].numpy(), axis=(1, 3)), -1, axes=(2, 3)),
+        np.rot90(np.flip(x_aug[6].numpy(), axis=(1, 3)), -2, axes=(2, 3)),
+        np.rot90(np.flip(x_aug[7].numpy(), axis=(1, 3)), -3, axes=(2, 3)),
+    ]
+    return np.mean(x_deaug, 0)

careamics/utils/__init__.py

@@ -0,0 +1,20 @@
+"""Utils module."""
+
+
+__all__ = [
+    "denormalize",
+    "normalize",
+    "get_device",
+    "check_axes_validity",
+    "add_axes",
+    "check_tiling_validity",
+    "cwd",
+    "MetricTracker",
+]
+
+
+from .context import cwd
+from .metrics import MetricTracker
+from .normalization import denormalize, normalize
+from .torch_utils import get_device
+from .validators import add_axes, check_axes_validity, check_tiling_validity

careamics/utils/ascii_logo.txt

@@ -0,0 +1,9 @@
+   ......       ......     ........     ........                                   ....
+ -+++----+-   -+++--+++-  :+++---+++:  :+++-----                                   .--:
+.+++     .:   +++.  .+++. :+++   :+++  :+++         :------.   .---:----..:----.   :---    :----:     :----:.
+.+++         .+++.  .+++. :+++   -++=  :+++        +=....=+++  :+++-..=+++-..=++=  -+++  .+++-..++   +++-..=+.
+.+++         .++++++++++. :++++++++=.  :++++++:          .+++. :+++   :+++   -+++  -+++  :+++       .+++=.
+.+++         .+++.  .+++. :+++   -+++  :+++        :=++==++++. :+++   :+++   -+++  -+++  :+++        .-=+++=:
+.+++     ..  .+++.  .+++. :+++   :+++  :+++       .+++.  .+++. :+++   :+++   -+++  -+++  :+++   ..   ..  :+++.
+ -++=-::-+=  .+++.  .+++. :+++   :+++  :+++-::::   =++=--=+++. :+++   :+++   -+++  -+++   =++=:-+=   =+-:=++=
+   ......     ...    ...   ...    ...   ........     .... ...   ...    ...   ....  ....     ....      .....

careamics/utils/augment.py

@@ -0,0 +1,65 @@
+"""Augmentation module."""
+from typing import Tuple
+
+import numpy as np
+
+
+# TODO: unused?
+def _flip_and_rotate(
+    image: np.ndarray, rotate_state: int, flip_state: int
+) -> np.ndarray:
+    """
+    Apply the given number of 90 degrees rotations and flip to an array.
+
+    Parameters
+    ----------
+    image : np.ndarray
+        Array containing single image or patch, 2D or 3D.
+    rotate_state : int
+        Number of 90 degree rotations to apply.
+    flip_state : int
+        0 or 1, whether to flip the array or not.
+
+    Returns
+    -------
+    np.ndarray
+        Flipped and rotated array.
+    """
+    rotated = np.rot90(image, k=rotate_state, axes=(-2, -1))
+    flipped = np.flip(rotated, axis=-1) if flip_state == 1 else rotated
+    return flipped.copy()
+
+
+def augment_batch(
+    patch: np.ndarray,
+    original_image: np.ndarray,
+    mask: np.ndarray,
+    seed: int = 42,
+) -> Tuple[np.ndarray, ...]:
+    """
+    Apply augmentation function to patches and masks.
+
+    Parameters
+    ----------
+    patch : np.ndarray
+        Array containing single image or patch, 2D or 3D with masked pixels.
+    original_image : np.ndarray
+        Array containing original image or patch, 2D or 3D.
+    mask : np.ndarray
+        Array containing only masked pixels, 2D or 3D.
+    seed : int, optional
+        Seed for random number generator, controls the rotation and falipping.
+
+    Returns
+    -------
+    Tuple[np.ndarray, ...]
+        Tuple of augmented arrays.
+    """
+    rng = np.random.default_rng(seed=seed)
+    rotate_state = rng.integers(0, 4)
+    flip_state = rng.integers(0, 2)
+    return (
+        _flip_and_rotate(patch, rotate_state, flip_state),
+        _flip_and_rotate(original_image, rotate_state, flip_state),
+        _flip_and_rotate(mask, rotate_state, flip_state),
+    )

careamics/utils/context.py

@@ -0,0 +1,45 @@
+"""
+Context submodule.
+
+A convenience function to change the working directory in order to save data.
+"""
+import os
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Iterator, Union
+
+
+@contextmanager
+def cwd(path: Union[str, Path]) -> Iterator[None]:
+    """
+    Change the current working directory to the given path.
+
+    This method can be used to generate files in a specific directory, once out of the
+    context, the working directory is set back to the original one.
+
+    Parameters
+    ----------
+    path : Union[str,Path]
+        New working directory path.
+
+    Returns
+    -------
+    Iterator[None]
+        None values.
+
+    Examples
+    --------
+    >>> with cwd(path):
+    ...     pass
+    """
+    path = Path(path)
+
+    if not path.exists():
+        path.mkdir(parents=True, exist_ok=True)
+
+    old_pwd = Path(".").absolute()
+    os.chdir(path)
+    try:
+        yield
+    finally:
+        os.chdir(old_pwd)