careamics 0.0.14__py3-none-any.whl → 0.0.16__py3-none-any.whl

careamics/lvae_training/dataset/multich_dataset.py

@@ -2,29 +2,35 @@
 A place for Datasets and Dataloaders.
 """
 
-from typing import Tuple, Union, Callable
+from pathlib import Path
+from typing import Any, Callable, Optional, Union
 
 import numpy as np
+import torch
+from torch.utils.data import Dataset
 
 from .utils.empty_patch_fetcher import EmptyPatchFetcher
 from .utils.index_manager import GridIndexManager
 from .utils.index_switcher import IndexSwitcher
-from .config import DatasetConfig
+from .config import MicroSplitDataConfig
 from .types import DataSplitType, TilingMode
 
 
-class MultiChDloader:
+class MultiChDloader(Dataset):
+    """Multi-channel dataset loader."""
+
     def __init__(
         self,
-        data_config: DatasetConfig,
-        fpath: str,
-        load_data_fn: Callable,
-        val_fraction: float = None,
-        test_fraction: float = None,
+        data_config: MicroSplitDataConfig,
+        datapath: Union[str, Path],
+        load_data_fn: Optional[Callable] = None,
+        val_fraction: float = 0.1,
+        test_fraction: float = 0.1,
+        allow_generation: bool = False,
     ):
         """ """
         self._data_type = data_config.data_type
-        self._fpath = fpath
+        self._fpath = datapath
         self._data = self._noise_data = None
         self.Z = 1
         self._5Ddata = False
@@ -395,7 +401,7 @@ class MultiChDloader:
         )
 
     def get_idx_manager_shapes(
-        self, patch_size: int, grid_size: Union[int, Tuple[int, int, int]]
+        self, patch_size: int, grid_size: Union[int, tuple[int, int, int]]
     ):
         numC = self._data.shape[-1]
         if self._5Ddata:
@@ -415,7 +421,7 @@ class MultiChDloader:
 
         return patch_shape, grid_shape
 
-    def set_img_sz(self, image_size, grid_size: Union[int, Tuple[int, int, int]]):
+    def set_img_sz(self, image_size, grid_size: Union[int, tuple[int, int, int]]):
         """
         If one wants to change the image size on the go, then this can be used.
         Args:
@@ -519,7 +525,7 @@ class MultiChDloader:
             },
         )
 
-    def _crop_img(self, img: np.ndarray, patch_start_loc: Tuple):
+    def _crop_img(self, img: np.ndarray, patch_start_loc: tuple):
         if self._tiling_mode in [TilingMode.TrimBoundary, TilingMode.ShiftBoundary]:
             # In training, this is used.
             # NOTE: It is my opinion that if I just use self._crop_img_with_padding, it will work perfectly fine.
@@ -600,7 +606,7 @@ class MultiChDloader:
         return new_img
 
     def _crop_flip_img(
-        self, img: np.ndarray, patch_start_loc: Tuple, h_flip: bool, w_flip: bool
+        self, img: np.ndarray, patch_start_loc: tuple, h_flip: bool, w_flip: bool
     ):
         new_img = self._crop_img(img, patch_start_loc)
         if h_flip:
@@ -611,8 +617,8 @@ class MultiChDloader:
         return new_img.astype(np.float32)
 
     def _load_img(
-        self, index: Union[int, Tuple[int, int]]
-    ) -> Tuple[np.ndarray, np.ndarray]:
+        self, index: Union[int, tuple[int, int]]
+    ) -> tuple[np.ndarray, np.ndarray]:
         """
         Returns the channels and also the respective noise channels.
         """
@@ -806,7 +812,7 @@ class MultiChDloader:
             w_start = 0
         return h_start, w_start
 
-    def _get_img(self, index: Union[int, Tuple[int, int]]):
+    def _get_img(self, index: Union[int, tuple[int, int]]):
         """
         Loads an image.
         Crops the image such that cropped image has content.
@@ -1056,8 +1062,8 @@ class MultiChDloader:
         return img_tuples, noise_tuples
 
     def __getitem__(
-        self, index: Union[int, Tuple[int, int]]
-    ) -> Tuple[np.ndarray, np.ndarray]:
+        self, index: Union[int, tuple[int, int]]
+    ) -> tuple[np.ndarray, np.ndarray]:
         # Vera: input can be both real microscopic image and two separate channels that are summed in the code
 
         if self._train_index_switcher is not None:

careamics/lvae_training/dataset/multifile_dataset.py

@@ -4,7 +4,7 @@ from typing import Callable, Union
 import numpy as np
 from numpy.typing import NDArray
 
-from .config import DatasetConfig
+from .config import MicroSplitDataConfig
 from .lc_dataset import LCMultiChDloader
 from .multich_dataset import MultiChDloader
 from .types import DataSplitType
@@ -82,7 +82,7 @@ class SingleFileLCDset(LCMultiChDloader):
     def __init__(
         self,
         preloaded_data: NDArray,
-        data_config: DatasetConfig,
+        data_config: MicroSplitDataConfig,
         fpath: str,
         load_data_fn: Callable,
         val_fraction=None,
@@ -106,7 +106,7 @@ class SingleFileLCDset(LCMultiChDloader):
 
     def load_data(
         self,
-        data_config: DatasetConfig,
+        data_config: MicroSplitDataConfig,
         datasplit_type: DataSplitType,
         load_data_fn: Callable,
         val_fraction=None,
@@ -124,7 +124,7 @@ class SingleFileDset(MultiChDloader):
     def __init__(
         self,
         preloaded_data: NDArray,
-        data_config: DatasetConfig,
+        data_config: MicroSplitDataConfig,
         fpath: str,
         load_data_fn: Callable,
         val_fraction=None,
@@ -148,7 +148,7 @@ class SingleFileDset(MultiChDloader):
 
     def load_data(
         self,
-        data_config: DatasetConfig,
+        data_config: MicroSplitDataConfig,
         datasplit_type: DataSplitType,
         load_data_fn: Callable[..., NDArray],
         val_fraction=None,
@@ -175,7 +175,7 @@ class MultiFileDset:
 
     def __init__(
         self,
-        data_config: DatasetConfig,
+        data_config: MicroSplitDataConfig,
         fpath: str,
         load_data_fn: Callable[..., Union[TwoChannelData, MultiChannelData]],
         val_fraction=None,

careamics/model_io/bioimage/model_description.py

@@ -1,10 +1,10 @@
 """Module use to build BMZ model description."""
 
 from pathlib import Path
-from typing import Optional, Union
+from typing import Union
 
 import numpy as np
-from bioimageio.spec._internal.io import resolve_and_extract
+from bioimageio.spec._internal.io import extract
 from bioimageio.spec.model.v0_5 import (
     ArchitectureFromLibraryDescr,
     Author,
@@ -12,7 +12,6 @@ from bioimageio.spec.model.v0_5 import (
     AxisId,
     BatchAxis,
     ChannelAxis,
-    EnvironmentFileDescr,
     FileDescr,
     FixedZeroMeanUnitVarianceAlongAxisKwargs,
     FixedZeroMeanUnitVarianceDescr,
@@ -36,7 +35,7 @@ from ._readme_factory import readme_factory
 def _create_axes(
     array: np.ndarray,
     data_config: DataConfig,
-    channel_names: Optional[list[str]] = None,
+    channel_names: list[str] | None = None,
     is_input: bool = True,
 ) -> list[AxisBase]:
     """Create axes description.
@@ -105,7 +104,7 @@ def _create_inputs_ouputs(
     data_config: DataConfig,
     input_path: Union[Path, str],
     output_path: Union[Path, str],
-    channel_names: Optional[list[str]] = None,
+    channel_names: list[str] | None = None,
 ) -> tuple[InputTensorDescr, OutputTensorDescr]:
     """Create input and output tensor description.
 
@@ -197,7 +196,7 @@ def create_model_description(
     config_path: Union[Path, str],
     env_path: Union[Path, str],
     covers: list[Union[Path, str]],
-    channel_names: Optional[list[str]] = None,
+    channel_names: list[str] | None = None,
     model_version: str = "0.1.0",
 ) -> ModelDescr:
     """Create model description.
@@ -269,7 +268,7 @@ def create_model_description(
             source=weights_path,
             architecture=architecture_descr,
             pytorch_version=Version(torch_version),
-            dependencies=EnvironmentFileDescr(source=env_path),
+            dependencies=FileDescr(source=Path(env_path)),
         ),
     )
 
@@ -322,9 +321,11 @@ def extract_model_path(model_desc: ModelDescr) -> tuple[Path, Path]:
     """
     if model_desc.weights.pytorch_state_dict is None:
         raise ValueError("No model weights found in model description.")
-    weights_path = resolve_and_extract(
-        model_desc.weights.pytorch_state_dict.source
-    ).path
+
+    # extract the zip model and return the directory
+    model_dir = extract(model_desc.root)
+
+    weights_path = model_dir.joinpath(model_desc.weights.pytorch_state_dict.source.path)
 
     for file in model_desc.attachments:
         file_path = file.source if isinstance(file.source, Path) else file.source.path
@@ -332,7 +333,7 @@ def extract_model_path(model_desc: ModelDescr) -> tuple[Path, Path]:
             continue
         file_path = Path(file_path)
         if file_path.name == "careamics.yaml":
-            config_path = resolve_and_extract(file.source).path
+            config_path = model_dir.joinpath(file.source.path)
             break
     else:
         raise ValueError("Configuration file not found.")

careamics/model_io/bmz_io.py

@@ -2,7 +2,7 @@
 
 import tempfile
 from pathlib import Path
-from typing import Optional, Union
+from typing import Union
 
 import numpy as np
 from bioimageio.core import load_model_description, test_model
@@ -90,8 +90,8 @@ def export_to_bmz(
     authors: list[dict],
     input_array: np.ndarray,
     output_array: np.ndarray,
-    covers: Optional[list[Union[Path, str]]] = None,
-    channel_names: Optional[list[str]] = None,
+    covers: list[Union[Path, str]] | None = None,
+    channel_names: list[str] | None = None,
     model_version: str = "0.1.0",
 ) -> None:
     """Export the model to BioImage Model Zoo format.
@@ -186,11 +186,15 @@ def export_to_bmz(
         )
 
         # test model description
-        test_kwargs = (
-            model_description.config.get("bioimageio", {})
-            .get("test_kwargs", {})
-            .get("pytorch_state_dict", {})
-        )
+        test_kwargs = {}
+        if hasattr(model_description, "config") and isinstance(
+            model_description.config, dict
+        ):
+            bioimageio_config = model_description.config.get("bioimageio", {})
+            test_kwargs = bioimageio_config.get("test_kwargs", {}).get(
+                "pytorch_state_dict", {}
+            )
+
         summary: ValidationSummary = test_model(model_description, **test_kwargs)
         if summary.status == "failed":
             raise ValueError(f"Model description test failed: {summary}")

careamics/models/layers.py

@@ -4,7 +4,7 @@ Layer module.
 This submodule contains layers used in the CAREamics models.
 """
 
-from typing import Optional, Union
+from typing import Union
 
 import torch
 import torch.nn as nn
@@ -207,8 +207,8 @@ def get_pascal_kernel_1d(
     kernel_size: int,
     norm: bool = False,
     *,
-    device: Optional[torch.device] = None,
-    dtype: Optional[torch.dtype] = None,
+    device: torch.device | None = None,
+    dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
     """Generate Yang Hui triangle (Pascal's triangle) for a given number.
 
@@ -270,8 +270,8 @@ def _get_pascal_kernel_nd(
     norm: bool = True,
     dim: int = 2,
     *,
-    device: Optional[torch.device] = None,
-    dtype: Optional[torch.dtype] = None,
+    device: torch.device | None = None,
+    dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
     """Generate pascal filter kernel by kernel size.
 

careamics/models/lvae/likelihoods.py

@@ -54,12 +54,8 @@ def likelihood_factory(
         )
     elif isinstance(config, NMLikelihoodConfig):
         return NoiseModelLikelihood(
-            data_mean=config.data_mean,
-            data_std=config.data_std,
             noise_model=noise_model,
         )
-    else:
-        raise ValueError(f"Invalid likelihood model type: {config.model_type}")
 
 
 # TODO: is it really worth to have this class? Or it just adds complexity? --> REFACTOR
@@ -290,27 +286,40 @@ class NoiseModelLikelihood(LikelihoodModule):
 
     def __init__(
         self,
-        data_mean: Union[np.ndarray, torch.Tensor],
-        data_std: Union[np.ndarray, torch.Tensor],
         noise_model: NoiseModel,
     ):
         """Constructor.
 
         Parameters
         ----------
-        data_mean: Union[np.ndarray, torch.Tensor]
-            The mean of the data, used to unnormalize data for noise model evaluation.
-        data_std: Union[np.ndarray, torch.Tensor]
-            The standard deviation of the data, used to unnormalize data for noise
-            model evaluation.
         noiseModel: NoiseModel
             The noise model instance used to compute the likelihood.
         """
         super().__init__()
-        self.data_mean = torch.Tensor(data_mean)
-        self.data_std = torch.Tensor(data_std)
+        self.data_mean = None
+        self.data_std = None
         self.noiseModel = noise_model
 
+    def set_data_stats(
+        self,
+        data_mean: Union[np.ndarray, torch.Tensor],
+        data_std: Union[np.ndarray, torch.Tensor],
+    ) -> None:
+        """Set the data mean and std for denormalization.
+        # TODO check this !!
+        Parameters
+        ----------
+        data_mean : Union[np.ndarray, torch.Tensor]
+            Mean values for each channel. Will be reshaped to (1, C, 1, 1, 1) for broadcasting.
+        data_std : Union[np.ndarray, torch.Tensor]
+            Standard deviation values for each channel. Will be reshaped to (1, C, 1, 1, 1) for broadcasting.
+        """
+        # Convert to tensor if needed
+        self.data_mean = torch.as_tensor(data_mean, dtype=torch.float32)
+        self.data_std = torch.as_tensor(data_std, dtype=torch.float32)
+
+        # TODO add extra dim for 3D ?
+
     def _set_params_to_same_device_as(
         self, correct_device_tensor: torch.Tensor
     ) -> None:
@@ -321,7 +330,10 @@ class NoiseModelLikelihood(LikelihoodModule):
         correct_device_tensor: torch.Tensor
             The tensor whose device is used to set the parameters.
         """
-        if self.data_mean.device != correct_device_tensor.device:
+        if (
+            self.data_mean is not None
+            and self.data_mean.device != correct_device_tensor.device
+        ):
             self.data_mean = self.data_mean.to(correct_device_tensor.device)
             self.data_std = self.data_std.to(correct_device_tensor.device)
         if correct_device_tensor.device != self.noiseModel.device:
@@ -367,6 +379,10 @@ class NoiseModelLikelihood(LikelihoodModule):
         torch.Tensor
             The log-likelihood tensor. Shape is (B, C, [Z], Y, X).
         """
+        if self.data_mean is None or self.data_std is None:
+            raise RuntimeError(
+                "NoiseModelLikelihood: data_mean and data_std must be set before calling log_likelihood."
+            )
         self._set_params_to_same_device_as(x)
         predicted_s_denormalized = params["mean"] * self.data_std + self.data_mean
         x_denormalized = x * self.data_std + self.data_mean

careamics/models/lvae/lvae.py

@@ -6,7 +6,7 @@ and Artefact Removal, Prakash et al."
 """
 
 from collections.abc import Iterable
-from typing import Optional, Union
+from typing import Union
 
 import numpy as np
 import torch
@@ -835,7 +835,7 @@ class LadderVAE(nn.Module):
             top_layer_shape = (n_imgs, mu_logvar, self._model_3D_depth, h, w)
         return top_layer_shape
 
-    def reset_for_inference(self, tile_size: Optional[tuple[int, int]] = None):
+    def reset_for_inference(self, tile_size: tuple[int, int] | None = None):
         """Should be called if we want to predict for a different input/output size."""
         self.mode_pred = True
         if tile_size is None:

careamics/models/lvae/noise_models.py

@@ -3,10 +3,10 @@ from __future__ import annotations
 import os
 from typing import TYPE_CHECKING, Optional
 
-from numpy.typing import NDArray
 import numpy as np
 import torch
 import torch.nn as nn
+from numpy.typing import NDArray
 
 if TYPE_CHECKING:
     from careamics.config import GaussianMixtureNMConfig, MultiChannelNMConfig
@@ -355,16 +355,16 @@ class GaussianMixtureNoiseModel(nn.Module):
 
         Parameters
         ----------
-        x: Tensor
-            Observations
-        mean: Tensor
-            Mean
-        std: Tensor
-            Standard-deviation
+        x: torch.Tensor
+            The ground-truth tensor. Shape is (batch, 1, dim1, dim2).
+        mean: torch.Tensor
+            The inferred mean of distribution. Shape is (batch, 1, dim1, dim2).
+        std: torch.Tensor
+            The inferred standard deviation of distribution. Shape is (batch, 1, dim1, dim2).
 
         Returns
         -------
-        tmp: Tensor
+        tmp: torch.Tensor
             Normal probability density of `x` given `mean` and `std`
         """
         tmp = -((x - mean) ** 2)
@@ -382,9 +382,9 @@ class GaussianMixtureNoiseModel(nn.Module):
         Parameters
         ----------
         observations : Tensor
-            Noisy observations
+            Noisy observations. Shape is (batch, 1, dim1, dim2).
         signals : Tensor
-            Underlying signals
+            Underlying signals. Shape is (batch, 1, dim1, dim2).
 
         Returns
         -------
@@ -392,15 +392,21 @@ class GaussianMixtureNoiseModel(nn.Module):
             Likelihood of observations given the signals and the GMM noise model
         """
         gaussian_parameters: list[torch.Tensor] = self.get_gaussian_parameters(signals)
-        p = 0
+        p = torch.zeros_like(observations)
         for gaussian in range(self.n_gaussian):
+            # Ensure all tensors have compatible shapes
+            mean = gaussian_parameters[gaussian]
+            std = gaussian_parameters[self.n_gaussian + gaussian]
+            weight = gaussian_parameters[2 * self.n_gaussian + gaussian]
+
+            # Compute normal density
             p += (
                 self.normal_density(
                     observations,
-                    gaussian_parameters[gaussian],
-                    gaussian_parameters[self.n_gaussian + gaussian],
+                    mean,
+                    std,
                 )
-                * gaussian_parameters[2 * self.n_gaussian + gaussian]
+                * weight
             )
         return p + self.tolerance
 

careamics/prediction_utils/__init__.py

@@ -2,9 +2,15 @@
 
 __all__ = [
     "convert_outputs",
+    "convert_outputs_microsplit",
     "stitch_prediction",
     "stitch_prediction_single",
+    "stitch_prediction_vae",
 ]
 
-from .prediction_outputs import convert_outputs
-from .stitch_prediction import stitch_prediction, stitch_prediction_single
+from .prediction_outputs import convert_outputs, convert_outputs_microsplit
+from .stitch_prediction import (
+    stitch_prediction,
+    stitch_prediction_single,
+    stitch_prediction_vae,
+)

careamics/prediction_utils/lvae_prediction.py

@@ -1,6 +1,6 @@
 """Module containing pytorch implementations for obtaining predictions from an LVAE."""
 
-from typing import Any, Optional
+from typing import Any
 
 import torch
 
@@ -18,7 +18,7 @@ def lvae_predict_single_sample(
     model: LVAE,
     likelihood_obj: LikelihoodModule,
     input: torch.Tensor,
-) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+) -> tuple[torch.Tensor, torch.Tensor | None]:
     """
     Generate a single sample prediction from an LVAE model, for a given input.
 
@@ -57,7 +57,7 @@ def lvae_predict_tiled_batch(
     model: LVAE,
     likelihood_obj: LikelihoodModule,
     input: tuple[Any],
-) -> tuple[tuple[Any], Optional[tuple[Any]]]:
+) -> tuple[tuple[Any], tuple[Any] | None]:
     # TODO: fix docstring return types, ... too many output options
     """
     Generate a single sample prediction from an LVAE model, for a given input.
@@ -98,7 +98,7 @@ def lvae_predict_mmse_tiled_batch(
     likelihood_obj: LikelihoodModule,
     input: tuple[Any],
     mmse_count: int,
-) -> tuple[tuple[Any], tuple[Any], Optional[tuple[Any]]]:
+) -> tuple[tuple[Any], tuple[Any], tuple[Any] | None]:
     # TODO: fix docstring return types, ... hard to make readable
     """
     Generate the MMSE (minimum mean squared error) prediction, for a given input.
@@ -137,7 +137,7 @@ def lvae_predict_mmse_tiled_batch(
 
     input_shape = x.shape
     output_shape = (input_shape[0], model.target_ch, *input_shape[2:])
-    log_var: Optional[torch.Tensor] = None
+    log_var: torch.Tensor | None = None
     # pre-declare empty array to fill with individual sample predictions
     sample_predictions = torch.zeros(size=(mmse_count, *output_shape))
     for mmse_idx in range(mmse_count):

careamics/prediction_utils/prediction_outputs.py

@@ -6,7 +6,7 @@ import numpy as np
 from numpy.typing import NDArray
 
 from ..config.tile_information import TileInformation
-from .stitch_prediction import stitch_prediction
+from .stitch_prediction import stitch_prediction, stitch_prediction_vae
 
 
 def convert_outputs(predictions: list[Any], tiled: bool) -> list[NDArray]:
@@ -41,6 +41,48 @@ def convert_outputs(predictions: list[Any], tiled: bool) -> list[NDArray]:
     return predictions_output
 
 
+def convert_outputs_microsplit(
+    predictions: list[tuple[NDArray, NDArray]], dataset
+) -> tuple[NDArray, NDArray]:
+    """
+    Convert microsplit Lightning trainer outputs using eval_utils stitching functions.
+
+    This function processes microsplit predictions that return (tile_prediction, tile_std) tuples
+    and stitches them back together using the same logic as get_single_file_mmse.
+
+    Parameters
+    ----------
+    predictions : list of tuple[NDArray, NDArray]
+        Predictions from Lightning trainer for microsplit. Each element is a tuple of
+        (tile_prediction, tile_std) where both are numpy arrays from predict_step.
+    dataset : Dataset
+        The dataset object used for prediction, needed for stitching function selection
+        and stitching process.
+
+    Returns
+    -------
+    tuple[NDArray, NDArray]
+        A tuple of (stitched_predictions, stitched_stds) representing the full
+        stitched predictions and standard deviations.
+    """
+    if len(predictions) == 0:
+        raise ValueError("No predictions provided")
+
+    # Separate predictions and stds from the list of tuples
+    tile_predictions = [pred for pred, _ in predictions]
+    tile_stds = [std for _, std in predictions]
+
+    # Concatenate all tiles exactly like get_single_file_mmse
+    tiles_arr = np.concatenate(tile_predictions, axis=0)
+    tile_stds_arr = np.concatenate(tile_stds, axis=0)
+
+    # Apply stitching using stitch_predictions_new
+    stitched_predictions = stitch_prediction_vae(tiles_arr, dataset)
+    stitched_stds = stitch_prediction_vae(tile_stds_arr, dataset)
+
+    return stitched_predictions, stitched_stds
+
+
 # for mypy
 @overload
 def combine_batches(  # numpydoc ignore=GL08
@@ -68,6 +110,8 @@ def combine_batches(
     """
     If predictions are in batches, they will be combined.
 
+    # TODO improve description!
+
     Parameters
     ----------
     predictions : list
@@ -107,11 +151,12 @@ def _combine_tiled_batches(
     """
     # turn list of lists into single list
     tile_infos = [
-        tile_info for _, tile_info_list in predictions for tile_info in tile_info_list
+        tile_info for *_, tile_info_list in predictions for tile_info in tile_info_list
     ]
     prediction_tiles: list[NDArray] = _combine_array_batches(
-        [preds for preds, _ in predictions]
+        [preds for preds, *_ in predictions]
     )
+
     return prediction_tiles, tile_infos