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

careamics/dataset/dataset_utils/__init__.py

@@ -1,6 +1,5 @@
 """Files and arrays utils used in the datasets."""
 
-
 __all__ = [
     "reshape_array",
     "get_files_size",

careamics/dataset/dataset_utils/dataset_utils.py

@@ -1,4 +1,5 @@
 """Convenience methods for datasets."""
+
 from typing import List, Tuple
 
 import numpy as np

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 ..config import DataModel, InferenceModel
+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,
@@ -29,7 +31,7 @@ class InMemoryDataset(Dataset):
 
     def __init__(
         self,
-        data_config: DataModel,
+        data_config: DataConfig,
         inputs: Union[np.ndarray, List[Path]],
         data_target: Optional[Union[np.ndarray, List[Path]]] = None,
         read_source_func: Callable = read_tiff,
@@ -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) "
@@ -279,7 +255,7 @@ class InMemoryPredictionDataset(Dataset):
 
     def __init__(
         self,
-        prediction_config: InferenceModel,
+        prediction_config: InferenceConfig,
         inputs: np.ndarray,
         data_target: Optional[np.ndarray] = None,
         read_source_func: Optional[Callable] = read_tiff,
@@ -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 ..config import DataModel, InferenceModel
+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
 
@@ -46,7 +45,7 @@ class PathIterableDataset(IterableDataset):
 
     def __init__(
         self,
-        data_config: Union[DataModel, InferenceModel],
+        data_config: Union[DataConfig, InferenceConfig],
         src_files: List[Path],
         target_files: Optional[List[Path]] = None,
         read_source_func: Callable = read_tiff,
@@ -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:
         """
@@ -346,7 +295,7 @@ class IterablePredictionDataset(PathIterableDataset):
 
     def __init__(
         self,
-        prediction_config: InferenceModel,
+        prediction_config: InferenceConfig,
         src_files: List[Path],
         read_source_func: Callable = read_tiff,
         **kwargs: Any,
@@ -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
+                    ),
                 )
             )