careamics 0.1.0rc6__py3-none-any.whl → 0.1.0rc8__py3-none-any.whl

careamics/models/lvae/utils.py

@@ -0,0 +1,395 @@
+"""
+Script for utility functions needed by the LVAE model.
+"""
+
+from typing import Iterable
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torchvision.transforms.functional as F
+from torch.distributions.normal import Normal
+
+
+def torch_nanmean(inp):
+    return torch.mean(inp[~inp.isnan()])
+
+
+def compute_batch_mean(x):
+    N = len(x)
+    return x.view(N, -1).mean(dim=1)
+
+
+def power_of_2(self, x):
+    assert isinstance(x, int)
+    if x == 1:
+        return True
+    if x == 0:
+        # happens with validation
+        return False
+    if x % 2 == 1:
+        return False
+    return self.power_of_2(x // 2)
+
+
+class Enum:
+    @classmethod
+    def name(cls, enum_type):
+        for key, value in cls.__dict__.items():
+            if enum_type == value:
+                return key
+
+    @classmethod
+    def contains(cls, enum_type):
+        for key, value in cls.__dict__.items():
+            if enum_type == value:
+                return True
+        return False
+
+    @classmethod
+    def from_name(cls, enum_type_str):
+        for key, value in cls.__dict__.items():
+            if key == enum_type_str:
+                return value
+        assert f"{cls.__name__}:{enum_type_str} doesnot exist."
+
+
+class LossType(Enum):
+    Elbo = 0
+    ElboWithCritic = 1
+    ElboMixedReconstruction = 2
+    MSE = 3
+    ElboWithNbrConsistency = 4
+    ElboSemiSupMixedReconstruction = 5
+    ElboCL = 6
+    ElboRestrictedReconstruction = 7
+    DenoiSplitMuSplit = 8
+
+
+class ModelType(Enum):
+    LadderVae = 3
+    LadderVaeTwinDecoder = 4
+    LadderVAECritic = 5
+    # Separate vampprior: two optimizers
+    LadderVaeSepVampprior = 6
+    # one encoder for mixed input, two for separate inputs.
+    LadderVaeSepEncoder = 7
+    LadderVAEMultiTarget = 8
+    LadderVaeSepEncoderSingleOptim = 9
+    UNet = 10
+    BraveNet = 11
+    LadderVaeStitch = 12
+    LadderVaeSemiSupervised = 13
+    LadderVaeStitch2Stage = 14  # Note that previously trained models will have issue.
+    # since earlier, LadderVaeStitch2Stage = 13, LadderVaeSemiSupervised = 14
+    LadderVaeMixedRecons = 15
+    LadderVaeCL = 16
+    LadderVaeTwoDataSet = (
+        17  # on one subdset, apply disentanglement, on other apply reconstruction
+    )
+    LadderVaeTwoDatasetMultiBranch = 18
+    LadderVaeTwoDatasetMultiOptim = 19
+    LVaeDeepEncoderIntensityAug = 20
+    AutoRegresiveLadderVAE = 21
+    LadderVAEInterleavedOptimization = 22
+    Denoiser = 23
+    DenoiserSplitter = 24
+    SplitterDenoiser = 25
+    LadderVAERestrictedReconstruction = 26
+    LadderVAETwoDataSetRestRecon = 27
+    LadderVAETwoDataSetFinetuning = 28
+
+
+def _pad_crop_img(x, size, mode) -> torch.Tensor:
+    """Pads or crops a tensor.
+    Pads or crops a tensor of shape (batch, channels, h, w) to new height
+    and width given by a tuple.
+    Args:
+        x (torch.Tensor): Input image
+        size (list or tuple): Desired size (height, width)
+        mode (str): Mode, either 'pad' or 'crop'
+    Returns:
+        The padded or cropped tensor
+    """
+    assert x.dim() == 4 and len(size) == 2
+    size = tuple(size)
+    x_size = x.size()[2:4]
+    if mode == "pad":
+        cond = x_size[0] > size[0] or x_size[1] > size[1]
+    elif mode == "crop":
+        cond = x_size[0] < size[0] or x_size[1] < size[1]
+    else:
+        raise ValueError(f"invalid mode '{mode}'")
+    if cond:
+        raise ValueError(f"trying to {mode} from size {x_size} to size {size}")
+    dr, dc = (abs(x_size[0] - size[0]), abs(x_size[1] - size[1]))
+    dr1, dr2 = dr // 2, dr - (dr // 2)
+    dc1, dc2 = dc // 2, dc - (dc // 2)
+    if mode == "pad":
+        return nn.functional.pad(x, [dc1, dc2, dr1, dr2, 0, 0, 0, 0])
+    elif mode == "crop":
+        return x[:, :, dr1 : x_size[0] - dr2, dc1 : x_size[1] - dc2]
+
+
+def pad_img_tensor(x, size) -> torch.Tensor:
+    """Pads a tensor.
+    Pads a tensor of shape (batch, channels, h, w) to a desired height and width.
+    Args:
+        x (torch.Tensor): Input image
+        size (list or tuple): Desired size (height, width)
+
+    Returns
+    -------
+        The padded tensor
+    """
+    return _pad_crop_img(x, size, "pad")
+
+
+def crop_img_tensor(x, size) -> torch.Tensor:
+    """Crops a tensor.
+    Crops a tensor of shape (batch, channels, h, w) to a desired height and width
+    given by a tuple.
+    Args:
+        x (torch.Tensor): Input image
+        size (list or tuple): Desired size (height, width)
+
+    Returns
+    -------
+        The cropped tensor
+    """
+    return _pad_crop_img(x, size, "crop")
+
+
+class StableExponential:
+    """
+    Class that redefines the definition of exp() to increase numerical stability.
+    Naturally, also the definition of log() must change accordingly.
+    However, it is worth noting that the two operations remain one the inverse of the other,
+    meaning that x = log(exp(x)) and x = exp(log(x)) are always true.
+
+    Definition:
+        exp(x) = {
+            exp(x) if x<=0
+            x+1    if x>0
+        }
+
+        log(x) = {
+            x        if x<=0
+            log(1+x) if x>0
+        }
+
+    NOTE 1:
+        Within the class everything is done on the tensor given as input to the constructor.
+        Therefore, when exp() is called, self._tensor.exp() is computed.
+        When log() is called, torch.log(self._tensor.exp()) is computed instead.
+
+    NOTE 2:
+        Given the output from exp(), torch.log() or the log() method of the class give identical results.
+    """
+
+    def __init__(self, tensor):
+        self._raw_tensor = tensor
+        posneg_dic = self.posneg_separation(self._raw_tensor)
+        self.pos_f, self.neg_f = posneg_dic["filter"]
+        self.pos_data, self.neg_data = posneg_dic["value"]
+
+    def posneg_separation(self, tensor):
+        pos = tensor > 0
+        pos_tensor = torch.clip(tensor, min=0)
+
+        neg = tensor <= 0
+        neg_tensor = torch.clip(tensor, max=0)
+
+        return {"filter": [pos, neg], "value": [pos_tensor, neg_tensor]}
+
+    def exp(self):
+        return torch.exp(self.neg_data) * self.neg_f + (1 + self.pos_data) * self.pos_f
+
+    def log(self):
+        return self.neg_data * self.neg_f + torch.log(1 + self.pos_data) * self.pos_f
+
+
+class StableLogVar:
+    """
+    Class that provides a numerically stable implementation of Log-Variance.
+    Specifically, it uses the exp() and log() formulas defined in `StableExponential` class.
+    """
+
+    def __init__(
+        self, logvar: torch.Tensor, enable_stable: bool = True, var_eps: float = 1e-6
+    ):
+        """
+        Contructor.
+
+        Parameters
+        ----------
+        logvar: torch.Tensor
+            The input (true) logvar vector, to be converted in the Stable version.
+        enable_stable: bool, optional
+            Whether to compute the stable version of log-variance. Default is `True`.
+        var_eps: float, optional
+            The minimum value attainable by the variance. Default is `1e-6`.
+        """
+        self._lv = logvar
+        self._enable_stable = enable_stable
+        self._eps = var_eps
+
+    def get(self) -> torch.Tensor:
+        if self._enable_stable is False:
+            return self._lv
+
+        return torch.log(self.get_var())
+
+    def get_var(self) -> torch.Tensor:
+        """
+        Get Variance from Log-Variance.
+        """
+        if self._enable_stable is False:
+            return torch.exp(self._lv)
+        return StableExponential(self._lv).exp() + self._eps
+
+    def get_std(self) -> torch.Tensor:
+        return torch.sqrt(self.get_var())
+
+    def centercrop_to_size(self, size: Iterable[int]) -> None:
+        """
+        Centercrop the log-variance tensor to the desired size.
+
+        Parameters
+        ----------
+        size: torch.Tensor
+            The desired size of the log-variance tensor.
+        """
+        if self._lv.shape[-1] == size:
+            return
+
+        diff = self._lv.shape[-1] - size
+        assert diff > 0 and diff % 2 == 0
+        self._lv = F.center_crop(self._lv, (size, size))
+
+
+class StableMean:
+
+    def __init__(self, mean):
+        self._mean = mean
+
+    def get(self) -> torch.Tensor:
+        return self._mean
+
+    def centercrop_to_size(self, size: Iterable[int]) -> None:
+        """
+        Centercrop the mean tensor to the desired size.
+
+        Parameters
+        ----------
+        size: torch.Tensor
+            The desired size of the log-variance tensor.
+        """
+        if self._mean.shape[-1] == size:
+            return
+
+        diff = self._mean.shape[-1] - size
+        assert diff > 0 and diff % 2 == 0
+        self._mean = F.center_crop(self._mean, (size, size))
+
+
+def allow_numpy(func):
+    """
+    All optional arguements are passed as is. positional arguments are checked. if they are numpy array,
+    they are converted to torch Tensor.
+    """
+
+    def numpy_wrapper(*args, **kwargs):
+        new_args = []
+        for arg in args:
+            if isinstance(arg, np.ndarray):
+                arg = torch.Tensor(arg)
+            new_args.append(arg)
+        new_args = tuple(new_args)
+
+        output = func(*new_args, **kwargs)
+        return output
+
+    return numpy_wrapper
+
+
+class Interpolate(nn.Module):
+    """Wrapper for torch.nn.functional.interpolate."""
+
+    def __init__(self, size=None, scale=None, mode="bilinear", align_corners=False):
+        super().__init__()
+        assert (size is None) == (scale is not None)
+        self.size = size
+        self.scale = scale
+        self.mode = mode
+        self.align_corners = align_corners
+
+    def forward(self, x):
+        out = F.interpolate(
+            x,
+            size=self.size,
+            scale_factor=self.scale,
+            mode=self.mode,
+            align_corners=self.align_corners,
+        )
+        return out
+
+
+def kl_normal_mc(z, p_mulv, q_mulv):
+    """
+    One-sample estimation of element-wise KL between two diagonal
+    multivariate normal distributions. Any number of dimensions,
+    broadcasting supported (be careful).
+    :param z:
+    :param p_mulv:
+    :param q_mulv:
+    :return:
+    """
+    assert isinstance(p_mulv, tuple)
+    assert isinstance(q_mulv, tuple)
+    p_mu, p_lv = p_mulv
+    q_mu, q_lv = q_mulv
+
+    p_std = p_lv.get_std()
+    q_std = q_lv.get_std()
+
+    p_distrib = Normal(p_mu.get(), p_std)
+    q_distrib = Normal(q_mu.get(), q_std)
+    return q_distrib.log_prob(z) - p_distrib.log_prob(z)
+
+
+def free_bits_kl(
+    kl: torch.Tensor, free_bits: float, batch_average: bool = False, eps: float = 1e-6
+) -> torch.Tensor:
+    """
+    Computes free-bits version of KL divergence.
+    Ensures that the KL doesn't go to zero for any latent dimension.
+    Hence, it contributes to use latent variables more efficiently,
+    leading to better representation learning.
+
+    NOTE:
+    Takes in the KL with shape (batch size, layers), returns the KL with
+    free bits (for optimization) with shape (layers,), which is the average
+    free-bits KL per layer in the current batch.
+    If batch_average is False (default), the free bits are per layer and
+    per batch element. Otherwise, the free bits are still per layer, but
+    are assigned on average to the whole batch. In both cases, the batch
+    average is returned, so it's simply a matter of doing mean(clamp(KL))
+    or clamp(mean(KL)).
+
+    Args:
+        kl (torch.Tensor)
+        free_bits (float)
+        batch_average (bool, optional))
+        eps (float, optional)
+
+    Returns
+    -------
+        The KL with free bits
+    """
+    assert kl.dim() == 2
+    if free_bits < eps:
+        return kl.mean(0)
+    if batch_average:
+        return kl.mean(0).clamp(min=free_bits)
+    return kl.clamp(min=free_bits).mean(0)

careamics/prediction_utils/__init__.py

@@ -0,0 +1,10 @@
+"""Package to house various prediction utilies."""
+
+__all__ = [
+    "stitch_prediction",
+    "stitch_prediction_single",
+    "convert_outputs",
+]
+
+from .prediction_outputs import convert_outputs
+from .stitch_prediction import stitch_prediction, stitch_prediction_single

careamics/prediction_utils/prediction_outputs.py

@@ -0,0 +1,137 @@
+"""Module containing functions to convert prediction outputs to desired form."""
+
+from typing import Any, List, Literal, Tuple, Union, overload
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ..config.tile_information import TileInformation
+from .stitch_prediction import stitch_prediction
+
+
+def convert_outputs(
+    predictions: List[Any], tiled: bool
+) -> Union[List[NDArray], NDArray]:
+    """
+    Convert the Lightning trainer outputs to the desired form.
+
+    This method allows stitching back together tiled predictions.
+
+    Parameters
+    ----------
+    predictions : list
+        Predictions that are output from `Trainer.predict`.
+    tiled : bool
+        Whether the predictions are tiled.
+
+    Returns
+    -------
+    list of numpy.ndarray or numpy.ndarray
+        List of arrays with the axes SC(Z)YX. If there is only 1 output it will not
+        be in a list.
+    """
+    if len(predictions) == 0:
+        return predictions
+
+    # this layout is to stop mypy complaining
+    if tiled:
+        predictions_comb = combine_batches(predictions, tiled)
+        predictions_output = stitch_prediction(*predictions_comb)
+    else:
+        predictions_output = combine_batches(predictions, tiled)
+
+    return predictions_output
+
+
+# for mypy
+@overload
+def combine_batches(  # numpydoc ignore=GL08
+    predictions: List[Any], tiled: Literal[True]
+) -> Tuple[List[NDArray], List[TileInformation]]: ...
+
+
+# for mypy
+@overload
+def combine_batches(  # numpydoc ignore=GL08
+    predictions: List[Any], tiled: Literal[False]
+) -> List[NDArray]: ...
+
+
+# for mypy
+@overload
+def combine_batches(  # numpydoc ignore=GL08
+    predictions: List[Any], tiled: Union[bool, Literal[True], Literal[False]]
+) -> Union[List[NDArray], Tuple[List[NDArray], List[TileInformation]]]: ...
+
+
+def combine_batches(
+    predictions: List[Any], tiled: bool
+) -> Union[List[NDArray], Tuple[List[NDArray], List[TileInformation]]]:
+    """
+    If predictions are in batches, they will be combined.
+
+    Parameters
+    ----------
+    predictions : list
+        Predictions that are output from `Trainer.predict`.
+    tiled : bool
+        Whether the predictions are tiled.
+
+    Returns
+    -------
+    (list of numpy.ndarray) or tuple of (list of numpy.ndarray, list of TileInformation)
+        Combined batches.
+    """
+    if tiled:
+        return _combine_tiled_batches(predictions)
+    else:
+        return _combine_array_batches(predictions)
+
+
+def _combine_tiled_batches(
+    predictions: List[Tuple[NDArray, List[TileInformation]]]
+) -> Tuple[List[NDArray], List[TileInformation]]:
+    """
+    Combine batches from tiled output.
+
+    Parameters
+    ----------
+    predictions : list of (numpy.ndarray, list of TileInformation)
+        Predictions that are output from `Trainer.predict`. For tiled batches, this is
+        a list of tuples. The first element of the tuples is the prediction output of
+        tiles with dimension (B, C, (Z), Y, X), where B is batch size. The second
+        element of the tuples is a list of TileInformation objects of length B.
+
+    Returns
+    -------
+    tuple of (list of numpy.ndarray, list of TileInformation)
+        Combined batches.
+    """
+    # turn list of lists into single list
+    tile_infos = [
+        tile_info for _, tile_info_list in predictions for tile_info in tile_info_list
+    ]
+    prediction_tiles: List[NDArray] = _combine_array_batches(
+        [preds for preds, _ in predictions]
+    )
+    return prediction_tiles, tile_infos
+
+
+def _combine_array_batches(predictions: List[NDArray]) -> List[NDArray]:
+    """
+    Combine batches of arrays.
+
+    Parameters
+    ----------
+    predictions : list
+        Prediction arrays that are output from `Trainer.predict`. A list of arrays that
+        have dimensions (B, C, (Z), Y, X), where B is batch size.
+
+    Returns
+    -------
+    list of numpy.ndarray
+        A list of arrays with dimensions (1, C, (Z), Y, X).
+    """
+    prediction_concat: NDArray = np.concatenate(predictions, axis=0)
+    prediction_split = np.split(prediction_concat, prediction_concat.shape[0], axis=0)
+    return prediction_split

careamics/prediction_utils/stitch_prediction.py

@@ -0,0 +1,103 @@
+"""Prediction utility functions."""
+
+import builtins
+from typing import List, Union
+
+import numpy as np
+from numpy.typing import NDArray
+
+from careamics.config.tile_information import TileInformation
+
+
+# TODO: why not allow input and output of torch.tensor ?
+def stitch_prediction(
+    tiles: List[np.ndarray],
+    tile_infos: List[TileInformation],
+) -> List[np.ndarray]:
+    """
+    Stitch tiles back together to form a full image(s).
+
+    Tiles are of dimensions SC(Z)YX, where C is the number of channels and can be a
+    singleton dimension.
+
+    Parameters
+    ----------
+    tiles : list of numpy.ndarray
+        Cropped tiles and their respective stitching coordinates. Can contain tiles
+        from multiple images.
+    tile_infos : list of TileInformation
+        List of information and coordinates obtained from
+        `dataset.tiled_patching.extract_tiles`.
+
+    Returns
+    -------
+    list of numpy.ndarray
+        Full image(s).
+    """
+    # Find where to split the lists so that only info from one image is contained.
+    # Do this by locating the last tiles of each image.
+    last_tiles = [tile_info.last_tile for tile_info in tile_infos]
+    last_tile_position = np.where(last_tiles)[0]
+    image_slices = [
+        slice(
+            None if i == 0 else last_tile_position[i - 1] + 1, last_tile_position[i] + 1
+        )
+        for i in range(len(last_tile_position))
+    ]
+    image_predictions = []
+    # slice the lists and apply stitch_prediction_single to each in turn.
+    for image_slice in image_slices:
+        image_predictions.append(
+            stitch_prediction_single(tiles[image_slice], tile_infos[image_slice])
+        )
+    return image_predictions
+
+
+def stitch_prediction_single(
+    tiles: List[NDArray],
+    tile_infos: List[TileInformation],
+) -> NDArray:
+    """
+    Stitch tiles back together to form a full image.
+
+    Tiles are of dimensions SC(Z)YX, where C is the number of channels and can be a
+    singleton dimension.
+
+    Parameters
+    ----------
+    tiles : list of numpy.ndarray
+        Cropped tiles and their respective stitching coordinates.
+    tile_infos : list of TileInformation
+        List of information and coordinates obtained from
+        `dataset.tiled_patching.extract_tiles`.
+
+    Returns
+    -------
+    numpy.ndarray
+        Full image, with dimensions SC(Z)YX.
+    """
+    # retrieve whole array size
+    input_shape = tile_infos[0].array_shape
+    predicted_image = np.zeros(input_shape, dtype=np.float32)
+
+    # reshape
+    # TODO: can be more elegantly solved if TileInformation allows singleton dims
+    singleton_dims = tuple(np.where(np.array(tiles[0].shape) == 1)[0])
+    predicted_image = np.expand_dims(predicted_image, singleton_dims)
+
+    for tile, tile_info in zip(tiles, tile_infos):
+
+        # Compute coordinates for cropping predicted tile
+        crop_slices: tuple[Union[builtins.ellipsis, slice], ...] = (
+            ...,
+            *[slice(c[0], c[1]) for c in tile_info.overlap_crop_coords],
+        )
+
+        # Crop predited tile according to overlap coordinates
+        cropped_tile = tile[crop_slices]
+
+        # Insert cropped tile into predicted image using stitch coordinates
+        image_slices = (..., *[slice(c[0], c[1]) for c in tile_info.stitch_coords])
+        predicted_image[image_slices] = cropped_tile.astype(np.float32)
+
+    return predicted_image

careamics/transforms/n2v_manipulate.py

@@ -60,7 +60,7 @@ class N2VManipulate(Transform):
         remove_center: bool = True,
         struct_mask_axis: Literal["horizontal", "vertical", "none"] = "none",
         struct_mask_span: int = 5,
-        seed: Optional[int] = None,  # TODO use in pixel manipulation
+        seed: Optional[int] = None,
     ):
         """Constructor.
 
@@ -127,6 +127,7 @@ class N2VManipulate(Transform):
                     subpatch_size=self.roi_size,
                     remove_center=self.remove_center,
                     struct_params=self.struct_mask,
+                    rng=self.rng,
                 )
         elif self.strategy == SupportedPixelManipulation.MEDIAN:
             # Iterate over the channels to apply manipulation separately
@@ -136,6 +137,7 @@ class N2VManipulate(Transform):
                     mask_pixel_percentage=self.masked_pixel_percentage,
                     subpatch_size=self.roi_size,
                     struct_params=self.struct_mask,
+                    rng=self.rng,
                 )
         else:
             raise ValueError(f"Unknown masking strategy ({self.strategy}).")