careamics 0.1.0rc5__py3-none-any.whl → 0.1.0rc7__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

@@ -4,47 +4,73 @@ from __future__ import annotations
 
 import copy
 from pathlib import Path
-from typing import Any, Callable, List, Optional, Tuple, Union
+from typing import Any, Callable, Optional, Union
 
 import numpy as np
 from torch.utils.data import Dataset
 
 from careamics.transforms import Compose
 
-from ..config import DataConfig, InferenceConfig
-from ..config.tile_information import TileInformation
+from ..config import DataConfig
+from ..config.transformations import NormalizeModel
 from ..utils.logging import get_logger
-from .dataset_utils import read_tiff, reshape_array
+from .dataset_utils import read_tiff
 from .patching.patching import (
+    PatchedOutput,
     prepare_patches_supervised,
     prepare_patches_supervised_array,
     prepare_patches_unsupervised,
     prepare_patches_unsupervised_array,
 )
-from .patching.tiled_patching import extract_tiles
 
 logger = get_logger(__name__)
 
 
 class InMemoryDataset(Dataset):
-    """Dataset storing data in memory and allowing generating patches from it."""
+    """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]],
-        data_target: Optional[Union[np.ndarray, List[Path]]] = None,
+        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.
 
-        # TODO
+        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.data_target = data_target
+        self.input_targets = input_target
         self.axes = self.data_config.axes
         self.patch_size = self.data_config.patch_size
 
@@ -52,30 +78,52 @@ class InMemoryDataset(Dataset):
         self.read_source_func = read_source_func
 
         # Generate patches
-        supervised = self.data_target is not None
-        patches = self._prepare_patches(supervised)
-
-        # Add results to members
-        self.data, self.data_targets, computed_mean, computed_std = patches
-
-        if not self.data_config.mean or not self.data_config.std:
-            self.mean, self.std = computed_mean, computed_std
-            logger.info(f"Computed dataset mean: {self.mean}, std: {self.std}")
-
-            # update mean and std in configuration
-            # the object is mutable and should then be recorded in the CAREamist obj
-            self.data_config.set_mean_and_std(self.mean, self.std)
+        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
+
+        if self.data_config.image_means is None:
+            self.image_means = patches_data.image_stats.means
+            self.image_stds = patches_data.image_stats.stds
+            logger.info(
+                f"Computed dataset mean: {self.image_means}, std: {self.image_stds}"
+            )
         else:
-            self.mean, self.std = self.data_config.mean, self.data_config.std
+            self.image_means = self.data_config.image_means
+            self.image_stds = self.data_config.image_stds
 
+        if self.data_config.target_means is None:
+            self.target_means = patches_data.target_stats.means
+            self.target_stds = patches_data.target_stats.stds
+        else:
+            self.target_means = self.data_config.target_means
+            self.target_stds = 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_mean_and_std(
+            image_means=self.image_means,
+            image_stds=self.image_stds,
+            target_means=self.target_means,
+            target_stds=self.target_stds,
+        )
         # get transforms
         self.patch_transform = Compose(
-            transform_list=self.data_config.transforms,
+            transform_list=[
+                NormalizeModel(
+                    image_means=self.image_means,
+                    image_stds=self.image_stds,
+                    target_means=self.target_means,
+                    target_stds=self.target_stds,
+                )
+            ]
+            + self.data_config.transforms,
         )
 
-    def _prepare_patches(
-        self, supervised: bool
-    ) -> Tuple[np.ndarray, Optional[np.ndarray], float, float]:
+    def _prepare_patches(self, supervised: bool) -> PatchedOutput:
         """
         Iterate over data source and create an array of patches.
 
@@ -86,23 +134,23 @@ class InMemoryDataset(Dataset):
 
         Returns
         -------
-        np.ndarray
+        numpy.ndarray
             Array of patches.
         """
         if supervised:
             if isinstance(self.inputs, np.ndarray) and isinstance(
-                self.data_target, np.ndarray
+                self.input_targets, np.ndarray
             ):
                 return prepare_patches_supervised_array(
                     self.inputs,
                     self.axes,
-                    self.data_target,
+                    self.input_targets,
                     self.patch_size,
                 )
-            elif isinstance(self.inputs, list) and isinstance(self.data_target, list):
+            elif isinstance(self.inputs, list) and isinstance(self.input_targets, list):
                 return prepare_patches_supervised(
                     self.inputs,
-                    self.data_target,
+                    self.input_targets,
                     self.axes,
                     self.patch_size,
                     self.read_source_func,
@@ -111,7 +159,7 @@ class InMemoryDataset(Dataset):
                 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.data_target)} (target)."
+                    f"and {type(self.input_targets)} (target)."
                 )
         else:
             if isinstance(self.inputs, np.ndarray):
@@ -137,9 +185,9 @@ class InMemoryDataset(Dataset):
         int
             Length of the dataset.
         """
-        return len(self.data)
+        return self.data.shape[0]
 
-    def __getitem__(self, index: int) -> Tuple[np.ndarray]:
+    def __getitem__(self, index: int) -> tuple[np.ndarray, ...]:
         """
         Return the patch corresponding to the provided index.
 
@@ -150,7 +198,7 @@ class InMemoryDataset(Dataset):
 
         Returns
         -------
-        Tuple[np.ndarray]
+        tuple of numpy.ndarray
             Patch.
 
         Raises
@@ -161,13 +209,13 @@ class InMemoryDataset(Dataset):
         patch = self.data[index]
 
         # if there is a target
-        if self.data_target is not None:
+        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():
+        elif self.data_config.has_n2v_manipulate():  # TODO not compatible with HDN
             return self.patch_transform(patch=patch)
         else:
             raise ValueError(
@@ -193,7 +241,7 @@ class InMemoryDataset(Dataset):
 
         Returns
         -------
-        InMemoryDataset
+        CAREamics InMemoryDataset
             New dataset with the extracted patches.
 
         Raises
@@ -244,117 +292,3 @@ class InMemoryDataset(Dataset):
             dataset.data_targets = val_targets
 
         return dataset
-
-
-class InMemoryPredictionDataset(Dataset):
-    """
-    Dataset storing data in memory and allowing generating patches from it.
-
-    # TODO
-    """
-
-    def __init__(
-        self,
-        prediction_config: InferenceConfig,
-        inputs: np.ndarray,
-        data_target: Optional[np.ndarray] = None,
-        read_source_func: Optional[Callable] = read_tiff,
-    ) -> None:
-        """Constructor.
-
-        Parameters
-        ----------
-        array : np.ndarray
-            Array containing the data.
-        axes : str
-            Description of axes in format STCZYX.
-
-        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.tile_size = self.pred_config.tile_size
-        self.tile_overlap = self.pred_config.tile_overlap
-        self.mean = self.pred_config.mean
-        self.std = self.pred_config.std
-        self.data_target = data_target
-
-        # tiling only if both tile size and overlap are provided
-        self.tiling = self.tile_size is not None and self.tile_overlap is not None
-
-        # read function
-        self.read_source_func = read_source_func
-
-        # Generate patches
-        self.data = self._prepare_tiles()
-        self.mean, self.std = self.pred_config.mean, self.pred_config.std
-
-        # get transforms
-        self.patch_transform = Compose(
-            transform_list=self.pred_config.transforms,
-        )
-
-    def _prepare_tiles(self) -> List[Tuple[np.ndarray, TileInformation]]:
-        """
-        Iterate over data source and create an array of patches.
-
-        Returns
-        -------
-        List[XArrayTile]
-            List of tiles.
-        """
-        # reshape array
-        reshaped_sample = reshape_array(self.input_array, self.axes)
-
-        if self.tiling:
-            # 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
-        else:
-            array_shape = reshaped_sample.squeeze().shape
-            return [(reshaped_sample, TileInformation(array_shape=array_shape))]
-
-    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[np.ndarray, TileInformation]:
-        """
-        Return the patch corresponding to the provided index.
-
-        Parameters
-        ----------
-        index : int
-            Index of the patch to return.
-
-        Returns
-        -------
-        Tuple[np.ndarray, 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

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