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

careamics/dataset/in_memory_dataset.py

@@ -1,4 +1,5 @@
 """In-memory dataset module."""
+
 from __future__ import annotations
 
 import copy
@@ -8,11 +9,12 @@ from typing import Any, Callable, List, Optional, Tuple, Union
 import numpy as np
 from torch.utils.data import Dataset
 
+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 .patching.patch_transform import get_patch_transform
 from .patching.patching import (
     prepare_patches_supervised,
     prepare_patches_supervised_array,
@@ -60,18 +62,15 @@ class InMemoryDataset(Dataset):
             self.mean, self.std = computed_mean, computed_std
             logger.info(f"Computed dataset mean: {self.mean}, std: {self.std}")
 
-            # if the transforms are not an instance of Compose
-            if self.data_config.has_transform_list():
-                # update mean and std in configuration
-                # the object is mutable and should then be recorded in the CAREamist obj
-                self.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 obj
+            self.data_config.set_mean_and_std(self.mean, self.std)
         else:
             self.mean, self.std = self.data_config.mean, self.data_config.std
 
         # get transforms
-        self.patch_transform = get_patch_transform(
-            patch_transforms=self.data_config.transforms,
-            with_target=self.data_target is not None,
+        self.patch_transform = Compose(
+            transform_list=self.data_config.transforms,
         )
 
     def _prepare_patches(
@@ -166,33 +165,10 @@ class InMemoryDataset(Dataset):
             # get target
             target = self.data_targets[index]
 
-            # Albumentations requires Channel last
-            c_patch = np.moveaxis(patch, 0, -1)
-            c_target = np.moveaxis(target, 0, -1)
-
-            # Apply transforms
-            transformed = self.patch_transform(image=c_patch, target=c_target)
-
-            # move axes back
-            patch = np.moveaxis(transformed["image"], -1, 0)
-            target = np.moveaxis(transformed["target"], -1, 0)
-
-            return patch, target
+            return self.patch_transform(patch=patch, target=target)
 
         elif self.data_config.has_n2v_manipulate():
-            # Albumentations requires Channel last
-            patch = np.moveaxis(patch, 0, -1)
-
-            # Apply transforms
-            transformed_patch = self.patch_transform(image=patch)["image"]
-            manip_patch, patch, mask = transformed_patch
-
-            # move C axes back
-            manip_patch = np.moveaxis(manip_patch, -1, 0)
-            patch = np.moveaxis(patch, -1, 0)
-            mask = np.moveaxis(mask, -1, 0)
-
-            return (manip_patch, patch, mask)
+            return self.patch_transform(patch=patch)
         else:
             raise ValueError(
                 "Something went wrong! No target provided (not supervised training) "
@@ -318,9 +294,8 @@ class InMemoryPredictionDataset(Dataset):
         self.mean, self.std = self.pred_config.mean, self.pred_config.std
 
         # get transforms
-        self.patch_transform = get_patch_transform(
-            patch_transforms=self.pred_config.transforms,
-            with_target=self.data_target is not None,
+        self.patch_transform = Compose(
+            transform_list=self.pred_config.transforms,
         )
 
     def _prepare_tiles(self) -> List[Tuple[np.ndarray, TileInformation]]:
@@ -379,13 +354,7 @@ class InMemoryPredictionDataset(Dataset):
         """
         tile_array, tile_info = self.data[index]
 
-        # Albumentations requires channel last, use the XArrayTile array
-        patch = np.moveaxis(tile_array, 0, -1)
-
         # Apply transforms
-        transformed_patch = self.patch_transform(image=patch)["image"]
-
-        # move C axes back
-        transformed_patch = np.moveaxis(transformed_patch, -1, 0)
+        transformed_tile, _ = self.patch_transform(patch=tile_array)
 
-        return transformed_patch, tile_info
+        return transformed_tile, tile_info

careamics/dataset/iterable_dataset.py

@@ -7,13 +7,12 @@ 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 ..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
 
@@ -61,26 +60,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]:
         """
@@ -192,49 +180,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:
         """
@@ -367,9 +316,8 @@ class IterablePredictionDataset(PathIterableDataset):
         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,
+        self.patch_transform = Compose(
+            transform_list=prediction_config.transforms,
         )
 
     def __iter__(
@@ -408,9 +356,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

@@ -3,6 +3,7 @@ Tiling submodule.
 
 These functions are used to tile images into patches or tiles.
 """
+
 from pathlib import Path
 from typing import Callable, List, Tuple, Union
 

careamics/dataset/patching/sequential_patching.py

@@ -135,15 +135,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
 
 
@@ -201,6 +198,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

careamics/dataset/patching/tiled_patching.py

@@ -43,9 +43,11 @@ def _compute_crop_and_stitch_coords_1d(
             stitch_coords.append(
                 (
                     i + overlap // 2 if i > 0 else 0,
-                    i + tile_size - overlap // 2
-                    if crop_coords[-1][1] < axis_size
-                    else axis_size,
+                    (
+                        i + tile_size - overlap // 2
+                        if crop_coords[-1][1] < axis_size
+                        else axis_size
+                    ),
                 )
             )
 
@@ -53,9 +55,11 @@ def _compute_crop_and_stitch_coords_1d(
             overlap_crop_coords.append(
                 (
                     overlap // 2 if i > 0 else 0,
-                    tile_size - overlap // 2
-                    if crop_coords[-1][1] < axis_size
-                    else tile_size,
+                    (
+                        tile_size - overlap // 2
+                        if crop_coords[-1][1] < axis_size
+                        else tile_size
+                    ),
                 )
             )
 

careamics/lightning_datamodule.py

@@ -1,10 +1,10 @@
 """Training and validation Lightning data modules."""
+
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Literal, Optional, Union
 
 import numpy as np
 import pytorch_lightning as L
-from albumentations import Compose
 from torch.utils.data import DataLoader
 
 from careamics.config import DataConfig
@@ -341,9 +341,9 @@ class CAREamicsTrainData(L.LightningDataModule):
                 self.train_dataset = InMemoryDataset(
                     data_config=self.data_config,
                     inputs=self.train_files,
-                    data_target=self.train_target_files
-                    if self.train_data_target
-                    else None,
+                    data_target=(
+                        self.train_target_files if self.train_data_target else None
+                    ),
                     read_source_func=self.read_source_func,
                 )
 
@@ -352,9 +352,9 @@ class CAREamicsTrainData(L.LightningDataModule):
                     self.val_dataset = InMemoryDataset(
                         data_config=self.data_config,
                         inputs=self.val_files,
-                        data_target=self.val_target_files
-                        if self.val_data_target
-                        else None,
+                        data_target=(
+                            self.val_target_files if self.val_data_target else None
+                        ),
                         read_source_func=self.read_source_func,
                     )
                 else:
@@ -370,9 +370,9 @@ class CAREamicsTrainData(L.LightningDataModule):
                 self.train_dataset = PathIterableDataset(
                     data_config=self.data_config,
                     src_files=self.train_files,
-                    target_files=self.train_target_files
-                    if self.train_data_target
-                    else None,
+                    target_files=(
+                        self.train_target_files if self.train_data_target else None
+                    ),
                     read_source_func=self.read_source_func,
                 )
 
@@ -382,9 +382,9 @@ class CAREamicsTrainData(L.LightningDataModule):
                     self.val_dataset = PathIterableDataset(
                         data_config=self.data_config,
                         src_files=self.val_files,
-                        target_files=self.val_target_files
-                        if self.val_data_target
-                        else None,
+                        target_files=(
+                            self.val_target_files if self.val_data_target else None
+                        ),
                         read_source_func=self.read_source_func,
                     )
                 elif len(self.train_files) <= self.val_minimum_split:
@@ -452,8 +452,7 @@ class TrainingDataWrapper(CAREamicsTrainData):
     In particular, N2V requires a specific transformation (N2V manipulates), which is
     not compatible with supervised training. The default transformations applied to the
     training patches are defined in `careamics.config.data_model`. To use different
-    transformations, pass a list of transforms or an albumentation `Compose` as
-    `transforms` parameter. See examples for more details.
+    transformations, pass a list of transforms. See examples for more details.
 
     By default, CAREamics only supports types defined in
     `careamics.config.support.SupportedData`. To read custom data types, you can set
@@ -488,7 +487,7 @@ class TrainingDataWrapper(CAREamicsTrainData):
         Batch size.
     val_data : Optional[Union[str, Path]], optional
         Validation data, by default None.
-    transforms : Optional[Union[List[TRANSFORMS_UNION], Compose]], optional
+    transforms : List[TRANSFORMS_UNION], optional
         List of transforms to apply to training patches. If None, default transforms
         are applied.
     train_target_data : Optional[Union[str, Path]], optional
@@ -584,7 +583,7 @@ class TrainingDataWrapper(CAREamicsTrainData):
         axes: str,
         batch_size: int,
         val_data: Optional[Union[str, Path]] = None,
-        transforms: Optional[Union[List[TRANSFORMS_UNION], Compose]] = None,
+        transforms: Optional[List[TRANSFORMS_UNION]] = None,
         train_target_data: Optional[Union[str, Path]] = None,
         val_target_data: Optional[Union[str, Path]] = None,
         read_source_func: Optional[Callable] = None,
@@ -617,8 +616,8 @@ class TrainingDataWrapper(CAREamicsTrainData):
         In particular, N2V requires a specific transformation (N2V manipulates), which
         is not compatible with supervised training. The default transformations applied
         to the training patches are defined in `careamics.config.data_model`. To use
-        different transformations, pass a list of transforms or an albumentation
-        `Compose` as `transforms` parameter. See examples for more details.
+        different transformations, pass a list of transforms. See examples for more
+        details.
 
         By default, CAREamics only supports types defined in
         `careamics.config.support.SupportedData`. To read custom data types, you can set
@@ -655,7 +654,7 @@ class TrainingDataWrapper(CAREamicsTrainData):
             Batch size.
         val_data : Optional[Union[str, Path]], optional
             Validation data, by default None.
-        transforms : Optional[Union[List[TRANSFORMS_UNION], Compose]], optional
+        transforms : Optional[List[TRANSFORMS_UNION]], optional
             List of transforms to apply to training patches. If None, default transforms
             are applied.
         train_target_data : Optional[Union[str, Path]], optional
@@ -709,10 +708,7 @@ class TrainingDataWrapper(CAREamicsTrainData):
         self.data_config = DataConfig(**data_dict)
 
         # N2V specific checks, N2V, structN2V, and transforms
-        if (
-            self.data_config.has_transform_list()
-            and self.data_config.has_n2v_manipulate()
-        ):
+        if self.data_config.has_n2v_manipulate():
             # there is not target, n2v2 and structN2V can be changed
             if train_target_data is None:
                 self.data_config.set_N2V2(use_n2v2)

careamics/lightning_module.py

@@ -162,7 +162,7 @@ class CAREamicsModule(L.LightningModule):
             mean=self._trainer.datamodule.predict_dataset.mean,
             std=self._trainer.datamodule.predict_dataset.std,
         )
-        denormalized_output = denorm(image=output)["image"]
+        denormalized_output, _ = denorm(patch=output)
 
         if len(aux) > 0:
             return denormalized_output, aux

careamics/lightning_prediction_datamodule.py

@@ -1,10 +1,10 @@
 """Prediction Lightning data modules."""
+
 from pathlib import Path
 from typing import Any, Callable, List, Literal, Optional, Tuple, Union
 
 import numpy as np
 import pytorch_lightning as L
-from albumentations import Compose
 from torch.utils.data import DataLoader
 from torch.utils.data.dataloader import default_collate
 
@@ -39,7 +39,7 @@ def _collate_tiles(batch: List[Tuple[np.ndarray, TileInformation]]) -> Any:
 
     Parameters
     ----------
-    batch : Tuple[Tuple[np.ndarray, TileInformation], ...]
+    batch : List[Tuple[np.ndarray, TileInformation], ...]
         Batch of tiles.
 
     Returns
@@ -257,14 +257,13 @@ class PredictDataWrapper(CAREamicsPredictData):
 
     The default transformations applied to the images are defined in
     `careamics.config.inference_model`. To use different transformations, pass a list
-    of transforms or an albumentation `Compose` as `transforms` parameter. See examples
+    of transforms. See examples
     for more details.
 
     The `mean` and `std` parameters are only used if Normalization is defined either
-    in the default transformations or in the `transforms` parameter, but not with
-    a `Compose` object. If you pass a `Normalization` transform in a list as
-    `transforms`, then the mean and std parameters will be overwritten by those passed
-    to this method.
+    in the default transformations or in the `transforms` parameter. If you pass a
+    `Normalization` transform in a list as `transforms`, then the mean and std
+    parameters will be overwritten by those passed to this method.
 
     By default, CAREamics only supports types defined in
     `careamics.config.support.SupportedData`. To read custom data types, you can set
@@ -276,6 +275,12 @@ class PredictDataWrapper(CAREamicsPredictData):
     dataloaders, except for `batch_size`, which is set by the `batch_size`
     parameter.
 
+    Note that if you are using a UNet model and tiling, the tile size must be
+    divisible in every dimension by 2**d, where d is the depth of the model. This
+    avoids artefacts arising from the broken shift invariance induced by the
+    pooling layers of the UNet. If your image has less dimensions, as it may
+    happen in the Z dimension, consider padding your image.
+
     Parameters
     ----------
     pred_data : Union[str, Path, np.ndarray]
@@ -298,7 +303,7 @@ class PredictDataWrapper(CAREamicsPredictData):
         Batch size.
     tta_transforms : bool, optional
         Use test time augmentation, by default True.
-    transforms : Optional[Union[List[TRANSFORMS_UNION], Compose]], optional
+    transforms : List, optional
         List of transforms to apply to prediction patches. If None, default
         transforms are applied.
     read_source_func : Optional[Callable], optional
@@ -321,7 +326,7 @@ class PredictDataWrapper(CAREamicsPredictData):
         axes: str = "YX",
         batch_size: int = 1,
         tta_transforms: bool = True,
-        transforms: Optional[Union[List, Compose]] = None,
+        transforms: Optional[List] = None,
         read_source_func: Optional[Callable] = None,
         extension_filter: str = "",
         dataloader_params: Optional[dict] = None,
@@ -351,7 +356,7 @@ class PredictDataWrapper(CAREamicsPredictData):
             Batch size.
         tta_transforms : bool, optional
             Use test time augmentation, by default True.
-        transforms : Optional[Union[List[TRANSFORMS_UNION], Compose]], optional
+        transforms : Optional[List], optional
             List of transforms to apply to prediction patches. If None, default
             transforms are applied.
         read_source_func : Optional[Callable], optional

careamics/losses/__init__.py

@@ -1,6 +1,5 @@
 """Losses module."""
 
-
 from .loss_factory import loss_factory
 
 # from .noise_model_factory import noise_model_factory as noise_model_factory

careamics/losses/loss_factory.py

@@ -3,6 +3,7 @@ Loss factory module.
 
 This module contains a factory function for creating loss functions.
 """
+
 from typing import Callable, Union
 
 from ..config.support import SupportedLoss

careamics/model_io/__init__.py

@@ -1,6 +1,5 @@
 """Model I/O utilities."""
 
-
 __all__ = ["load_pretrained", "export_to_bmz"]
 
 

careamics/model_io/bioimage/_readme_factory.py

@@ -1,4 +1,5 @@
 """Functions used to create a README.md file for BMZ export."""
+
 from pathlib import Path
 from typing import Optional
 
@@ -117,4 +118,4 @@ def readme_factory(
 
         readme.write_text("".join(description))
 
-    return readme
+        return readme.absolute()

careamics/model_io/bioimage/bioimage_utils.py

@@ -1,4 +1,5 @@
 """Bioimage.io utils."""
+
 from pathlib import Path
 from typing import Union
 

careamics/model_io/bioimage/model_description.py

@@ -1,4 +1,5 @@
 """Module use to build BMZ model description."""
+
 from pathlib import Path
 from typing import List, Optional, Tuple, Union
 

careamics/model_io/bmz_io.py

@@ -1,4 +1,5 @@
 """Function to export to the BioImage Model Zoo format."""
+
 import tempfile
 from pathlib import Path
 from typing import List, Optional, Tuple, Union
@@ -177,7 +178,7 @@ def export_to_bmz(
         )
 
         # test model description
-        summary: ValidationSummary = test_model(model_description)
+        summary: ValidationSummary = test_model(model_description, decimal=0)
         if summary.status == "failed":
             raise ValueError(f"Model description test failed: {summary}")
 

careamics/models/layers.py

@@ -3,6 +3,7 @@ Layer module.
 
 This submodule contains layers used in the CAREamics models.
 """
+
 from typing import List, Optional, Tuple, Union
 
 import torch

careamics/models/model_factory.py

@@ -3,6 +3,7 @@ Model factory.
 
 Model creation factory functions.
 """
+
 from typing import Union
 
 import torch