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

careamics/dataset/patching/validate_patch_dimension.py

@@ -0,0 +1,64 @@
+"""Patch validation functions."""
+
+from typing import List, Tuple, Union
+
+import numpy as np
+
+
+def validate_patch_dimensions(
+    arr: np.ndarray,
+    patch_size: Union[List[int], Tuple[int, ...]],
+    is_3d_patch: bool,
+) -> None:
+    """
+    Check patch size and array compatibility.
+
+    This method validates the patch sizes with respect to the array dimensions:
+
+    - Patch must have two dimensions fewer than the array (S and C).
+    - Patch sizes are smaller than the corresponding array dimensions.
+
+    If one of these conditions is not met, a ValueError is raised.
+
+    This method should be called after inputs have been resized.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Input array.
+    patch_size : Union[List[int], Tuple[int, ...]]
+        Size of the patches along each dimension of the array, except the first.
+    is_3d_patch : bool
+        Whether the patch is 3D or not.
+
+    Raises
+    ------
+    ValueError
+        If the patch size is not consistent with the array shape (one more array
+        dimension).
+    ValueError
+        If the patch size in Z is larger than the array dimension.
+    ValueError
+        If either of the patch sizes in X or Y is larger than the corresponding array
+        dimension.
+    """
+    if len(patch_size) != len(arr.shape[2:]):
+        raise ValueError(
+            f"There must be a patch size for each spatial dimensions "
+            f"(got {patch_size} patches for dims {arr.shape}). Check the axes order."
+        )
+
+    # Sanity checks on patch sizes versus array dimension
+    if is_3d_patch and patch_size[0] > arr.shape[-3]:
+        raise ValueError(
+            f"Z patch size is inconsistent with image shape "
+            f"(got {patch_size[0]} patches for dim {arr.shape[1]}). Check the axes "
+            f"order."
+        )
+
+    if patch_size[-2] > arr.shape[-2] or patch_size[-1] > arr.shape[-1]:
+        raise ValueError(
+            f"At least one of YX patch dimensions is larger than the corresponding "
+            f"image dimension (got {patch_size} patches for dims {arr.shape[-2:]}). "
+            f"Check the axes order."
+        )

careamics/dataset/tiling/__init__.py

@@ -0,0 +1,10 @@
+"""Tiling functions."""
+
+__all__ = [
+    "stitch_prediction",
+    "extract_tiles",
+    "collate_tiles",
+]
+
+from .collate_tiles import collate_tiles
+from .tiled_patching import extract_tiles

careamics/dataset/tiling/collate_tiles.py

@@ -0,0 +1,33 @@
+"""Collate function for tiling."""
+
+from typing import Any, List, Tuple
+
+import numpy as np
+from torch.utils.data.dataloader import default_collate
+
+from careamics.config.tile_information import TileInformation
+
+
+def collate_tiles(batch: List[Tuple[np.ndarray, TileInformation]]) -> Any:
+    """
+    Collate tiles received from CAREamics prediction dataloader.
+
+    CAREamics prediction dataloader returns tuples of arrays and TileInformation. In
+    case of non-tiled data, this function will return the arrays. In case of tiled data,
+    it will return the arrays, the last tile flag, the overlap crop coordinates and the
+    stitch coordinates.
+
+    Parameters
+    ----------
+    batch : List[Tuple[np.ndarray, TileInformation], ...]
+        Batch of tiles.
+
+    Returns
+    -------
+    Any
+        Collated batch.
+    """
+    new_batch = [tile for tile, _ in batch]
+    tiles_batch = [tile_info for _, tile_info in batch]
+
+    return default_collate(new_batch), tiles_batch

careamics/dataset/tiling/tiled_patching.py

@@ -0,0 +1,164 @@
+"""Tiled patching utilities."""
+
+import itertools
+from typing import Generator, List, Tuple, Union
+
+import numpy as np
+
+from careamics.config.tile_information import TileInformation
+
+
+def _compute_crop_and_stitch_coords_1d(
+    axis_size: int, tile_size: int, overlap: int
+) -> Tuple[List[Tuple[int, int]], List[Tuple[int, int]], List[Tuple[int, int]]]:
+    """
+    Compute the coordinates of each tile along an axis, given the overlap.
+
+    Parameters
+    ----------
+    axis_size : int
+        Length of the axis.
+    tile_size : int
+        Size of the tile for the given axis.
+    overlap : int
+        Size of the overlap for the given axis.
+
+    Returns
+    -------
+    Tuple[Tuple[int, ...], ...]
+        Tuple of all coordinates for given axis.
+    """
+    # Compute the step between tiles
+    step = tile_size - overlap
+    crop_coords = []
+    stitch_coords = []
+    overlap_crop_coords = []
+
+    # Iterate over the axis with step
+    for i in range(0, max(1, axis_size - overlap), step):
+        # Check if the tile fits within the axis
+        if i + tile_size <= axis_size:
+            # Add the coordinates to crop one tile
+            crop_coords.append((i, i + tile_size))
+
+            # Add the pixel coordinates of the cropped tile in the original image space
+            stitch_coords.append(
+                (
+                    i + overlap // 2 if i > 0 else 0,
+                    (
+                        i + tile_size - overlap // 2
+                        if crop_coords[-1][1] < axis_size
+                        else axis_size
+                    ),
+                )
+            )
+
+            # Add the coordinates to crop the overlap from the prediction.
+            overlap_crop_coords.append(
+                (
+                    overlap // 2 if i > 0 else 0,
+                    (
+                        tile_size - overlap // 2
+                        if crop_coords[-1][1] < axis_size
+                        else tile_size
+                    ),
+                )
+            )
+
+        # If the tile does not fit within the axis, perform the abovementioned
+        # operations starting from the end of the axis
+        else:
+            # if (axis_size - tile_size, axis_size) not in crop_coords:
+            crop_coords.append((max(0, axis_size - tile_size), axis_size))
+            last_tile_end_coord = stitch_coords[-1][1] if stitch_coords else 1
+            stitch_coords.append((last_tile_end_coord, axis_size))
+            overlap_crop_coords.append(
+                (tile_size - (axis_size - last_tile_end_coord), tile_size)
+            )
+            break
+    return crop_coords, stitch_coords, overlap_crop_coords
+
+
+def extract_tiles(
+    arr: np.ndarray,
+    tile_size: Union[List[int], Tuple[int, ...]],
+    overlaps: Union[List[int], Tuple[int, ...]],
+) -> Generator[Tuple[np.ndarray, TileInformation], None, None]:
+    """Generate tiles from the input array with specified overlap.
+
+    The tiles cover the whole array. The method returns a generator that yields
+    tuples of array and tile information, the latter includes whether
+    the tile is the last one, the coordinates of the overlap crop, and the coordinates
+    of the stitched tile.
+
+    Input array should have shape SC(Z)YX, while the returned tiles have shape C(Z)YX,
+    where C can be a singleton.
+
+    Parameters
+    ----------
+    arr : np.ndarray
+        Array of shape (S, C, (Z), Y, X).
+    tile_size : Union[List[int], Tuple[int]]
+        Tile sizes in each dimension, of length 2 or 3.
+    overlaps : Union[List[int], Tuple[int]]
+        Overlap values in each dimension, of length 2 or 3.
+
+    Yields
+    ------
+    Generator[Tuple[np.ndarray, TileInformation], None, None]
+        Tile generator, yields the tile and additional information.
+    """
+    # Iterate over num samples (S)
+    for sample_idx in range(arr.shape[0]):
+        sample: np.ndarray = arr[sample_idx, ...]
+
+        # Create a list of coordinates for cropping and stitching all axes.
+        # [crop coordinates, stitching coordinates, overlap crop coordinates]
+        # For axis of size 35 and patch size of 32 compute_crop_and_stitch_coords_1d
+        # will output ([(0, 32), (3, 35)], [(0, 20), (20, 35)], [(0, 20), (17, 32)])
+        crop_and_stitch_coords_list = [
+            _compute_crop_and_stitch_coords_1d(
+                sample.shape[i + 1], tile_size[i], overlaps[i]
+            )
+            for i in range(len(tile_size))
+        ]
+
+        # Rearrange crop coordinates from a list of coordinate pairs per axis to a list
+        # grouped by type.
+        all_crop_coords, all_stitch_coords, all_overlap_crop_coords = zip(
+            *crop_and_stitch_coords_list
+        )
+
+        # Maximum tile index
+        max_tile_idx = np.prod([len(axis) for axis in all_crop_coords]) - 1
+
+        # Iterate over generated coordinate pairs:
+        for tile_idx, (crop_coords, stitch_coords, overlap_crop_coords) in enumerate(
+            zip(
+                itertools.product(*all_crop_coords),
+                itertools.product(*all_stitch_coords),
+                itertools.product(*all_overlap_crop_coords),
+            )
+        ):
+            # Extract tile from the sample
+            tile: np.ndarray = sample[
+                (..., *[slice(c[0], c[1]) for c in list(crop_coords)])  # type: ignore
+            ]
+
+            # Check if we are at the end of the sample by computing the length of the
+            # array that contains all the tiles
+            if tile_idx == max_tile_idx:
+                last_tile = True
+            else:
+                last_tile = False
+
+            # create tile information
+            tile_info = TileInformation(
+                array_shape=sample.shape,
+                last_tile=last_tile,
+                overlap_crop_coords=overlap_crop_coords,
+                stitch_coords=stitch_coords,
+                sample_id=sample_idx,
+            )
+
+            yield tile, tile_info

careamics/dataset/zarr_dataset.py

@@ -0,0 +1,151 @@
+"""Zarr dataset."""
+
+# from itertools import islice
+# from typing import Callable, Dict, List, Optional, Tuple, Union
+
+# import numpy as np
+# import torch
+# import zarr
+
+# from careamics.utils import RunningStats
+# from careamics.utils.logging import get_logger
+
+# from ..utils import normalize
+# from .dataset_utils import read_zarr
+# from .patching.patching import (
+#     generate_patches_unsupervised,
+# )
+
+# logger = get_logger(__name__)
+
+
+# class ZarrDataset(torch.utils.data.IterableDataset):
+#     """Dataset to extract patches from a zarr storage.
+
+#     Parameters
+#     ----------
+#     data_source : Union[zarr.Group, zarr.Array]
+#         Zarr storage.
+#     axes : str
+#         Description of axes in format STCZYX.
+#     patch_extraction_method : Union[ExtractionStrategies, None]
+#         Patch extraction strategy, as defined in extraction_strategy.
+#     patch_size : Optional[Union[List[int], Tuple[int]]], optional
+#         Size of the patches in each dimension, by default None.
+#     num_patches : Optional[int], optional
+#         Number of patches to extract, by default None.
+#     mean : Optional[float], optional
+#         Expected mean of the dataset, by default None.
+#     std : Optional[float], optional
+#         Expected standard deviation of the dataset, by default None.
+#     patch_transform : Optional[Callable], optional
+#         Patch transform callable, by default None.
+#     patch_transform_params : Optional[Dict], optional
+#         Patch transform parameters, by default None.
+#     running_stats_window_perc : float, optional
+#         Percentage of the dataset to use for calculating the initial mean and standard
+#         deviation, by default 0.01.
+#     mode : str, optional
+#         train/predict, controls running stats calculation.
+#     """
+
+#     def __init__(
+#         self,
+#         data_source: Union[zarr.Group, zarr.Array],
+#         axes: str,
+#         patch_extraction_method: Union[SupportedExtractionStrategy, None],
+#         patch_size: Optional[Union[List[int], Tuple[int]]] = None,
+#         num_patches: Optional[int] = None,
+#         mean: Optional[float] = None,
+#         std: Optional[float] = None,
+#         patch_transform: Optional[Callable] = None,
+#         patch_transform_params: Optional[Dict] = None,
+#         running_stats_window_perc: float = 0.01,
+#         mode: str = "train",
+#     ) -> None:
+#         self.data_source = data_source
+#         self.axes = axes
+#         self.patch_extraction_method = patch_extraction_method
+#         self.patch_size = patch_size
+#         self.num_patches = num_patches
+#         self.mean = mean
+#         self.std = std
+#         self.patch_transform = patch_transform
+#         self.patch_transform_params = patch_transform_params
+#         self.sample = read_zarr(self.data_source, self.axes)
+#         self.running_stats_window = int(
+#             np.prod(self.sample._cdata_shape) * running_stats_window_perc
+#         )
+#         self.mode = mode
+#         self.running_stats = RunningStats()
+
+#         self._calculate_initial_mean_std()
+
+#     def _calculate_initial_mean_std(self):
+#         """Calculate initial mean and std of the dataset."""
+#         if self.mean is None and self.std is None:
+#             idxs = np.random.randint(
+#                 0,
+#                 np.prod(self.sample._cdata_shape),
+#                 size=max(1, self.running_stats_window),
+#             )
+#             random_chunks = self.sample[idxs]
+#             self.running_stats.init(random_chunks.mean(), random_chunks.std())
+
+#     def _generate_patches(self):
+#         """Generate patches from the dataset and calculates running stats.
+
+#         Yields
+#         ------
+#         np.ndarray
+#             Patch.
+#         """
+#         patches = generate_patches_unsupervised(
+#             self.sample,
+#             self.patch_extraction_method,
+#             self.patch_size,
+#         )
+
+#         # num_patches = np.ceil(
+#         #     np.prod(self.sample.chunks)
+#         #     / (np.prod(self.patch_size) * self.running_stats_window)
+#         # ).astype(int)
+
+#         for idx, patch in enumerate(patches):
+#             if self.mode != "predict":
+#                 self.running_stats.update(patch.mean())
+#             if isinstance(patch, tuple):
+#                 normalized_patch = normalize(
+#                     img=patch[0],
+#                     mean=self.running_stats.avg_mean.value,
+#                     std=self.running_stats.avg_std.value,
+#                 )
+#                 patch = (normalized_patch, *patch[1:])
+#             else:
+#                 patch = normalize(
+#                     img=patch,
+#                     mean=self.running_stats.avg_mean.value,
+#                     std=self.running_stats.avg_std.value,
+#                 )
+
+#             if self.patch_transform is not None:
+#                 assert self.patch_transform_params is not None
+#                 patch = self.patch_transform(patch, **self.patch_transform_params)
+#             if self.num_patches is not None and idx >= self.num_patches:
+#                 return
+#             else:
+#                 yield patch
+#         self.mean = self.running_stats.avg_mean.value
+#         self.std = self.running_stats.avg_std.value
+
+#     def __iter__(self):
+#         """
+#         Iterate over data source and yield single patch.
+
+#         Yields
+#         ------
+#         np.ndarray
+#         """
+#         worker_info = torch.utils.data.get_worker_info()
+#         num_workers = worker_info.num_workers if worker_info is not None else 1
+#         yield from islice(self._generate_patches(), 0, None, num_workers)

careamics/file_io/__init__.py

@@ -0,0 +1,15 @@
+"""Functions relating reading and writing image files."""
+
+__all__ = [
+    "read",
+    "write",
+    "get_read_func",
+    "get_write_func",
+    "ReadFunc",
+    "WriteFunc",
+    "SupportedWriteType",
+]
+
+from . import read, write
+from .read import ReadFunc, get_read_func
+from .write import SupportedWriteType, WriteFunc, get_write_func

careamics/file_io/read/__init__.py

@@ -0,0 +1,12 @@
+"""Functions relating to reading image files of different formats."""
+
+__all__ = [
+    "get_read_func",
+    "read_tiff",
+    "read_zarr",
+    "ReadFunc",
+]
+
+from .get_func import ReadFunc, get_read_func
+from .tiff import read_tiff
+from .zarr import read_zarr

careamics/file_io/read/get_func.py

@@ -0,0 +1,56 @@
+"""Module to get read functions."""
+
+from pathlib import Path
+from typing import Callable, Dict, Protocol, Union
+
+from numpy.typing import NDArray
+
+from careamics.config.support import SupportedData
+
+from .tiff import read_tiff
+
+
+# This is very strict, function signature has to match including arg names
+# See WriteFunc notes
+class ReadFunc(Protocol):
+    """Protocol for type hinting read functions."""
+
+    def __call__(self, file_path: Path, *args, **kwargs) -> NDArray:
+        """
+        Type hinted callables must match this function signature (not including self).
+
+        Parameters
+        ----------
+        file_path : pathlib.Path
+            Path to file.
+        *args
+            Other positional arguments.
+        **kwargs
+            Other keyword arguments.
+        """
+
+
+READ_FUNCS: Dict[SupportedData, ReadFunc] = {
+    SupportedData.TIFF: read_tiff,
+}
+
+
+def get_read_func(data_type: Union[str, SupportedData]) -> Callable:
+    """
+    Get the read function for the data type.
+
+    Parameters
+    ----------
+    data_type : SupportedData
+        Data type.
+
+    Returns
+    -------
+    callable
+        Read function.
+    """
+    if data_type in READ_FUNCS:
+        data_type = SupportedData(data_type)  # mypy complaining about dict key type
+        return READ_FUNCS[data_type]
+    else:
+        raise NotImplementedError(f"Data type '{data_type}' is not supported.")

careamics/file_io/read/tiff.py

@@ -0,0 +1,58 @@
+"""Funtions to read tiff images."""
+
+import logging
+from fnmatch import fnmatch
+from pathlib import Path
+
+import numpy as np
+import tifffile
+
+from careamics.config.support import SupportedData
+from careamics.utils.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def read_tiff(file_path: Path, *args: list, **kwargs: dict) -> np.ndarray:
+    """
+    Read a tiff file and return a numpy array.
+
+    Parameters
+    ----------
+    file_path : Path
+        Path to a file.
+    *args : list
+        Additional arguments.
+    **kwargs : dict
+        Additional keyword arguments.
+
+    Returns
+    -------
+    np.ndarray
+        Resulting array.
+
+    Raises
+    ------
+    ValueError
+        If the file failed to open.
+    OSError
+        If the file failed to open.
+    ValueError
+        If the file is not a valid tiff.
+    ValueError
+        If the data dimensions are incorrect.
+    ValueError
+        If the axes length is incorrect.
+    """
+    if fnmatch(
+        file_path.suffix, SupportedData.get_extension_pattern(SupportedData.TIFF)
+    ):
+        try:
+            array = tifffile.imread(file_path)
+        except (ValueError, OSError) as e:
+            logging.exception(f"Exception in file {file_path}: {e}, skipping it.")
+            raise e
+    else:
+        raise ValueError(f"File {file_path} is not a valid tiff.")
+
+    return array

careamics/file_io/read/zarr.py

@@ -0,0 +1,60 @@
+"""Function to read zarr images."""
+
+from typing import Union
+
+from zarr import Group, core, hierarchy, storage
+
+
+def read_zarr(
+    zarr_source: Group, axes: str
+) -> Union[core.Array, storage.DirectoryStore, hierarchy.Group]:
+    """Read a file and returns a pointer.
+
+    Parameters
+    ----------
+    zarr_source : Group
+        Zarr storage.
+    axes : str
+        Axes of the data.
+
+    Returns
+    -------
+    np.ndarray
+        Pointer to zarr storage.
+
+    Raises
+    ------
+    ValueError, OSError
+        if a file is not a valid tiff or damaged.
+    ValueError
+        if data dimensions are not 2, 3 or 4.
+    ValueError
+        if axes parameter from config is not consistent with data dimensions.
+    """
+    if isinstance(zarr_source, hierarchy.Group):
+        array = zarr_source[0]
+
+    elif isinstance(zarr_source, storage.DirectoryStore):
+        raise NotImplementedError("DirectoryStore not supported yet")
+
+    elif isinstance(zarr_source, core.Array):
+        # array should be of shape (S, (C), (Z), Y, X), iterating over S ?
+        if zarr_source.dtype == "O":
+            raise NotImplementedError("Object type not supported yet")
+        else:
+            array = zarr_source
+    else:
+        raise ValueError(f"Unsupported zarr object type {type(zarr_source)}")
+
+    # sanity check on dimensions
+    if len(array.shape) < 2 or len(array.shape) > 4:
+        raise ValueError(
+            f"Incorrect data dimensions. Must be 2, 3 or 4 (got {array.shape})."
+        )
+
+    # sanity check on axes length
+    if len(axes) != len(array.shape):
+        raise ValueError(f"Incorrect axes length (got {axes}).")
+
+    # arr = fix_axes(arr, axes)
+    return array

careamics/file_io/write/__init__.py

@@ -0,0 +1,15 @@
+"""Functions relating to writing image files of different formats."""
+
+__all__ = [
+    "get_write_func",
+    "write_tiff",
+    "WriteFunc",
+    "SupportedWriteType",
+]
+
+from .get_func import (
+    SupportedWriteType,
+    WriteFunc,
+    get_write_func,
+)
+from .tiff import write_tiff