careamics 0.1.0rc2__py3-none-any.whl → 0.1.0rc3__py3-none-any.whl

careamics/dataset/iterable_dataset.py

@@ -0,0 +1,416 @@
+from __future__ import annotations
+
+import copy
+from pathlib import Path
+from typing import Any, Callable, Generator, List, Optional, Tuple, Union
+
+import numpy as np
+from torch.utils.data import IterableDataset, get_worker_info
+
+from ..config import DataModel, InferenceModel
+from ..config.tile_information import TileInformation
+from ..utils.logging import get_logger
+from .dataset_utils import read_tiff, reshape_array
+from .patching import (
+    get_patch_transform,
+)
+from .patching.random_patching import extract_patches_random
+from .patching.tiled_patching import extract_tiles
+
+logger = get_logger(__name__)
+
+
+class PathIterableDataset(IterableDataset):
+    """
+    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.
+    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[DataModel, InferenceModel],
+        src_files: List[Path],
+        target_files: Optional[List[Path]] = None,
+        read_source_func: Callable = read_tiff,
+    ) -> None:
+        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()
+
+            # if the transforms are not an instance of Compose
+            # Check if the data_config is an instance of DataModel or InferenceModel
+            # isinstance isn't working properly here
+            if hasattr(data_config, "has_transform_list"):
+                if data_config.has_transform_list():
+                    # 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:
+                data_config.set_mean_and_std(self.mean, self.std)
+
+        else:
+            self.mean = data_config.mean
+            self.std = data_config.std
+
+        # get transforms
+        self.patch_transform = get_patch_transform(
+            patch_transforms=data_config.transforms,
+            with_target=target_files is not None,
+        )
+
+    def _calculate_mean_and_std(self) -> Tuple[float, float]:
+        """
+        Calculate mean and std of the dataset.
+
+        Returns
+        -------
+        Tuple[float, float]
+            Tuple containing mean and standard deviation.
+        """
+        means, stds = 0, 0
+        num_samples = 0
+
+        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
+
+    def _iterate_over_files(
+        self,
+    ) -> Generator[Tuple[np.ndarray, Optional[np.ndarray]], None, None]:
+        """
+        Iterate over data source and yield whole image.
+
+        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}")
+
+    def __iter__(
+        self,
+    ) -> Generator[Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]], 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"
+
+        # 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)
+            )
+
+            patches = extract_patches_random(
+                arr=reshaped_sample,
+                patch_size=self.data_config.patch_size,
+                target=reshaped_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:
+                # if there is a target
+                if self.target_files is not None:
+                    # Albumentations expects the channel dimension to be last
+                    # Taking the first element because patch_data can include target
+                    c_patch = np.moveaxis(patch_data[0], 0, -1)
+                    c_target = np.moveaxis(patch_data[1], 0, -1)
+
+                    # apply the transform to the patch and the target
+                    transformed = self.patch_transform(
+                        image=c_patch,
+                        target=c_target,
+                    )
+
+                    # move the axes back to the original position
+                    c_patch = np.moveaxis(transformed["image"], -1, 0)
+                    c_target = np.moveaxis(transformed["target"], -1, 0)
+
+                    yield (c_patch, c_target)
+                elif self.data_config.has_n2v_manipulate():
+                    # Albumentations expects the channel dimension to be last
+                    # Taking the first element because patch_data can include target
+                    patch = np.moveaxis(patch_data[0], 0, -1)
+
+                    # apply transform
+                    transformed = self.patch_transform(image=patch)
+
+                    # retrieve the output of ManipulateN2V
+                    results = transformed["image"]
+                    masked_patch: np.ndarray = results[0]
+                    original_patch: np.ndarray = results[1]
+                    mask: np.ndarray = results[2]
+
+                    # move C axes back
+                    masked_patch = np.moveaxis(masked_patch, -1, 0)
+                    original_patch = np.moveaxis(original_patch, -1, 0)
+                    mask = np.moveaxis(mask, -1, 0)
+
+                    yield (masked_patch, original_patch, mask)
+                else:
+                    raise ValueError(
+                        "Something went wrong! Not target file (no supervised "
+                        "training) and no N2V transform (no n2v training either)."
+                    )
+
+    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
+
+
+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: InferenceModel,
+        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 = get_patch_transform(
+            patch_transforms=prediction_config.transforms,
+            with_target=False,
+        )
+
+    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:
+                # albumentations expects the channel dimension to be last
+                patch = np.moveaxis(patch_array, 0, -1)
+                transformed_patch = self.patch_transform(image=patch)
+                transformed_patch = np.moveaxis(transformed_patch["image"], -1, 0)
+
+                yield transformed_patch, tile_info

careamics/dataset/patching/__init__.py

@@ -0,0 +1,8 @@
+"""Patching and tiling functions."""
+
+
+__all__ = [
+    "get_patch_transform",
+]
+
+from .patch_transform import get_patch_transform

careamics/dataset/patching/patch_transform.py

@@ -0,0 +1,44 @@
+from typing import List, Union
+
+import albumentations as Aug
+
+from careamics.config.data_model import TRANSFORMS_UNION
+from careamics.transforms import get_all_transforms
+
+
+# TODO add some explanations on how the additional_targets is used
+def get_patch_transform(
+    patch_transforms: Union[List[TRANSFORMS_UNION], Aug.Compose],
+    with_target: bool,
+    normalize_mask: bool = True,
+) -> Aug.Compose:
+    """Return a pixel manipulation function."""
+    # if we passed a Compose, we just return it
+    if isinstance(patch_transforms, Aug.Compose):
+        return patch_transforms
+
+    # empty list of transforms is a NoOp
+    elif len(patch_transforms) == 0:
+        return Aug.Compose(
+            [Aug.NoOp()],
+            additional_targets={},  # TODO this part need be checked (wrt segmentation)
+        )
+
+    # else we have a list of transforms
+    else:
+        # retrieve all transforms
+        all_transforms = get_all_transforms()
+
+        # instantiate all transforms
+        transforms = [
+            all_transforms[transform.name](**transform.model_dump())
+            for transform in patch_transforms
+        ]
+
+        return Aug.Compose(
+            transforms,
+            # apply image aug to "target"
+            additional_targets={"target": "image"}
+            if (with_target and normalize_mask)  # TODO check this
+            else {},
+        )

careamics/dataset/patching/patching.py

@@ -0,0 +1,212 @@
+"""
+Tiling submodule.
+
+These functions are used to tile images into patches or tiles.
+"""
+from pathlib import Path
+from typing import Callable, List, Tuple, Union
+
+import numpy as np
+
+from ...utils.logging import get_logger
+from ..dataset_utils import reshape_array
+from .sequential_patching import extract_patches_sequential
+
+logger = get_logger(__name__)
+
+
+# called by in memory dataset
+def prepare_patches_supervised(
+    train_files: List[Path],
+    target_files: List[Path],
+    axes: str,
+    patch_size: Union[List[int], Tuple[int]],
+    read_source_func: Callable,
+) -> Tuple[np.ndarray, np.ndarray, float, float]:
+    """
+    Iterate over data source and create an array of patches and corresponding targets.
+
+    Returns
+    -------
+    np.ndarray
+        Array of patches.
+    """
+    train_files.sort()
+    target_files.sort()
+
+    means, stds, num_samples = 0, 0, 0
+    all_patches, all_targets = [], []
+    for train_filename, target_filename in zip(train_files, target_files):
+        try:
+            sample: np.ndarray = read_source_func(train_filename, axes)
+            target: np.ndarray = read_source_func(target_filename, axes)
+            means += sample.mean()
+            stds += sample.std()
+            num_samples += 1
+
+            # reshape array
+            sample = reshape_array(sample, axes)
+            target = reshape_array(target, axes)
+
+            # generate patches, return a generator
+            patches, targets = extract_patches_sequential(
+                sample, patch_size=patch_size, target=target
+            )
+
+            # convert generator to list and add to all_patches
+            all_patches.append(patches)
+
+            # ensure targets are not None (type checking)
+            if targets is not None:
+                all_targets.append(targets)
+            else:
+                raise ValueError(f"No target found for {target_filename}.")
+
+        except Exception as e:
+            # emit warning and continue
+            logger.error(f"Failed to read {train_filename} or {target_filename}: {e}")
+
+    # raise error if no valid samples found
+    if num_samples == 0 or len(all_patches) == 0:
+        raise ValueError(
+            f"No valid samples found in the input data: {train_files} and "
+            f"{target_files}."
+        )
+
+    result_mean, result_std = means / num_samples, stds / num_samples
+
+    patch_array: np.ndarray = np.concatenate(all_patches, axis=0)
+    target_array: np.ndarray = np.concatenate(all_targets, axis=0)
+    logger.info(f"Extracted {patch_array.shape[0]} patches from input array.")
+
+    return (
+        patch_array,
+        target_array,
+        result_mean,
+        result_std,
+    )
+
+
+# called by in_memory_dataset
+def prepare_patches_unsupervised(
+    train_files: List[Path],
+    axes: str,
+    patch_size: Union[List[int], Tuple[int]],
+    read_source_func: Callable,
+) -> Tuple[np.ndarray, None, float, float]:
+    """
+    Iterate over data source and create an array of patches.
+
+    Returns
+    -------
+    np.ndarray
+        Array of patches.
+    """
+    means, stds, num_samples = 0, 0, 0
+    all_patches = []
+    for filename in train_files:
+        try:
+            sample: np.ndarray = read_source_func(filename, axes)
+            means += sample.mean()
+            stds += sample.std()
+            num_samples += 1
+
+            # reshape array
+            sample = reshape_array(sample, axes)
+
+            # generate patches, return a generator
+            patches, _ = extract_patches_sequential(sample, patch_size=patch_size)
+
+            # convert generator to list and add to all_patches
+            all_patches.append(patches)
+        except Exception as e:
+            # emit warning and continue
+            logger.error(f"Failed to read {filename}: {e}")
+
+    # raise error if no valid samples found
+    if num_samples == 0:
+        raise ValueError(f"No valid samples found in the input data: {train_files}.")
+
+    result_mean, result_std = means / num_samples, stds / num_samples
+
+    patch_array: np.ndarray = np.concatenate(all_patches)
+    logger.info(f"Extracted {patch_array.shape[0]} patches from input array.")
+
+    return patch_array, _, result_mean, result_std  # TODO return object?
+
+
+# called on arrays by in memory dataset
+def prepare_patches_supervised_array(
+    data: np.ndarray,
+    axes: str,
+    data_target: np.ndarray,
+    patch_size: Union[List[int], Tuple[int]],
+) -> Tuple[np.ndarray, np.ndarray, float, float]:
+    """Iterate over data source and create an array of patches.
+
+    This method expects an array of shape SC(Z)YX, where S and C can be singleton
+    dimensions.
+
+    Patches returned are of shape SC(Z)YX, where S is now the patches dimension.
+
+    Returns
+    -------
+    np.ndarray
+        Array of patches.
+    """
+    # compute statistics
+    mean = data.mean()
+    std = data.std()
+
+    # reshape array
+    reshaped_sample = reshape_array(data, axes)
+    reshaped_target = reshape_array(data_target, axes)
+
+    # generate patches, return a generator
+    patches, patch_targets = extract_patches_sequential(
+        reshaped_sample, patch_size=patch_size, target=reshaped_target
+    )
+
+    if patch_targets is None:
+        raise ValueError("No target extracted.")
+
+    logger.info(f"Extracted {patches.shape[0]} patches from input array.")
+
+    return (
+        patches,
+        patch_targets,
+        mean,
+        std,
+    )
+
+
+# called by in memory dataset
+def prepare_patches_unsupervised_array(
+    data: np.ndarray,
+    axes: str,
+    patch_size: Union[List[int], Tuple[int]],
+) -> Tuple[np.ndarray, None, float, float]:
+    """
+    Iterate over data source and create an array of patches.
+
+    This method expects an array of shape SC(Z)YX, where S and C can be singleton
+    dimensions.
+
+    Patches returned are of shape SC(Z)YX, where S is now the patches dimension.
+
+    Returns
+    -------
+    np.ndarray
+        Array of patches.
+    """
+    # calculate mean and std
+    mean = data.mean()
+    std = data.std()
+
+    # reshape array
+    reshaped_sample = reshape_array(data, axes)
+
+    # generate patches, return a generator
+    patches, _ = extract_patches_sequential(reshaped_sample, patch_size=patch_size)
+
+    return patches, _, mean, std  # TODO inelegant, replace  by dataclass?