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

careamics/prediction_utils/prediction_outputs.py

@@ -0,0 +1,165 @@
+"""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 outputs to the desired form.
+
+    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)
+        # remove sample dimension (always 1) `stitch_predict` func expects no S dim
+        tiles = [pred[0] for pred in predictions_comb[0]]
+        tile_infos = predictions_comb[1]
+        predictions_output = stitch_prediction(tiles, tile_infos)
+    else:
+        predictions_output = combine_batches(predictions, tiled)
+
+    # TODO: add this in? Returns output with same axes as input
+    # Won't work with tiling rn because stitch_prediction func removes S axis
+    # predictions = reshape(predictions, axes)
+    # At least make sure stitched prediction and non-tiled prediction have matching axes
+
+    # TODO: might want to remove this
+    if len(predictions_output) == 1:
+        return predictions_output[0]
+    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_untiled_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
+        Predictions that are output from `Trainer.predict`.
+
+    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_untiled_batches(
+        [preds for preds, _ in predictions]
+    )
+    return prediction_tiles, tile_infos
+
+
+def _combine_untiled_batches(predictions: List[NDArray]) -> List[NDArray]:
+    """
+    Combine batches from un-tiled output.
+
+    Parameters
+    ----------
+        predictions : list
+            Predictions that are output from `Trainer.predict`.
+
+    Returns
+    -------
+        list of nunpy.ndarray
+            Combined batches.
+    """
+    prediction_concat: NDArray = np.concatenate(predictions, axis=0)
+    prediction_split = np.split(prediction_concat, prediction_concat.shape[0], axis=0)
+    return prediction_split
+
+
+def reshape(predictions: List[NDArray], axes: str) -> List[NDArray]:
+    """
+    Reshape predictions to have dimensions of input.
+
+    Parameters
+    ----------
+    predictions : list
+        Predictions that are output from `Trainer.predict`.
+    axes : str
+        Axes SC(Z)YX.
+
+    Returns
+    -------
+    List[NDArray]
+        Reshaped predicitions.
+    """
+    if "C" not in axes:
+        predictions = [pred[:, 0] for pred in predictions]
+    if "S" not in axes:
+        predictions = [pred[0] for pred in predictions]
+    return predictions

careamics/prediction_utils/stitch_prediction.py

@@ -0,0 +1,100 @@
+"""Prediction utility functions."""
+
+from typing import List
+
+import numpy as np
+
+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[np.ndarray],
+    tile_infos: List[TileInformation],
+) -> np.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.
+    """
+    # retrieve whole array size
+    input_shape = tile_infos[0].array_shape
+    predicted_image = np.zeros(input_shape, dtype=np.float32)
+
+    for tile, tile_info in zip(tiles, tile_infos):
+        n_channels = tile.shape[0]
+
+        # Compute coordinates for cropping predicted tile
+        slices = (slice(0, n_channels),) + tuple(
+            [slice(c[0], c[1]) for c in tile_info.overlap_crop_coords]
+        )
+
+        # Crop predited tile according to overlap coordinates
+        cropped_tile = tile[slices]
+
+        # Insert cropped tile into predicted image using stitch coordinates
+        predicted_image[
+            (
+                ...,
+                *[slice(c[0], c[1]) for c in tile_info.stitch_coords],
+            )
+        ] = 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}).")

careamics/transforms/normalize.py

@@ -1,12 +1,34 @@
 """Normalization and denormalization transforms for image patches."""
 
-from typing import Optional, Tuple
+from typing import Optional
 
 import numpy as np
+from numpy.typing import NDArray
 
 from careamics.transforms.transform import Transform
 
 
+def _reshape_stats(stats: list[float], ndim: int) -> NDArray:
+    """Reshape stats to match the number of dimensions of the input image.
+
+    This allows to broadcast the stats (mean or std) to the image dimensions, and
+    thus directly perform a vectorial calculation.
+
+    Parameters
+    ----------
+    stats : list of float
+        List of stats, mean or standard deviation.
+    ndim : int
+        Number of dimensions of the image, including the C channel.
+
+    Returns
+    -------
+    NDArray
+        Reshaped stats.
+    """
+    return np.array(stats)[(..., *[np.newaxis] * (ndim - 1))]
+
+
 class Normalize(Transform):
     """
     Normalize an image or image patch.
@@ -19,154 +41,203 @@ class Normalize(Transform):
 
     Parameters
     ----------
-    mean : float
-        Mean value.
-    std : float
-        Standard deviation value.
+    image_means : list of float
+        Mean value per channel.
+    image_stds : list of float
+        Standard deviation value per channel.
+    target_means : list of float, optional
+        Target mean value per channel, by default None.
+    target_stds : list of float, optional
+        Target standard deviation value per channel, by default None.
 
     Attributes
     ----------
-    mean : float
-        Mean value.
-    std : float
-        Standard deviation value.
+    image_means : list of float
+        Mean value per channel.
+    image_stds : list of float
+        Standard deviation value per channel.
+    target_means :list of float, optional
+        Target mean value per channel, by default None.
+    target_stds : list of float, optional
+        Target standard deviation value per channel, by default None.
     """
 
     def __init__(
         self,
-        mean: float,
-        std: float,
+        image_means: list[float],
+        image_stds: list[float],
+        target_means: Optional[list[float]] = None,
+        target_stds: Optional[list[float]] = None,
     ):
         """Constructor.
 
         Parameters
         ----------
-        mean : float
-            Mean value.
-        std : float
-            Standard deviation value.
+        image_means : list of float
+            Mean value per channel.
+        image_stds : list of float
+            Standard deviation value per channel.
+        target_means : list of float, optional
+            Target mean value per channel, by default None.
+        target_stds : list of float, optional
+            Target standard deviation value per channel, by default None.
         """
-        self.mean = mean
-        self.std = std
+        self.image_means = image_means
+        self.image_stds = image_stds
+        self.target_means = target_means
+        self.target_stds = target_stds
+
         self.eps = 1e-6
 
     def __call__(
-        self, patch: np.ndarray, target: Optional[np.ndarray] = None
-    ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+        self, patch: np.ndarray, target: Optional[NDArray] = None
+    ) -> tuple[NDArray, Optional[NDArray]]:
         """Apply the transform to the source patch and the target (optional).
 
         Parameters
         ----------
-        patch : np.ndarray
+        patch : NDArray
             Patch, 2D or 3D, shape C(Z)YX.
-        target : Optional[np.ndarray], optional
+        target : NDArray, optional
             Target for the patch, by default None.
 
         Returns
         -------
-        Tuple[np.ndarray, Optional[np.ndarray]]
-            Transformed patch and target.
+        tuple of NDArray
+            Transformed patch and target, the target can be returned as `None`.
         """
-        norm_patch = self._apply(patch)
-        norm_target = self._apply(target) if target is not None else None
+        if len(self.image_means) != patch.shape[0]:
+            raise ValueError(
+                f"Number of means (got a list of size {len(self.image_means)}) and "
+                f"number of channels (got shape {patch.shape} for C(Z)YX) do not match."
+            )
+
+        # reshape mean and std and apply the normalization to the patch
+        means = _reshape_stats(self.image_means, patch.ndim)
+        stds = _reshape_stats(self.image_stds, patch.ndim)
+        norm_patch = self._apply(patch, means, stds)
+
+        # same for the target patch
+        if (
+            target is not None
+            and self.target_means is not None
+            and self.target_stds is not None
+        ):
+            target_means = _reshape_stats(self.target_means, target.ndim)
+            target_stds = _reshape_stats(self.target_stds, target.ndim)
+            norm_target = self._apply(target, target_means, target_stds)
+        else:
+            norm_target = None
 
         return norm_patch, norm_target
 
-    def _apply(self, patch: np.ndarray) -> np.ndarray:
+    def _apply(self, patch: NDArray, mean: NDArray, std: NDArray) -> NDArray:
         """
         Apply the transform to the image.
 
         Parameters
         ----------
-        patch : np.ndarray
+        patch : NDArray
             Image patch, 2D or 3D, shape C(Z)YX.
+        mean : NDArray
+            Mean values.
+        std : NDArray
+            Standard deviations.
 
         Returns
         -------
-        np.ndarray
-            Normalizedimage patch.
+        NDArray
+            Normalized image patch.
         """
-        return ((patch - self.mean) / (self.std + self.eps)).astype(np.float32)
+        return ((patch - mean) / (std + self.eps)).astype(np.float32)
 
 
 class Denormalize:
     """
-    Denormalize an image or image patch.
+    Denormalize an image.
 
     Denormalization is performed expecting a zero mean and unit variance input. This
     transform expects C(Z)YX dimensions.
 
-    Not that an epsilon value of 1e-6 is added to the standard deviation to avoid
+    Note 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.
 
     Parameters
     ----------
-    mean : float
-        Mean value.
-    std : float
-        Standard deviation value.
+    image_means : list or tuple of float
+        Mean value per channel.
+    image_stds : list or tuple of float
+        Standard deviation value per channel.
 
-    Attributes
-    ----------
-    mean : float
-        Mean value.
-    std : float
-        Standard deviation value.
     """
 
     def __init__(
         self,
-        mean: float,
-        std: float,
+        image_means: list[float],
+        image_stds: list[float],
     ):
         """Constructor.
 
         Parameters
         ----------
-        mean : float
-            Mean.
-        std : float
-            Standard deviation.
+        image_means : list of float
+            Mean value per channel.
+        image_stds : list of float
+            Standard deviation value per channel.
         """
-        self.mean = mean
-        self.std = std
+        self.image_means = image_means
+        self.image_stds = image_stds
+
         self.eps = 1e-6
 
-    def __call__(
-        self, patch: np.ndarray, target: Optional[np.ndarray] = None
-    ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
-        """Apply the transform to the source patch and the target (optional).
+    def __call__(self, patch: NDArray) -> NDArray:
+        """Reverse the normalization operation for a batch of patches.
 
         Parameters
         ----------
-        patch : np.ndarray
-            Patch, 2D or 3D, shape C(Z)YX.
-        target : Optional[np.ndarray], optional
-            Target for the patch, by default None.
+        patch : NDArray
+            Patch, 2D or 3D, shape BC(Z)YX.
 
         Returns
         -------
-        Tuple[np.ndarray, Optional[np.ndarray]]
-            Transformed patch and target.
+        NDArray
+            Transformed array.
         """
-        norm_patch = self._apply(patch)
-        norm_target = self._apply(target) if target is not None else None
+        if len(self.image_means) != patch.shape[1]:
+            raise ValueError(
+                f"Number of means (got a list of size {len(self.image_means)}) and "
+                f"number of channels (got shape {patch.shape} for BC(Z)YX) do not "
+                f"match."
+            )
 
-        return norm_patch, norm_target
+        means = _reshape_stats(self.image_means, patch.ndim)
+        stds = _reshape_stats(self.image_stds, patch.ndim)
+
+        denorm_array = self._apply(
+            patch,
+            np.swapaxes(means, 0, 1),  # swap axes as C channel is axis 1
+            np.swapaxes(stds, 0, 1),
+        )
+
+        return denorm_array.astype(np.float32)
 
-    def _apply(self, patch: np.ndarray) -> np.ndarray:
+    def _apply(self, array: NDArray, mean: NDArray, std: NDArray) -> NDArray:
         """
         Apply the transform to the image.
 
         Parameters
         ----------
-        patch : np.ndarray
+        array : NDArray
             Image patch, 2D or 3D, shape C(Z)YX.
+        mean : NDArray
+            Mean values.
+        std : NDArray
+            Standard deviations.
 
         Returns
         -------
-        np.ndarray
-            Denormalized image patch.
+        NDArray
+            Denormalized image array.
         """
-        return patch * (self.std + self.eps) + self.mean
+        return array * (std + self.eps) + mean