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

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/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
 
 

careamics/prediction_utils/stitch_prediction.py

@@ -9,6 +9,77 @@ from numpy.typing import NDArray
 from careamics.config.tile_information import TileInformation
 
 
+class TilingMode:
+    """Enum for the tiling mode."""
+
+    TrimBoundary = 0
+    PadBoundary = 1
+    ShiftBoundary = 2
+
+
+def stitch_prediction_vae(predictions, dset):
+    """
+    Stitch predictions back together using dataset's index manager.
+
+    Args:
+        predictions: Array of predictions with shape (n_tiles, channels, height, width)
+        dset: Dataset object with idx_manager containing tiling information
+    """
+    mng = dset.idx_manager
+
+    # if there are more channels, use all of them.
+    shape = list(dset.get_data_shape())
+    shape[-1] = max(shape[-1], predictions.shape[1])
+
+    output = np.zeros(shape, dtype=predictions.dtype)
+    # frame_shape = dset.get_data_shape()[:-1]
+    for dset_idx in range(predictions.shape[0]):
+        # loc = get_location_from_idx(dset, dset_idx, predictions.shape[-2], predictions.shape[-1])
+        # grid start, grid end
+        gs = np.array(mng.get_location_from_dataset_idx(dset_idx), dtype=int)
+        ge = gs + mng.grid_shape
+
+        # patch start, patch end
+        ps = gs - mng.patch_offset()
+        pe = ps + mng.patch_shape
+
+        # valid grid start, valid grid end
+        vgs = np.array([max(0, x) for x in gs], dtype=int)
+        vge = np.array(
+            [min(x, y) for x, y in zip(ge, mng.data_shape, strict=False)], dtype=int
+        )
+
+        if mng.tiling_mode == TilingMode.ShiftBoundary:
+            for dim in range(len(vgs)):
+                if ps[dim] == 0:
+                    vgs[dim] = 0
+                if pe[dim] == mng.data_shape[dim]:
+                    vge[dim] = mng.data_shape[dim]
+
+        # relative start, relative end. This will be used on pred_tiled
+        rs = vgs - ps
+        re = rs + (vge - vgs)
+
+        for ch_idx in range(predictions.shape[1]):
+            if len(output.shape) == 4:
+                # channel dimension is the last one.
+                output[vgs[0] : vge[0], vgs[1] : vge[1], vgs[2] : vge[2], ch_idx] = (
+                    predictions[dset_idx][ch_idx, rs[1] : re[1], rs[2] : re[2]]
+                )
+            elif len(output.shape) == 5:
+                # channel dimension is the last one.
+                assert vge[0] - vgs[0] == 1, "Only one frame is supported"
+                output[
+                    vgs[0], vgs[1] : vge[1], vgs[2] : vge[2], vgs[3] : vge[3], ch_idx
+                ] = predictions[dset_idx][
+                    ch_idx, rs[1] : re[1], rs[2] : re[2], rs[3] : re[3]
+                ]
+            else:
+                raise ValueError(f"Unsupported shape {output.shape}")
+
+    return output
+
+
 # TODO: why not allow input and output of torch.tensor ?
 def stitch_prediction(
     tiles: list[np.ndarray],

careamics/transforms/xy_random_rotate90.py

@@ -74,7 +74,7 @@ class XYRandomRotate90(Transform):
             return patch, target, additional_arrays
 
         # number of rotations
-        n_rot = self.rng.integers(1, 4)
+        n_rot = int(self.rng.integers(1, 4))
 
         axes = (-2, -1)
         patch_transformed = self._apply(patch, n_rot, axes)

{careamics-0.0.15.dist-info → careamics-0.0.16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: careamics
-Version: 0.0.15
+Version: 0.0.16
 Summary: Toolbox for running N2V and friends.
 Project-URL: homepage, https://careamics.github.io/
 Project-URL: repository, https://github.com/CAREamics/careamics
@@ -10,28 +10,31 @@ License-File: LICENSE
 Classifier: Development Status :: 3 - Alpha
 Classifier: License :: OSI Approved :: BSD License
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: Typing :: Typed
-Requires-Python: >=3.10
-Requires-Dist: bioimageio-core==0.9.0
-Requires-Dist: matplotlib<=3.10.3
-Requires-Dist: numpy<2.0.0
-Requires-Dist: pillow<=11.2.1
-Requires-Dist: psutil<=7.0.0
+Requires-Python: >=3.11
+Requires-Dist: bioimageio-core>=0.9.0
+Requires-Dist: matplotlib<=3.10.6
+Requires-Dist: numpy>=1.21
+Requires-Dist: numpy>=2.1.0; python_version >= '3.13'
+Requires-Dist: pillow<=11.3.0
+Requires-Dist: psutil<=7.1.0
 Requires-Dist: pydantic<=2.12,>=2.11
-Requires-Dist: pytorch-lightning<=2.5.2,>=2.2
-Requires-Dist: pyyaml!=6.0.0,<=6.0.2
+Requires-Dist: pytorch-lightning<=2.5.5,>=2.2
+Requires-Dist: pyyaml!=6.0.0,<=6.0.3
 Requires-Dist: scikit-image<=0.25.2
-Requires-Dist: tifffile<=2025.5.10
-Requires-Dist: torch<=2.7.1,>=2.0
-Requires-Dist: torchvision<=0.22.1
-Requires-Dist: typer<=0.16.0,>=0.12.3
-Requires-Dist: zarr<3.0.0
+Requires-Dist: tifffile<=2025.9.30
+Requires-Dist: torch<=2.8.0,>=2.0
+Requires-Dist: torchvision<=0.23.0
+Requires-Dist: typer<=0.19.2,>=0.12.3
+Requires-Dist: validators<=0.35.0
+Requires-Dist: zarr<4.0.0,>=3.0.0
 Provides-Extra: czi
 Requires-Dist: pylibczirw<6.0.0,>=4.1.2; extra == 'czi'
 Provides-Extra: dev
+Requires-Dist: ml-dtypes>=0.5.0; extra == 'dev'
 Requires-Dist: onnx; extra == 'dev'
 Requires-Dist: pre-commit; extra == 'dev'
 Requires-Dist: pytest; extra == 'dev'