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

careamics/lightning/train_data_module.py

@@ -8,7 +8,7 @@ import pytorch_lightning as L
 from numpy.typing import NDArray
 from torch.utils.data import DataLoader, IterableDataset
 
-from careamics.config.data import DataConfig, GeneralDataConfig, N2VDataConfig
+from careamics.config.data import DataConfig
 from careamics.config.support import SupportedData
 from careamics.config.transformations import TransformModel
 from careamics.dataset.dataset_utils import (
@@ -118,7 +118,7 @@ class TrainDataModule(L.LightningDataModule):
 
     def __init__(
         self,
-        data_config: GeneralDataConfig,
+        data_config: DataConfig,
         train_data: Union[Path, str, NDArray],
         val_data: Optional[Union[Path, str, NDArray]] = None,
         train_data_target: Optional[Union[Path, str, NDArray]] = None,
@@ -218,7 +218,7 @@ class TrainDataModule(L.LightningDataModule):
             )
 
         # configuration
-        self.data_config: GeneralDataConfig = data_config
+        self.data_config: DataConfig = data_config
         self.data_type: str = data_config.data_type
         self.batch_size: int = data_config.batch_size
         self.use_in_memory: bool = use_in_memory
@@ -486,9 +486,6 @@ def create_train_datamodule(
     val_minimum_patches: int = 5,
     dataloader_params: Optional[dict] = None,
     use_in_memory: bool = True,
-    use_n2v2: bool = False,
-    struct_n2v_axis: Literal["horizontal", "vertical", "none"] = "none",
-    struct_n2v_span: int = 5,
 ) -> TrainDataModule:
     """Create a TrainDataModule.
 
@@ -499,16 +496,8 @@ def create_train_datamodule(
     parameters passed to the datamodule are consistent with the model's requirements and
     are coherent.
 
-    By default, the train DataModule will be set for Noise2Void if no target data is
-    provided. That means that it will add a `N2VManipulateModel` transformation to the
-    list of augmentations. The default augmentations are XY flip, XY rotation, and N2V
-    pixel manipulation. If you pass a training target data, the default behaviour is to
-    train a supervised model. It will use the default XY flip and rotation
-    augmentations.
-
-    To use a different set of transformations, you can pass a list of transforms to
-    `transforms`. Note that if you intend to use Noise2Void, you should add
-    `N2VManipulateModel` as the last transform in the list of transformations.
+    The default augmentations are XY flip and XY rotation. To use a different set of
+    transformations, you can pass a list of transforms to `transforms`.
 
     The data module can be used with Path, str or numpy arrays. In the case of
     numpy arrays, it loads and computes all the patches in memory. For Path and str
@@ -534,11 +523,6 @@ def create_train_datamodule(
     In `dataloader_params`, you can pass any parameter accepted by PyTorch dataloaders,
     except for `batch_size`, which is set by the `batch_size` parameter.
 
-    Finally, if you intend to use N2V family of algorithms, you can set `use_n2v2` to
-    use N2V2, and set the `struct_n2v_axis` and `struct_n2v_span` parameters to define
-    the axis and span of the structN2V mask. These parameters are without effect if
-    a `train_target_data` or if `transforms` are provided.
-
     Parameters
     ----------
     train_data : pathlib.Path or str or numpy.ndarray
@@ -575,13 +559,6 @@ def create_train_datamodule(
         Pytorch dataloader parameters, by default {}.
     use_in_memory : bool, optional
         Use in memory dataset if possible, by default True.
-    use_n2v2 : bool, optional
-        Use N2V2 transformation during training, by default False.
-    struct_n2v_axis : {"horizontal", "vertical", "none"}, optional
-        Axis for the structN2V mask, only applied if `struct_n2v_axis` is `none`, by
-        default "none".
-    struct_n2v_span : int, optional
-        Span for the structN2V mask, by default 5.
 
     Returns
     -------
@@ -624,12 +601,11 @@ def create_train_datamodule(
     transforms:
     >>> import numpy as np
     >>> from careamics.lightning import create_train_datamodule
-    >>> from careamics.config.transformations import XYFlipModel, N2VManipulateModel
+    >>> from careamics.config.transformations import XYFlipModel
     >>> from careamics.config.support import SupportedTransform
     >>> my_array = np.arange(256).reshape(16, 16)
     >>> my_transforms = [
     ...     XYFlipModel(flip_y=False),
-    ...     N2VManipulateModel()
     ... ]
     >>> data_module = create_train_datamodule(
     ...     train_data=my_array,
@@ -656,15 +632,8 @@ def create_train_datamodule(
     if transforms is not None:
         data_dict["transforms"] = transforms
 
-    # TODO not compatible with HDN, consider adding an argument for n2v/hdn
-    if train_target_data is None:
-        data_config: GeneralDataConfig = N2VDataConfig(**data_dict)
-        assert isinstance(data_config, N2VDataConfig)
-
-        data_config.set_n2v2(use_n2v2)
-        data_config.set_structN2V_mask(struct_n2v_axis, struct_n2v_span)
-    else:
-        data_config = DataConfig(**data_dict)
+    # instantiate data configuration
+    data_config = DataConfig(**data_dict)
 
     # sanity check on the dataloader parameters
     if "batch_size" in dataloader_params:

careamics/losses/fcn/losses.py

@@ -8,7 +8,7 @@ import torch
 from torch.nn import L1Loss, MSELoss
 
 
-def mse_loss(source: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+def mse_loss(source: torch.Tensor, target: torch.Tensor, *args) -> torch.Tensor:
     """
     Mean squared error loss.
 
@@ -18,6 +18,8 @@ def mse_loss(source: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         Source patches.
     target : torch.Tensor
         Target patches.
+    *args : Any
+        Additional arguments.
 
     Returns
     -------
@@ -29,34 +31,37 @@ def mse_loss(source: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
 
 
 def n2v_loss(
-    manipulated_patches: torch.Tensor,
-    original_patches: torch.Tensor,
+    manipulated_batch: torch.Tensor,
+    original_batch: torch.Tensor,
     masks: torch.Tensor,
+    *args,
 ) -> torch.Tensor:
     """
     N2V Loss function described in A Krull et al 2018.
 
     Parameters
     ----------
-    manipulated_patches : torch.Tensor
-        Patches with manipulated pixels.
-    original_patches : torch.Tensor
-        Noisy patches.
+    manipulated_batch : torch.Tensor
+        Batch after manipulation function applied.
+    original_batch : torch.Tensor
+        Original images.
     masks : torch.Tensor
-        Array containing masked pixel locations.
+        Coordinates of changed pixels.
+    *args : Any
+        Additional arguments.
 
     Returns
     -------
     torch.Tensor
         Loss value.
     """
-    errors = (original_patches - manipulated_patches) ** 2
+    errors = (original_batch - manipulated_batch) ** 2
     # Average over pixels and batch
     loss = torch.sum(errors * masks) / torch.sum(masks)
     return loss  # TODO change output to dict ?
 
 
-def mae_loss(samples: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+def mae_loss(samples: torch.Tensor, labels: torch.Tensor, *args) -> torch.Tensor:
     """
     N2N Loss function described in to J Lehtinen et al 2018.
 
@@ -66,6 +71,8 @@ def mae_loss(samples: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
         Raw patches.
     labels : torch.Tensor
         Different subset of noisy patches.
+    *args : Any
+        Additional arguments.
 
     Returns
     -------

careamics/lvae_training/eval_utils.py

@@ -107,6 +107,15 @@ def get_first_index(bin_count, quantile):
     return None
 
 
+def get_device():
+    if torch.cuda.is_available():
+        return "cuda"
+    elif torch.backends.mps.is_available():
+        return "mps"
+    else:
+        return "cpu"
+
+
 def show_for_one(
     idx,
     val_dset,
@@ -516,7 +525,7 @@ def get_predictions(
                 num_workers=num_workers,
             )
             # get filename without extension and path
-            filename = str(d._fpath).split("/")[-1].split(".")[0]
+            filename = d._fpath.name
             multifile_stitched_predictions[filename] = stitched_predictions
             multifile_stitched_stds[filename] = stitched_stds
         return (
@@ -534,7 +543,7 @@ def get_predictions(
             num_workers=num_workers,
         )
         # get filename without extension and path
-        filename = str(dset._fpath).split("/")[-1].split(".")[0]
+        filename = dset._fpath.name
         return (
             {filename: stitched_predictions},
             {filename: stitched_stds},
@@ -553,6 +562,8 @@ def get_single_file_predictions(
     if tile_size and grid_size:
         dset.set_img_sz(tile_size, grid_size)
 
+    device = get_device()
+
     dloader = DataLoader(
         dset,
         pin_memory=False,
@@ -561,14 +572,14 @@ def get_single_file_predictions(
         batch_size=batch_size,
     )
     model.eval()
-    model.cuda()
+    model.to(device)
     tiles = []
     logvar_arr = []
     with torch.no_grad():
         for batch in tqdm(dloader, desc="Predicting tiles"):
             inp, tar = batch
-            inp = inp.cuda()
-            tar = tar.cuda()
+            inp = inp.to(device)
+            tar = tar.to(device)
 
             # get model output
             rec, _ = model(inp)
@@ -597,6 +608,8 @@ def get_single_file_mmse(
     num_workers: int = 4,
 ) -> tuple[np.ndarray, np.ndarray]:
     """Get patch-wise predictions from a model for a single file dataset."""
+    device = get_device()
+
     dloader = DataLoader(
         dset,
         pin_memory=False,
@@ -608,15 +621,15 @@ def get_single_file_mmse(
         dset.set_img_sz(tile_size, grid_size)
 
     model.eval()
-    model.cuda()
+    model.to(device)
     tile_mmse = []
     tile_stds = []
     logvar_arr = []
     with torch.no_grad():
         for batch in tqdm(dloader, desc="Predicting tiles"):
             inp, tar = batch
-            inp = inp.cuda()
-            tar = tar.cuda()
+            inp = inp.to(device)
+            tar = tar.to(device)
 
             rec_img_list = []
             for _ in range(mmse_count):

careamics/model_io/bioimage/bioimage_utils.py

@@ -3,6 +3,8 @@
 from pathlib import Path
 from typing import Union
 
+from careamics.utils.version import get_careamics_version
+
 
 def get_unzip_path(zip_path: Union[Path, str]) -> Path:
     """Generate unzipped folder path from the bioimage.io model path.
@@ -44,11 +46,11 @@ def create_env_text(pytorch_version: str, torchvision_version: str) -> str:
         f"name: careamics\n"
         f"dependencies:\n"
         f"  - python=3.10\n"
-        f"  - pytorch={pytorch_version}\n"
-        f"  - torchvision={torchvision_version}\n"
         f"  - pip\n"
         f"  - pip:\n"
-        f"    - git+https://github.com/CAREamics/careamics.git\n"
+        f"    - torch=={pytorch_version}\n"
+        f"    - torchvision=={torchvision_version}\n"
+        f"    - careamics=={get_careamics_version()}"
     )
 
     return env

careamics/model_io/bioimage/model_description.py

@@ -28,14 +28,14 @@ from bioimageio.spec.model.v0_5 import (
     WeightsDescr,
 )
 
-from careamics.config import Configuration, GeneralDataConfig
+from careamics.config import Configuration, DataConfig
 
 from ._readme_factory import readme_factory
 
 
 def _create_axes(
     array: np.ndarray,
-    data_config: GeneralDataConfig,
+    data_config: DataConfig,
     channel_names: Optional[list[str]] = None,
     is_input: bool = True,
 ) -> list[AxisBase]:
@@ -102,7 +102,7 @@ def _create_axes(
 def _create_inputs_ouputs(
     input_array: np.ndarray,
     output_array: np.ndarray,
-    data_config: GeneralDataConfig,
+    data_config: DataConfig,
     input_path: Union[Path, str],
     output_path: Union[Path, str],
     channel_names: Optional[list[str]] = None,

careamics/model_io/bmz_io.py

@@ -5,7 +5,6 @@ from pathlib import Path
 from typing import Optional, Union
 
 import numpy as np
-import pkg_resources
 from bioimageio.core import load_model_description, test_model
 from bioimageio.spec import ValidationSummary, save_bioimageio_package
 from pydantic import HttpUrl
@@ -16,6 +15,7 @@ from torchvision import __version__ as TORCHVISION_VERSION
 from careamics.config import Configuration, load_configuration, save_configuration
 from careamics.config.support import SupportedArchitecture
 from careamics.lightning.lightning_module import FCNModule, VAEModule
+from careamics.utils.version import get_careamics_version
 
 from .bioimage import (
     create_env_text,
@@ -139,7 +139,7 @@ def export_to_bmz(
         path_to_archive.parent.mkdir(parents=True, exist_ok=True)
 
     # versions
-    careamics_version = pkg_resources.get_distribution("careamics").version
+    careamics_version = get_careamics_version()
 
     # save files in temporary folder
     with tempfile.TemporaryDirectory() as tmpdirname:

careamics/model_io/model_io_utils.py

@@ -5,7 +5,7 @@ from typing import Union
 
 import torch
 
-from careamics.config import Configuration, configuration_factory
+from careamics.config import Configuration
 from careamics.lightning.lightning_module import FCNModule, VAEModule
 from careamics.model_io.bmz_io import load_from_bmz
 from careamics.utils import check_path_exists
@@ -92,4 +92,4 @@ def _load_checkpoint(
             f"{cfg_dict['algorithm_config']['model']['architecture']}"
         )
 
-    return model, configuration_factory(cfg_dict)
+    return model, Configuration(**cfg_dict)

careamics/transforms/__init__.py

@@ -5,15 +5,16 @@ __all__ = [
     "Denormalize",
     "ImageRestorationTTA",
     "N2VManipulate",
+    "N2VManipulateTorch",
     "Normalize",
     "XYFlip",
     "XYRandomRotate90",
     "get_all_transforms",
 ]
 
-
 from .compose import Compose, get_all_transforms
 from .n2v_manipulate import N2VManipulate
+from .n2v_manipulate_torch import N2VManipulateTorch
 from .normalize import Denormalize, Normalize
 from .tta import ImageRestorationTTA
 from .xy_flip import XYFlip

careamics/transforms/compose.py

@@ -6,7 +6,6 @@ from numpy.typing import NDArray
 
 from careamics.config.transformations import NORM_AND_SPATIAL_UNION
 
-from .n2v_manipulate import N2VManipulate
 from .normalize import Normalize
 from .transform import Transform
 from .xy_flip import XYFlip
@@ -14,7 +13,6 @@ from .xy_random_rotate90 import XYRandomRotate90
 
 ALL_TRANSFORMS = {
     "Normalize": Normalize,
-    "N2VManipulate": N2VManipulate,
     "XYFlip": XYFlip,
     "XYRandomRotate90": XYRandomRotate90,
 }
@@ -82,21 +80,13 @@ class Compose:
         tuple[np.ndarray, Optional[np.ndarray]]
             The output of the transformations.
         """
-        params: Union[
-            tuple[NDArray, Optional[NDArray]],
-            tuple[NDArray, NDArray, NDArray],  # N2VManiupulate output
-        ] = (patch, target)
+        params: Union[tuple[NDArray, Optional[NDArray]],] = (patch, target)
 
         for t in self.transforms:
-            # N2VManipulate returns tuple of 3 arrays
-            #   - Other transoforms return tuple of (patch, target, additional_arrays)
-            if isinstance(t, N2VManipulate):
-                patch, *_ = params
-                params = t(patch=patch)
-            else:
-                *params, _ = t(*params)  # ignore additional_arrays dict
-
-        return params
+            *params, _ = t(*params)  # ignore additional_arrays dict
+
+        # avoid None values that create problems for collating
+        return tuple(p for p in params if p is not None)
 
     def _chain_transforms_additional_arrays(
         self,

careamics/transforms/n2v_manipulate_torch.py

@@ -0,0 +1,143 @@
+"""N2V manipulation transform for PyTorch."""
+
+import platform
+from typing import Any, Optional
+
+import torch
+
+from careamics.config.support import SupportedPixelManipulation, SupportedStructAxis
+from careamics.config.transformations import N2VManipulateModel
+
+from .pixel_manipulation_torch import (
+    median_manipulate_torch,
+    uniform_manipulate_torch,
+)
+from .struct_mask_parameters import StructMaskParameters
+
+
+class N2VManipulateTorch:
+    """
+    Default augmentation for the N2V model.
+
+    This transform expects C(Z)YX dimensions.
+
+    Parameters
+    ----------
+    n2v_manipulate_config : N2VManipulateConfig
+        N2V manipulation configuration.
+    seed : Optional[int], optional
+        Random seed, by default None.
+
+    Attributes
+    ----------
+    masked_pixel_percentage : float
+        Percentage of pixels to mask.
+    roi_size : int
+        Size of the replacement area.
+    strategy : Literal[ "uniform", "median" ]
+        Replacement strategy, uniform or median.
+    remove_center : bool
+        Whether to remove central pixel from patch.
+    struct_mask : Optional[StructMaskParameters]
+        StructN2V mask parameters.
+    rng : Generator
+        Random number generator.
+    """
+
+    def __init__(
+        self,
+        n2v_manipulate_config: N2VManipulateModel,
+        seed: Optional[int] = None,
+    ):
+        """Constructor.
+
+        Parameters
+        ----------
+        n2v_manipulate_config : N2VManipulateModel
+            N2V manipulation configuration.
+        seed : Optional[int], optional
+            Random seed, by default None.
+        """
+        self.masked_pixel_percentage = n2v_manipulate_config.masked_pixel_percentage
+        self.roi_size = n2v_manipulate_config.roi_size
+        self.strategy = n2v_manipulate_config.strategy
+        self.remove_center = n2v_manipulate_config.remove_center
+
+        if n2v_manipulate_config.struct_mask_axis == SupportedStructAxis.NONE:
+            self.struct_mask: Optional[StructMaskParameters] = None
+        else:
+            self.struct_mask = StructMaskParameters(
+                axis=(
+                    0
+                    if n2v_manipulate_config.struct_mask_axis
+                    == SupportedStructAxis.HORIZONTAL
+                    else 1
+                ),
+                span=n2v_manipulate_config.struct_mask_span,
+            )
+
+        # PyTorch random generator
+        # TODO refactor into careamics.utils.torch_utils.get_device
+        if torch.cuda.is_available():
+            device = "cuda"
+        elif torch.backends.mps.is_available() and platform.processor() in (
+            "arm",
+            "arm64",
+        ):
+            device = "mps"
+        else:
+            device = "cpu"
+
+        self.rng = (
+            torch.Generator(device=device).manual_seed(seed)
+            if seed is not None
+            else torch.Generator(device=device)
+        )
+
+    def __call__(
+        self, batch: torch.Tensor, *args: Any, **kwargs: Any
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Apply the transform to the image.
+
+        Parameters
+        ----------
+        batch : torch.Tensor
+            Batch if image patches, 2D or 3D, shape BC(Z)YX.
+        *args : Any
+            Additional arguments, unused.
+        **kwargs : Any
+            Additional keyword arguments, unused.
+
+        Returns
+        -------
+        tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+            Masked patch, original patch, and mask.
+        """
+        masked = torch.zeros_like(batch)
+        mask = torch.zeros_like(batch, dtype=torch.uint8)
+
+        if self.strategy == SupportedPixelManipulation.UNIFORM:
+            # Iterate over the channels to apply manipulation separately
+            for c in range(batch.shape[1]):
+                masked[:, c, ...], mask[:, c, ...] = uniform_manipulate_torch(
+                    patch=batch[:, c, ...],
+                    mask_pixel_percentage=self.masked_pixel_percentage,
+                    subpatch_size=self.roi_size,
+                    remove_center=self.remove_center,
+                    struct_params=self.struct_mask,
+                    rng=self.rng,
+                )
+        elif self.strategy == SupportedPixelManipulation.MEDIAN:
+            # Iterate over the channels to apply manipulation separately
+            for c in range(batch.shape[1]):
+                masked[:, c, ...], mask[:, c, ...] = median_manipulate_torch(
+                    batch=batch[:, c, ...],
+                    mask_pixel_percentage=self.masked_pixel_percentage,
+                    subpatch_size=self.roi_size,
+                    struct_params=self.struct_mask,
+                    rng=self.rng,
+                )
+        else:
+            raise ValueError(f"Unknown masking strategy ({self.strategy}).")
+
+        return masked, batch, mask

careamics/transforms/pixel_manipulation.py

@@ -213,6 +213,7 @@ def _create_subpatch_struct_mask(
     np.ndarray
         StructN2V mask for the subpatch.
     """
+    # TODO no test for this function!
     # Create a mask with the center of the subpatch masked
     mask_placeholder = np.ones(subpatch.shape)