careamics 0.1.0rc4__py3-none-any.whl → 0.1.0rc6__py3-none-any.whl

careamics/dataset/iterable_dataset.py

@@ -1,3 +1,5 @@
+"""Iterable dataset used to load data file by file."""
+
 from __future__ import annotations
 
 import copy
@@ -7,26 +9,98 @@ from typing import Any, Callable, Generator, List, Optional, Tuple, Union
 import numpy as np
 from torch.utils.data import IterableDataset, get_worker_info
 
+from careamics.transforms import Compose
+
 from ..config import DataConfig, InferenceConfig
 from ..config.tile_information import TileInformation
+from ..config.transformations import NormalizeModel
 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__)
 
 
+def _iterate_over_files(
+    data_config: Union[DataConfig, InferenceConfig],
+    data_files: List[Path],
+    target_files: Optional[List[Path]] = None,
+    read_source_func: Callable = read_tiff,
+) -> Generator[Tuple[np.ndarray, Optional[np.ndarray]], None, None]:
+    """
+    Iterate over data source and yield whole image.
+
+    Parameters
+    ----------
+    data_config : Union[DataConfig, InferenceConfig]
+        Data configuration.
+    data_files : List[Path]
+        List of data files.
+    target_files : Optional[List[Path]]
+        List of target files, by default None.
+    read_source_func : Optional[Callable]
+        Function to read the source, by default read_tiff.
+
+    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(data_files):
+        # retrieve file corresponding to the worker id
+        if i % num_workers == worker_id:
+            try:
+                # read data
+                sample = read_source_func(filename, data_config.axes)
+
+                # read target, if available
+                if target_files is not None:
+                    if filename.name != target_files[i].name:
+                        raise ValueError(
+                            f"File {filename} does not match target file "
+                            f"{target_files[i]}. Have you passed sorted "
+                            f"arrays?"
+                        )
+
+                    # read target
+                    target = read_source_func(target_files[i], data_config.axes)
+
+                    yield sample, target
+                else:
+                    yield sample, None
+
+            except Exception as e:
+                logger.error(f"Error reading file {filename}: {e}")
+
+
 class PathIterableDataset(IterableDataset):
     """
     Dataset allowing extracting patches w/o loading whole data into memory.
 
     Parameters
     ----------
-    data_path : Union[str, Path]
+    data_config : DataConfig
+        Data configuration.
+    src_files : List[Path]
+        List of data files.
+    target_files : Optional[List[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[Path]
         Path to the data, must be a directory.
     axes : str
         Description of axes in format STCZYX.
@@ -46,11 +120,24 @@ class PathIterableDataset(IterableDataset):
 
     def __init__(
         self,
-        data_config: Union[DataConfig, InferenceConfig],
+        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 : Optional[List[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.
+        """
         self.data_config = data_config
         self.data_files = src_files
         self.target_files = target_files
@@ -61,26 +148,15 @@ class PathIterableDataset(IterableDataset):
         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)
-
+            # 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
 
         # get transforms
-        self.patch_transform = get_patch_transform(
-            patch_transforms=data_config.transforms,
-            with_target=target_files is not None,
-        )
+        self.patch_transform = Compose(transform_list=data_config.transforms)
 
     def _calculate_mean_and_std(self) -> Tuple[float, float]:
         """
@@ -94,7 +170,9 @@ class PathIterableDataset(IterableDataset):
         means, stds = 0, 0
         num_samples = 0
 
-        for sample, _ in self._iterate_over_files():
+        for sample, _ in _iterate_over_files(
+            self.data_config, self.data_files, self.target_files, self.read_source_func
+        ):
             means += sample.mean()
             stds += sample.std()
             num_samples += 1
@@ -109,57 +187,9 @@ class PathIterableDataset(IterableDataset):
         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]:
+    ) -> Generator[Tuple[np.ndarray, ...], None, None]:
         """
         Iterate over data source and yield single patch.
 
@@ -173,7 +203,9 @@ class PathIterableDataset(IterableDataset):
         ), "Mean and std must be provided"
 
         # iterate over files
-        for sample_input, sample_target in self._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
+        ):
             reshaped_sample = reshape_array(sample_input, self.data_config.axes)
             reshaped_target = (
                 None
@@ -192,49 +224,10 @@ class PathIterableDataset(IterableDataset):
             # 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)."
-                    )
+                yield self.patch_transform(
+                    patch=patch_data[0],
+                    target=patch_data[1],
+                )
 
     def get_number_of_files(self) -> int:
         """
@@ -260,9 +253,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
         -------
@@ -326,12 +319,23 @@ class PathIterableDataset(IterableDataset):
         return dataset
 
 
-class IterablePredictionDataset(PathIterableDataset):
+class IterablePredictionDataset(IterableDataset):
     """
-    Dataset allowing extracting patches w/o loading whole data into memory.
+    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
@@ -351,13 +355,26 @@ class IterablePredictionDataset(PathIterableDataset):
         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,
-        )
+        """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.
+        """
         self.prediction_config = prediction_config
+        self.data_files = src_files
         self.axes = prediction_config.axes
         self.tile_size = self.prediction_config.tile_size
         self.tile_overlap = self.prediction_config.tile_overlap
@@ -366,11 +383,21 @@ class IterablePredictionDataset(PathIterableDataset):
         # 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,
-        )
+        # check mean and std and create normalize transform
+        if self.prediction_config.mean is None or self.prediction_config.std is None:
+            raise ValueError("Mean and std must be provided for prediction.")
+        else:
+            self.mean = self.prediction_config.mean
+            self.std = self.prediction_config.std
+
+            # instantiate normalize transform
+            self.patch_transform = Compose(
+                transform_list=[
+                    NormalizeModel(
+                        mean=prediction_config.mean, std=prediction_config.std
+                    )
+                ],
+            )
 
     def __iter__(
         self,
@@ -387,11 +414,19 @@ class IterablePredictionDataset(PathIterableDataset):
             self.mean is not None and self.std is not None
         ), "Mean and std must be provided"
 
-        for sample, _ in self._iterate_over_files():
+        for sample, _ in _iterate_over_files(
+            self.prediction_config,
+            self.data_files,
+            read_source_func=self.read_source_func,
+        ):
             # reshape array
             reshaped_sample = reshape_array(sample, self.axes)
 
-            if self.tile:
+            if (
+                self.tile
+                and self.tile_size is not None
+                and self.tile_overlap is not None
+            ):
                 # generate patches, return a generator
                 patch_gen = extract_tiles(
                     arr=reshaped_sample,
@@ -408,9 +443,6 @@ class IterablePredictionDataset(PathIterableDataset):
 
             # 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)
+                transformed_patch, _ = self.patch_transform(patch=patch_array)
 
                 yield transformed_patch, tile_info

careamics/dataset/patching/__init__.py

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

careamics/dataset/patching/patching.py

@@ -1,8 +1,5 @@
-"""
-Tiling submodule.
+"""Patching functions."""
 
-These functions are used to tile images into patches or tiles.
-"""
 from pathlib import Path
 from typing import Callable, List, Tuple, Union
 
@@ -20,12 +17,25 @@ def prepare_patches_supervised(
     train_files: List[Path],
     target_files: List[Path],
     axes: str,
-    patch_size: Union[List[int], Tuple[int]],
+    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.
 
+    Parameters
+    ----------
+    train_files : List[Path]
+        List of paths to training data.
+    target_files : List[Path]
+        List of paths to target data.
+    axes : str
+        Axes of the data.
+    patch_size : Union[List[int], Tuple[int]]
+        Size of the patches.
+    read_source_func : Callable
+        Function to read the data.
+
     Returns
     -------
     np.ndarray
@@ -94,13 +104,25 @@ def prepare_patches_unsupervised(
     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.
+    """Iterate over data source and create an array of patches.
+
+    This method returns the mean and standard deviation of the image.
+
+    Parameters
+    ----------
+    train_files : List[Path]
+        List of paths to training data.
+    axes : str
+        Axes of the data.
+    patch_size : Union[List[int], Tuple[int]]
+        Size of the patches.
+    read_source_func : Callable
+        Function to read the data.
 
     Returns
     -------
-    np.ndarray
-        Array of patches.
+    Tuple[np.ndarray, None, float, float]
+        Source and target patches, mean and standard deviation.
     """
     means, stds, num_samples = 0, 0, 0
     all_patches = []
@@ -149,10 +171,21 @@ def prepare_patches_supervised_array(
 
     Patches returned are of shape SC(Z)YX, where S is now the patches dimension.
 
+    Parameters
+    ----------
+    data : np.ndarray
+        Input data array.
+    axes : str
+        Axes of the data.
+    data_target : np.ndarray
+        Target data array.
+    patch_size : Union[List[int], Tuple[int]]
+        Size of the patches.
+
     Returns
     -------
-    np.ndarray
-        Array of patches.
+    Tuple[np.ndarray, np.ndarray, float, float]
+        Source and target patches, mean and standard deviation.
     """
     # compute statistics
     mean = data.mean()
@@ -194,10 +227,19 @@ def prepare_patches_unsupervised_array(
 
     Patches returned are of shape SC(Z)YX, where S is now the patches dimension.
 
+    Parameters
+    ----------
+    data : np.ndarray
+        Input data array.
+    axes : str
+        Axes of the data.
+    patch_size : Union[List[int], Tuple[int]]
+        Size of the patches.
+
     Returns
     -------
-    np.ndarray
-        Array of patches.
+    Tuple[np.ndarray, None, float, float]
+        Source patches, mean and standard deviation.
     """
     # calculate mean and std
     mean = data.mean()
@@ -209,4 +251,4 @@ def prepare_patches_unsupervised_array(
     # generate patches, return a generator
     patches, _ = extract_patches_sequential(reshaped_sample, patch_size=patch_size)
 
-    return patches, _, mean, std  # TODO inelegant, replace  by dataclass?
+    return patches, _, mean, std  # TODO inelegant, replace by dataclass?

careamics/dataset/patching/random_patching.py

@@ -1,3 +1,5 @@
+"""Random patching utilities."""
+
 from typing import Generator, List, Optional, Tuple, Union
 
 import numpy as np
@@ -30,6 +32,8 @@ def extract_patches_random(
         Input image array.
     patch_size : Tuple[int]
         Patch sizes in each dimension.
+    target : Optional[np.ndarray], optional
+        Target array, by default None.
 
     Yields
     ------
@@ -120,10 +124,12 @@ def extract_patches_random_from_chunks(
     ----------
     arr : np.ndarray
         Input image array.
-    patch_size : Tuple[int]
+    patch_size : Union[List[int], Tuple[int, ...]]
         Patch sizes in each dimension.
-    chunk_size : Tuple[int]
+    chunk_size : Union[List[int], Tuple[int, ...]]
         Chunk sizes to load from the.
+    chunk_limit : Optional[int], optional
+        Number of chunks to load, by default None.
 
     Yields
     ------

careamics/dataset/patching/sequential_patching.py

@@ -1,3 +1,5 @@
+"""Sequential patching functions."""
+
 from typing import List, Optional, Tuple, Union
 
 import numpy as np
@@ -14,14 +16,14 @@ def _compute_number_of_patches(
 
     Parameters
     ----------
-    arr : Tuple[int, ...]
+    arr_shape : Tuple[int, ...]
         Shape of the input array.
-    patch_sizes : Tuple[int]
+    patch_sizes : Union[List[int], Tuple[int, ...]
         Shape of the patches.
 
     Returns
     -------
-    Tuple[int]
+    Tuple[int, ...]
         Number of patches in each dimension.
     """
     if len(arr_shape) != len(patch_sizes):
@@ -55,14 +57,14 @@ def _compute_overlap(
 
     Parameters
     ----------
-    arr : Tuple[int, ...]
+    arr_shape : Tuple[int, ...]
         Input array shape.
-    patch_sizes : Tuple[int]
+    patch_sizes : Union[List[int], Tuple[int, ...]]
         Size of the patches.
 
     Returns
     -------
-    Tuple[int]
+    Tuple[int, ...]
         Overlap between patches in each dimension.
     """
     n_patches = _compute_number_of_patches(arr_shape, patch_sizes)
@@ -123,6 +125,8 @@ def _compute_patch_views(
         Steps between views.
     output_shape : Tuple[int]
         Shape of the output array.
+    target : Optional[np.ndarray], optional
+        Target array, by default None.
 
     Returns
     -------
@@ -135,15 +139,12 @@ def _compute_patch_views(
         arr = np.stack([arr, target], axis=0)
         window_shape = [arr.shape[0], *window_shape]
         step = (arr.shape[0], *step)
-        output_shape = [arr.shape[0], -1, arr.shape[2], *output_shape[2:]]
+        output_shape = [-1, arr.shape[0], arr.shape[2], *output_shape[2:]]
 
     patches = view_as_windows(arr, window_shape=window_shape, step=step).reshape(
         *output_shape
     )
-    if target is not None:
-        rng.shuffle(patches, axis=1)
-    else:
-        rng.shuffle(patches, axis=0)
+    rng.shuffle(patches, axis=0)
     return patches
 
 
@@ -164,11 +165,13 @@ def extract_patches_sequential(
         Input image array.
     patch_size : Tuple[int]
         Patch sizes in each dimension.
+    target : Optional[np.ndarray], optional
+        Target array, by default None.
 
     Returns
     -------
-    Generator[Tuple[np.ndarray, ...], None, None]
-        Generator of patches.
+    Tuple[np.ndarray, Optional[np.ndarray]]
+        Patches.
     """
     is_3d_patch = len(patch_size) == 3
 
@@ -201,6 +204,9 @@ def extract_patches_sequential(
 
     if target is not None:
         # target was concatenated to patches in _compute_reshaped_view
-        return (patches[0, ...], patches[1, ...])  # TODO  in _compute_reshaped_view?
+        return (
+            patches[:, 0, ...],
+            patches[:, 1, ...],
+        )  # TODO  in _compute_reshaped_view?
     else:
         return patches, None