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

careamics/dataset/patching/patching.py

@@ -1,5 +1,6 @@
 """Patching functions."""
 
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Callable, List, Tuple, Union
 
@@ -7,11 +8,38 @@ import numpy as np
 
 from ...utils.logging import get_logger
 from ..dataset_utils import reshape_array
+from ..dataset_utils.running_stats import compute_normalization_stats
 from .sequential_patching import extract_patches_sequential
 
 logger = get_logger(__name__)
 
 
+@dataclass
+class Stats:
+    """Dataclass to store statistics."""
+
+    means: Union[np.ndarray, tuple, list, None]
+    stds: Union[np.ndarray, tuple, list, None]
+
+
+@dataclass
+class PatchedOutput:
+    """Dataclass to store patches and statistics."""
+
+    patches: Union[np.ndarray]
+    targets: Union[np.ndarray, None]
+    image_stats: Stats
+    target_stats: Stats
+
+
+@dataclass
+class StatsOutput:
+    """Dataclass to store patches and statistics."""
+
+    image_stats: Stats
+    target_stats: Stats
+
+
 # called by in memory dataset
 def prepare_patches_supervised(
     train_files: List[Path],
@@ -19,10 +47,12 @@ def prepare_patches_supervised(
     axes: str,
     patch_size: Union[List[int], Tuple[int, ...]],
     read_source_func: Callable,
-) -> Tuple[np.ndarray, np.ndarray, float, float]:
+) -> PatchedOutput:
     """
     Iterate over data source and create an array of patches and corresponding targets.
 
+    The lists of Paths should be pre-sorted.
+
     Parameters
     ----------
     train_files : List[Path]
@@ -41,9 +71,6 @@ def prepare_patches_supervised(
     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):
@@ -83,17 +110,18 @@ def prepare_patches_supervised(
             f"{target_files}."
         )
 
-    result_mean, result_std = means / num_samples, stds / num_samples
+    image_means, image_stds = compute_normalization_stats(np.concatenate(all_patches))
+    target_means, target_stds = compute_normalization_stats(np.concatenate(all_targets))
 
     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 (
+    return PatchedOutput(
         patch_array,
         target_array,
-        result_mean,
-        result_std,
+        Stats(image_means, image_stds),
+        Stats(target_means, target_stds),
     )
 
 
@@ -103,7 +131,7 @@ def prepare_patches_unsupervised(
     axes: str,
     patch_size: Union[List[int], Tuple[int]],
     read_source_func: Callable,
-) -> Tuple[np.ndarray, None, float, float]:
+) -> PatchedOutput:
     """Iterate over data source and create an array of patches.
 
     This method returns the mean and standard deviation of the image.
@@ -149,12 +177,14 @@ def prepare_patches_unsupervised(
     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
+    image_means, image_stds = compute_normalization_stats(np.concatenate(all_patches))
 
     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?
+    return PatchedOutput(
+        patch_array, None, Stats(image_means, image_stds), Stats((), ())
+    )
 
 
 # called on arrays by in memory dataset
@@ -163,7 +193,7 @@ def prepare_patches_supervised_array(
     axes: str,
     data_target: np.ndarray,
     patch_size: Union[List[int], Tuple[int]],
-) -> Tuple[np.ndarray, np.ndarray, float, float]:
+) -> PatchedOutput:
     """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
@@ -187,14 +217,14 @@ def prepare_patches_supervised_array(
     Tuple[np.ndarray, np.ndarray, float, float]
         Source and target patches, mean and standard deviation.
     """
-    # compute statistics
-    mean = data.mean()
-    std = data.std()
-
     # reshape array
     reshaped_sample = reshape_array(data, axes)
     reshaped_target = reshape_array(data_target, axes)
 
+    # compute statistics
+    image_means, image_stds = compute_normalization_stats(reshaped_sample)
+    target_means, target_stds = compute_normalization_stats(reshaped_target)
+
     # generate patches, return a generator
     patches, patch_targets = extract_patches_sequential(
         reshaped_sample, patch_size=patch_size, target=reshaped_target
@@ -205,11 +235,11 @@ def prepare_patches_supervised_array(
 
     logger.info(f"Extracted {patches.shape[0]} patches from input array.")
 
-    return (
+    return PatchedOutput(
         patches,
         patch_targets,
-        mean,
-        std,
+        Stats(image_means, image_stds),
+        Stats(target_means, target_stds),
     )
 
 
@@ -218,7 +248,7 @@ def prepare_patches_unsupervised_array(
     data: np.ndarray,
     axes: str,
     patch_size: Union[List[int], Tuple[int]],
-) -> Tuple[np.ndarray, None, float, float]:
+) -> PatchedOutput:
     """
     Iterate over data source and create an array of patches.
 
@@ -241,14 +271,13 @@ def prepare_patches_unsupervised_array(
     Tuple[np.ndarray, None, float, float]
         Source patches, mean and standard deviation.
     """
-    # calculate mean and std
-    mean = data.mean()
-    std = data.std()
-
     # reshape array
     reshaped_sample = reshape_array(data, axes)
 
+    # calculate mean and std
+    means, stds = compute_normalization_stats(reshaped_sample)
+
     # 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 PatchedOutput(patches, None, Stats(means, stds), Stats((), ()))

careamics/dataset/patching/random_patching.py

@@ -13,6 +13,7 @@ def extract_patches_random(
     arr: np.ndarray,
     patch_size: Union[List[int], Tuple[int, ...]],
     target: Optional[np.ndarray] = None,
+    seed: Optional[int] = None,
 ) -> Generator[Tuple[np.ndarray, Optional[np.ndarray]], None, None]:
     """
     Generate patches from an array in a random manner.
@@ -34,12 +35,16 @@ def extract_patches_random(
         Patch sizes in each dimension.
     target : Optional[np.ndarray], optional
         Target array, by default None.
+    seed : Optional[int], optional
+        Random seed, by default None.
 
     Yields
     ------
     Generator[np.ndarray, None, None]
         Generator of patches.
     """
+    rng = np.random.default_rng(seed=seed)
+
     is_3d_patch = len(patch_size) == 3
 
     # patches sanity check
@@ -48,9 +53,6 @@ def extract_patches_random(
     # Update patch size to encompass S and C dimensions
     patch_size = [1, arr.shape[1], *patch_size]
 
-    # random generator
-    rng = np.random.default_rng()
-
     # iterate over the number of samples (S or T)
     for sample_idx in range(arr.shape[0]):
         # get sample array
@@ -113,6 +115,7 @@ def extract_patches_random_from_chunks(
     patch_size: Union[List[int], Tuple[int, ...]],
     chunk_size: Union[List[int], Tuple[int, ...]],
     chunk_limit: Optional[int] = None,
+    seed: Optional[int] = None,
 ) -> Generator[np.ndarray, None, None]:
     """
     Generate patches from an array in a random manner.
@@ -130,6 +133,8 @@ def extract_patches_random_from_chunks(
         Chunk sizes to load from the.
     chunk_limit : Optional[int], optional
         Number of chunks to load, by default None.
+    seed : Optional[int], optional
+        Random seed, by default None.
 
     Yields
     ------
@@ -141,7 +146,7 @@ def extract_patches_random_from_chunks(
     # Patches sanity check
     validate_patch_dimensions(arr, patch_size, is_3d_patch)
 
-    rng = np.random.default_rng()
+    rng = np.random.default_rng(seed=seed)
     num_chunks = chunk_limit if chunk_limit else np.prod(arr._cdata_shape)
 
     # Iterate over num chunks in the array

careamics/dataset/patching/validate_patch_dimension.py

@@ -45,18 +45,20 @@ def validate_patch_dimensions(
     if len(patch_size) != len(arr.shape[2:]):
         raise ValueError(
             f"There must be a patch size for each spatial dimensions "
-            f"(got {patch_size} patches for dims {arr.shape})."
+            f"(got {patch_size} patches for dims {arr.shape}). Check the axes order."
         )
 
     # Sanity checks on patch sizes versus array dimension
     if is_3d_patch and patch_size[0] > arr.shape[-3]:
         raise ValueError(
             f"Z patch size is inconsistent with image shape "
-            f"(got {patch_size[0]} patches for dim {arr.shape[1]})."
+            f"(got {patch_size[0]} patches for dim {arr.shape[1]}). Check the axes "
+            f"order."
         )
 
     if patch_size[-2] > arr.shape[-2] or patch_size[-1] > arr.shape[-1]:
         raise ValueError(
             f"At least one of YX patch dimensions is larger than the corresponding "
-            f"image dimension (got {patch_size} patches for dims {arr.shape[-2:]})."
+            f"image dimension (got {patch_size} patches for dims {arr.shape[-2:]}). "
+            f"Check the axes order."
         )

careamics/dataset/tiling/__init__.py

@@ -0,0 +1,10 @@
+"""Tiling functions."""
+
+__all__ = [
+    "stitch_prediction",
+    "extract_tiles",
+    "collate_tiles",
+]
+
+from .collate_tiles import collate_tiles
+from .tiled_patching import extract_tiles

careamics/dataset/tiling/collate_tiles.py

@@ -0,0 +1,33 @@
+"""Collate function for tiling."""
+
+from typing import Any, List, Tuple
+
+import numpy as np
+from torch.utils.data.dataloader import default_collate
+
+from careamics.config.tile_information import TileInformation
+
+
+def collate_tiles(batch: List[Tuple[np.ndarray, TileInformation]]) -> Any:
+    """
+    Collate tiles received from CAREamics prediction dataloader.
+
+    CAREamics prediction dataloader returns tuples of arrays and TileInformation. In
+    case of non-tiled data, this function will return the arrays. In case of tiled data,
+    it will return the arrays, the last tile flag, the overlap crop coordinates and the
+    stitch coordinates.
+
+    Parameters
+    ----------
+    batch : List[Tuple[np.ndarray, TileInformation], ...]
+        Batch of tiles.
+
+    Returns
+    -------
+    Any
+        Collated batch.
+    """
+    new_batch = [tile for tile, _ in batch]
+    tiles_batch = [tile_info for _, tile_info in batch]
+
+    return default_collate(new_batch), tiles_batch

careamics/dataset/{patching → tiling}/tiled_patching.py

@@ -84,15 +84,15 @@ def extract_tiles(
     tile_size: Union[List[int], Tuple[int, ...]],
     overlaps: Union[List[int], Tuple[int, ...]],
 ) -> Generator[Tuple[np.ndarray, TileInformation], None, None]:
-    """
-    Generate tiles from the input array with specified overlap.
+    """Generate tiles from the input array with specified overlap.
 
     The tiles cover the whole array. The method returns a generator that yields
     tuples of array and tile information, the latter includes whether
     the tile is the last one, the coordinates of the overlap crop, and the coordinates
     of the stitched tile.
 
-    The array has shape C(Z)YX, where C can be a singleton.
+    Input array should have shape SC(Z)YX, while the returned tiles have shape C(Z)YX,
+    where C can be a singleton.
 
     Parameters
     ----------
@@ -155,10 +155,10 @@ def extract_tiles(
             # create tile information
             tile_info = TileInformation(
                 array_shape=sample.squeeze().shape,
-                tiled=True,
                 last_tile=last_tile,
                 overlap_crop_coords=overlap_crop_coords,
                 stitch_coords=stitch_coords,
+                sample_id=sample_idx,
             )
 
             yield tile, tile_info

careamics/lightning_datamodule.py

@@ -583,12 +583,7 @@ class TrainingDataWrapper(CAREamicsTrainData):
     >>> my_array = np.arange(256).reshape(16, 16)
     >>> my_transforms = [
     ...     {
-    ...         "name": SupportedTransform.NORMALIZE.value,
-    ...         "mean": 0,
-    ...         "std": 1,
-    ...     },
-    ...     {
-    ...         "name": SupportedTransform.N2V_MANIPULATE.value,
+    ...         "name": SupportedTransform.XY_FLIP.value,
     ...     }
     ... ]
     >>> data_module = TrainingDataWrapper(

careamics/lightning_module.py

@@ -148,13 +148,17 @@ class CAREamicsModule(L.LightningModule):
         Any
             Model output.
         """
-        x, *aux = batch
+        if self._trainer.datamodule.tiled:
+            x, *aux = batch
+        else:
+            x = batch
+            aux = []
 
         # apply test-time augmentation if available
         # TODO: probably wont work with batch size > 1
         if self._trainer.datamodule.prediction_config.tta_transforms:
             tta = ImageRestorationTTA()
-            augmented_batch = tta.forward(batch[0])  # list of augmented tensors
+            augmented_batch = tta.forward(x)  # list of augmented tensors
             augmented_output = []
             for augmented in augmented_batch:
                 augmented_pred = self.model(augmented)
@@ -165,13 +169,13 @@ class CAREamicsModule(L.LightningModule):
 
         # Denormalize the output
         denorm = Denormalize(
-            mean=self._trainer.datamodule.predict_dataset.mean,
-            std=self._trainer.datamodule.predict_dataset.std,
+            image_means=self._trainer.datamodule.predict_dataset.image_means,
+            image_stds=self._trainer.datamodule.predict_dataset.image_stds,
         )
-        denormalized_output, _ = denorm(patch=output)
+        denormalized_output = denorm(patch=output.cpu().numpy())
 
-        if len(aux) > 0:
-            return denormalized_output, aux
+        if len(aux) > 0:  # aux can be tiling information
+            return denormalized_output, *aux
         else:
             return denormalized_output
 

careamics/lightning_prediction_datamodule.py

@@ -1,68 +1,37 @@
 """Prediction Lightning data modules."""
 
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Literal, Optional, Tuple, Union
 
 import numpy as np
 import pytorch_lightning as L
 from torch.utils.data import DataLoader
-from torch.utils.data.dataloader import default_collate
 
 from careamics.config import InferenceConfig
 from careamics.config.support import SupportedData
-from careamics.config.tile_information import TileInformation
+from careamics.dataset import (
+    InMemoryPredDataset,
+    InMemoryTiledPredDataset,
+    IterablePredDataset,
+    IterableTiledPredDataset,
+)
 from careamics.dataset.dataset_utils import (
     get_read_func,
     list_files,
 )
-from careamics.dataset.in_memory_dataset import (
-    InMemoryPredictionDataset,
-)
-from careamics.dataset.iterable_dataset import (
-    IterablePredictionDataset,
-)
+from careamics.dataset.tiling.collate_tiles import collate_tiles
 from careamics.utils import get_logger
 
-PredictDatasetType = Union[InMemoryPredictionDataset, IterablePredictionDataset]
+PredictDatasetType = Union[
+    InMemoryPredDataset,
+    InMemoryTiledPredDataset,
+    IterablePredDataset,
+    IterableTiledPredDataset,
+]
 
 logger = get_logger(__name__)
 
 
-def _collate_tiles(batch: List[Tuple[np.ndarray, TileInformation]]) -> Any:
-    """
-    Collate tiles received from CAREamics prediction dataloader.
-
-    CAREamics prediction dataloader returns tuples of arrays and TileInformation. In
-    case of non-tiled data, this function will return the arrays. In case of tiled data,
-    it will return the arrays, the last tile flag, the overlap crop coordinates and the
-    stitch coordinates.
-
-    Parameters
-    ----------
-    batch : List[Tuple[np.ndarray, TileInformation], ...]
-        Batch of tiles.
-
-    Returns
-    -------
-    Any
-        Collated batch.
-    """
-    first_tile_info: TileInformation = batch[0][1]
-    # if not tiled, then return arrays
-    if not first_tile_info.tiled:
-        arrays, _ = zip(*batch)
-
-        return default_collate(arrays)
-    # else we explicit the last_tile flag and coordinates
-    else:
-        new_batch = [
-            (tile, t.last_tile, t.array_shape, t.overlap_crop_coords, t.stitch_coords)
-            for tile, t in batch
-        ]
-
-        return default_collate(new_batch)
-
-
 class CAREamicsPredictData(L.LightningDataModule):
     """
     CAREamics Lightning prediction data module.
@@ -182,6 +151,9 @@ class CAREamicsPredictData(L.LightningDataModule):
         self.tile_size = pred_config.tile_size
         self.tile_overlap = pred_config.tile_overlap
 
+        # check if it is tiled
+        self.tiled = self.tile_size is not None and self.tile_overlap is not None
+
         # read source function
         if pred_config.data_type == SupportedData.CUSTOM:
             # mypy check
@@ -212,17 +184,29 @@ class CAREamicsPredictData(L.LightningDataModule):
         """
         # if numpy array
         if self.data_type == SupportedData.ARRAY:
-            # prediction dataset
-            self.predict_dataset: PredictDatasetType = InMemoryPredictionDataset(
-                prediction_config=self.prediction_config,
-                inputs=self.pred_data,
-            )
+            if self.tiled:
+                self.predict_dataset: PredictDatasetType = InMemoryTiledPredDataset(
+                    prediction_config=self.prediction_config,
+                    inputs=self.pred_data,
+                )
+            else:
+                self.predict_dataset = InMemoryPredDataset(
+                    prediction_config=self.prediction_config,
+                    inputs=self.pred_data,
+                )
         else:
-            self.predict_dataset = IterablePredictionDataset(
-                prediction_config=self.prediction_config,
-                src_files=self.pred_files,
-                read_source_func=self.read_source_func,
-            )
+            if self.tiled:
+                self.predict_dataset = IterableTiledPredDataset(
+                    prediction_config=self.prediction_config,
+                    src_files=self.pred_files,
+                    read_source_func=self.read_source_func,
+                )
+            else:
+                self.predict_dataset = IterablePredDataset(
+                    prediction_config=self.prediction_config,
+                    src_files=self.pred_files,
+                    read_source_func=self.read_source_func,
+                )
 
     def predict_dataloader(self) -> DataLoader:
         """
@@ -236,7 +220,7 @@ class CAREamicsPredictData(L.LightningDataModule):
         return DataLoader(
             self.predict_dataset,
             batch_size=self.batch_size,
-            collate_fn=_collate_tiles,
+            collate_fn=collate_tiles if self.tiled else None,
             **self.dataloader_params,
         )  # TODO check workers are used
 
@@ -287,12 +271,10 @@ class PredictDataWrapper(CAREamicsPredictData):
         Prediction data.
     data_type : Union[Literal["array", "tiff", "custom"], SupportedData]
         Data type, see `SupportedData` for available options.
-    mean : float
-        Mean value for normalization, only used if Normalization is defined in the
-        transforms.
-    std : float
-        Standard deviation value for normalization, only used if Normalization is
-        defined in the transform.
+    image_means : list of float
+        Mean values for normalization, only used if Normalization is defined.
+    image_stds : list of float
+        Std values for normalization, only used if Normalization is defined.
     tile_size : Tuple[int, ...]
         Tile size, 2D or 3D tile size.
     tile_overlap : Tuple[int, ...]
@@ -316,8 +298,8 @@ class PredictDataWrapper(CAREamicsPredictData):
         self,
         pred_data: Union[str, Path, np.ndarray],
         data_type: Union[Literal["array", "tiff", "custom"], SupportedData],
-        mean: float,
-        std: float,
+        image_means=list[float],
+        image_stds=list[float],
         tile_size: Optional[Tuple[int, ...]] = None,
         tile_overlap: Optional[Tuple[int, ...]] = None,
         axes: str = "YX",
@@ -336,12 +318,10 @@ class PredictDataWrapper(CAREamicsPredictData):
             Prediction data.
         data_type : Union[Literal["array", "tiff", "custom"], SupportedData]
             Data type, see `SupportedData` for available options.
-        mean : float
-            Mean value for normalization, only used if Normalization is defined in the
-            transforms.
-        std : float
-            Standard deviation value for normalization, only used if Normalization is
-            defined in the transform.
+        image_means : list of float
+            Mean values for normalization, only used if Normalization is defined.
+        image_stds : list of float
+            Std values for normalization, only used if Normalization is defined.
         tile_size : List[int]
             Tile size, 2D or 3D tile size.
         tile_overlap : List[int]
@@ -367,8 +347,8 @@ class PredictDataWrapper(CAREamicsPredictData):
             "tile_size": tile_size,
             "tile_overlap": tile_overlap,
             "axes": axes,
-            "mean": mean,
-            "std": std,
+            "image_means": image_means,
+            "image_stds": image_stds,
             "tta": tta_transforms,
             "batch_size": batch_size,
             "transforms": [],