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

careamics/dataset_ng/patching_strategies/random_patching.py

@@ -0,0 +1,338 @@
+"""A module for random patching strategies."""
+
+from collections.abc import Sequence
+from typing import Optional
+
+import numpy as np
+
+from .patching_strategy_protocol import PatchSpecs
+
+
+class RandomPatchingStrategy:
+    """
+    A patching strategy for sampling random patches, it implements the
+    `PatchingStrategy` `Protocol`.
+
+    The output of `get_patch_spec` will be random, i.e. if the same index is given
+    twice the two outputs can be different.
+
+    However the strategy still ensures that there will be a known number of patches for
+    each sample in each image stack. This is achieved through defining a set of bins
+    that map to each sample in each image stack. Whichever bin an `index` passed to
+    `get_patch_spec` falls into, determines the `"data_idx"` and `"sample_idx"` in
+    the returned `PatchSpecs`, but the `"coords"` will be random.
+
+    The number of patches in each sample is based on the number of patches that would
+    fit if they were sampled sequentially, non-overlapping, and covering the entire
+    array.
+    """
+
+    def __init__(
+        self,
+        data_shapes: Sequence[Sequence[int]],
+        patch_size: Sequence[int],
+        seed: Optional[int] = None,
+    ):
+        """
+        A patching strategy for sampling random patches.
+
+        Parameters
+        ----------
+        data_shapes : sequence of (sequence of int)
+            The shapes of the underlying data. Each element is the dimension of the
+            axes SC(Z)YX.
+        patch_size : sequence of int
+            The size of the patch. The sequence will have length 2 or 3, for 2D and 3D
+            data respectively.
+        seed : int, optional
+            An optional seed to ensure the reproducibility of the random patches.
+        """
+        self.rng = np.random.default_rng(seed=seed)
+        self.patch_size = patch_size
+        self.data_shapes = data_shapes
+
+        # these bins will determine which image stack and sample a patch comes from
+        # the image_stack_cumulative_patches map a patch index to each image stack
+        # the sample_cumulative_patches map a patch index to each sample
+        # the image_stack_cumulative_samples map a sample index to each image stack
+        (
+            self.image_stack_cumulative_patches,
+            self.sample_cumulative_patches,
+            self.image_stack_cumulative_samples,
+        ) = self._calc_bins(self.data_shapes, self.patch_size)
+
+    @property
+    def n_patches(self) -> int:
+        """
+        The number of patches that this patching strategy will return.
+
+        It also determines the maximum index that can be given to `get_patch_spec`.
+        """
+        # last bin boundary will be total patches
+        return self.image_stack_cumulative_patches[-1]
+
+    def get_patch_spec(self, index: int) -> PatchSpecs:
+        """Return the patch specs for a given index.
+
+        Parameters
+        ----------
+        index : int
+            A patch index.
+
+        Returns
+        -------
+        PatchSpecs
+            A dictionary that specifies a single patch in a series of `ImageStacks`.
+        """
+        # TODO: break into smaller testable functions?
+        if index >= self.n_patches:
+            raise IndexError(
+                f"Index {index} out of bounds for RandomPatchingStrategy with number "
+                f"of patches {self.n_patches}"
+            )
+        # digitize returns the bin that `index` belongs to
+        data_index = np.digitize(index, bins=self.image_stack_cumulative_patches).item()
+        # maps to a particular sample within the whole series of image stacks
+        #   (not just a single image stack)
+        total_samples_index = np.digitize(
+            index, bins=self.sample_cumulative_patches
+        ).item()
+
+        data_shape = self.data_shapes[data_index]
+        spatial_shape = data_shape[2:]
+
+        # calculate sample index relative to image stack:
+        #   subtract the total number of samples in the previous image stacks
+        if data_index == 0:
+            n_previous_samples = 0
+        else:
+            n_previous_samples = self.image_stack_cumulative_samples[data_index - 1]
+        sample_index = total_samples_index - n_previous_samples
+        coords = _generate_random_coords(spatial_shape, self.patch_size, self.rng)
+        return {
+            "data_idx": data_index,
+            "sample_idx": sample_index,
+            "coords": coords,
+            "patch_size": self.patch_size,
+        }
+
+    @staticmethod
+    def _calc_bins(
+        data_shapes: Sequence[Sequence[int]], patch_size: Sequence[int]
+    ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
+        """Calculate bins used to map an index to an image_stack and a sample.
+
+        The number of patches in each sample is based on the number of patches that
+        would fit if they were sampled sequentially.
+
+        Parameters
+        ----------
+        data_shapes : sequence of (sequence of int)
+            The shapes of the underlying data. Each element is the dimension of the
+            axes SC(Z)YX.
+        patch_size : sequence of int
+            The size of the patch. The sequence will have length 2 or 3, for 2D and 3D
+            data respectively.
+
+        Returns
+        -------
+        image_stack_cumulative_patches: tuple of int
+            The bins that map a patch index to an image stack. E.g. if a patch index
+            falls below the first bin boundary it belongs to the first image stack, if
+            a patch index falls between the first bin boundary and the second bin
+            boundary it belongs to the second image stack, and so on.
+        sample_cumulative_patches: tuple of int
+            The bins that map a patch index to a sample. E.g. if a patch index
+            falls below the first bin boundary it belongs to the first sample, if
+            a patch index falls between the first bin boundary and the second bin
+            boundary it belongs to the second sample, and so on.
+        image_stack_cumulative_samples: tuple of int
+            The bins that map a sample index to an image stack. E.g. if a sample index
+            falls below the first bin boundary it belongs to the first image stack, if
+            a patch index falls between the first bin boundary and the second bin
+            boundary it belongs to the second image stack, and so on.
+        """
+        patches_per_image_stack: list[int] = []
+        patches_per_sample: list[int] = []
+        samples_per_image_stack: list[int] = []
+        for data_shape in data_shapes:
+            spatial_shape = data_shape[2:]
+            n_single_sample_patches = _calc_n_patches(spatial_shape, patch_size)
+            # multiply by number of samples in image_stack
+            patches_per_image_stack.append(n_single_sample_patches * data_shape[0])
+            # list of length `sample` filled with `n_single_sample_patches`
+            patches_per_sample.extend([n_single_sample_patches] * data_shape[0])
+            # number of samples in each image stack
+            samples_per_image_stack.append(data_shape[0])
+
+        # cumulative sum creates the bins
+        image_stack_cumulative_patches = np.cumsum(patches_per_image_stack)
+        sample_cumulative_patches = np.cumsum(patches_per_sample)
+        image_stack_cumulative_samples = np.cumsum(samples_per_image_stack)
+        return (
+            tuple(image_stack_cumulative_patches),
+            tuple(sample_cumulative_patches),
+            tuple(image_stack_cumulative_samples),
+        )
+
+
+class FixedRandomPatchingStrategy:
+    """
+    A patching strategy for sampling random patches it implements the `PatchingStrategy`
+    `Protocol`.
+
+    The output of `get_patch_spec` will be deterministic, i.e. if the same index is
+    given twice the two outputs will be the same.
+
+    The number of patches in each sample is based on the number of patches that would
+    fit if they were sampled sequentially, non-overlapping, and covering the entire
+    array.
+    """
+
+    def __init__(
+        self,
+        data_shapes: Sequence[Sequence[int]],
+        patch_size: Sequence[int],
+        seed: Optional[int] = None,
+    ):
+        """A patching strategy for sampling random patches.
+
+        Parameters
+        ----------
+        data_shapes : sequence of (sequence of int)
+            The shapes of the underlying data. Each element is the dimension of the
+            axes SC(Z)YX.
+        patch_size : sequence of int
+            The size of the patch. The sequence will have length 2 or 3, for 2D and 3D
+            data respectively.
+        seed : int, optional
+            An optional seed to ensure the reproducibility of the random patches.
+        """
+        self.rng = np.random.default_rng(seed=seed)
+        self.patch_size = patch_size
+        self.data_shapes = data_shapes
+
+        # simply generate all the patches at initialisation, so they will be fixed
+        self.fixed_patch_specs: list[PatchSpecs] = []
+        for data_idx, data_shape in enumerate(self.data_shapes):
+            spatial_shape = data_shape[2:]
+            n_patches = _calc_n_patches(spatial_shape, self.patch_size)
+            for sample_idx in range(data_shape[0]):
+                for _ in range(n_patches):
+                    random_coords = _generate_random_coords(
+                        spatial_shape, self.patch_size, self.rng
+                    )
+                    patch_specs: PatchSpecs = {
+                        "data_idx": data_idx,
+                        "sample_idx": sample_idx,
+                        "coords": random_coords,
+                        "patch_size": self.patch_size,
+                    }
+                    self.fixed_patch_specs.append(patch_specs)
+
+    @property
+    def n_patches(self):
+        """
+        The number of patches that this patching strategy will return.
+
+        It also determines the maximum index that can be given to `get_patch_spec`.
+        """
+        return len(self.fixed_patch_specs)
+
+    def get_patch_spec(self, index: int) -> PatchSpecs:
+        """Return the patch specs for a given index.
+
+        Parameters
+        ----------
+        index : int
+            A patch index.
+
+        Returns
+        -------
+        PatchSpecs
+            A dictionary that specifies a single patch in a series of `ImageStacks`.
+        """
+        if index >= self.n_patches:
+            raise IndexError(
+                f"Index {index} out of bounds for FixedRandomPatchingStrategy with "
+                f"number of patches, {self.n_patches}"
+            )
+        # simply index the pre-generated patches to get the correct patch
+        return self.fixed_patch_specs[index]
+
+
+def _generate_random_coords(
+    spatial_shape: Sequence[int], patch_size: Sequence[int], rng: np.random.Generator
+) -> tuple[int, ...]:
+    """Generate random patch coordinates for a given `spatial_shape` and `patch_size`.
+
+    The coords are the top-left (and first z-slice for 3D data) of a patch. The
+    sequence will have length 2 or 3, for 2D and 3D data respectively.
+
+    Parameters
+    ----------
+    spatial_shape : sequence of int
+        The dimension of the axes (Z)YX, a sequence of length 2 or 3, for 2D and 3D
+        data respectively.
+    patch_size : sequence of int
+        The size of the patch. The sequence will have length 2 or 3, for 2D and 3D
+        data respectively.
+    rng : numpy.random.Generator
+        A numpy generator to ensure the reproducibility of the random patches.
+
+    Returns
+    -------
+    coords: tuple of int
+        The top-left (and first z-slice for 3D data) coords of a patch. The tuple will
+        have length 2 or 3, for 2D and 3D data respectively.
+
+    Raises
+    ------
+    ValueError
+        Raises if the number of spatial dimensions do not match the number of patch
+        dimensions.
+    """
+    if len(patch_size) != len(spatial_shape):
+        raise ValueError(
+            f"Number of patch dimension {len(patch_size)}, do not match the number of "
+            f"spatial dimensions {len(spatial_shape)}, for `patch_size={patch_size}` "
+            f"and `spatial_shape={spatial_shape}`."
+        )
+    return tuple(
+        rng.integers(
+            np.zeros(len(patch_size), dtype=int),
+            np.array(spatial_shape) - np.array(patch_size),
+            endpoint=False,
+            dtype=int,
+        ).tolist()
+    )
+
+
+def _calc_n_patches(spatial_shape: Sequence[int], patch_size: Sequence[int]) -> int:
+    """
+    Calculates the number of patches for a given `spatial_shape` and `patch_size`.
+
+    This is based on the number of patches that would fit if they were sampled
+    sequentially.
+
+    Parameters
+    ----------
+    spatial_shape : sequence of int
+        The dimension of the axes (Z)YX, a sequence of length 2 or 3, for 2D and 3D
+        data respectively.
+    patch_size : sequence of int
+        The size of the patch. The sequence will have length 2 or 3, for 2D and 3D
+        data respectively.
+
+    Returns
+    -------
+    int
+        The number of patches.
+    """
+    if len(patch_size) != len(spatial_shape):
+        raise ValueError(
+            f"Number of patch dimension {len(patch_size)}, do not match the number of "
+            f"spatial dimensions {len(spatial_shape)}, for `patch_size={patch_size}` "
+            f"and `spatial_shape={spatial_shape}`."
+        )
+    return int(np.ceil(np.prod(spatial_shape) / np.prod(patch_size)))

careamics/dataset_ng/patching_strategies/sequential_patching.py

@@ -0,0 +1,75 @@
+import itertools
+from collections.abc import Sequence
+from typing import Optional
+
+import numpy as np
+from typing_extensions import ParamSpec
+
+from .patching_strategy_protocol import PatchSpecs
+
+P = ParamSpec("P")
+
+
+# TODO: this is an unfinished prototype based on current tiling implementation
+#  not guaranteed to work!
+class SequentialPatchingStrategy:
+    # TODO: docs
+    def __init__(
+        self,
+        data_shapes: Sequence[Sequence[int]],
+        patch_size: Sequence[int],
+        overlap: Optional[Sequence[int]] = None,
+    ):
+        self.data_shapes = data_shapes
+        self.patch_size = patch_size
+        if overlap is None:
+            overlap = [0] * len(patch_size)
+        self.overlap = np.asarray(overlap)
+
+        self.patch_specs: list[PatchSpecs] = self._initialize_patch_specs()
+
+    @property
+    def n_patches(self) -> int:
+        return len(self.patch_specs)
+
+    def get_patch_spec(self, index: int) -> PatchSpecs:
+        return self.patch_specs[index]
+
+    def _compute_coords_1d(
+        self, patch_size: int, spatial_shape: int, overlap: int
+    ) -> list[tuple[int, int]]:
+        step = patch_size - overlap
+        crop_coords = []
+
+        current_pos = 0
+        while current_pos <= spatial_shape - patch_size:
+            crop_coords.append((current_pos, current_pos + patch_size))
+            current_pos += step
+
+        if crop_coords[-1][1] < spatial_shape:
+            crop_coords.append((spatial_shape - patch_size, spatial_shape))
+
+        return crop_coords
+
+    def _initialize_patch_specs(self) -> list[PatchSpecs]:
+        patch_specs: list[PatchSpecs] = []
+        for data_idx, data_shape in enumerate(self.data_shapes):
+
+            data_spatial_shape = data_shape[-len(self.patch_size) :]
+            coords_list = [
+                self._compute_coords_1d(
+                    self.patch_size[i], data_spatial_shape[i], self.overlap[i]
+                )
+                for i in range(len(self.patch_size))
+            ]
+            for sample_idx in range(data_shape[0]):
+                for crop_coord in itertools.product(*coords_list):
+                    patch_specs.append(
+                        PatchSpecs(
+                            data_idx=data_idx,
+                            sample_idx=sample_idx,
+                            coords=tuple(coord[0] for coord in crop_coord),
+                            patch_size=self.patch_size,
+                        )
+                    )
+        return patch_specs

careamics/lightning/lightning_module.py

@@ -1,12 +1,17 @@
 """CAREamics Lightning module."""
 
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable, Literal, Optional, Union
 
 import numpy as np
 import pytorch_lightning as L
 from torch import Tensor, nn
 
-from careamics.config import UNetBasedAlgorithm, VAEBasedAlgorithm
+from careamics.config import (
+    N2VAlgorithm,
+    UNetBasedAlgorithm,
+    VAEBasedAlgorithm,
+    algorithm_factory,
+)
 from careamics.config.support import (
     SupportedAlgorithm,
     SupportedArchitecture,
@@ -27,7 +32,11 @@ from careamics.models.lvae.noise_models import (
     noise_model_factory,
 )
 from careamics.models.model_factory import model_factory
-from careamics.transforms import Denormalize, ImageRestorationTTA
+from careamics.transforms import (
+    Denormalize,
+    ImageRestorationTTA,
+    N2VManipulateTorch,
+)
 from careamics.utils.metrics import RunningPSNR, scale_invariant_psnr
 from careamics.utils.torch_utils import get_optimizer, get_scheduler
 
@@ -73,13 +82,21 @@ class FCNModule(L.LightningModule):
             Algorithm configuration.
         """
         super().__init__()
-        # if loading from a checkpoint, AlgorithmModel needs to be instantiated
+
         if isinstance(algorithm_config, dict):
-            algorithm_config = UNetBasedAlgorithm(
-                **algorithm_config
-            )  # TODO this needs to be updated using the algorithm-specific class
+            algorithm_config = algorithm_factory(algorithm_config)
+
+        # create preprocessing, model and loss function
+        if isinstance(algorithm_config, N2VAlgorithm):
+            self.use_n2v = True
+            self.n2v_preprocess: Optional[N2VManipulateTorch] = N2VManipulateTorch(
+                n2v_manipulate_config=algorithm_config.n2v_config
+            )
+        else:
+            self.use_n2v = False
+            self.n2v_preprocess = None
 
-        # create model and loss function
+        self.algorithm = algorithm_config.algorithm
         self.model: nn.Module = model_factory(algorithm_config.model)
         self.loss_func = loss_factory(algorithm_config.loss)
 
@@ -119,10 +136,15 @@ class FCNModule(L.LightningModule):
         Any
             Loss value.
         """
-        # TODO can N2V be simplified by returning mask*original_patch
-        x, *aux = batch
-        out = self.model(x)
-        loss = self.loss_func(out, *aux)
+        x, *targets = batch
+        if self.use_n2v and self.n2v_preprocess is not None:
+            x_preprocessed, *aux = self.n2v_preprocess(x)
+        else:
+            x_preprocessed = x
+            aux = []
+
+        out = self.model(x_preprocessed)
+        loss = self.loss_func(out, *aux, *targets)
         self.log(
             "train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, logger=True
         )
@@ -138,9 +160,15 @@ class FCNModule(L.LightningModule):
         batch_idx : Any
             Batch index.
         """
-        x, *aux = batch
-        out = self.model(x)
-        val_loss = self.loss_func(out, *aux)
+        x, *targets = batch
+        if self.use_n2v and self.n2v_preprocess is not None:
+            x_preprocessed, *aux = self.n2v_preprocess(x)
+        else:
+            x_preprocessed = x
+            aux = []
+
+        out = self.model(x_preprocessed)
+        val_loss = self.loss_func(out, *aux, *targets)
 
         # log validation loss
         self.log(
@@ -177,10 +205,16 @@ class FCNModule(L.LightningModule):
             and isinstance(batch[1][0], TileInformation)
         )
 
+        # TODO add explanations for what is happening here
         if is_tiled:
             x, *aux = batch
+            if type(x) in [list, tuple]:
+                x = x[0]
         else:
-            x = batch
+            if type(batch) in [list, tuple]:
+                x = batch[0]  # TODO change, ugly way to deal with n2v refac
+            else:
+                x = batch
             aux = []
 
         # apply test-time augmentation if available
@@ -593,6 +627,9 @@ def create_careamics_module(
     algorithm: Union[SupportedAlgorithm, str],
     loss: Union[SupportedLoss, str],
     architecture: Union[SupportedArchitecture, str],
+    use_n2v2: bool = False,
+    struct_n2v_axis: Literal["horizontal", "vertical", "none"] = "none",
+    struct_n2v_span: int = 5,
     model_parameters: Optional[dict] = None,
     optimizer: Union[SupportedOptimizer, str] = "Adam",
     optimizer_parameters: Optional[dict] = None,
@@ -612,6 +649,12 @@ def create_careamics_module(
         Loss function to use for training (see SupportedLoss).
     architecture : SupportedArchitecture or str
         Model architecture to use for training (see SupportedArchitecture).
+    use_n2v2 : bool, default=False
+        Whether to use N2V2 or Noise2Void.
+    struct_n2v_axis : "horizontal", "vertical", or "none", default="none"
+        Axis of the StructN2V mask.
+    struct_n2v_span : int, default=5
+        Span of the StructN2V mask.
     model_parameters : dict, optional
         Model parameters to use for training, by default {}. Model parameters are
         defined in the relevant `torch.nn.Module` class, or Pyddantic model (see
@@ -633,14 +676,15 @@ def create_careamics_module(
     CAREamicsModule
         CAREamics Lightning module.
     """
-    # create a AlgorithmModel compatible dictionary
+    # TODO should use the same functions are in configuration_factory.py
+    # create an AlgorithmModel compatible dictionary
     if lr_scheduler_parameters is None:
         lr_scheduler_parameters = {}
     if optimizer_parameters is None:
         optimizer_parameters = {}
     if model_parameters is None:
         model_parameters = {}
-    algorithm_configuration: dict[str, Any] = {
+    algorithm_dict: dict[str, Any] = {
         "algorithm": algorithm,
         "loss": loss,
         "optimizer": {
@@ -652,18 +696,25 @@ def create_careamics_module(
             "parameters": lr_scheduler_parameters,
         },
     }
-    model_configuration = {"architecture": architecture}
-    model_configuration.update(model_parameters)
+
+    model_dict = {"architecture": architecture}
+    model_dict.update(model_parameters)
 
     # add model parameters to algorithm configuration
-    algorithm_configuration["model"] = model_configuration
+    algorithm_dict["model"] = model_dict
+
+    which_algo = algorithm_dict["algorithm"]
+    if which_algo in UNetBasedAlgorithm.get_compatible_algorithms():
+        algorithm_cfg = algorithm_factory(algorithm_dict)
+
+        # if use N2V
+        if isinstance(algorithm_cfg, N2VAlgorithm):
+            algorithm_cfg.n2v_config.struct_mask_axis = struct_n2v_axis
+            algorithm_cfg.n2v_config.struct_mask_span = struct_n2v_span
+            algorithm_cfg.set_n2v2(use_n2v2)
 
-    # call the parent init using an AlgorithmModel instance
-    # TODO broken by new configutations!
-    algorithm_str = algorithm_configuration["algorithm"]
-    if algorithm_str in UNetBasedAlgorithm.get_compatible_algorithms():
-        return FCNModule(UNetBasedAlgorithm(**algorithm_configuration))
+        return FCNModule(algorithm_cfg)
     else:
         raise NotImplementedError(
-            f"Model {algorithm_str} is not implemented or unknown."
+            f"Algorithm {which_algo} is not implemented or unknown."
         )