careamics 0.1.0rc1__py3-none-any.whl → 0.1.0rc3__py3-none-any.whl

careamics/models/unet.py

@@ -3,12 +3,14 @@ UNet model.
 
 A UNet encoder, decoder and complete model.
 """
-from typing import Callable, List, Optional
+from typing import Any, List, Union
 
 import torch
 import torch.nn as nn
 
-from .layers import Conv_Block
+from ..config.support import SupportedActivation
+from .activation import get_activation
+from .layers import Conv_Block, MaxBlurPool
 
 
 class UnetEncoder(nn.Module):
@@ -42,6 +44,7 @@ class UnetEncoder(nn.Module):
         use_batch_norm: bool = True,
         dropout: float = 0.0,
         pool_kernel: int = 2,
+        n2v2: bool = False,
     ) -> None:
         """
         Constructor.
@@ -65,7 +68,11 @@ class UnetEncoder(nn.Module):
         """
         super().__init__()
 
-        self.pooling = getattr(nn, f"MaxPool{conv_dim}d")(kernel_size=pool_kernel)
+        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 = []
 
@@ -82,7 +89,6 @@ class UnetEncoder(nn.Module):
                 )
             )
             encoder_blocks.append(self.pooling)
-
         self.encoder_blocks = nn.ModuleList(encoder_blocks)
 
     def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
@@ -134,6 +140,7 @@ class UnetDecoder(nn.Module):
         num_channels_init: int = 64,
         use_batch_norm: bool = True,
         dropout: float = 0.0,
+        n2v2: bool = False,
     ) -> None:
         """
         Constructor.
@@ -157,6 +164,9 @@ class UnetDecoder(nn.Module):
             scale_factor=2, mode="bilinear" if conv_dim == 2 else "trilinear"
         )
         in_channels = out_channels = num_channels_init * 2 ** (depth - 1)
+
+        self.n2v2 = n2v2
+
         self.bottleneck = Conv_Block(
             conv_dim,
             in_channels=in_channels,
@@ -169,12 +179,18 @@ class UnetDecoder(nn.Module):
         decoder_blocks = []
         for n in range(depth):
             decoder_blocks.append(upsampling)
-            in_channels = num_channels_init * 2 ** (depth - n)
-            out_channels = num_channels_init
+            in_channels = (
+                num_channels_init ** (depth - n)
+                if (self.n2v2 and n == depth - 1)
+                else num_channels_init * 2 ** (depth - n)
+            )
+            out_channels = in_channels // 2
             decoder_blocks.append(
                 Conv_Block(
                     conv_dim,
-                    in_channels=in_channels,
+                    in_channels=in_channels + in_channels // 2
+                    if n > 0
+                    else in_channels,
                     out_channels=out_channels,
                     intermediate_channel_multiplier=2,
                     dropout_perc=dropout,
@@ -200,13 +216,19 @@ class UnetDecoder(nn.Module):
         torch.Tensor
             Output of the decoder.
         """
-        x = features[0]
-        skip_connections = features[1:][::-1]
+        x: torch.Tensor = features[0]
+        skip_connections: torch.Tensor = 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)
+                if self.n2v2:
+                    if x.shape != skip_connections[-1].shape:
+                        x = torch.cat([x, skip_connections[i // 2]], axis=1)
+                else:
+                    x = torch.cat([x, skip_connections[i // 2]], axis=1)
         return x
 
 
@@ -214,12 +236,12 @@ class UNet(nn.Module):
     """
     UNet model.
 
-    Adapted for PyTorch from
+    Adapted for PyTorch from:
     https://github.com/juglab/n2v/blob/main/n2v/nets/unet_blocks.py.
 
     Parameters
     ----------
-    conv_dim : int
+    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.
@@ -241,7 +263,7 @@ class UNet(nn.Module):
 
     def __init__(
         self,
-        conv_dim: int,
+        conv_dims: int,
         num_classes: int = 1,
         in_channels: int = 1,
         depth: int = 3,
@@ -249,14 +271,16 @@ class UNet(nn.Module):
         use_batch_norm: bool = True,
         dropout: float = 0.0,
         pool_kernel: int = 2,
-        last_activation: Optional[Callable] = None,
+        final_activation: Union[SupportedActivation, str] = SupportedActivation.NONE,
+        n2v2: bool = False,
+        **kwargs: Any,
     ) -> None:
         """
         Constructor.
 
         Parameters
         ----------
-        conv_dim : int
+        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.
@@ -278,28 +302,30 @@ class UNet(nn.Module):
         super().__init__()
 
         self.encoder = UnetEncoder(
-            conv_dim,
+            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,
         )
 
         self.decoder = UnetDecoder(
-            conv_dim,
+            conv_dims,
             depth=depth,
             num_channels_init=num_channels_init,
             use_batch_norm=use_batch_norm,
             dropout=dropout,
+            n2v2=n2v2,
         )
-        self.final_conv = getattr(nn, f"Conv{conv_dim}d")(
+        self.final_conv = getattr(nn, f"Conv{conv_dims}d")(
             in_channels=num_channels_init,
             out_channels=num_classes,
             kernel_size=1,
         )
-        self.last_activation = last_activation if last_activation else nn.Identity()
+        self.final_activation = get_activation(final_activation)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
@@ -318,5 +344,5 @@ class UNet(nn.Module):
         encoder_features = self.encoder(x)
         x = self.decoder(*encoder_features)
         x = self.final_conv(x)
-        x = self.last_activation(x)
+        x = self.final_activation(x)
         return x

careamics/prediction/__init__.py

@@ -2,8 +2,6 @@
 
 __all__ = [
     "stitch_prediction",
-    "tta_backward",
-    "tta_forward",
 ]
 
-from .prediction_utils import stitch_prediction, tta_backward, tta_forward
+from .stitch_prediction import stitch_prediction

careamics/prediction/stitch_prediction.py

@@ -0,0 +1,73 @@
+"""
+Prediction convenience functions.
+
+These functions are used during prediction.
+"""
+from typing import List
+
+import numpy as np
+import torch
+
+
+def stitch_prediction(
+    tiles: List[torch.Tensor],
+    stitching_data: List[List[torch.Tensor]],
+) -> torch.Tensor:
+    """
+    Stitch tiles back together to form a full image.
+
+    Parameters
+    ----------
+    tiles : List[torch.Tensor]
+        Cropped tiles and their respective stitching coordinates.
+    stitching_coords : List
+        List of information and coordinates obtained from
+        `dataset.tiled_patching.extract_tiles`.
+
+    Returns
+    -------
+    np.ndarray
+        Full image.
+    """
+    # retrieve whole array size, there is two cases to consider:
+    # 1. the tiles are stored in a list
+    # 2. the tiles are stored in a list with batches along the first dim
+    if tiles[0].shape[0] > 1:
+        input_shape = np.array(
+            [el.numpy() for el in stitching_data[0][0][0]], dtype=int
+        ).squeeze()
+    else:
+        input_shape = np.array(
+            [el.numpy() for el in stitching_data[0][0]], dtype=int
+        ).squeeze()
+
+    # TODO should use torch.zeros instead of np.zeros
+    predicted_image = torch.Tensor(np.zeros(input_shape, dtype=np.float32))
+
+    for tile_batch, (_, overlap_crop_coords_batch, stitch_coords_batch) in zip(
+        tiles, stitching_data
+    ):
+        for batch_idx in range(tile_batch.shape[0]):
+            # Compute coordinates for cropping predicted tile
+            slices = tuple(
+                [
+                    slice(c[0][batch_idx], c[1][batch_idx])
+                    for c in overlap_crop_coords_batch
+                ]
+            )
+
+            # Crop predited tile according to overlap coordinates
+            cropped_tile = tile_batch[batch_idx].squeeze()[slices]
+
+            # Insert cropped tile into predicted image using stitch coordinates
+            predicted_image[
+                (
+                    ...,
+                    *[
+                        slice(c[0][batch_idx], c[1][batch_idx])
+                        for c in stitch_coords_batch
+                    ],
+                )
+            ] = cropped_tile.to(torch.float32)
+
+    return predicted_image

careamics/transforms/__init__.py

@@ -0,0 +1,41 @@
+"""Transforms that are used to augment the data."""
+
+__all__ = [
+    "get_all_transforms",
+    "N2VManipulate",
+    "NDFlip",
+    "XYRandomRotate90",
+    "ImageRestorationTTA",
+    "Denormalize",
+    "Normalize",
+]
+
+
+from .n2v_manipulate import N2VManipulate
+from .nd_flip import NDFlip
+from .normalize import Denormalize, Normalize
+from .tta import ImageRestorationTTA
+from .xy_random_rotate90 import XYRandomRotate90
+
+ALL_TRANSFORMS = {
+    "Normalize": Normalize,
+    "N2VManipulate": N2VManipulate,
+    "NDFlip": NDFlip,
+    "XYRandomRotate90": XYRandomRotate90,
+}
+
+
+def get_all_transforms() -> dict:
+    """Return all the transforms accepted by CAREamics.
+
+    Note that while CAREamics accepts any `Compose` transforms from Albumentations (see
+    https://albumentations.ai/), only a few transformations are explicitely supported
+    (see `SupportedTransform`).
+
+    Returns
+    -------
+    dict
+        A dictionary with all the transforms accepted by CAREamics, where the keys are
+        the transform names and the values are the transform classes.
+    """
+    return ALL_TRANSFORMS

careamics/transforms/n2v_manipulate.py

@@ -0,0 +1,113 @@
+from typing import Any, Literal, Optional, Tuple
+
+import numpy as np
+from albumentations import ImageOnlyTransform
+
+from careamics.config.support import SupportedPixelManipulation, SupportedStructAxis
+
+from .pixel_manipulation import median_manipulate, uniform_manipulate
+from .struct_mask_parameters import StructMaskParameters
+
+
+class N2VManipulate(ImageOnlyTransform):
+    """
+    Default augmentation for the N2V model.
+
+    This transform expects (Z)YXC dimensions.
+
+    Parameters
+    ----------
+    mask_pixel_percentage : float
+        Approximate percentage of pixels to be masked.
+    roi_size : int
+        Size of the ROI the new pixel value is sampled from, by default 11.
+    """
+
+    def __init__(
+        self,
+        roi_size: int = 11,
+        masked_pixel_percentage: float = 0.2,
+        strategy: Literal[
+            "uniform", "median"
+        ] = SupportedPixelManipulation.UNIFORM.value,
+        remove_center: bool = True,
+        struct_mask_axis: Literal["horizontal", "vertical", "none"] = "none",
+        struct_mask_span: int = 5,
+    ):
+        """Constructor.
+
+        Parameters
+        ----------
+        roi_size : int, optional
+            Size of the replacement area, by default 11
+        masked_pixel_percentage : float, optional
+            Percentage of pixels to mask, by default 0.2
+        strategy : Literal[ "uniform", "median" ], optional
+            Replaccement strategy, uniform or median, by default uniform
+        remove_center : bool, optional
+            Whether to remove central pixel from patch, by default True
+        struct_mask_axis : Literal["horizontal", "vertical", "none"], optional
+            StructN2V mask axis, by default "none"
+        struct_mask_span : int, optional
+            StructN2V mask span, by default 5
+        """
+        super().__init__(p=1)
+        self.masked_pixel_percentage = masked_pixel_percentage
+        self.roi_size = roi_size
+        self.strategy = strategy
+        self.remove_center = remove_center
+
+        if struct_mask_axis == SupportedStructAxis.NONE:
+            self.struct_mask: Optional[StructMaskParameters] = None
+        else:
+            self.struct_mask = StructMaskParameters(
+                axis=0 if struct_mask_axis == SupportedStructAxis.HORIZONTAL else 1,
+                span=struct_mask_span,
+            )
+
+    def apply(
+        self, patch: np.ndarray, **kwargs: Any
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Apply the transform to the image.
+
+        Parameters
+        ----------
+        image : np.ndarray
+            Image or image patch, 2D or 3D, shape (y, x, c) or (z, y, x, c).
+        """
+        masked = np.zeros_like(patch)
+        mask = np.zeros_like(patch)
+        if self.strategy == SupportedPixelManipulation.UNIFORM:
+            # Iterate over the channels to apply manipulation separately
+            for c in range(patch.shape[-1]):
+                masked[..., c], mask[..., c] = uniform_manipulate(
+                    patch=patch[..., c],
+                    mask_pixel_percentage=self.masked_pixel_percentage,
+                    subpatch_size=self.roi_size,
+                    remove_center=self.remove_center,
+                    struct_params=self.struct_mask,
+                )
+        elif self.strategy == SupportedPixelManipulation.MEDIAN:
+            # Iterate over the channels to apply manipulation separately
+            for c in range(patch.shape[-1]):
+                masked[..., c], mask[..., c] = median_manipulate(
+                    patch=patch[..., c],
+                    mask_pixel_percentage=self.masked_pixel_percentage,
+                    subpatch_size=self.roi_size,
+                    struct_params=self.struct_mask,
+                )
+        else:
+            raise ValueError(f"Unknown masking strategy ({self.strategy}).")
+
+        # TODO why return patch?
+        return masked, patch, mask
+
+    def get_transform_init_args_names(self) -> Tuple[str, ...]:
+        """Get the transform parameters.
+
+        Returns
+        -------
+        Tuple[str, ...]
+            Transform parameters.
+        """
+        return ("roi_size", "masked_pixel_percentage", "strategy", "struct_mask")

careamics/transforms/nd_flip.py

@@ -0,0 +1,93 @@
+from typing import Any, Dict, Tuple
+
+import numpy as np
+from albumentations import DualTransform
+
+
+class NDFlip(DualTransform):
+    """Flip ND arrays on a single axis.
+
+    This transform ignores singleton axes and randomly flips one of the other
+    axes, to the exception of the first and last axes (sample and channels).
+
+    This transform expects (Z)YXC dimensions.
+    """
+
+    def __init__(self, p: float = 0.5, is_3D: bool = False, flip_z: bool = True):
+        """Constructor.
+
+        Parameters
+        ----------
+        p : float, optional
+            Probability to apply the transform, by default 0.5
+        is_3D : bool, optional
+            Whether the data is 3D, by default False
+        flip_z : bool, optional
+            Whether to flip Z dimension, by default True
+        """
+        super().__init__(p=p)
+
+        self.is_3D = is_3D
+        self.flip_z = flip_z
+
+        # "flippable" axes
+        if is_3D:
+            self.axis_indices = [0, 1, 2] if flip_z else [1, 2]
+        else:
+            self.axis_indices = [0, 1]
+
+    def get_params(self, **kwargs: Any) -> Dict[str, int]:
+        """Get the transform parameters.
+
+        Returns
+        -------
+        Dict[str, int]
+            Transform parameters.
+        """
+        return {"flip_axis": np.random.choice(self.axis_indices)}
+
+    def apply(self, patch: np.ndarray, flip_axis: int, **kwargs: Any) -> np.ndarray:
+        """Apply the transform to the image.
+
+        Parameters
+        ----------
+        patch : np.ndarray
+            Image or image patch, 2D or 3D, shape (y, x, c) or (z, y, x, c).
+        flip_axis : int
+            Axis along which to flip the patch.
+        """
+        if len(patch.shape) == 3 and self.is_3D:
+            raise ValueError(
+                "Incompatible patch shape and dimensionality. ZYXC patch shape "
+                "expected, but got YXC shape."
+            )
+
+        return np.ascontiguousarray(np.flip(patch, axis=flip_axis))
+
+    def apply_to_mask(
+        self, mask: np.ndarray, flip_axis: int, **kwargs: Any
+    ) -> np.ndarray:
+        """Apply the transform to the mask.
+
+        Parameters
+        ----------
+        mask : np.ndarray
+            Mask or mask patch, 2D or 3D, shape (y, x, c) or (z, y, x, c).
+        """
+        if len(mask.shape) == 3 and self.is_3D:
+            raise ValueError(
+                "Incompatible mask shape and dimensionality. ZYXC patch shape "
+                "expected, but got YXC shape."
+            )
+
+        return np.ascontiguousarray(np.flip(mask, axis=flip_axis))
+
+    def get_transform_init_args_names(self, **kwargs: Any) -> Tuple[str, ...]:
+        """Get the transform arguments names.
+
+        Returns
+        -------
+        Tuple[str, ...]
+            Transform arguments names.
+        """
+        return ("is_3D", "flip_z")

careamics/transforms/normalize.py

@@ -0,0 +1,109 @@
+from typing import Any
+
+import numpy as np
+from albumentations import DualTransform
+
+
+class Normalize(DualTransform):
+    """
+    Normalize an image or image patch.
+
+    Normalization is a zero mean and unit variance. This transform expects (Z)YXC
+    dimensions.
+
+    Not that an epsilon value of 1e-6 is added to the standard deviation to avoid
+    division by zero and that it returns a float32 image.
+
+    Attributes
+    ----------
+    mean : float
+        Mean value.
+    std : float
+        Standard deviation value.
+    eps : float
+        Epsilon value to avoid division by zero.
+    """
+
+    def __init__(
+        self,
+        mean: float,
+        std: float,
+    ):
+        super().__init__(always_apply=True, p=1)
+
+        self.mean = mean
+        self.std = std
+        self.eps = 1e-6
+
+    def apply(self, patch: np.ndarray, **kwargs: Any) -> np.ndarray:
+        """
+        Apply the transform to the image.
+
+        Parameters
+        ----------
+        patch : np.ndarray
+            Image or image patch, 2D or 3D, shape (y, x, c) or (z, y, x, c).
+
+        Returns
+        -------
+        np.ndarray
+            Normalized image or image patch.
+        """
+        return ((patch - self.mean) / (self.std + self.eps)).astype(np.float32)
+
+    def apply_to_mask(self, mask: np.ndarray, **kwargs: Any) -> np.ndarray:
+        """
+        Apply the transform to the mask.
+
+        The mask is returned as is.
+
+        Parameters
+        ----------
+        mask : np.ndarray
+            Mask or mask patch, 2D or 3D, shape (y, x, c) or (z, y, x, c).
+        """
+        return mask
+
+
+class Denormalize(DualTransform):
+    """
+    Denormalize an image or image patch.
+
+    Denormalization is performed expecting a zero mean and unit variance input. This
+    transform expects (Z)YXC dimensions.
+
+    Not that an epsilon value of 1e-6 is added to the standard deviation to avoid
+    division by zero during the normalization step, which is taken into account during
+    denormalization.
+
+    Attributes
+    ----------
+    mean : float
+        Mean value.
+    std : float
+        Standard deviation value.
+    eps : float
+        Epsilon value to avoid division by zero.
+    """
+
+    def __init__(
+        self,
+        mean: float,
+        std: float,
+    ):
+        super().__init__(always_apply=True, p=1)
+
+        self.mean = mean
+        self.std = std
+        self.eps = 1e-6
+
+    def apply(self, patch: np.ndarray, **kwargs: Any) -> np.ndarray:
+        """
+        Apply the transform to the image.
+
+        Parameters
+        ----------
+        patch : np.ndarray
+            Image or image patch, 2D or 3D, shape (y, x, c) or (z, y, x, c).
+        """
+        return patch * (self.std + self.eps) + self.mean