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

careamics/prediction_utils/create_pred_datamodule.py

@@ -0,0 +1,185 @@
+"""Module containing functions to create `CAREamicsPredictData`."""
+
+from pathlib import Path
+from typing import Callable, Dict, Literal, Optional, Tuple, Union
+
+import numpy as np
+from numpy.typing import NDArray
+
+from careamics.config import Configuration, create_inference_configuration
+from careamics.utils import check_path_exists
+
+from ..lightning_prediction_datamodule import CAREamicsPredictData
+
+
+def create_pred_datamodule(
+    source: Union[CAREamicsPredictData, Path, str, NDArray],
+    config: Configuration,
+    batch_size: Optional[int] = None,
+    tile_size: Optional[Tuple[int, ...]] = None,
+    tile_overlap: Tuple[int, ...] = (48, 48),
+    axes: Optional[str] = None,
+    data_type: Optional[Literal["array", "tiff", "custom"]] = None,
+    tta_transforms: bool = True,
+    dataloader_params: Optional[Dict] = None,
+    read_source_func: Optional[Callable] = None,
+    extension_filter: str = "",
+) -> CAREamicsPredictData:
+    """
+    Create a `CAREamicsPredictData` module.
+
+    Parameters
+    ----------
+    source : CAREamicsPredData, pathlib.Path, str or numpy.ndarray
+        Data to predict on.
+    config : Configuration
+        Global configuration.
+    batch_size : int, default=1
+        Batch size for prediction.
+    tile_size : tuple of int, optional
+        Size of the tiles to use for prediction.
+    tile_overlap : tuple of int, default=(48, 48)
+        Overlap between tiles.
+    axes : str, optional
+        Axes of the input data, by default None.
+    data_type : {"array", "tiff", "custom"}, optional
+        Type of the input data.
+    tta_transforms : bool, default=True
+        Whether to apply test-time augmentation.
+    dataloader_params : dict, optional
+        Parameters to pass to the dataloader.
+    read_source_func : Callable, optional
+        Function to read the source data.
+    extension_filter : str, default=""
+        Filter for the file extension.
+
+    Returns
+    -------
+    prediction datamodule: CAREamicsPredictData
+        Subclass of `pytorch_lightning.LightningDataModule` for creating predictions.
+
+    Raises
+    ------
+    ValueError
+        If the input is not a CAREamicsPredData instance, a path or a numpy array.
+    """
+    # Reuse batch size if not provided explicitly
+    if batch_size is None:
+        batch_size = config.data_config.batch_size
+
+    # create predict config, reuse training config if parameters missing
+    prediction_config = create_inference_configuration(
+        configuration=config,
+        tile_size=tile_size,
+        tile_overlap=tile_overlap,
+        data_type=data_type,
+        axes=axes,
+        tta_transforms=tta_transforms,
+        batch_size=batch_size,
+    )
+
+    # remove batch from dataloader parameters (priority given to config)
+    if dataloader_params is None:
+        dataloader_params = {}
+    if "batch_size" in dataloader_params:
+        del dataloader_params["batch_size"]
+
+    if isinstance(source, CAREamicsPredictData):
+        pred_datamodule = source
+    elif isinstance(source, Path) or isinstance(source, str):
+        pred_datamodule = _create_from_path(
+            source=source,
+            pred_config=prediction_config,
+            read_source_func=read_source_func,
+            extension_filter=extension_filter,
+            dataloader_params=dataloader_params,
+        )
+    elif isinstance(source, np.ndarray):
+        pred_datamodule = _create_from_array(
+            source=source,
+            pred_config=prediction_config,
+            dataloader_params=dataloader_params,
+        )
+    else:
+        raise ValueError(
+            f"Invalid input. Expected a CAREamicsPredData instance, paths or "
+            f"NDArray (got {type(source)})."
+        )
+
+    return pred_datamodule
+
+
+def _create_from_path(
+    source: Union[Path, str],
+    pred_config: Configuration,
+    read_source_func: Optional[Callable] = None,
+    extension_filter: str = "",
+    dataloader_params: Optional[Dict] = None,
+    **kwargs,
+) -> CAREamicsPredictData:
+    """
+    Create `CAREamicsPredictData` from path.
+
+    Parameters
+    ----------
+    source : Path or str
+        _Data to predict on.
+    pred_config : Configuration
+        Prediction configuration.
+    read_source_func : Callable, optional
+        Function to read the source data.
+    extension_filter : str, default=""
+        Function to read the source data.
+    dataloader_params : Optional[Dict], optional
+        Parameters to pass to the dataloader.
+    **kwargs
+        Unused.
+
+    Returns
+    -------
+    prediction datamodule: CAREamicsPredictData
+        Subclass of `pytorch_lightning.LightningDataModule` for creating predictions.
+    """
+    source_path = check_path_exists(source)
+
+    datamodule = CAREamicsPredictData(
+        pred_config=pred_config,
+        pred_data=source_path,
+        read_source_func=read_source_func,
+        extension_filter=extension_filter,
+        dataloader_params=dataloader_params,
+    )
+    return datamodule
+
+
+def _create_from_array(
+    source: NDArray,
+    pred_config: Configuration,
+    dataloader_params: Optional[Dict] = None,
+    **kwargs,
+) -> CAREamicsPredictData:
+    """
+    Create `CAREamicsPredictData` from array.
+
+    Parameters
+    ----------
+    source : Path or str
+        _Data to predict on.
+    pred_config : Configuration
+        Prediction configuration.
+    dataloader_params : Optional[Dict], optional
+        Parameters to pass to the dataloader.
+    **kwargs
+        Unused. Added for compatible function signature with `_create_from_path`.
+
+    Returns
+    -------
+    prediction datamodule: CAREamicsPredictData
+        Subclass of `pytorch_lightning.LightningDataModule` for creating predictions.
+    """
+    datamodule = CAREamicsPredictData(
+        pred_config=pred_config,
+        pred_data=source,
+        dataloader_params=dataloader_params,
+    )
+    return datamodule

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/__init__.py

@@ -3,7 +3,7 @@
 __all__ = [
     "get_all_transforms",
     "N2VManipulate",
-    "NDFlip",
+    "XYFlip",
     "XYRandomRotate90",
     "ImageRestorationTTA",
     "Denormalize",
@@ -14,7 +14,7 @@ __all__ = [
 
 from .compose import Compose, get_all_transforms
 from .n2v_manipulate import N2VManipulate
-from .nd_flip import NDFlip
 from .normalize import Denormalize, Normalize
 from .tta import ImageRestorationTTA
+from .xy_flip import XYFlip
 from .xy_random_rotate90 import XYRandomRotate90

careamics/transforms/compose.py

@@ -1,26 +1,26 @@
 """A class chaining transforms together."""
 
-from typing import Callable, List, Optional, Tuple
+from typing import Callable, Dict, List, Optional, Tuple
 
 import numpy as np
 
 from careamics.config.data_model import TRANSFORMS_UNION
 
 from .n2v_manipulate import N2VManipulate
-from .nd_flip import NDFlip
 from .normalize import Normalize
 from .transform import Transform
+from .xy_flip import XYFlip
 from .xy_random_rotate90 import XYRandomRotate90
 
 ALL_TRANSFORMS = {
     "Normalize": Normalize,
     "N2VManipulate": N2VManipulate,
-    "NDFlip": NDFlip,
+    "XYFlip": XYFlip,
     "XYRandomRotate90": XYRandomRotate90,
 }
 
 
-def get_all_transforms() -> dict:
+def get_all_transforms() -> Dict[str, type]:
     """Return all the transforms accepted by CAREamics.
 
     Returns
@@ -33,7 +33,19 @@ def get_all_transforms() -> dict:
 
 
 class Compose:
-    """A class chaining transforms together."""
+    """A class chaining transforms together.
+
+    Parameters
+    ----------
+    transform_list : List[TRANSFORMS_UNION]
+        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[TRANSFORMS_UNION]) -> None:
         """Instantiate a Compose object.
@@ -68,7 +80,21 @@ class Compose:
 
         def _chain(
             patch: np.ndarray, target: Optional[np.ndarray]
-        ) -> Tuple[np.ndarray, ...]:
+        ) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+            """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 = (patch, target)
 
             for t in transforms:
@@ -88,7 +114,7 @@ class Compose:
         patch : np.ndarray
             The input data.
         target : Optional[np.ndarray], optional
-            Target data, by default None
+            Target data, by default None.
 
         Returns
         -------

careamics/transforms/n2v_manipulate.py

@@ -1,3 +1,5 @@
+"""N2V manipulation transform."""
+
 from typing import Any, Literal, Optional, Tuple
 
 import numpy as np
@@ -17,10 +19,35 @@ class N2VManipulate(Transform):
 
     Parameters
     ----------
-    mask_pixel_percentage : float
-        Approximate percentage of pixels to be masked.
+    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.
+    seed : Optional[int], optional
+        Random seed, by default None.
+
+    Attributes
+    ----------
+    masked_pixel_percentage : float
+        Percentage of pixels to mask.
     roi_size : int
-        Size of the ROI the new pixel value is sampled from, by default 11.
+        Size of the replacement area.
+    strategy : Literal[ "uniform", "median" ]
+        Replaccement strategy, uniform or median.
+    remove_center : bool
+        Whether to remove central pixel from patch.
+    struct_mask : Optional[StructMaskParameters]
+        StructN2V mask parameters.
+    rng : Generator
+        Random number generator.
     """
 
     def __init__(
@@ -33,31 +60,31 @@ 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.
 
         Parameters
         ----------
         roi_size : int, optional
-            Size of the replacement area, by default 11
+            Size of the replacement area, by default 11.
         masked_pixel_percentage : float, optional
-            Percentage of pixels to mask, by default 0.2
+            Percentage of pixels to mask, by default 0.2.
         strategy : Literal[ "uniform", "median" ], optional
-            Replaccement strategy, uniform or median, by default uniform
+            Replaccement strategy, uniform or median, by default uniform.
         remove_center : bool, optional
-            Whether to remove central pixel from patch, by default True
+            Whether to remove central pixel from patch, by default True.
         struct_mask_axis : Literal["horizontal", "vertical", "none"], optional
-            StructN2V mask axis, by default "none"
+            StructN2V mask axis, by default "none".
         struct_mask_span : int, optional
-            StructN2V mask span, by default 5
+            StructN2V mask span, by default 5.
         seed : Optional[int], optional
-            Random seed, by default None
+            Random seed, by default None.
         """
         self.masked_pixel_percentage = masked_pixel_percentage
         self.roi_size = roi_size
         self.strategy = strategy
-        self.remove_center = remove_center
+        self.remove_center = remove_center  # TODO is this ever used?
 
         if struct_mask_axis == SupportedStructAxis.NONE:
             self.struct_mask: Optional[StructMaskParameters] = None
@@ -77,8 +104,17 @@ class N2VManipulate(Transform):
 
         Parameters
         ----------
-        image : np.ndarray
-            Image or image patch, 2D or 3D, shape C(Z)YX.
+        patch : np.ndarray
+            Image patch, 2D or 3D, shape C(Z)YX.
+        *args : Any
+            Additional arguments, unused.
+        **kwargs : Any
+            Additional keyword arguments, unused.
+
+        Returns
+        -------
+        Tuple[np.ndarray, np.ndarray, np.ndarray]
+            Masked patch, original patch, and mask.
         """
         masked = np.zeros_like(patch)
         mask = np.zeros_like(patch)
@@ -91,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
@@ -100,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}).")