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

careamics/dataset/iterable_dataset.py

@@ -0,0 +1,295 @@
+"""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 Callable, Optional
+
+import numpy as np
+from torch.utils.data import IterableDataset
+
+from careamics.config import DataConfig
+from careamics.config.transformations import NormalizeModel
+from careamics.file_io.read import read_tiff
+from careamics.transforms import Compose
+
+from ..utils.logging import get_logger
+from .dataset_utils import iterate_over_files
+from .dataset_utils.running_stats import WelfordStatistics
+from .patching.patching import Stats
+from .patching.random_patching import extract_patches_random
+
+logger = get_logger(__name__)
+
+
+class PathIterableDataset(IterableDataset):
+    """
+    Dataset allowing extracting patches w/o loading whole data into memory.
+
+    Parameters
+    ----------
+    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.
+    """
+
+    def __init__(
+        self,
+        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.read_source_func = read_source_func
+
+        # compute mean and std over the dataset
+        # 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.image_stats, self.target_stats = self._calculate_mean_and_std()
+            logger.info(
+                f"Computed dataset mean: {self.image_stats.means},"
+                f"std: {self.image_stats.stds}"
+            )
+
+            # update the mean in the config
+            self.data_config.set_means_and_stds(
+                image_means=self.image_stats.means,
+                image_stds=self.image_stats.stds,
+                target_means=(
+                    list(self.target_stats.means)
+                    if self.target_stats.means is not None
+                    else None
+                ),
+                target_stds=(
+                    list(self.target_stats.stds)
+                    if self.target_stats.stds is not None
+                    else None
+                ),
+            )
+
+        else:
+            # if mean and std are provided in the config, use them
+            self.image_stats, self.target_stats = (
+                Stats(self.data_config.image_means, self.data_config.image_stds),
+                Stats(self.data_config.target_means, self.data_config.target_stds),
+            )
+
+        # create transform composed of normalization and other 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,
+                )
+            ]
+            + data_config.transforms
+        )
+
+    def _calculate_mean_and_std(self) -> tuple[Stats, Stats]:
+        """
+        Calculate mean and std of the dataset.
+
+        Returns
+        -------
+        tuple of Stats and optional Stats
+            Data classes containing the image and target statistics.
+        """
+        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)
+
+            num_samples += 1
+
+        if num_samples == 0:
+            raise ValueError("No samples found in the dataset.")
+
+        # Average the means and stds per sample
+        image_means, image_stds = image_stats.finalize()
+
+        if target is not None:
+            target_means, target_stds = target_stats.finalize()
+
+            return (
+                Stats(image_means, image_stds),
+                Stats(np.array(target_means), np.array(target_stds)),
+            )
+        else:
+            return Stats(image_means, image_stds), Stats(None, None)
+
+    def __iter__(
+        self,
+    ) -> Generator[tuple[np.ndarray, ...], None, None]:
+        """
+        Iterate over data source and yield single patch.
+
+        Yields
+        ------
+        np.ndarray
+            Single patch.
+        """
+        assert (
+            self.image_stats.means is not None and self.image_stats.stds is not None
+        ), "Mean and std must be provided"
+
+        # iterate over files
+        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=sample_input,
+                patch_size=self.data_config.patch_size,
+                target=sample_target,
+            )
+
+            # iterate over patches
+            # patches are tuples of (patch, target) if target is available
+            # or (patch, None) only if no target is available
+            # patch is of dimensions (C)ZYX
+            for patch_data in patches:
+                yield self.patch_transform(
+                    patch=patch_data[0],
+                    target=patch_data[1],
+                )
+
+    def get_data_statistics(self) -> tuple[list[float], list[float]]:
+        """Return training data statistics.
+
+        Returns
+        -------
+        tuple of list of floats
+            Means and standard deviations across channels of the training data.
+        """
+        return self.image_stats.get_statistics()
+
+    def get_number_of_files(self) -> int:
+        """
+        Return the number of files in the dataset.
+
+        Returns
+        -------
+        int
+            Number of files in the dataset.
+        """
+        return len(self.data_files)
+
+    def split_dataset(
+        self,
+        percentage: float = 0.1,
+        minimum_number: int = 5,
+    ) -> PathIterableDataset:
+        """Split up dataset in two.
+
+        Splits the datest sing a percentage of the data (files) to extract, or the
+        minimum number of the percentage is less than the minimum number.
+
+        Parameters
+        ----------
+        percentage : float, optional
+            Percentage of files to split up, by default 0.1.
+        minimum_number : int, optional
+            Minimum number of files to split up, by default 5.
+
+        Returns
+        -------
+        IterableDataset
+            Dataset containing the split data.
+
+        Raises
+        ------
+        ValueError
+            If the percentage is smaller than 0 or larger than 1.
+        ValueError
+            If the minimum number is smaller than 1 or larger than the number of files.
+        """
+        if percentage < 0 or percentage > 1:
+            raise ValueError(f"Percentage must be between 0 and 1, got {percentage}.")
+
+        if minimum_number < 1 or minimum_number > self.get_number_of_files():
+            raise ValueError(
+                f"Minimum number of files must be between 1 and "
+                f"{self.get_number_of_files()} (number of files), got "
+                f"{minimum_number}."
+            )
+
+        # compute number of files
+        total_files = self.get_number_of_files()
+        n_files = max(round(percentage * total_files), minimum_number)
+
+        # get random indices
+        indices = np.random.choice(total_files, n_files, replace=False)
+
+        # extract files
+        val_files = [self.data_files[i] for i in indices]
+
+        # remove patches from self.patch
+        data_files = []
+        for i, file in enumerate(self.data_files):
+            if i not in indices:
+                data_files.append(file)
+        self.data_files = data_files
+
+        # same for targets
+        if self.target_files is not None:
+            val_target_files = [self.target_files[i] for i in indices]
+
+            data_target_files = []
+            for i, file in enumerate(self.target_files):
+                if i not in indices:
+                    data_target_files.append(file)
+            self.target_files = data_target_files
+
+        # clone the dataset
+        dataset = copy.deepcopy(self)
+
+        # reassign patches
+        dataset.data_files = val_files
+
+        # reassign targets
+        if self.target_files is not None:
+            dataset.target_files = val_target_files
+
+        return dataset

careamics/dataset/iterable_pred_dataset.py

@@ -0,0 +1,122 @@
+"""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.file_io.read import read_tiff
+from careamics.transforms import Compose
+
+from ..config import InferenceConfig
+from ..config.transformations import NormalizeModel
+from .dataset_utils import iterate_over_files
+
+
+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

careamics/dataset/iterable_tiled_pred_dataset.py

@@ -0,0 +1,140 @@
+"""Iterable tiled 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.file_io.read import read_tiff
+from careamics.transforms import Compose
+
+from ..config import InferenceConfig
+from ..config.tile_information import TileInformation
+from ..config.transformations import NormalizeModel
+from .dataset_utils import iterate_over_files
+from .tiling import extract_tiles
+
+
+class IterableTiledPredDataset(IterableDataset):
+    """Tiled prediction dataset.
+
+    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.
+
+    Attributes
+    ----------
+    data_path : str or pathlib.Path
+        Path to the data, must be a directory.
+    axes : str
+        Description of axes in format STCZYX.
+    mean : float, optional
+        Expected mean of the dataset, by default None.
+    std : float, optional
+        Expected standard deviation of the dataset, by default None.
+    patch_transform : 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[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.
+        """
+        if (
+            prediction_config.tile_size is None
+            or prediction_config.tile_overlap is None
+        ):
+            raise ValueError(
+                "Tile size and overlap must be provided for tiled prediction."
+            )
+
+        self.prediction_config = prediction_config
+        self.data_files = src_files
+        self.axes = prediction_config.axes
+        self.tile_size = prediction_config.tile_size
+        self.tile_overlap = prediction_config.tile_overlap
+        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[tuple[NDArray, TileInformation], None, None]:
+        """
+        Iterate over data source and yield single patch.
+
+        Yields
+        ------
+        Generator of NDArray and TileInformation tuple
+            Generator of single tiles.
+        """
+        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,
+        ):
+            # generate patches, return a generator of single tiles
+            patch_gen = extract_tiles(
+                arr=sample,
+                tile_size=self.tile_size,
+                overlaps=self.tile_overlap,
+            )
+
+            # 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/patching/__init__.py

@@ -0,0 +1 @@
+"""Patching and tiling functions."""