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

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/lvae_training/eval_utils.py

@@ -32,7 +32,7 @@ class TilingMode:
     ShiftBoundary = 2
 
 
-# ------------------------------------------------------------------------------------------------
+# ------------------------------------------------------------------------------------
 # Function of plotting: TODO -> moved them to another file, plot_utils.py
 def clean_ax(ax):
     """
@@ -68,7 +68,9 @@ def get_psnr_str(tar_hsnr, pred, col_idx):
     """
     Compute PSNR between the ground truth (`tar_hsnr`) and the predicted image (`pred`).
     """
-    return f"{scale_invariant_psnr(tar_hsnr[col_idx][None], pred[col_idx][None]).item():.1f}"
+    psnr = scale_invariant_psnr(tar_hsnr[col_idx][None], pred[col_idx][None]).item()
+
+    return f"{psnr:.1f}"
 
 
 def add_psnr_str(ax_, psnr):
@@ -129,8 +131,10 @@ def show_for_one(
     baseline_preds=None,
 ):
     """
-    Given an index, it plots the input, target, reconstructed images and the difference image.
-    Note the the difference image is computed with respect to a ground truth image, obtained from the high SNR dataset.
+    Given an index, it plots the input, target, reconstructed images and the difference
+    image.
+    Note the the difference image is computed with respect to a ground truth image,
+    obtained from the high SNR dataset.
     """
     highsnr_val_dset.set_img_sz(patch_size, 64)
     highsnr_val_dset.disable_noise()
@@ -164,7 +168,8 @@ def plot_crops(
         for i in range(len(baseline_preds)):
             if baseline_preds[i].shape != tar_hsnr.shape:
                 print(
-                    f"Baseline prediction {i} shape {baseline_preds[i].shape} does not match target shape {tar_hsnr.shape}"
+                    f"Baseline prediction {i} shape {baseline_preds[i].shape} does not "
+                    f"match target shape {tar_hsnr.shape}"
                 )
                 print("This happens when we want to predict the edges of the image.")
                 return
@@ -333,14 +338,21 @@ def plot_crops(
     ax_temp.imshow(inp[0, 0].cpu().numpy(), cmap="magma")
     clean_ax(ax_temp)
 
-    # line_ch1 = mlines.Line2D([0, 1], [0, 1], color=color_ch_list[0], linestyle='-', label='$C_1$')
-    # line_ch2 = mlines.Line2D([0, 1], [0, 1], color=color_ch_list[1], linestyle='-', label='$C_2$')
-    # line_pred = mlines.Line2D([0, 1], [0, 1], color=color_pred, linestyle='-', label='Pred')
-    # line_noisych1 = mlines.Line2D([0, 1], [0, 1], color=color_ch_list[0], linestyle='--', label='$C^N_1$')
-    # line_noisych2 = mlines.Line2D([0, 1], [0, 1], color=color_ch_list[1], linestyle='--', label='$C^N_2$')
-    # legend_ch1 = legend_ch1_ax.legend(handles=[line_ch1, line_noisych1, line_pred], loc='upper right', frameon=False, labelcolor='white',
+    # line_ch1 = mlines.Line2D([0, 1], [0, 1], color=color_ch_list[0], linestyle='-',
+    # label='$C_1$')
+    # line_ch2 = mlines.Line2D([0, 1], [0, 1], color=color_ch_list[1], linestyle='-',
+    # label='$C_2$')
+    # line_pred = mlines.Line2D([0, 1], [0, 1], color=color_pred, linestyle='-',
+    # label='Pred')
+    # line_noisych1 = mlines.Line2D([0, 1], [0, 1], color=color_ch_list[0],
+    # linestyle='--', label='$C^N_1$')
+    # line_noisych2 = mlines.Line2D([0, 1], [0, 1], color=color_ch_list[1],
+    # linestyle='--', label='$C^N_2$')
+    # legend_ch1 = legend_ch1_ax.legend(handles=[line_ch1, line_noisych1, line_pred],
+    # loc='upper right', frameon=False, labelcolor='white',
     #                         prop={'size': 11})
-    # legend_ch2 = legend_ch2_ax.legend(handles=[line_ch2, line_noisych2, line_pred], loc='upper right', frameon=False, labelcolor='white',
+    # legend_ch2 = legend_ch2_ax.legend(handles=[line_ch2, line_noisych2, line_pred],
+    # loc='upper right', frameon=False, labelcolor='white',
     #                             prop={'size': 11})
 
     if calibration_stats is not None:
@@ -383,7 +395,9 @@ def plot_calibration(ax, calibration_stats):
 
 def shiftedColorMap(cmap, start=0, midpoint=0.5, stop=1.0, name="shiftedcmap"):
     """
-    Adapted from https://stackoverflow.com/questions/7404116/defining-the-midpoint-of-a-colormap-in-matplotlib
+    Adapted from
+    https://stackoverflow.com/questions/7404116/defining-the-midpoint-of-a-colormap-in-
+    matplotlib
 
     Function to offset the "center" of a colormap. Useful for
     data with a negative min and positive max and you want the
@@ -444,7 +458,8 @@ def get_fractional_change(target, prediction, max_val=None):
 
 def get_zero_centered_midval(error):
     """
-    When done this way, the midval ensures that the colorbar is centered at 0. (Don't know how, but it works ;))
+    When done this way, the midval ensures that the colorbar is centered at 0. (Don't
+    know how, but it works ;))
     """
     vmax = error.max()
     vmin = error.min()
@@ -670,7 +685,7 @@ def get_single_file_mmse(
     return stitched_predictions, stitched_stds
 
 
-# ------------------------------------------------------------------------------------------
+# ---------------------------------------------------------------------------------
 ### Classes and Functions used to stitch predictions
 class PatchLocation:
     """
@@ -698,9 +713,10 @@ def _get_location(extra_padding, hwt, pred_h, pred_w):
 
 def get_location_from_idx(dset, dset_input_idx, pred_h, pred_w):
     """
-    For a given idx of the dataset, it returns where exactly in the dataset, does this prediction lies.
-    Note that this prediction also has padded pixels and so a subset of it will be used in the final prediction.
-    Which time frame, which spatial location (h_start, h_end, w_start,w_end)
+    For a given idx of the dataset, it returns where exactly in the dataset, does this
+    prediction lies. Note that this prediction also has padded pixels and so a subset of
+    it will be used in the final prediction. Which time frame, which spatial location
+    (h_start, h_end, w_start,w_end)
     Args:
         dset:
         dset_input_idx:
@@ -785,9 +801,11 @@ def stitch_predictions(predictions, dset, smoothening_pixelcount=0):
                 # NOTE: don't need to compute it for every patch.
                 assert (
                     smoothening_pixelcount == 0
-                ), "For smoothing,enable the get_smoothing_mask. It is disabled since I don't use it and it needs modification to work with non-square images"
+                ), "For smoothing,enable the get_smoothing_mask. It is disabled since I"
+                "don't use it and it needs modification to work with non-square images"
                 mask = 1
-                # mask = _get_smoothing_mask(cropped_pred_i.shape, smoothening_pixelcount, loc, frame_size)
+                # mask = _get_smoothing_mask(cropped_pred_i.shape,
+                # smoothening_pixelcount, loc, frame_size)
 
             cropped_pred_list.append(cropped_pred_i)
 
@@ -827,7 +845,8 @@ def stitch_predictions_new(predictions, dset):
     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])
+        # 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
@@ -843,7 +862,8 @@ def stitch_predictions_new(predictions, dset):
         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)], dtype=int)
         # assert np.all(vgs == gs)
-        # assert np.all(vge == ge) # TODO comented out this shit cuz I have no interest to dig why it's failing at this point !
+        # assert np.all(vge == ge) # TODO comented out this shit cuz I have no interest
+        # to dig why it's failing at this point !
         # print('VGS')
         # print(gs)
         # print(ge)
@@ -898,7 +918,8 @@ def stitch_predictions_general(predictions, dset):
     # frame_shape = dset.get_data_shape()[:-1]
     for patch_idx in range(predictions.shape[0]):
         # grid start, grid end
-        # channel_idx is 0 because during prediction we're only use one channel. # TODO revisit this
+        # channel_idx is 0 because during prediction we're only use one channel.
+        # # TODO revisit this
         # 0th dimension is sample index in the output list
         grid_coords = np.array(
             mng.get_location_from_patch_idx(channel_idx=0, patch_idx=patch_idx),
@@ -906,7 +927,8 @@ def stitch_predictions_general(predictions, dset):
         )
         sample_idx = grid_coords[0]
         grid_start = grid_coords[1:]
-        # from here on, coordinates are relative to the sample(file in the list of inputs)
+        # from here on, coordinates are relative to the sample(file in the list of
+        # inputs)
         grid_end = grid_start + mng.grid_shape
 
         # patch start, patch end

careamics/model_io/bmz_io.py

@@ -186,11 +186,15 @@ def export_to_bmz(
         )
 
         # test model description
-        test_kwargs = (
-            model_description.config.bioimageio.model_dump()
-            .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/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,11 @@ 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"
+                "callinglog_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,49 @@ 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 +111,8 @@ def combine_batches(
     """
     If predictions are in batches, they will be combined.
 
+    # TODO improve description!
+
     Parameters
     ----------
     predictions : list
@@ -107,11 +152,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,86 @@ 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) -> NDArray:
+    """
+    Stitch predictions back together using dataset's index manager.
+
+    Parameters
+    ----------
+    predictions : np.ndarray
+        Array of predictions with shape (n_tiles, channels, height, width).
+    dset : Dataset
+        Dataset object with idx_manager containing tiling information.
+
+    Returns
+    -------
+    np.ndarray
+        Stitched array with shape matching the original data shape.
+    """
+    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],
@@ -107,6 +187,8 @@ def stitch_prediction_single(
 
         # Insert cropped tile into predicted image using stitch coordinates
         image_slices = (..., *[slice(c[0], c[1]) for c in tile_info.stitch_coords])
-        predicted_image[image_slices] = cropped_tile.astype(np.float32)
+
+        # TODO fix mypy error here, potentially due to numpy 2
+        predicted_image[image_slices] = cropped_tile.astype(np.float32)  # type: ignore
 
     return predicted_image

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/utils/version.py

@@ -21,9 +21,9 @@ def get_careamics_version() -> str:
     parts = __version__.split(".")
 
     # for local installs that do not detect the latest versions via tags
-    # (typically our CI will install `0.1.devX<hash>` versions)
-    if "dev" in parts[-1]:
-        parts[-1] = "*"
+    # (typically our CI will install `0.X.devX<hash>.<other hash>` versions)
+    if "dev" in parts[2]:
+        parts[2] = "*"
         clean_version = ".".join(parts[:3])
 
         logger.warning(
@@ -34,5 +34,5 @@ def get_careamics_version() -> str:
             f"closest CAREamics version from PyPI or conda-forge."
         )
 
-    # Remove any local version identifier)
+    # Remove any local version identifier
     return ".".join(parts[:3])