careamics 0.0.19__py3-none-any.whl

careamics/prediction_utils/prediction_outputs.py

@@ -0,0 +1,238 @@
+"""Module containing functions to convert prediction outputs to desired form."""
+
+from typing import Any, Literal, Union, overload
+
+import numpy as np
+from numpy.typing import NDArray
+
+from ..config.data.tile_information import TileInformation
+from .stitch_prediction import stitch_prediction, stitch_prediction_vae
+
+
+def convert_outputs(predictions: list[Any], tiled: bool) -> list[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
+
+
+def convert_outputs_pn2v(
+    predictions: list[Any], tiled: bool
+) -> tuple[list[NDArray], list[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`. Length of list the total
+        number of tiles divided by the batch size. Each element consists of a tuple of
+        ((prediction, mse), tile_info_list). 1st dimension of each tensor is the bs.
+        Length of tile info list is the batch size.
+
+    tiled : bool
+        Whether the predictions are tiled.
+
+    Returns
+    -------
+    tuple[list[NDArray], list[NDArray]]
+        Tuple of (predictions, mmse) where each is a list of arrays with axes SC(Z)YX.
+    """
+    if len(predictions) == 0:
+        return [], []
+    # TODO test with multi_channel predictions
+    if tiled:
+        # Separate predictions and mmse, keeping tile info for each
+        pred_with_tiles = [
+            (pred, tile_info_list) for (pred, _), tile_info_list in predictions
+        ]
+        mse_with_tiles = [
+            (mse, tile_info_list) for (_, mse), tile_info_list in predictions
+        ]
+
+        # Process predictions
+        pred_comb = combine_batches(pred_with_tiles, tiled)
+        predictions_output = stitch_prediction(*pred_comb)
+
+        # Process mmse
+        mse_comb = combine_batches(mse_with_tiles, tiled)
+        mse_output = stitch_prediction(*mse_comb)
+
+        return predictions_output, mse_output
+    else:
+        # Separate predictions and mmse for non-tiled case
+        pred_only_tuple, mse_only_tuple = zip(*predictions, strict=False)
+        pred_only_list: list[NDArray] = list(pred_only_tuple)
+        mse_only_list: list[NDArray] = list(mse_only_tuple)
+
+        predictions_output = combine_batches(pred_only_list, tiled=False)
+        mse_output = combine_batches(mse_only_list, tiled=False)
+
+        return predictions_output, mse_output
+
+
+def convert_outputs_microsplit(
+    predictions: list[tuple[NDArray, NDArray]], dataset
+) -> tuple[NDArray, NDArray]:
+    """
+    Convert microsplit Lightning trainer outputs using eval_utils stitching functions.
+
+    This function processes microsplit predictions that return
+    (tile_prediction, tile_std) tuples and stitches them back together using the same
+    logic as get_single_file_mmse.
+
+    Parameters
+    ----------
+    predictions : list of tuple[NDArray, NDArray]
+        Predictions from Lightning trainer for microsplit. Each element is a tuple of
+        (tile_prediction, tile_std) where both are numpy arrays from predict_step.
+    dataset : Dataset
+        The dataset object used for prediction, needed for stitching function selection
+        and stitching process.
+
+    Returns
+    -------
+    tuple[NDArray, NDArray]
+        A tuple of (stitched_predictions, stitched_stds) representing the full
+        stitched predictions and standard deviations.
+    """
+    if len(predictions) == 0:
+        raise ValueError("No predictions provided")
+
+    # Separate predictions and stds from the list of tuples
+    tile_predictions = [pred for pred, _ in predictions]
+    tile_stds = [std for _, std in predictions]
+
+    # Concatenate all tiles exactly like get_single_file_mmse
+    tiles_arr = np.concatenate(tile_predictions, axis=0)
+    tile_stds_arr = np.concatenate(tile_stds, axis=0)
+
+    # Apply stitching using stitch_predictions_new
+    stitched_predictions = stitch_prediction_vae(tiles_arr, dataset)
+    stitched_stds = stitch_prediction_vae(tile_stds_arr, dataset)
+
+    return stitched_predictions, stitched_stds
+
+
+# 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.
+
+    # TODO improve description!
+
+    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,193 @@
+"""Prediction utility functions."""
+
+import builtins
+from typing import Union
+
+import numpy as np
+from numpy.typing import NDArray
+
+from careamics.config.data.tile_information import TileInformation
+
+
+class TilingMode:
+    """Enum for the tiling mode."""
+
+    TrimBoundary = 0
+    PadBoundary = 1
+    ShiftBoundary = 2
+
+
+def stitch_prediction_vae(predictions, dset):
+    """Stitch predictions back together using dataset's index manager.
+
+    Parameters
+    ----------
+    predictions : numpy.ndarray
+        Array of predictions with shape (n_tiles, channels, height, width).
+    dset : Dataset
+        Dataset object with idx_manager containing tiling information.
+
+    Returns
+    -------
+    numpy.ndarray
+        Stitched predictions.
+    """
+    mng = dset.idx_manager
+
+    # if there are more channels, use all of them.
+    shape = list(dset.get_data_shape())
+    shape[-1] = max(shape[-1], predictions.shape[1])
+
+    output = np.zeros(shape, dtype=predictions.dtype)
+    # frame_shape = dset.get_data_shape()[:-1]
+    for dset_idx in range(predictions.shape[0]):
+        # loc = get_location_from_idx(dset, dset_idx, predictions.shape[-2],
+        # predictions.shape[-1])
+        # grid start, grid end
+        gs = np.array(mng.get_location_from_dataset_idx(dset_idx), dtype=int)
+        ge = gs + mng.grid_shape
+
+        # patch start, patch end
+        ps = gs - mng.patch_offset()
+        pe = ps + mng.patch_shape
+
+        # valid grid start, valid grid end
+        vgs = np.array([max(0, x) for x in gs], dtype=int)
+        vge = np.array(
+            [min(x, y) for x, y in zip(ge, mng.data_shape, strict=False)], dtype=int
+        )
+
+        if mng.tiling_mode == TilingMode.ShiftBoundary:
+            for dim in range(len(vgs)):
+                if ps[dim] == 0:
+                    vgs[dim] = 0
+                if pe[dim] == mng.data_shape[dim]:
+                    vge[dim] = mng.data_shape[dim]
+
+        # relative start, relative end. This will be used on pred_tiled
+        rs = vgs - ps
+        re = rs + (vge - vgs)
+
+        for ch_idx in range(predictions.shape[1]):
+            if len(output.shape) == 4:
+                # channel dimension is the last one.
+                output[vgs[0] : vge[0], vgs[1] : vge[1], vgs[2] : vge[2], ch_idx] = (
+                    predictions[dset_idx][ch_idx, rs[1] : re[1], rs[2] : re[2]]
+                )
+            elif len(output.shape) == 5:
+                # channel dimension is the last one.
+                assert vge[0] - vgs[0] == 1, "Only one frame is supported"
+                output[
+                    vgs[0], vgs[1] : vge[1], vgs[2] : vge[2], vgs[3] : vge[3], ch_idx
+                ] = predictions[dset_idx][
+                    ch_idx, rs[1] : re[1], rs[2] : re[2], rs[3] : re[3]
+                ]
+            else:
+                raise ValueError(f"Unsupported shape {output.shape}")
+
+    return output
+
+
+# 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.
+    """
+    # TODO: this is hacky... need a better way to deal with when input channels and
+    #   target channels do not match
+    if len(tile_infos[0].array_shape) == 4:
+        # 4 dimensions => 3 spatial dimensions so -4 is channel dimension
+        tile_channels = tiles[0].shape[-4]
+    elif len(tile_infos[0].array_shape) == 3:
+        # 3 dimensions => 2 spatial dimensions so -3 is channel dimension
+        tile_channels = tiles[0].shape[-3]
+    else:
+        # Note pretty sure this is unreachable because array shape is already
+        #   validated by TileInformation
+        raise ValueError(
+            f"Unsupported number of output dimension {len(tile_infos[0].array_shape)}"
+        )
+    # retrieve whole array size, add S dim and use number of channels in tile
+    input_shape = (1, tile_channels, *tile_infos[0].array_shape[1:])
+    predicted_image = np.zeros(input_shape, dtype=np.float32)
+
+    for tile, tile_info in zip(tiles, tile_infos, strict=False):
+
+        # 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])
+
+        # TODO fix mypy error here, potentially due to numpy 2
+        predicted_image[image_slices] = cropped_tile.astype(np.float32)  # type: ignore
+
+    return predicted_image

careamics/py.typed

@@ -0,0 +1,5 @@
+You may remove this file if you don't intend to add types to your package
+
+Details at:
+
+https://mypy.readthedocs.io/en/stable/installed_packages.html#creating-pep-561-compatible-packages

careamics/transforms/__init__.py

@@ -0,0 +1,22 @@
+"""Transforms that are used to augment the data."""
+
+__all__ = [
+    "Compose",
+    "Denormalize",
+    "ImageRestorationTTA",
+    "N2VManipulate",
+    "N2VManipulateTorch",
+    "Normalize",
+    "TrainDenormalize",
+    "XYFlip",
+    "XYRandomRotate90",
+    "get_all_transforms",
+]
+
+from .compose import Compose, get_all_transforms
+from .n2v_manipulate import N2VManipulate
+from .n2v_manipulate_torch import N2VManipulateTorch
+from .normalize import Denormalize, Normalize, TrainDenormalize
+from .tta import ImageRestorationTTA
+from .xy_flip import XYFlip
+from .xy_random_rotate90 import XYRandomRotate90

careamics/transforms/compose.py

@@ -0,0 +1,173 @@
+"""A class chaining transforms together."""
+
+from typing import Union, cast
+
+from numpy.typing import NDArray
+
+from careamics.config.transformations import NORM_AND_SPATIAL_UNION
+
+from .normalize import Normalize
+from .transform import Transform
+from .xy_flip import XYFlip
+from .xy_random_rotate90 import XYRandomRotate90
+
+ALL_TRANSFORMS = {
+    "Normalize": Normalize,
+    "XYFlip": XYFlip,
+    "XYRandomRotate90": XYRandomRotate90,
+}
+
+
+def get_all_transforms() -> dict[str, type]:
+    """Return all the transforms accepted by CAREamics.
+
+    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
+
+
+class Compose:
+    """A class chaining transforms together.
+
+    Parameters
+    ----------
+    transform_list : list[TransformConfig]
+        A list of dictionaries where each dictionary contains the name of a
+        transform and its parameters.
+
+    Attributes
+    ----------
+    _callable_transforms : Callable
+        A callable that applies the transforms to the input data.
+    """
+
+    def __init__(self, transform_list: list[NORM_AND_SPATIAL_UNION]) -> None:
+        """Instantiate a Compose object.
+
+        Parameters
+        ----------
+        transform_list : list[NORM_AND_SPATIAL_UNION]
+            A list of dictionaries where each dictionary contains the name of a
+            transform and its parameters.
+        """
+        # retrieve all available transforms
+        # TODO: correctly type hint get_all_transforms function output
+        all_transforms: dict[str, type[Transform]] = get_all_transforms()
+
+        # instantiate all transforms
+        self.transforms: list[Transform] = [
+            all_transforms[t.name](**t.model_dump()) for t in transform_list
+        ]
+
+    def _chain_transforms(
+        self, patch: NDArray, target: NDArray | None
+    ) -> tuple[NDArray | None, ...]:
+        """Chain transforms on the input data.
+
+        Parameters
+        ----------
+        patch : np.ndarray
+            Input data.
+        target : Optional[np.ndarray]
+            Target data, by default None.
+
+        Returns
+        -------
+        tuple[np.ndarray, Optional[np.ndarray]]
+            The output of the transformations.
+        """
+        params: Union[tuple[NDArray, NDArray | None],] = (patch, target)
+
+        for t in self.transforms:
+            *params, _ = t(*params)  # ignore additional_arrays dict
+
+        # avoid None values that create problems for collating
+        # TODO: removing None should be handled in dataset, not here
+        return tuple(p for p in params if p is not None)
+
+    def _chain_transforms_additional_arrays(
+        self,
+        patch: NDArray,
+        target: NDArray | None,
+        **additional_arrays: NDArray,
+    ) -> tuple[NDArray, NDArray | None, dict[str, NDArray]]:
+        """Chain transforms on the input data, with additional arrays.
+
+        Parameters
+        ----------
+        patch : np.ndarray
+            Input data.
+        target : Optional[np.ndarray]
+            Target data, by default None.
+        **additional_arrays : NDArray
+            Additional arrays that will be transformed identically to `patch` and
+            `target`.
+
+        Returns
+        -------
+        tuple[np.ndarray, Optional[np.ndarray]]
+            The output of the transformations.
+        """
+        params = {"patch": patch, "target": target, **additional_arrays}
+
+        for t in self.transforms:
+            patch, target, additional_arrays = t(**params)
+            params = {"patch": patch, "target": target, **additional_arrays}
+
+        return patch, target, additional_arrays
+
+    def __call__(
+        self, patch: NDArray, target: NDArray | None = None
+    ) -> tuple[NDArray, ...]:
+        """Apply the transforms to the input data.
+
+        Parameters
+        ----------
+        patch : np.ndarray
+            The input data.
+        target : Optional[np.ndarray], optional
+            Target data, by default None.
+
+        Returns
+        -------
+        tuple[np.ndarray, ...]
+            The output of the transformations.
+        """
+        # TODO: solve casting Compose.__call__ ouput
+        return cast(tuple[NDArray, ...], self._chain_transforms(patch, target))
+
+    def transform_with_additional_arrays(
+        self,
+        patch: NDArray,
+        target: NDArray | None = None,
+        **additional_arrays: NDArray,
+    ) -> tuple[NDArray, NDArray | None, dict[str, NDArray]]:
+        """Apply the transforms to the input data, including additional arrays.
+
+        Parameters
+        ----------
+        patch : np.ndarray
+            The input data.
+        target : Optional[np.ndarray], optional
+            Target data, by default None.
+        **additional_arrays : NDArray
+            Additional arrays that will be transformed identically to `patch` and
+            `target`.
+
+        Returns
+        -------
+        NDArray
+            The transformed patch.
+        NDArray | None
+            The transformed target.
+        dict of {str, NDArray}
+            Transformed additional arrays. Keys correspond to the keyword argument
+            names.
+        """
+        return self._chain_transforms_additional_arrays(
+            patch, target, **additional_arrays
+        )