careamics 0.1.0rc1__py3-none-any.whl → 0.1.0rc2__py3-none-any.whl

careamics/engine.py

@@ -10,12 +10,11 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import torch
-from bioimageio.spec.model.raw_nodes import Model as BioimageModel
 from torch.utils.data import DataLoader, TensorDataset
 
 from .bioimage import (
-    build_zip_model,
     get_default_model_specs,
+    save_bioimage_model,
 )
 from .config import Configuration, load_configuration
 from .dataset.prepare_dataset import (
@@ -32,7 +31,7 @@ from .prediction import (
 )
 from .utils import (
     MetricTracker,
-    check_array_validity,
+    add_axes,
     denormalize,
     get_device,
     normalize,
@@ -40,7 +39,6 @@ from .utils import (
 from .utils.logging import ProgressBar, get_logger
 
 
-# TODO: refactor private methods and bioimage.io to other modules
 class Engine:
     """
     Class allowing training of a model and subsequent prediction.
@@ -165,47 +163,52 @@ class Engine:
             self.scaler,
             self.cfg,
         ) = create_model(config=self.cfg, model_path=model_path, device=self.device)
+        assert self.cfg is not None
 
         # create loss function
-        if self.cfg is not None:
-            self.loss_func = create_loss_function(self.cfg)
+        self.loss_func = create_loss_function(self.cfg)
+
+        # Set logging
+        log_path = self.cfg.working_directory / "log.txt"
+        self.logger = get_logger(__name__, log_path=log_path)
 
-            # Set logging
-            log_path = self.cfg.working_directory / "log.txt"
-            self.logger = get_logger(__name__, log_path=log_path)
+        # wandb
+        self.use_wandb = self.cfg.training.use_wandb
 
-            # wandb
-            self.use_wandb = self.cfg.training.use_wandb
+        if self.use_wandb:
+            try:
+                from wandb.errors import UsageError
+
+                from careamics.utils.wandb import WandBLogging
 
-            if self.use_wandb:
                 try:
-                    from wandb.errors import UsageError
-
-                    from careamics.utils.wandb import WandBLogging
-
-                    try:
-                        self.wandb = WandBLogging(
-                            experiment_name=self.cfg.experiment_name,
-                            log_path=self.cfg.working_directory,
-                            config=self.cfg,
-                            model_to_watch=self.model,
-                        )
-                    except UsageError as e:
-                        self.logger.warning(
-                            f"Wandb usage error, using default logger. Check whether "
-                            f"wandb correctly configured:\n"
-                            f"{e}"
-                        )
-                        self.use_wandb = False
-
-                except ModuleNotFoundError:
+                    self.wandb = WandBLogging(
+                        experiment_name=self.cfg.experiment_name,
+                        log_path=self.cfg.working_directory,
+                        config=self.cfg,
+                        model_to_watch=self.model,
+                    )
+                except UsageError as e:
                     self.logger.warning(
-                        "Wandb not installed, using default logger. Try pip install "
-                        "wandb"
+                        f"Wandb usage error, using default logger. Check whether "
+                        f"wandb correctly configured:\n"
+                        f"{e}"
                     )
                     self.use_wandb = False
-        else:
-            raise ValueError("Configuration is not defined.")
+
+            except ModuleNotFoundError:
+                self.logger.warning(
+                    "Wandb not installed, using default logger. Try pip install "
+                    "wandb"
+                )
+                self.use_wandb = False
+
+        # BMZ inputs/outputs placeholders, filled during validation
+        self._input = None
+        self._outputs = None
+
+        # torch version
+        self.torch_version = torch.__version__
 
     def train(
         self,
@@ -387,6 +390,12 @@ class Engine:
 
         with torch.no_grad():
             for patch, *auxillary in val_loader:
+                # if inputs is None, record a single patch
+                if self._input is None:
+                    # patch has dimension SC(Z)YX
+                    self._input = patch.clone().detach().cpu().numpy()
+
+                # evaluate
                 outputs = self.model(patch.to(self.device))
                 loss = self.loss_func(
                     outputs, *[a.to(self.device) for a in auxillary], self.device
@@ -409,10 +418,18 @@ class Engine:
         The Engine must have previously been trained and mean/std be specified in
         its configuration.
 
+        Data should be compatible with the axes, either from the configuration or
+        as passed using the `axes` parameter. If the batch and channel dimensions are
+        missing, then singleton dimensions are added.
+
         To use tiling, both `tile_shape` and `overlaps` must be specified, have same
         length, be divisible by 2 and greater than 0. Finally, the overlaps must be
         smaller than the tiles.
 
+        By setting `tta` to `True`, the prediction is performed using test time
+        augmentation, meaning that the input is augmented and the prediction is averaged
+        over the augmentations.
+
         Parameters
         ----------
         input : Union[np.ndarra, str, Path]
@@ -497,6 +514,18 @@ class Engine:
         UserWarning
             If the samples have different shapes, the prediction then returns a list.
         """
+        # checks are done here to satisfy mypy
+        # check that configuration exists
+        if self.cfg is None:
+            raise ValueError("Configuration is not defined, cannot predict.")
+
+        # Check that the mean and std are there (= has been trained)
+        if not self.cfg.data.mean or not self.cfg.data.std:
+            raise ValueError(
+                "Mean or std are not specified in the configuration, prediction cannot "
+                "be performed."
+            )
+
         prediction = []
         tiles = []
         stitching_data = []
@@ -513,9 +542,7 @@ class Engine:
                     predicted_augments = []
                     for augmented_tile in augmented_tiles:
                         augmented_pred = self.model(augmented_tile.to(self.device))
-                        predicted_augments.append(
-                            augmented_pred.squeeze().cpu().numpy()
-                        )
+                        predicted_augments.append(augmented_pred.cpu())
                     tiles.append(tta_backward(predicted_augments).squeeze())
                 else:
                     tiles.append(
@@ -529,8 +556,8 @@ class Engine:
                     predicted_sample = stitch_prediction(tiles, stitching_data)
                     predicted_sample = denormalize(
                         predicted_sample,
-                        float(self.cfg.data.mean),  # type: ignore
-                        float(self.cfg.data.std),  # type: ignore
+                        float(self.cfg.data.mean),
+                        float(self.cfg.data.std),
                     )
                     prediction.append(predicted_sample)
                     tiles.clear()
@@ -566,6 +593,18 @@ class Engine:
         np.ndarray
             Predicted image.
         """
+        # checks are done here to satisfy mypy
+        # check that configuration exists
+        if self.cfg is None:
+            raise ValueError("Configuration is not defined, cannot predict.")
+
+        # Check that the mean and std are there (= has been trained)
+        if not self.cfg.data.mean or not self.cfg.data.std:
+            raise ValueError(
+                "Mean or std are not specified in the configuration, prediction cannot "
+                "be performed."
+            )
+
         prediction = []
         with torch.no_grad():
             for i, sample in enumerate(pred_loader):
@@ -574,9 +613,7 @@ class Engine:
                     predicted_augments = []
                     for augmented_pred in augmented_preds:
                         augmented_pred = self.model(augmented_pred.to(self.device))
-                        predicted_augments.append(
-                            augmented_pred.squeeze().cpu().numpy()
-                        )
+                        predicted_augments.append(augmented_pred.cpu())
                     prediction.append(tta_backward(predicted_augments).squeeze())
                 else:
                     prediction.append(
@@ -584,7 +621,9 @@ class Engine:
                     )
                 progress_bar.update(i, 1)
         output = denormalize(
-            np.stack(prediction), float(self.cfg.data.mean), float(self.cfg.data.std)  # type: ignore
+            np.stack(prediction).squeeze(),
+            float(self.cfg.data.mean),
+            float(self.cfg.data.std),
         )
         return output
 
@@ -696,14 +735,13 @@ class Engine:
             )
 
         if input is None:
-            raise ValueError("Ccannot predict on None input.")
+            raise ValueError("Input cannot be None.")
 
         # Create dataset
         if isinstance(input, np.ndarray):  # np.ndarray
-            # Check that the axes fit the input
+            # Validate axes and add missing dimensions (S)C if necessary
             img_axes = self.cfg.data.axes if axes is None else axes
-            # TODO are self.cfg.data.axes and axes compatible (same spatial dim)?
-            check_array_validity(input, img_axes)
+            input_expanded = add_axes(input, img_axes)
 
             # Check if tiling requested
             tiled = tile_shape is not None and overlaps is not None
@@ -714,11 +752,9 @@ class Engine:
                     "Tiling with in memory array is currently not implemented."
                 )
 
-                # check_tiling_validity(tile_shape, overlaps)
-
             # Normalize input and cast to float32
             normalized_input = normalize(
-                img=input, mean=self.cfg.data.mean, std=self.cfg.data.std
+                img=input_expanded, mean=self.cfg.data.mean, std=self.cfg.data.std
             )
             normalized_input = normalized_input.astype(np.float32)
 
@@ -812,7 +848,66 @@ class Engine:
                     self.logger.removeHandler(handler)
                     handler.close()
 
-    def _generate_rdf(self, model_specs: Optional[dict] = None) -> dict:
+    def _get_sample_io_files(
+        self,
+        input_array: Optional[np.ndarray] = None,
+        axes: Optional[str] = None,
+    ) -> Tuple[List[str], List[str]]:
+        """
+        Create numpy format for use as inputs and outputs in the bioimage.io archive.
+
+        Parameters
+        ----------
+        input_array : Optional[np.ndarray], optional
+            Input array to use for the bioimage.io model zoo, by default None.
+        axes : Optional[str], optional
+            Axes from the configuration.
+
+        Returns
+        -------
+        Tuple[List[str], List[str]]
+            Tuple of input and output file paths.
+
+        Raises
+        ------
+        ValueError
+            If the configuration is not defined.
+        """
+        if self.cfg is not None and self._input is not None:
+            # use the input array if provided, otherwise use the first validation sample
+            if input_array is not None:
+                array_in = input_array
+
+                # add axes to be compatible with the axes declared in the RDF specs
+                add_axes(array_in, axes)
+            else:
+                array_in = self._input
+
+            # predict (no tta since BMZ does not apply it)
+            array_out = self.predict(array_in, tta=False)
+
+            # add singleton dimensions (for compatibility with model axes)
+            # indeed, BMZ applies the model but CAREamics function are meant
+            # to work on user data (potentially with no S or C axe)
+            array_out = array_out[np.newaxis, np.newaxis, ...]
+
+            # save numpy files
+            workdir = self.cfg.working_directory
+            in_file = workdir.joinpath("test_inputs.npy")
+            np.save(in_file, array_in)
+            out_file = workdir.joinpath("test_outputs.npy")
+            np.save(out_file, array_out)
+
+            return [str(in_file.absolute())], [str(out_file.absolute())]
+        else:
+            raise ValueError("Configuration is not defined or model was not trained.")
+
+    def _generate_rdf(
+        self,
+        *,
+        model_specs: Optional[dict] = None,
+        input_array: Optional[np.ndarray] = None,
+    ) -> dict:
         """
         Generate rdf data for bioimage.io format export.
 
@@ -820,6 +915,8 @@ class Engine:
         ----------
         model_specs : Optional[dict], optional
             Custom specs if different than the default ones, by default None.
+        input_array : Optional[np.ndarray], optional
+            Input array to use for the bioimage.io model zoo, by default None.
 
         Returns
         -------
@@ -848,7 +945,9 @@ class Engine:
                 axes = "b" + axes
 
             # get in/out samples' files
-            test_inputs, test_outputs = self._get_sample_io_files(axes)
+            test_inputs, test_outputs = self._get_sample_io_files(
+                input_array, self.cfg.data.axes
+            )
 
             specs = get_default_model_specs(
                 "Noise2Void",
@@ -861,7 +960,6 @@ class Engine:
 
             specs.update(
                 {
-                    "architecture": "careamics.models.unet",
                     "test_inputs": test_inputs,
                     "test_outputs": test_outputs,
                     "input_axes": [axes],
@@ -870,14 +968,21 @@ class Engine:
             )
             return specs
         else:
-            raise ValueError("Configuration is not defined.")
+            raise ValueError("Configuration is not defined or model was not trained.")
 
     def save_as_bioimage(
-        self, output_zip: Union[Path, str], model_specs: Optional[dict] = None
-    ) -> BioimageModel:
+        self,
+        output_zip: Union[Path, str],
+        model_specs: Optional[dict] = None,
+        input_array: Optional[np.ndarray] = None,
+    ) -> None:
         """
         Export the current model to BioImage.io model zoo format.
 
+        Custom specs can be passed in `model_specs (e.g. maintainers). For a description
+        of the model RDF, refer to
+        github.com/bioimage-io/spec-bioimage-io/blob/gh-pages/model_spec_latest.md.
+
         Parameters
         ----------
         output_zip : Union[Path, str]
@@ -885,11 +990,10 @@ class Engine:
         model_specs : Optional[dict]
             A dictionary with keys being the bioimage-core build_model parameters. If
             None then it will be populated by the model default specs.
-
-        Returns
-        -------
-        BioimageModel
-            Bioimage.io model object.
+        input_array : Optional[np.ndarray]
+            An array to use as input for the bioimage.io model zoo. If None then the
+            first validation sample will be used. Note that the array must have S and
+            C dimensions (e.g. SCYX), even if only singleton dimensions.
 
         Raises
         ------
@@ -898,57 +1002,13 @@ class Engine:
         """
         if self.cfg is not None:
             # Generate specs
-            specs = self._generate_rdf(model_specs)
+            specs = self._generate_rdf(model_specs=model_specs, input_array=input_array)
 
             # Build model
-            raw_model = build_zip_model(
+            save_bioimage_model(
                 path=output_zip,
                 config=self.cfg,
-                model_specs=specs,
+                specs=specs,
             )
-
-            return raw_model
-        else:
-            raise ValueError("Configuration is not defined.")
-
-    def _get_sample_io_files(self, axes: str) -> Tuple[List[str], List[str]]:
-        """
-        Create numpy format for use as inputs and outputs in the bioimage.io archive.
-
-        Parameters
-        ----------
-        axes : str
-            Input and output axes.
-
-        Returns
-        -------
-        Tuple[List[str], List[str]]
-            Tuple of input and output file paths.
-
-        Raises
-        ------
-        ValueError
-            If the configuration is not defined.
-        """
-        # input:
-        if self.cfg is not None:
-            sample_input = np.random.randn(*self.cfg.training.patch_size)
-            # if there are more input axes (like channel, ...),
-            # then expand the sample dimensions.
-            len_diff = len(axes) - len(self.cfg.training.patch_size)
-            if len_diff > 0:
-                sample_input = np.expand_dims(
-                    sample_input, axis=tuple(i for i in range(len_diff))
-                )
-            sample_output = np.random.randn(*sample_input.shape)
-
-            # save numpy files
-            workdir = self.cfg.working_directory
-            in_file = workdir.joinpath("test_inputs.npy")
-            np.save(in_file, sample_input)
-            out_file = workdir.joinpath("test_outputs.npy")
-            np.save(out_file, sample_output)
-
-            return [str(in_file.absolute())], [str(out_file.absolute())]
         else:
             raise ValueError("Configuration is not defined.")

careamics/losses/loss_factory.py

@@ -5,8 +5,9 @@ This module contains a factory function for creating loss functions.
 """
 from typing import Callable
 
-from ..config import Configuration
-from ..config.algorithm import Loss
+from careamics.config import Configuration
+from careamics.config.algorithm import Loss
+
 from .losses import n2v_loss
 
 

careamics/models/model_factory.py

@@ -8,10 +8,11 @@ from typing import Dict, Optional, Tuple, Union
 
 import torch
 
-from ..bioimage import import_bioimage_model
-from ..config import Configuration
-from ..config.algorithm import Models
-from ..utils.logging import get_logger
+from careamics.bioimage import import_bioimage_model
+from careamics.config import Configuration
+from careamics.config.algorithm import Models
+from careamics.utils.logging import get_logger
+
 from .unet import UNet
 
 logger = get_logger(__name__)
@@ -49,16 +50,27 @@ def create_model(
     model_path: Optional[Union[str, Path]] = None,
     config: Optional[Configuration] = None,
     device: Optional[torch.device] = None,
-) -> torch.nn.Module:
+) -> Tuple[
+    torch.nn.Module,
+    torch.optim.Optimizer,
+    Union[
+        torch.optim.lr_scheduler.LRScheduler,
+        torch.optim.lr_scheduler.ReduceLROnPlateau,  # not a subclass of LRScheduler
+    ],
+    torch.cuda.amp.GradScaler,
+    Configuration,
+]:
     """
-    Instantiate a model from a checkpoint or configuration.
+    Instantiate a model from a model path or configuration.
 
-    If both checkpoint and configuration are provided, the checkpoint is used.
+    If both path and configuration are provided, the model path is used. The model
+    path should point to either a checkpoint (created during training) or a model
+    exported to the bioimage.io format.
 
     Parameters
     ----------
     model_path : Optional[Union[str, Path]], optional
-        Path to a checkpoint, by default None.
+        Path to a checkpoint or bioimage.io archive, by default None.
     config : Optional[Configuration], optional
         Configuration, by default None.
     device : Optional[torch.device], optional
@@ -103,12 +115,13 @@ def create_model(
             raise ValueError("Invalid checkpoint format, no configuration found.")
 
         # Create model
-        model = model_registry(model_name)(
+        model: torch.nn.Module = model_registry(model_name)(
             depth=model_config.depth,
             conv_dim=algo_config.get_conv_dim(),
             num_channels_init=model_config.num_channels_init,
         )
         model.to(device)
+
         # Load the model state dict
         if "model_state_dict" in checkpoint:
             model.load_state_dict(checkpoint["model_state_dict"])
@@ -141,13 +154,19 @@ def create_model(
 
     else:
         raise ValueError("Either config or model_path must be provided")
-    # model = compile_model(model)
+
     return model, optimizer, scheduler, scaler, config
 
 
 def get_optimizer_and_scheduler(
     config: Configuration, model: torch.nn.Module, state_dict: Optional[Dict] = None
-) -> Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LRScheduler]:
+) -> Tuple[
+    torch.optim.Optimizer,
+    Union[
+        torch.optim.lr_scheduler.LRScheduler,
+        torch.optim.lr_scheduler.ReduceLROnPlateau,  # not a subclass of LRScheduler
+    ],
+]:
     """
     Create optimizer and learning rate schedulers.
 

careamics/prediction/prediction_utils.py

@@ -46,16 +46,18 @@ def stitch_prediction(
     return predicted_image
 
 
-def tta_forward(x: np.ndarray) -> List:
+def tta_forward(x: torch.Tensor) -> List[torch.Tensor]:
     """
     Augment 8-fold an array.
 
     The augmentation is performed using all 90 deg rotations and their flipped version,
     as well as the original image flipped.
 
+    Tensors should be of shape SC(Z)YX, with S and C potentially singleton dimensions.
+
     Parameters
     ----------
-    x : torch.tensor
+    x : torch.Tensor
         Data to augment.
 
     Returns
@@ -75,13 +77,15 @@ def tta_forward(x: np.ndarray) -> List:
     return x_aug_flip
 
 
-def tta_backward(x_aug: List) -> np.ndarray:
+def tta_backward(x_aug: List[torch.Tensor]) -> np.ndarray:
     """
     Invert `tta_forward` and average the 8 images.
 
+    The function takes a list of torch tensors and returns a numpy array.
+
     Parameters
     ----------
-    x_aug : List
+    x_aug : List[torch.Tensor]
         Stack of 8-fold augmented images.
 
     Returns
@@ -90,13 +94,13 @@ def tta_backward(x_aug: List) -> np.ndarray:
         Average of de-augmented x_aug.
     """
     x_deaug = [
-        x_aug[0],
-        np.rot90(x_aug[1], -1),
-        np.rot90(x_aug[2], -2),
-        np.rot90(x_aug[3], -3),
-        np.fliplr(x_aug[4]),
-        np.rot90(np.fliplr(x_aug[5]), -1),
-        np.rot90(np.fliplr(x_aug[6]), -2),
-        np.rot90(np.fliplr(x_aug[7]), -3),
+        x_aug[0].numpy(),
+        np.rot90(x_aug[1], -1, axes=(2, 3)),
+        np.rot90(x_aug[2], -2, axes=(2, 3)),
+        np.rot90(x_aug[3], -3, axes=(2, 3)),
+        np.flip(x_aug[4].numpy(), axis=(1, 3)),
+        np.rot90(np.flip(x_aug[5].numpy(), axis=(1, 3)), -1, axes=(2, 3)),
+        np.rot90(np.flip(x_aug[6].numpy(), axis=(1, 3)), -2, axes=(2, 3)),
+        np.rot90(np.flip(x_aug[7].numpy(), axis=(1, 3)), -3, axes=(2, 3)),
     ]
     return np.mean(x_deaug, 0)

careamics/utils/__init__.py

@@ -5,11 +5,10 @@ __all__ = [
     "denormalize",
     "normalize",
     "get_device",
-    "check_array_validity",
     "check_axes_validity",
+    "add_axes",
     "check_tiling_validity",
     "cwd",
-    "compile_model",
     "MetricTracker",
 ]
 
@@ -17,5 +16,5 @@ __all__ = [
 from .context import cwd
 from .metrics import MetricTracker
 from .normalization import denormalize, normalize
-from .torch_utils import compile_model, get_device
-from .validators import check_array_validity, check_axes_validity, check_tiling_validity
+from .torch_utils import get_device
+from .validators import add_axes, check_axes_validity, check_tiling_validity