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

careamics/dataset/dataset_utils/running_stats.py

@@ -0,0 +1,186 @@
+"""Computing data statistics."""
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+def compute_normalization_stats(image: NDArray) -> tuple[NDArray, NDArray]:
+    """
+    Compute mean and standard deviation of an array.
+
+    Expected input shape is (S, C, (Z), Y, X). The mean and standard deviation are
+    computed per channel.
+
+    Parameters
+    ----------
+    image : NDArray
+        Input array.
+
+    Returns
+    -------
+    tuple of (list of floats, list of floats)
+        Lists of mean and standard deviation values per channel.
+    """
+    # Define the list of axes excluding the channel axis
+    axes = tuple(np.delete(np.arange(image.ndim), 1))
+    return np.mean(image, axis=axes), np.std(image, axis=axes)
+
+
+def update_iterative_stats(
+    count: NDArray, mean: NDArray, m2: NDArray, new_values: NDArray
+) -> tuple[NDArray, NDArray, NDArray]:
+    """Update the mean and variance of an array iteratively.
+
+    Parameters
+    ----------
+    count : NDArray
+        Number of elements in the array.
+    mean : NDArray
+        Mean of the array.
+    m2 : NDArray
+        Variance of the array.
+    new_values : NDArray
+        New values to add to the mean and variance.
+
+    Returns
+    -------
+    tuple[NDArray, NDArray, NDArray]
+        Updated count, mean, and variance.
+    """
+    count += np.array([np.prod(channel.shape) for channel in new_values])
+    # newvalues - oldMean
+    delta = [
+        np.subtract(v.flatten(), [m] * len(v.flatten()))
+        for v, m in zip(new_values, mean)
+    ]
+
+    mean += np.array([np.sum(d / c) for d, c in zip(delta, count)])
+    # newvalues - newMeant
+    delta2 = [
+        np.subtract(v.flatten(), [m] * len(v.flatten()))
+        for v, m in zip(new_values, mean)
+    ]
+
+    m2 += np.array([np.sum(d * d2) for d, d2 in zip(delta, delta2)])
+
+    return (count, mean, m2)
+
+
+def finalize_iterative_stats(
+    count: NDArray, mean: NDArray, m2: NDArray
+) -> tuple[NDArray, NDArray]:
+    """Finalize the mean and variance computation.
+
+    Parameters
+    ----------
+    count : NDArray
+        Number of elements in the array.
+    mean : NDArray
+        Mean of the array.
+    m2 : NDArray
+        Variance of the array.
+
+    Returns
+    -------
+    tuple[NDArray, NDArray]
+        Final mean and standard deviation.
+    """
+    std = np.array([np.sqrt(m / c) for m, c in zip(m2, count)])
+    if any(c < 2 for c in count):
+        return np.full(mean.shape, np.nan), np.full(std.shape, np.nan)
+    else:
+        return mean, std
+
+
+class WelfordStatistics:
+    """Compute Welford statistics iteratively.
+
+    The Welford algorithm is used to compute the mean and variance of an array
+    iteratively. Based on the implementation from:
+    https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_online_algorithm
+    """
+
+    def update(self, array: NDArray, sample_idx: int) -> None:
+        """Update the Welford statistics.
+
+        Parameters
+        ----------
+        array : NDArray
+            Input array.
+        sample_idx : int
+            Current sample number.
+        """
+        self.sample_idx = sample_idx
+        sample_channels = np.array(np.split(array, array.shape[1], axis=1))
+
+        # Initialize the statistics
+        if self.sample_idx == 0:
+            # Compute the mean and standard deviation
+            self.mean, _ = compute_normalization_stats(array)
+            # Initialize the count and m2 with zero-valued arrays of shape (C,)
+            self.count, self.mean, self.m2 = update_iterative_stats(
+                count=np.zeros(array.shape[1]),
+                mean=self.mean,
+                m2=np.zeros(array.shape[1]),
+                new_values=sample_channels,
+            )
+        else:
+            # Update the statistics
+            self.count, self.mean, self.m2 = update_iterative_stats(
+                count=self.count, mean=self.mean, m2=self.m2, new_values=sample_channels
+            )
+
+        self.sample_idx += 1
+
+    def finalize(self) -> tuple[NDArray, NDArray]:
+        """Finalize the Welford statistics.
+
+        Returns
+        -------
+        tuple or numpy arrays
+            Final mean and standard deviation.
+        """
+        return finalize_iterative_stats(self.count, self.mean, self.m2)
+
+
+# from multiprocessing import Value
+# from typing import tuple
+
+# import numpy as np
+
+
+# class RunningStats:
+#     """Calculates running mean and std."""
+
+#     def __init__(self) -> None:
+#         self.reset()
+
+#     def reset(self) -> None:
+#         """Reset the running stats."""
+#         self.avg_mean = Value("d", 0)
+#         self.avg_std = Value("d", 0)
+#         self.m2 = Value("d", 0)
+#         self.count = Value("i", 0)
+
+#     def init(self, mean: float, std: float) -> None:
+#         """Initialize running stats."""
+#         with self.avg_mean.get_lock():
+#             self.avg_mean.value += mean
+#         with self.avg_std.get_lock():
+#             self.avg_std.value = std
+
+#     def compute_std(self) -> tuple[float, float]:
+#         """Compute std."""
+#         if self.count.value >= 2:
+#             self.avg_std.value = np.sqrt(self.m2.value / self.count.value)
+
+#     def update(self, value: float) -> None:
+#         """Update running stats."""
+#         with self.count.get_lock():
+#             self.count.value += 1
+#         delta = value - self.avg_mean.value
+#         with self.avg_mean.get_lock():
+#             self.avg_mean.value += delta / self.count.value
+#         delta2 = value - self.avg_mean.value
+#         with self.m2.get_lock():
+#             self.m2.value += delta * delta2

careamics/dataset/in_memory_dataset.py

@@ -0,0 +1,310 @@
+"""In-memory dataset module."""
+
+from __future__ import annotations
+
+import copy
+from pathlib import Path
+from typing import Any, Callable, Optional, Union
+
+import numpy as np
+from torch.utils.data import Dataset
+
+from careamics.file_io.read import read_tiff
+from careamics.transforms import Compose
+
+from ..config import DataConfig
+from ..config.transformations import NormalizeModel
+from ..utils.logging import get_logger
+from .patching.patching import (
+    PatchedOutput,
+    Stats,
+    prepare_patches_supervised,
+    prepare_patches_supervised_array,
+    prepare_patches_unsupervised,
+    prepare_patches_unsupervised_array,
+)
+
+logger = get_logger(__name__)
+
+
+class InMemoryDataset(Dataset):
+    """Dataset storing data in memory and allowing generating patches from it.
+
+    Parameters
+    ----------
+    data_config : CAREamics DataConfig
+        (see careamics.config.data_model.DataConfig)
+        Data configuration.
+    inputs : numpy.ndarray or list[pathlib.Path]
+        Input data.
+    input_target : numpy.ndarray or list[pathlib.Path], optional
+        Target data, by default None.
+    read_source_func : Callable, optional
+        Read source function for custom types, by default read_tiff.
+    **kwargs : Any
+        Additional keyword arguments, unused.
+    """
+
+    def __init__(
+        self,
+        data_config: DataConfig,
+        inputs: Union[np.ndarray, list[Path]],
+        input_target: Optional[Union[np.ndarray, list[Path]]] = None,
+        read_source_func: Callable = read_tiff,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        data_config : DataConfig
+            Data configuration.
+        inputs : numpy.ndarray or list[pathlib.Path]
+            Input data.
+        input_target : numpy.ndarray or list[pathlib.Path], optional
+            Target data, by default None.
+        read_source_func : Callable, optional
+            Read source function for custom types, by default read_tiff.
+        **kwargs : Any
+            Additional keyword arguments, unused.
+        """
+        self.data_config = data_config
+        self.inputs = inputs
+        self.input_targets = input_target
+        self.axes = self.data_config.axes
+        self.patch_size = self.data_config.patch_size
+
+        # read function
+        self.read_source_func = read_source_func
+
+        # generate patches
+        supervised = self.input_targets is not None
+        patches_data = self._prepare_patches(supervised)
+
+        # unpack the dataclass
+        self.data = patches_data.patches
+        self.data_targets = patches_data.targets
+
+        # set image statistics
+        if self.data_config.image_means is None:
+            self.image_stats = patches_data.image_stats
+            logger.info(
+                f"Computed dataset mean: {self.image_stats.means}, "
+                f"std: {self.image_stats.stds}"
+            )
+        else:
+            self.image_stats = Stats(
+                self.data_config.image_means, self.data_config.image_stds
+            )
+
+        # set target statistics
+        if self.data_config.target_means is None:
+            self.target_stats = patches_data.target_stats
+        else:
+            self.target_stats = Stats(
+                self.data_config.target_means, self.data_config.target_stds
+            )
+
+        # update mean and std in configuration
+        # the object is mutable and should then be recorded in the CAREamist obj
+        self.data_config.set_means_and_stds(
+            image_means=self.image_stats.means,
+            image_stds=self.image_stats.stds,
+            target_means=self.target_stats.means,
+            target_stds=self.target_stats.stds,
+        )
+        # get transforms
+        self.patch_transform = Compose(
+            transform_list=[
+                NormalizeModel(
+                    image_means=self.image_stats.means,
+                    image_stds=self.image_stats.stds,
+                    target_means=self.target_stats.means,
+                    target_stds=self.target_stats.stds,
+                )
+            ]
+            + self.data_config.transforms,
+        )
+
+    def _prepare_patches(self, supervised: bool) -> PatchedOutput:
+        """
+        Iterate over data source and create an array of patches.
+
+        Parameters
+        ----------
+        supervised : bool
+            Whether the dataset is supervised or not.
+
+        Returns
+        -------
+        numpy.ndarray
+            Array of patches.
+        """
+        if supervised:
+            if isinstance(self.inputs, np.ndarray) and isinstance(
+                self.input_targets, np.ndarray
+            ):
+                return prepare_patches_supervised_array(
+                    self.inputs,
+                    self.axes,
+                    self.input_targets,
+                    self.patch_size,
+                )
+            elif isinstance(self.inputs, list) and isinstance(self.input_targets, list):
+                return prepare_patches_supervised(
+                    self.inputs,
+                    self.input_targets,
+                    self.axes,
+                    self.patch_size,
+                    self.read_source_func,
+                )
+            else:
+                raise ValueError(
+                    f"Data and target must be of the same type, either both numpy "
+                    f"arrays or both lists of paths, got {type(self.inputs)} (data) "
+                    f"and {type(self.input_targets)} (target)."
+                )
+        else:
+            if isinstance(self.inputs, np.ndarray):
+                return prepare_patches_unsupervised_array(
+                    self.inputs,
+                    self.axes,
+                    self.patch_size,
+                )
+            else:
+                return prepare_patches_unsupervised(
+                    self.inputs,
+                    self.axes,
+                    self.patch_size,
+                    self.read_source_func,
+                )
+
+    def __len__(self) -> int:
+        """
+        Return the length of the dataset.
+
+        Returns
+        -------
+        int
+            Length of the dataset.
+        """
+        return self.data.shape[0]
+
+    def __getitem__(self, index: int) -> tuple[np.ndarray, ...]:
+        """
+        Return the patch corresponding to the provided index.
+
+        Parameters
+        ----------
+        index : int
+            Index of the patch to return.
+
+        Returns
+        -------
+        tuple of numpy.ndarray
+            Patch.
+
+        Raises
+        ------
+        ValueError
+            If dataset mean and std are not set.
+        """
+        patch = self.data[index]
+
+        # if there is a target
+        if self.data_targets is not None:
+            # get target
+            target = self.data_targets[index]
+
+            return self.patch_transform(patch=patch, target=target)
+
+        elif self.data_config.has_n2v_manipulate():  # TODO not compatible with HDN
+            return self.patch_transform(patch=patch)
+        else:
+            raise ValueError(
+                "Something went wrong! No target provided (not supervised training) "
+                "and no N2V manipulation (no N2V training)."
+            )
+
+    def get_data_statistics(self) -> tuple[list[float], list[float]]:
+        """Return training data statistics.
+
+        This does not return the target data statistics, only those of the input.
+
+        Returns
+        -------
+        tuple of list of floats
+            Means and standard deviations across channels of the training data.
+        """
+        return self.image_stats.get_statistics()
+
+    def split_dataset(
+        self,
+        percentage: float = 0.1,
+        minimum_patches: int = 1,
+    ) -> InMemoryDataset:
+        """Split a new dataset away from the current one.
+
+        This method is used to extract random validation patches from the dataset.
+
+        Parameters
+        ----------
+        percentage : float, optional
+            Percentage of patches to extract, by default 0.1.
+        minimum_patches : int, optional
+            Minimum number of patches to extract, by default 5.
+
+        Returns
+        -------
+        CAREamics InMemoryDataset
+            New dataset with the extracted patches.
+
+        Raises
+        ------
+        ValueError
+            If `percentage` is not between 0 and 1.
+        ValueError
+            If `minimum_number` is not between 1 and the number of patches.
+        """
+        if percentage < 0 or percentage > 1:
+            raise ValueError(f"Percentage must be between 0 and 1, got {percentage}.")
+
+        if minimum_patches < 1 or minimum_patches > len(self):
+            raise ValueError(
+                f"Minimum number of patches must be between 1 and "
+                f"{len(self)} (number of patches), got "
+                f"{minimum_patches}. Adjust the patch size or the minimum number of "
+                f"patches."
+            )
+
+        total_patches = len(self)
+
+        # number of patches to extract (either percentage rounded or minimum number)
+        n_patches = max(round(total_patches * percentage), minimum_patches)
+
+        # get random indices
+        indices = np.random.choice(total_patches, n_patches, replace=False)
+
+        # extract patches
+        val_patches = self.data[indices]
+
+        # remove patches from self.patch
+        self.data = np.delete(self.data, indices, axis=0)
+
+        # same for targets
+        if self.data_targets is not None:
+            val_targets = self.data_targets[indices]
+            self.data_targets = np.delete(self.data_targets, indices, axis=0)
+
+        # clone the dataset
+        dataset = copy.deepcopy(self)
+
+        # reassign patches
+        dataset.data = val_patches
+
+        # reassign targets
+        if self.data_targets is not None:
+            dataset.data_targets = val_targets
+
+        return dataset

careamics/dataset/in_memory_pred_dataset.py

@@ -0,0 +1,88 @@
+"""In-memory prediction dataset."""
+
+from __future__ import annotations
+
+from numpy.typing import NDArray
+from torch.utils.data import Dataset
+
+from careamics.transforms import Compose
+
+from ..config import InferenceConfig
+from ..config.transformations import NormalizeModel
+from .dataset_utils import reshape_array
+
+
+class InMemoryPredDataset(Dataset):
+    """Simple prediction dataset returning images along the sample axis.
+
+    Parameters
+    ----------
+    prediction_config : InferenceConfig
+        Prediction configuration.
+    inputs : NDArray
+        Input data.
+    """
+
+    def __init__(
+        self,
+        prediction_config: InferenceConfig,
+        inputs: NDArray,
+    ) -> None:
+        """Constructor.
+
+        Parameters
+        ----------
+        prediction_config : InferenceConfig
+            Prediction configuration.
+        inputs : NDArray
+            Input data.
+
+        Raises
+        ------
+        ValueError
+            If data_path is not a directory.
+        """
+        self.pred_config = prediction_config
+        self.input_array = inputs
+        self.axes = self.pred_config.axes
+        self.image_means = self.pred_config.image_means
+        self.image_stds = self.pred_config.image_stds
+
+        # Reshape data
+        self.data = reshape_array(self.input_array, self.axes)
+
+        # get transforms
+        self.patch_transform = Compose(
+            transform_list=[
+                NormalizeModel(image_means=self.image_means, image_stds=self.image_stds)
+            ],
+        )
+
+    def __len__(self) -> int:
+        """
+        Return the length of the dataset.
+
+        Returns
+        -------
+        int
+            Length of the dataset.
+        """
+        return len(self.data)
+
+    def __getitem__(self, index: int) -> NDArray:
+        """
+        Return the patch corresponding to the provided index.
+
+        Parameters
+        ----------
+        index : int
+            Index of the patch to return.
+
+        Returns
+        -------
+        NDArray
+            Transformed patch.
+        """
+        transformed_patch, _ = self.patch_transform(patch=self.data[index])
+
+        return transformed_patch

careamics/dataset/in_memory_tiled_pred_dataset.py

@@ -0,0 +1,129 @@
+"""In-memory tiled prediction dataset."""
+
+from __future__ import annotations
+
+from numpy.typing import NDArray
+from torch.utils.data import Dataset
+
+from careamics.transforms import Compose
+
+from ..config import InferenceConfig
+from ..config.tile_information import TileInformation
+from ..config.transformations import NormalizeModel
+from .dataset_utils import reshape_array
+from .tiling import extract_tiles
+
+
+class InMemoryTiledPredDataset(Dataset):
+    """Prediction dataset storing data in memory and returning tiles of each image.
+
+    Parameters
+    ----------
+    prediction_config : InferenceConfig
+        Prediction configuration.
+    inputs : NDArray
+        Input data.
+    """
+
+    def __init__(
+        self,
+        prediction_config: InferenceConfig,
+        inputs: NDArray,
+    ) -> None:
+        """Constructor.
+
+        Parameters
+        ----------
+        prediction_config : InferenceConfig
+            Prediction configuration.
+        inputs : NDArray
+            Input data.
+
+        Raises
+        ------
+        ValueError
+            If data_path is not a directory.
+        """
+        if (
+            prediction_config.tile_size is None
+            or prediction_config.tile_overlap is None
+        ):
+            raise ValueError(
+                "Tile size and overlap must be provided to use the tiled prediction "
+                "dataset."
+            )
+
+        self.pred_config = prediction_config
+        self.input_array = inputs
+        self.axes = self.pred_config.axes
+        self.tile_size = prediction_config.tile_size
+        self.tile_overlap = prediction_config.tile_overlap
+        self.image_means = self.pred_config.image_means
+        self.image_stds = self.pred_config.image_stds
+
+        # Generate patches
+        self.data = self._prepare_tiles()
+
+        # get transforms
+        self.patch_transform = Compose(
+            transform_list=[
+                NormalizeModel(image_means=self.image_means, image_stds=self.image_stds)
+            ],
+        )
+
+    def _prepare_tiles(self) -> list[tuple[NDArray, TileInformation]]:
+        """
+        Iterate over data source and create an array of patches.
+
+        Returns
+        -------
+        list of tuples of NDArray and TileInformation
+            List of tiles and tile information.
+        """
+        # reshape array
+        reshaped_sample = reshape_array(self.input_array, self.axes)
+
+        # generate patches, which returns a generator
+        patch_generator = extract_tiles(
+            arr=reshaped_sample,
+            tile_size=self.tile_size,
+            overlaps=self.tile_overlap,
+        )
+        patches_list = list(patch_generator)
+
+        if len(patches_list) == 0:
+            raise ValueError("No tiles generated, ")
+
+        return patches_list
+
+    def __len__(self) -> int:
+        """
+        Return the length of the dataset.
+
+        Returns
+        -------
+        int
+            Length of the dataset.
+        """
+        return len(self.data)
+
+    def __getitem__(self, index: int) -> tuple[NDArray, TileInformation]:
+        """
+        Return the patch corresponding to the provided index.
+
+        Parameters
+        ----------
+        index : int
+            Index of the patch to return.
+
+        Returns
+        -------
+        tuple of NDArray and TileInformation
+            Transformed patch.
+        """
+        tile_array, tile_info = self.data[index]
+
+        # Apply transforms
+        transformed_tile, _ = self.patch_transform(patch=tile_array)
+
+        return transformed_tile, tile_info