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

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

careamics/dataset/iterable_dataset.py

@@ -1,20 +1,27 @@
+"""Iterable dataset used to load data file by file."""
+
 from __future__ import annotations
 
 import copy
+from collections.abc import Generator
 from pathlib import Path
-from typing import Any, Callable, Generator, List, Optional, Tuple, Union
+from typing import Callable, Optional
 
 import numpy as np
-from torch.utils.data import IterableDataset, get_worker_info
+from torch.utils.data import IterableDataset
 
+from careamics.config import DataConfig
+from careamics.config.transformations import NormalizeModel
 from careamics.transforms import Compose
 
-from ..config import DataConfig, InferenceConfig
-from ..config.tile_information import TileInformation
 from ..utils.logging import get_logger
-from .dataset_utils import read_tiff, reshape_array
+from .dataset_utils import (
+    iterate_over_files,
+    read_tiff,
+)
+from .dataset_utils.running_stats import WelfordStatistics
+from .patching.patching import Stats, StatsOutput
 from .patching.random_patching import extract_patches_random
-from .patching.tiled_patching import extract_tiles
 
 logger = get_logger(__name__)
 
@@ -25,129 +32,142 @@ class PathIterableDataset(IterableDataset):
 
     Parameters
     ----------
-    data_path : Union[str, Path]
+    data_config : DataConfig
+        Data configuration.
+    src_files : list of pathlib.Path
+        List of data files.
+    target_files : list of pathlib.Path, optional
+        Optional list of target files, by default None.
+    read_source_func : Callable, optional
+        Read source function for custom types, by default read_tiff.
+
+    Attributes
+    ----------
+    data_path : list of pathlib.Path
         Path to the data, must be a directory.
     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.
-    patch_overlap : Optional[Union[List[int], Tuple[int]]], optional
-        Overlap of the patches in each dimension, 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.
     """
 
     def __init__(
         self,
-        data_config: Union[DataConfig, InferenceConfig],
-        src_files: List[Path],
-        target_files: Optional[List[Path]] = None,
+        data_config: DataConfig,
+        src_files: list[Path],
+        target_files: Optional[list[Path]] = None,
         read_source_func: Callable = read_tiff,
     ) -> None:
+        """Constructors.
+
+        Parameters
+        ----------
+        data_config : DataConfig
+            Data configuration.
+        src_files : list[Path]
+            List of data files.
+        target_files : list[Path] or None, optional
+            Optional list of target files, by default None.
+        read_source_func : Callable, optional
+            Read source function for custom types, by default read_tiff.
+        """
         self.data_config = data_config
         self.data_files = src_files
         self.target_files = target_files
-        self.data_config = data_config
         self.read_source_func = read_source_func
 
         # compute mean and std over the dataset
-        if not data_config.mean or not data_config.std:
-            self.mean, self.std = self._calculate_mean_and_std()
+        # only checking the image_mean because the DataConfig class ensures that
+        # if image_mean is provided, image_std is also provided
+        if not self.data_config.image_means:
+            self.data_stats = self._calculate_mean_and_std()
+            logger.info(
+                f"Computed dataset mean: {self.data_stats.image_stats.means},"
+                f"std: {self.data_stats.image_stats.stds}"
+            )
+
+            # update the mean in the config
+            self.data_config.set_mean_and_std(
+                image_means=self.data_stats.image_stats.means,
+                image_stds=self.data_stats.image_stats.stds,
+                target_means=(
+                    list(self.data_stats.target_stats.means)
+                    if self.data_stats.target_stats.means is not None
+                    else None
+                ),
+                target_stds=(
+                    list(self.data_stats.target_stats.stds)
+                    if self.data_stats.target_stats.stds is not None
+                    else None
+                ),
+            )
 
-            # update mean and std in configuration
-            # the object is mutable and should then be recorded in the CAREamist
-            data_config.set_mean_and_std(self.mean, self.std)
         else:
-            self.mean = data_config.mean
-            self.std = data_config.std
+            # if mean and std are provided in the config, use them
+            self.data_stats = StatsOutput(
+                Stats(self.data_config.image_means, self.data_config.image_stds),
+                Stats(self.data_config.target_means, self.data_config.target_stds),
+            )
 
-        # get transforms
-        self.patch_transform = Compose(transform_list=data_config.transforms)
+        # create transform composed of normalization and other transforms
+        self.patch_transform = Compose(
+            transform_list=[
+                NormalizeModel(
+                    image_means=self.data_stats.image_stats.means,
+                    image_stds=self.data_stats.image_stats.stds,
+                    target_means=self.data_stats.target_stats.means,
+                    target_stds=self.data_stats.target_stats.stds,
+                )
+            ]
+            + data_config.transforms
+        )
 
-    def _calculate_mean_and_std(self) -> Tuple[float, float]:
+    def _calculate_mean_and_std(self) -> StatsOutput:
         """
         Calculate mean and std of the dataset.
 
         Returns
         -------
-        Tuple[float, float]
-            Tuple containing mean and standard deviation.
+        PatchedOutput
+            Data class containing the image statistics.
         """
-        means, stds = 0, 0
         num_samples = 0
+        image_stats = WelfordStatistics()
+        if self.target_files is not None:
+            target_stats = WelfordStatistics()
+
+        for sample, target in iterate_over_files(
+            self.data_config, self.data_files, self.target_files, self.read_source_func
+        ):
+            # update the image statistics
+            image_stats.update(sample, num_samples)
+
+            # update the target statistics if target is available
+            if target is not None:
+                target_stats.update(target, num_samples)
 
-        for sample, _ in self._iterate_over_files():
-            means += sample.mean()
-            stds += sample.std()
             num_samples += 1
 
         if num_samples == 0:
             raise ValueError("No samples found in the dataset.")
 
-        result_mean = means / num_samples
-        result_std = stds / num_samples
-
-        logger.info(f"Calculated mean and std for {num_samples} images")
-        logger.info(f"Mean: {result_mean}, std: {result_std}")
-        return result_mean, result_std
+        # Average the means and stds per sample
+        image_means, image_stds = image_stats.finalize()
 
-    def _iterate_over_files(
-        self,
-    ) -> Generator[Tuple[np.ndarray, Optional[np.ndarray]], None, None]:
-        """
-        Iterate over data source and yield whole image.
+        if target is not None:
+            target_means, target_stds = target_stats.finalize()
 
-        Yields
-        ------
-        np.ndarray
-            Image.
-        """
-        # When num_workers > 0, each worker process will have a different copy of the
-        # dataset object
-        # Configuring each copy independently to avoid having duplicate data returned
-        # from the workers
-        worker_info = get_worker_info()
-        worker_id = worker_info.id if worker_info is not None else 0
-        num_workers = worker_info.num_workers if worker_info is not None else 1
-
-        # iterate over the files
-        for i, filename in enumerate(self.data_files):
-            # retrieve file corresponding to the worker id
-            if i % num_workers == worker_id:
-                try:
-                    # read data
-                    sample = self.read_source_func(filename, self.data_config.axes)
-
-                    # read target, if available
-                    if self.target_files is not None:
-                        if filename.name != self.target_files[i].name:
-                            raise ValueError(
-                                f"File {filename} does not match target file "
-                                f"{self.target_files[i]}. Have you passed sorted "
-                                f"arrays?"
-                            )
-
-                        # read target
-                        target = self.read_source_func(
-                            self.target_files[i], self.data_config.axes
-                        )
-
-                        yield sample, target
-                    else:
-                        yield sample, None
-
-                except Exception as e:
-                    logger.error(f"Error reading file {filename}: {e}")
+        logger.info(f"Calculated mean and std for {num_samples} images")
+        logger.info(f"Mean: {image_means}, std: {image_stds}")
+        return StatsOutput(
+            Stats(image_means, image_stds),
+            Stats(
+                np.array(target_means) if target is not None else None,
+                np.array(target_stds) if target is not None else None,
+            ),
+        )
 
     def __iter__(
         self,
-    ) -> Generator[Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]], None, None]:
+    ) -> Generator[tuple[np.ndarray, ...], None, None]:
         """
         Iterate over data source and yield single patch.
 
@@ -157,22 +177,18 @@ class PathIterableDataset(IterableDataset):
             Single patch.
         """
         assert (
-            self.mean is not None and self.std is not None
+            self.data_stats.image_stats.means is not None
+            and self.data_stats.image_stats.stds is not None
         ), "Mean and std must be provided"
 
         # iterate over files
-        for sample_input, sample_target in self._iterate_over_files():
-            reshaped_sample = reshape_array(sample_input, self.data_config.axes)
-            reshaped_target = (
-                None
-                if sample_target is None
-                else reshape_array(sample_target, self.data_config.axes)
-            )
-
+        for sample_input, sample_target in iterate_over_files(
+            self.data_config, self.data_files, self.target_files, self.read_source_func
+        ):
             patches = extract_patches_random(
-                arr=reshaped_sample,
+                arr=sample_input,
                 patch_size=self.data_config.patch_size,
-                target=reshaped_target,
+                target=sample_target,
             )
 
             # iterate over patches
@@ -209,9 +225,9 @@ class PathIterableDataset(IterableDataset):
         Parameters
         ----------
         percentage : float, optional
-            Percentage of files to split up, by default 0.1
+            Percentage of files to split up, by default 0.1.
         minimum_number : int, optional
-            Minimum number of files to split up, by default 5
+            Minimum number of files to split up, by default 5.
 
         Returns
         -------
@@ -273,89 +289,3 @@ class PathIterableDataset(IterableDataset):
             dataset.target_files = val_target_files
 
         return dataset
-
-
-class IterablePredictionDataset(PathIterableDataset):
-    """
-    Dataset allowing extracting patches w/o loading whole data into memory.
-
-    Parameters
-    ----------
-    data_path : Union[str, Path]
-        Path to the data, must be a directory.
-    axes : str
-        Description of axes in format STCZYX.
-    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.
-    """
-
-    def __init__(
-        self,
-        prediction_config: InferenceConfig,
-        src_files: List[Path],
-        read_source_func: Callable = read_tiff,
-        **kwargs: Any,
-    ) -> None:
-        super().__init__(
-            data_config=prediction_config,
-            src_files=src_files,
-            read_source_func=read_source_func,
-        )
-
-        self.prediction_config = prediction_config
-        self.axes = prediction_config.axes
-        self.tile_size = self.prediction_config.tile_size
-        self.tile_overlap = self.prediction_config.tile_overlap
-        self.read_source_func = read_source_func
-
-        # tile only if both tile size and overlaps are provided
-        self.tile = self.tile_size is not None and self.tile_overlap is not None
-
-        # get tta transforms
-        self.patch_transform = Compose(
-            transform_list=prediction_config.transforms,
-        )
-
-    def __iter__(
-        self,
-    ) -> Generator[Tuple[np.ndarray, TileInformation], None, None]:
-        """
-        Iterate over data source and yield single patch.
-
-        Yields
-        ------
-        np.ndarray
-            Single patch.
-        """
-        assert (
-            self.mean is not None and self.std is not None
-        ), "Mean and std must be provided"
-
-        for sample, _ in self._iterate_over_files():
-            # reshape array
-            reshaped_sample = reshape_array(sample, self.axes)
-
-            if self.tile:
-                # generate patches, return a generator
-                patch_gen = extract_tiles(
-                    arr=reshaped_sample,
-                    tile_size=self.tile_size,
-                    overlaps=self.tile_overlap,
-                )
-            else:
-                # just wrap the sample in a generator with default tiling info
-                array_shape = reshaped_sample.squeeze().shape
-                patch_gen = (
-                    (reshaped_sample, TileInformation(array_shape=array_shape))
-                    for _ in range(1)
-                )
-
-            # apply transform to patches
-            for patch_array, tile_info in patch_gen:
-                transformed_patch, _ = self.patch_transform(patch=patch_array)
-
-                yield transformed_patch, tile_info

careamics/dataset/iterable_pred_dataset.py

@@ -0,0 +1,121 @@
+"""Iterable prediction dataset used to load data file by file."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any, Callable, Generator
+
+from numpy.typing import NDArray
+from torch.utils.data import IterableDataset
+
+from careamics.transforms import Compose
+
+from ..config import InferenceConfig
+from ..config.transformations import NormalizeModel
+from .dataset_utils import iterate_over_files, read_tiff
+
+
+class IterablePredDataset(IterableDataset):
+    """Simple iterable prediction dataset.
+
+    Parameters
+    ----------
+    prediction_config : InferenceConfig
+        Inference configuration.
+    src_files : List[Path]
+        List of data files.
+    read_source_func : Callable, optional
+        Read source function for custom types, by default read_tiff.
+    **kwargs : Any
+        Additional keyword arguments, unused.
+
+    Attributes
+    ----------
+    data_path : Union[str, Path]
+        Path to the data, must be a directory.
+    axes : str
+        Description of axes in format STCZYX.
+    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.
+    """
+
+    def __init__(
+        self,
+        prediction_config: InferenceConfig,
+        src_files: list[Path],
+        read_source_func: Callable = read_tiff,
+        **kwargs: Any,
+    ) -> None:
+        """Constructor.
+
+        Parameters
+        ----------
+        prediction_config : InferenceConfig
+            Inference configuration.
+        src_files : list of pathlib.Path
+            List of data files.
+        read_source_func : Callable, optional
+            Read source function for custom types, by default read_tiff.
+        **kwargs : Any
+            Additional keyword arguments, unused.
+
+        Raises
+        ------
+        ValueError
+            If mean and std are not provided in the inference configuration.
+        """
+        self.prediction_config = prediction_config
+        self.data_files = src_files
+        self.axes = prediction_config.axes
+        self.read_source_func = read_source_func
+
+        # check mean and std and create normalize transform
+        if (
+            self.prediction_config.image_means is None
+            or self.prediction_config.image_stds is None
+        ):
+            raise ValueError("Mean and std must be provided for prediction.")
+        else:
+            self.image_means = self.prediction_config.image_means
+            self.image_stds = self.prediction_config.image_stds
+
+        # instantiate normalize transform
+        self.patch_transform = Compose(
+            transform_list=[
+                NormalizeModel(
+                    image_means=self.image_means,
+                    image_stds=self.image_stds,
+                )
+            ],
+        )
+
+    def __iter__(
+        self,
+    ) -> Generator[NDArray, None, None]:
+        """
+        Iterate over data source and yield single patch.
+
+        Yields
+        ------
+        NDArray
+            Single patch.
+        """
+        assert (
+            self.image_means is not None and self.image_stds is not None
+        ), "Mean and std must be provided"
+
+        for sample, _ in iterate_over_files(
+            self.prediction_config,
+            self.data_files,
+            read_source_func=self.read_source_func,
+        ):
+            # sample has S dimension
+            for i in range(sample.shape[0]):
+
+                transformed_sample, _ = self.patch_transform(patch=sample[i])
+
+                yield transformed_sample