careamics 0.0.9__py3-none-any.whl → 0.0.10__py3-none-any.whl

careamics/config/transformations/__init__.py

@@ -1,7 +1,6 @@
 """CAREamics transformation Pydantic models."""
 
 __all__ = [
-    "N2V_TRANSFORMS_UNION",
     "NORM_AND_SPATIAL_UNION",
     "SPATIAL_TRANSFORMS_UNION",
     "N2VManipulateModel",
@@ -16,7 +15,6 @@ from .n2v_manipulate_model import N2VManipulateModel
 from .normalize_model import NormalizeModel
 from .transform_model import TransformModel
 from .transform_unions import (
-    N2V_TRANSFORMS_UNION,
     NORM_AND_SPATIAL_UNION,
     SPATIAL_TRANSFORMS_UNION,
 )

careamics/config/transformations/n2v_manipulate_model.py

@@ -7,6 +7,8 @@ from pydantic import ConfigDict, Field, field_validator
 from .transform_model import TransformModel
 
 
+# TODO should probably not be a TransformModel anymore, no reason for it
+# `name` is used as a discriminator field in the transforms
 class N2VManipulateModel(TransformModel):
     """
     Pydantic model used to represent N2V manipulation.
@@ -32,11 +34,24 @@ class N2VManipulateModel(TransformModel):
     )
 
     name: Literal["N2VManipulate"] = "N2VManipulate"
+
     roi_size: int = Field(default=11, ge=3, le=21)
+    """Size of the region where the pixel manipulation is applied."""
+
     masked_pixel_percentage: float = Field(default=0.2, ge=0.05, le=10.0)
+    """Percentage of masked pixels per image."""
+
+    remove_center: bool = Field(default=True)  # TODO remove it
+    """Exclude center pixel from average calculation."""  # TODO rephrase this
+
     strategy: Literal["uniform", "median"] = Field(default="uniform")
+    """Strategy for pixel value replacement."""
+
     struct_mask_axis: Literal["horizontal", "vertical", "none"] = Field(default="none")
+    """Orientation of the structN2V mask. Set to `\"non\"` to not apply StructN2V."""
+
     struct_mask_span: int = Field(default=5, ge=3, le=15)
+    """Size of the structN2V mask."""
 
     @field_validator("roi_size", "struct_mask_span")
     @classmethod

careamics/config/transformations/transform_unions.py

@@ -4,7 +4,6 @@ from typing import Annotated, Union
 
 from pydantic import Discriminator
 
-from .n2v_manipulate_model import N2VManipulateModel
 from .normalize_model import NormalizeModel
 from .xy_flip_model import XYFlipModel
 from .xy_random_rotate90_model import XYRandomRotate90Model
@@ -14,7 +13,6 @@ NORM_AND_SPATIAL_UNION = Annotated[
         NormalizeModel,
         XYFlipModel,
         XYRandomRotate90Model,
-        N2VManipulateModel,
     ],
     Discriminator("name"),  # used to tell the different transform models apart
 ]
@@ -29,14 +27,3 @@ SPATIAL_TRANSFORMS_UNION = Annotated[
     Discriminator("name"),  # used to tell the different transform models apart
 ]
 """Available spatial transforms in CAREamics."""
-
-
-N2V_TRANSFORMS_UNION = Annotated[
-    Union[
-        XYFlipModel,
-        XYRandomRotate90Model,
-        N2VManipulateModel,
-    ],
-    Discriminator("name"),  # used to tell the different transform models apart
-]
-"""Available N2V-compatible transforms in CAREamics."""

careamics/dataset/dataset_utils/iterate_over_files.py

@@ -9,7 +9,7 @@ from typing import Callable, Optional, Union
 from numpy.typing import NDArray
 from torch.utils.data import get_worker_info
 
-from careamics.config import GeneralDataConfig, InferenceConfig
+from careamics.config import DataConfig, InferenceConfig
 from careamics.file_io.read import read_tiff
 from careamics.utils.logging import get_logger
 
@@ -19,7 +19,7 @@ logger = get_logger(__name__)
 
 
 def iterate_over_files(
-    data_config: Union[GeneralDataConfig, InferenceConfig],
+    data_config: Union[DataConfig, InferenceConfig],
     data_files: list[Path],
     target_files: Optional[list[Path]] = None,
     read_source_func: Callable = read_tiff,

careamics/dataset/in_memory_dataset.py

@@ -9,7 +9,7 @@ from typing import Any, Callable, Optional, Union
 import numpy as np
 from torch.utils.data import Dataset
 
-from careamics.config import GeneralDataConfig, N2VDataConfig
+from careamics.config import DataConfig
 from careamics.config.transformations import NormalizeModel
 from careamics.dataset.patching.patching import (
     PatchedOutput,
@@ -46,7 +46,7 @@ class InMemoryDataset(Dataset):
 
     def __init__(
         self,
-        data_config: GeneralDataConfig,
+        data_config: DataConfig,
         inputs: Union[np.ndarray, list[Path]],
         input_target: Optional[Union[np.ndarray, list[Path]]] = None,
         read_source_func: Callable = read_tiff,
@@ -215,16 +215,9 @@ class InMemoryDataset(Dataset):
         if self.data_targets is not None:
             # get target
             target = self.data_targets[index]
-
             return self.patch_transform(patch=patch, target=target)
 
-        elif isinstance(self.data_config, N2VDataConfig):
-            return self.patch_transform(patch=patch)
-        else:
-            raise ValueError(
-                "Something went wrong! No target provided (not supervised training) "
-                "while the algorithm is not Noise2Void."
-            )
+        return self.patch_transform(patch=patch)
 
     def get_data_statistics(self) -> tuple[list[float], list[float]]:
         """Return training data statistics.

careamics/dataset/in_memory_pred_dataset.py

@@ -69,7 +69,7 @@ class InMemoryPredDataset(Dataset):
         """
         return len(self.data)
 
-    def __getitem__(self, index: int) -> NDArray:
+    def __getitem__(self, index: int) -> tuple[NDArray, ...]:
         """
         Return the patch corresponding to the provided index.
 
@@ -80,9 +80,7 @@ class InMemoryPredDataset(Dataset):
 
         Returns
         -------
-        NDArray
+        tuple(numpy.ndarray, ...)
             Transformed patch.
         """
-        transformed_patch, _ = self.patch_transform(patch=self.data[index])
-
-        return transformed_patch
+        return self.patch_transform(patch=self.data[index])

careamics/dataset/in_memory_tiled_pred_dataset.py

@@ -107,7 +107,7 @@ class InMemoryTiledPredDataset(Dataset):
         """
         return len(self.data)
 
-    def __getitem__(self, index: int) -> tuple[NDArray, TileInformation]:
+    def __getitem__(self, index: int) -> tuple[tuple[NDArray, ...], TileInformation]:
         """
         Return the patch corresponding to the provided index.
 
@@ -124,6 +124,6 @@ class InMemoryTiledPredDataset(Dataset):
         tile_array, tile_info = self.data[index]
 
         # Apply transforms
-        transformed_tile, _ = self.patch_transform(patch=tile_array)
+        transformed_tile = self.patch_transform(patch=tile_array)
 
         return transformed_tile, tile_info

careamics/dataset/iterable_dataset.py

@@ -10,7 +10,7 @@ from typing import Callable, Optional
 import numpy as np
 from torch.utils.data import IterableDataset
 
-from careamics.config import GeneralDataConfig
+from careamics.config import DataConfig
 from careamics.config.transformations import NormalizeModel
 from careamics.file_io.read import read_tiff
 from careamics.transforms import Compose
@@ -49,7 +49,7 @@ class PathIterableDataset(IterableDataset):
 
     def __init__(
         self,
-        data_config: GeneralDataConfig,
+        data_config: DataConfig,
         src_files: list[Path],
         target_files: Optional[list[Path]] = None,
         read_source_func: Callable = read_tiff,

careamics/dataset/iterable_pred_dataset.py

@@ -97,13 +97,13 @@ class IterablePredDataset(IterableDataset):
 
     def __iter__(
         self,
-    ) -> Generator[NDArray, None, None]:
+    ) -> Generator[tuple[NDArray, ...], None, None]:
         """
         Iterate over data source and yield single patch.
 
         Yields
         ------
-        NDArray
+        (numpy.ndarray, numpy.ndarray or None)
             Single patch.
         """
         assert (
@@ -118,6 +118,4 @@ class IterablePredDataset(IterableDataset):
             # sample has S dimension
             for i in range(sample.shape[0]):
 
-                transformed_sample, _ = self.patch_transform(patch=sample[i])
-
-                yield transformed_sample
+                yield self.patch_transform(patch=sample[i])

careamics/dataset/iterable_tiled_pred_dataset.py

@@ -109,13 +109,13 @@ class IterableTiledPredDataset(IterableDataset):
 
     def __iter__(
         self,
-    ) -> Generator[tuple[NDArray, TileInformation], None, None]:
+    ) -> Generator[tuple[tuple[NDArray, ...], TileInformation], None, None]:
         """
         Iterate over data source and yield single patch.
 
         Yields
         ------
-        Generator of NDArray and TileInformation tuple
+        Generator of (np.ndarray, np.ndarray or None) and TileInformation tuple
             Generator of single tiles.
         """
         assert (
@@ -136,6 +136,6 @@ class IterableTiledPredDataset(IterableDataset):
 
             # apply transform to patches
             for patch_array, tile_info in patch_gen:
-                transformed_patch, _ = self.patch_transform(patch=patch_array)
+                transformed_patch = self.patch_transform(patch=patch_array)
 
                 yield transformed_patch, tile_info

careamics/dataset_ng/dataset/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ["CareamicsDataset"]
+
+from .dataset import CareamicsDataset

careamics/dataset_ng/dataset/dataset.py

@@ -0,0 +1,184 @@
+from collections.abc import Sequence
+from enum import Enum
+from pathlib import Path
+from typing import Literal, NamedTuple, Optional, Union
+
+import numpy as np
+from numpy.typing import NDArray
+from torch.utils.data import Dataset
+from typing_extensions import ParamSpec
+
+from careamics.config import DataConfig, InferenceConfig
+from careamics.config.support import SupportedData
+from careamics.dataset.patching.patching import Stats
+from careamics.dataset_ng.patch_extractor import (
+    ImageStackLoader,
+    PatchExtractor,
+    create_patch_extractor,
+)
+from careamics.dataset_ng.patching_strategies import (
+    FixedRandomPatchingStrategy,
+    PatchingStrategy,
+    PatchSpecs,
+    RandomPatchingStrategy,
+)
+from careamics.transforms import Compose
+
+P = ParamSpec("P")
+
+
+class Mode(str, Enum):
+    TRAINING = "training"
+    VALIDATING = "validating"
+    PREDICTING = "predicting"
+
+
+class ImageRegionData(NamedTuple):
+    data: NDArray
+    source: Union[Path, Literal["array"]]
+    data_shape: Sequence[int]
+    dtype: str  # dtype should be str for collate
+    axes: str
+    region_spec: PatchSpecs
+
+
+InputType = Union[Sequence[np.ndarray], Sequence[Path]]
+
+
+class CareamicsDataset(Dataset):
+    def __init__(
+        self,
+        data_config: Union[DataConfig, InferenceConfig],
+        mode: Mode,
+        inputs: InputType,
+        targets: Optional[InputType] = None,
+        image_stack_loader: Optional[ImageStackLoader[P]] = None,
+        *args: P.args,
+        **kwargs: P.kwargs,
+    ):
+        self.config = data_config
+        self.mode = mode
+
+        data_type_enum = SupportedData(self.config.data_type)
+        self.input_extractor = create_patch_extractor(
+            inputs,
+            self.config.axes,
+            data_type_enum,
+            image_stack_loader,
+            *args,
+            **kwargs,
+        )
+        if targets is not None:
+            self.target_extractor: Optional[PatchExtractor] = create_patch_extractor(
+                targets,
+                self.config.axes,
+                data_type_enum,
+                image_stack_loader,
+                *args,
+                **kwargs,
+            )
+        else:
+            self.target_extractor = None
+
+        self.patching_strategy = self._initialize_patching_strategy()
+
+        self.input_stats, self.target_stats = self._initialize_statistics()
+
+        self.transforms = self._initialize_transforms()
+
+    def _initialize_patching_strategy(self) -> PatchingStrategy:
+        patching_strategy: PatchingStrategy
+        if self.mode == Mode.TRAINING:
+            if isinstance(self.config, InferenceConfig):
+                raise ValueError("Inference config cannot be used for training.")
+            patching_strategy = RandomPatchingStrategy(
+                data_shapes=self.input_extractor.shape,
+                patch_size=self.config.patch_size,
+                # TODO: Add random seed to dataconfig
+                seed=getattr(self.config, "random_seed", None),
+            )
+        elif self.mode == Mode.VALIDATING:
+            if isinstance(self.config, InferenceConfig):
+                raise ValueError("Inference config cannot be used for validating.")
+            patching_strategy = FixedRandomPatchingStrategy(
+                data_shapes=self.input_extractor.shape,
+                patch_size=self.config.patch_size,
+                # TODO: Add random seed to dataconfig
+                seed=getattr(self.config, "random_seed", None),
+            )
+        elif self.mode == Mode.PREDICTING:
+            # TODO: patching strategy will be tilingStrategy in upcoming PR
+            raise NotImplementedError(
+                "Prediction mode for the CAREamicsDataset has not been implemented yet."
+            )
+        else:
+            raise ValueError(f"Unrecognised dataset mode {self.mode}.")
+
+        return patching_strategy
+
+    def _initialize_transforms(self) -> Optional[Compose]:
+        if isinstance(self.config, DataConfig):
+            return Compose(
+                transform_list=list(self.config.transforms),
+            )
+        # TODO: add TTA
+        return None
+
+    def _initialize_statistics(self) -> tuple[Stats, Optional[Stats]]:
+        # TODO: add running stats
+        # Currently assume that stats are provided in the configuration
+        input_stats = Stats(self.config.image_means, self.config.image_stds)
+        target_stats = None
+        if isinstance(self.config, DataConfig):
+            target_means = self.config.target_means
+            target_stds = self.config.target_stds
+            if target_means is not None and target_stds is not None:
+                target_stats = Stats(target_means, target_stds)
+        return input_stats, target_stats
+
+    def __len__(self):
+        return self.patching_strategy.n_patches
+
+    def _create_image_region(
+        self, patch: np.ndarray, patch_spec: PatchSpecs, extractor: PatchExtractor
+    ) -> ImageRegionData:
+        data_idx = patch_spec["data_idx"]
+        return ImageRegionData(
+            data=patch,
+            source=extractor.image_stacks[data_idx].source,
+            dtype=str(extractor.image_stacks[data_idx].data_dtype),
+            data_shape=extractor.image_stacks[data_idx].data_shape,
+            # TODO: should it be axes of the original image instead?
+            axes=self.config.axes,
+            region_spec=patch_spec,
+        )
+
+    def __getitem__(
+        self, index: int
+    ) -> tuple[ImageRegionData, Optional[ImageRegionData]]:
+        patch_spec = self.patching_strategy.get_patch_spec(index)
+        input_patch = self.input_extractor.extract_patch(**patch_spec)
+
+        target_patch = (
+            self.target_extractor.extract_patch(**patch_spec)
+            if self.target_extractor is not None
+            else None
+        )
+
+        if self.transforms is not None:
+            input_patch, target_patch = self.transforms(input_patch, target_patch)
+
+        input_data = self._create_image_region(
+            patch=input_patch, patch_spec=patch_spec, extractor=self.input_extractor
+        )
+
+        if target_patch is not None and self.target_extractor is not None:
+            target_data = self._create_image_region(
+                patch=target_patch,
+                patch_spec=patch_spec,
+                extractor=self.target_extractor,
+            )
+        else:
+            target_data = None
+
+        return input_data, target_data