careamics 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

careamics/prediction_utils/lvae_tiling_manager.py

@@ -0,0 +1,362 @@
+"""Module contiaing tiling manager class."""
+
+# # TODO: remove this file, left as a reference for now.
+
+# from typing import Any, Optional
+
+# import numpy as np
+# from numpy.typing import NDArray
+
+# from careamics.config.tile_information import TileInformation
+# from careamics.config.validators import check_axes_validity
+
+
+# def calculate_padding(
+#     patch_start_location: NDArray,
+#     patch_size: NDArray,
+#     data_shape: NDArray,
+# ) -> NDArray:
+#     patch_end_location = patch_start_location + patch_size
+
+#     pad_before = np.zeros_like(patch_start_location)
+#     start_out_of_bounds = patch_start_location < 0
+#     pad_before[start_out_of_bounds] = -patch_start_location[start_out_of_bounds]
+
+#     pad_after = np.zeros_like(patch_start_location)
+#     end_out_of_bounds = patch_end_location > data_shape
+#     pad_after[end_out_of_bounds] = (
+#         patch_end_location - data_shape
+#     )[end_out_of_bounds]
+
+#     return np.stack([pad_before, pad_after], axis=1)
+
+
+# def extract_tile(
+#     img: np.ndarray,
+#     grid_start_loc: tuple[int, ...],
+#     patch_size: tuple[int, ...],
+#     overlap: tuple[int, ...],
+#     padding: bool,
+#     padding_kwargs: Optional[dict[str, Any]] = None,
+# ) -> NDArray:
+#     if padding_kwargs is None:
+#         padding_kwargs = {}
+
+#     data_shape = img.shape
+#     patch_start_loc = np.array(grid_start_loc) - np.array(overlap) // 2
+#     crop_slices = tuple(
+#         slice(max(0, start), min(start + size, dim_shape))
+#         for start, size, dim_shape in zip(patch_start_loc, patch_size, data_shape)
+#     )
+#     crop = img[crop_slices]
+#     if padding:
+#         pad = calculate_padding(
+#             patch_start_location=patch_start_loc,
+#             patch_size=patch_size,
+#             data_shape=data_shape,
+#         )
+#         crop = np.pad(crop, pad, **padding_kwargs)
+
+#     return crop
+
+
+# class TilingManager:
+
+#     def __init__(
+#         self,
+#         data_shape: tuple[int, ...],
+#         tile_size: tuple[int, ...],
+#         overlaps: tuple[int, ...],
+#         trim_boundary: tuple[int, ...],
+#     ):
+#         # --- validation
+#         if len(data_shape) != len(tile_size):
+#             raise ValueError(
+#                 f"Data shape:{data_shape} and tile size:{tile_size} must have the "
+#                 "same dimension"
+#             )
+#         if len(data_shape) != len(overlaps):
+#             raise ValueError(
+#                 f"Data shape:{data_shape} and tile overlaps:{overlaps} must have the "
+#                 "same dimension"
+#             )
+#         # overlaps = np.array(tile_size) - np.array(grid_shape)
+#         if (np.array(overlaps) < 0).any():
+#             raise ValueError(
+#                 "Tile overlap must be positive or zero in all dimension."
+#             )
+#         if ((np.array(overlaps) % 2) != 0).any():
+#             # TODO: currently not required by CAREamics tiling,
+#             #   -> because floor divide is used.
+#             raise ValueError("Tile overlaps must be even.")
+
+#         # initialize attributes
+#         self.data_shape = data_shape
+#         self.overlaps = overlaps
+#         self.grid_shape = tuple(np.array(tile_size) - np.array(overlaps))
+#         self.patch_shape = tile_size
+#         self.trim_boundary = trim_boundary
+
+#     def compute_tile_info(self, index: int, axes: str):
+
+#         # TODO: better axis validation, data should already be in the form SC(Z)YX
+
+#         # validate axes
+#         check_axes_validity(axes)
+#         # z will be -1 if not present
+#         spatial_axes = [axes.find("Z"), axes.find("Y"), axes.find("X")]
+
+#         # convert to numpy for convenience
+#         data_shape = np.array(self.data_shape)
+#         patch_shape = np.array(self.patch_shape)
+
+#         # --- calculate stitch coords
+#         stitch_coords_start = np.array(self.get_location_from_dataset_idx(index))
+#         stitch_coords_end = stitch_coords_start + np.array(self.grid_shape)
+
+#         # --- patch coords
+#         patch_coords_start = stitch_coords_start - np.array(self.overlaps) // 2
+#         patch_coords_end = patch_coords_start + patch_shape
+
+#         # --- replace out of bounds indices
+
+#         out_of_lower_bound = stitch_coords_start < 0
+#         out_of_upper_bound = stitch_coords_end > data_shape
+
+#         stitch_coords_start[out_of_lower_bound] = 0
+#         stitch_coords_end[out_of_upper_bound] = data_shape[out_of_upper_bound]
+
+#         # --- calculate overlap crop coords
+#         overlap_crop_coords_start = stitch_coords_start - patch_coords_start
+#         overlap_crop_coords_end = overlap_crop_coords_start + (
+#             stitch_coords_end - stitch_coords_start
+#         )
+
+#         # --- combine start and end
+#         stitch_coords = tuple(
+#             (stitch_coords_start[axis], stitch_coords_end[axis])
+#             for axis in spatial_axes
+#             if axis != -1
+#         )
+#         overlap_crop_coords = tuple(
+#             (overlap_crop_coords_start[axis], overlap_crop_coords_end[axis])
+#             for axis in spatial_axes
+#             if axis != -1
+#         )
+
+#         channel_axis = axes.find("C")
+#         array_shape_processed = tuple(
+#             data_shape[axis] for axis in [channel_axis, *spatial_axes] if axis != -1
+#         )
+
+#         tile_info = TileInformation(
+#             array_shape=array_shape_processed,
+#             last_tile=index == self.total_grid_count() - 1,
+#             overlap_crop_coords=overlap_crop_coords,
+#             stitch_coords=stitch_coords,
+#             sample_id=0, # TODO: in iterable dataset this is also always 0 pretty sure
+#         )
+#         return tile_info
+
+#     def patch_offset(self):
+#         return (np.array(self.patch_shape) - np.array(self.grid_shape)) // 2
+
+#     def get_individual_dim_grid_count(self, dim: int):
+#         """
+#         Returns the number of the grid in the specified dimension, ignoring all other
+#         dimensions.
+#         """
+#         assert dim < len(
+#             self.data_shape
+#         ), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+#         assert dim >= 0, "Dimension must be greater than or equal to 0"
+
+#         if self.grid_shape[dim] == 1 and self.patch_shape[dim] == 1:
+#             return self.data_shape[dim]
+#         elif self.trim_boundary is False:
+#             return int(np.ceil(self.data_shape[dim] / self.grid_shape[dim]))
+#         else:
+#             excess_size = self.patch_shape[dim] - self.grid_shape[dim]
+#             return int(
+#                 np.floor((self.data_shape[dim] - excess_size) / self.grid_shape[dim])
+#             )
+
+#     def total_grid_count(self):
+#         """
+#         Returns the total number of grids in the dataset.
+#         """
+#         return self.grid_count(0) * self.get_individual_dim_grid_count(0)
+
+#     def grid_count(self, dim: int):
+#         """
+#         Returns the total number of grids for one value in the specified dimension.
+#         """
+#         assert dim < len(
+#             self.data_shape
+#         ), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+#         assert dim >= 0, "Dimension must be greater than or equal to 0"
+#         if dim == len(self.data_shape) - 1:
+#             return 1
+
+#         return self.get_individual_dim_grid_count(dim + 1) * self.grid_count(dim + 1)
+
+#     def get_grid_index(self, dim: int, coordinate: int):
+#         """
+#         Returns the index of the grid in the specified dimension.
+#         """
+#         assert dim < len(
+#             self.data_shape
+#         ), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+#         assert dim >= 0, "Dimension must be greater than or equal to 0"
+#         assert (
+#             coordinate < self.data_shape[dim]
+#         ), (
+#             f"Coordinate {coordinate} is out of bounds for data "
+#             f"shape {self.data_shape}"
+#         )
+#         if self.grid_shape[dim] == 1 and self.patch_shape[dim] == 1:
+#             return coordinate
+#         elif self.trim_boundary is False:
+#             return np.floor(coordinate / self.grid_shape[dim])
+#         else:
+#             excess_size = (self.patch_shape[dim] - self.grid_shape[dim]) // 2
+#             # can be <0 if coordinate is in [0,grid_shape[dim]]
+#             return max(0, np.floor((coordinate - excess_size) / self.grid_shape[dim]))
+
+#     def dataset_idx_from_grid_idx(self, grid_idx: tuple):
+#         """
+#         Returns the index of the grid in the dataset.
+#         """
+#         assert len(grid_idx) == len(
+#             self.data_shape
+#         ), (
+#             f"Dimension indices {grid_idx} must have the same dimension as data "
+#             f"shape {self.data_shape}"
+#         )
+#         index = 0
+#         for dim in range(len(grid_idx)):
+#             index += grid_idx[dim] * self.grid_count(dim)
+#         return index
+
+#     def get_patch_location_from_dataset_idx(self, dataset_idx: int):
+#         """
+#         Returns the patch location of the grid in the dataset.
+#         """
+#         location = self.get_location_from_dataset_idx(dataset_idx)
+#         offset = self.patch_offset()
+#         return tuple(np.array(location) - np.array(offset))
+
+#     def get_dataset_idx_from_grid_location(self, location: tuple):
+#         assert len(location) == len(
+#             self.data_shape
+#         ), (
+#             f"Location {location} must have the same dimension as data shape "
+#             f"{self.data_shape}"
+#         )
+#         grid_idx = [
+#             self.get_grid_index(dim, location[dim]) for dim in range(len(location))
+#         ]
+#         return self.dataset_idx_from_grid_idx(tuple(grid_idx))
+
+#     def get_gridstart_location_from_dim_index(self, dim: int, dim_index: int):
+#         """
+#         Returns the grid-start coordinate of the grid in the specified dimension.
+#         """
+#         assert dim < len(
+#             self.data_shape
+#         ), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+#         assert dim >= 0, "Dimension must be greater than or equal to 0"
+#         assert dim_index < self.get_individual_dim_grid_count(
+#             dim
+#         ), (
+#             f"Dimension index {dim_index} is out of bounds for data shape "
+#             f"{self.data_shape}"
+#         )
+
+#         if self.grid_shape[dim] == 1 and self.patch_shape[dim] == 1:
+#             return dim_index
+#         elif self.trim_boundary is False:
+#             return dim_index * self.grid_shape[dim]
+#         else:
+#             excess_size = (self.patch_shape[dim] - self.grid_shape[dim]) // 2
+#             return dim_index * self.grid_shape[dim] + excess_size
+
+#     def get_location_from_dataset_idx(self, dataset_idx: int):
+#         grid_idx = []
+#         for dim in range(len(self.data_shape)):
+#             grid_idx.append(dataset_idx // self.grid_count(dim))
+#             dataset_idx = dataset_idx % self.grid_count(dim)
+#         location = [
+#             self.get_gridstart_location_from_dim_index(dim, grid_idx[dim])
+#             for dim in range(len(self.data_shape))
+#         ]
+#         return tuple(location)
+
+#     def on_boundary(self, dataset_idx: int, dim: int):
+#         """
+#         Returns True if the grid is on the boundary in the specified dimension.
+#         """
+#         assert dim < len(
+#             self.data_shape
+#         ), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+#         assert dim >= 0, "Dimension must be greater than or equal to 0"
+
+#         if dim > 0:
+#             dataset_idx = dataset_idx % self.grid_count(dim - 1)
+
+#         dim_index = dataset_idx // self.grid_count(dim)
+#         return (
+#             dim_index == 0 or dim_index == self.get_individual_dim_grid_count(dim) - 1
+#         )
+
+#     def next_grid_along_dim(self, dataset_idx: int, dim: int):
+#         """
+#         Returns the index of the grid in the specified dimension in the specified "
+#         "direction.
+#         """
+#         assert dim < len(
+#             self.data_shape
+#         ), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+#         assert dim >= 0, "Dimension must be greater than or equal to 0"
+#         new_idx = dataset_idx + self.grid_count(dim)
+#         if new_idx >= self.total_grid_count():
+#             return None
+#         return new_idx
+
+#     def prev_grid_along_dim(self, dataset_idx: int, dim: int):
+#         """
+#         Returns the index of the grid in the specified dimension in the specified "
+#         "direction.
+#         """
+#         assert dim < len(
+#             self.data_shape
+#         ), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+#         assert dim >= 0, "Dimension must be greater than or equal to 0"
+#         new_idx = dataset_idx - self.grid_count(dim)
+#         if new_idx < 0:
+#             return None
+
+
+# if __name__ == "__main__":
+#     data_shape = (1, 1, 103, 103, 2)
+#     grid_shape = (1, 1, 16, 16, 2)
+#     patch_shape = (1, 1, 32, 32, 2)
+#     overlap = tuple(np.array(patch_shape) - np.array(grid_shape))
+
+#     trim_boundary = False
+#     manager = TilingManager(
+#         data_shape=data_shape,
+#         tile_size=patch_shape,
+#         overlaps=overlap,
+#         trim_boundary=trim_boundary,
+#     )
+#     gc = manager.total_grid_count()
+#     print("Grid count", gc)
+#     for i in range(gc):
+#         loc = manager.get_location_from_dataset_idx(i)
+#         print(i, loc)
+#         inferred_i = manager.get_dataset_idx_from_grid_location(loc)
+#         assert i == inferred_i, f"Index mismatch: {i} != {inferred_i}"
+
+#     for i in range(5):
+#         print(manager.on_boundary(40, i))

careamics/prediction_utils/prediction_outputs.py

@@ -0,0 +1,135 @@
+"""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) -> 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
+
+
+# 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,112 @@
+"""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.
+    """
+    # 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):
+
+        # 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/__init__.py

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