careamics 0.0.4.2__py3-none-any.whl → 0.0.5__py3-none-any.whl

careamics/lvae_training/eval_utils.py

@@ -6,13 +6,13 @@ It includes functions to:
     - create plots to visualize the results.
 """
 
-import math
 import os
-from typing import Dict, List, Literal, Union
+from typing import List, Literal, Union
 
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
+from scipy import stats
 import torch
 from torch import nn
 from torch.utils.data import Dataset
@@ -21,14 +21,21 @@ from torch.utils.data import DataLoader
 from tqdm import tqdm
 
 from careamics.lightning import VAEModule
-from careamics.losses.lvae.losses import (
-    get_reconstruction_loss,
-    reconstruction_loss_musplit_denoisplit,
-)
+
 from careamics.models.lvae.utils import ModelType
 from careamics.utils.metrics import scale_invariant_psnr, RunningPSNR
 
 
+class TilingMode:
+    """
+    Enum for the tiling mode.
+    """
+
+    TrimBoundary = 0
+    PadBoundary = 1
+    ShiftBoundary = 2
+
+
 # ------------------------------------------------------------------------------------------------
 # Function of plotting: TODO -> moved them to another file, plot_utils.py
 def clean_ax(ax):
@@ -596,51 +603,49 @@ def get_dset_predictions(
                 logvar_arr.append(logvar.cpu().numpy())
 
                 # compute reconstruction loss
-                if loss_type == "musplit":
-                    rec_loss = get_reconstruction_loss(
-                        reconstruction=rec, target=tar, likelihood_obj=gauss_likelihood
-                    )
-                elif loss_type == "denoisplit":
-                    rec_loss = get_reconstruction_loss(
-                        reconstruction=rec, target=tar, likelihood_obj=nm_likelihood
-                    )
-                elif loss_type == "denoisplit_musplit":
-                    rec_loss = reconstruction_loss_musplit_denoisplit(
-                        predictions=rec,
-                        targets=tar,
-                        gaussian_likelihood=gauss_likelihood,
-                        nm_likelihood=nm_likelihood,
-                        nm_weight=model.loss_parameters.denoisplit_weight,
-                        gaussian_weight=model.loss_parameters.musplit_weight,
-                    )
-                    rec_loss = {"loss": rec_loss}  # hacky, but ok for now
-
-                # store rec loss values for first pred
-                if mmse_idx == 0:
-                    try:
-                        losses.append(rec_loss["loss"].cpu().numpy())
-                    except:
-                        losses.append(rec_loss["loss"])
+                # if loss_type == "musplit":
+                #     rec_loss = get_reconstruction_loss(
+                #         reconstruction=rec, target=tar, likelihood_obj=gauss_likelihood
+                #     )
+                # elif loss_type == "denoisplit":
+                #     rec_loss = get_reconstruction_loss(
+                #         reconstruction=rec, target=tar, likelihood_obj=nm_likelihood
+                #     )
+                # elif loss_type == "denoisplit_musplit":
+                #     rec_loss = reconstruction_loss_musplit_denoisplit(
+                #         predictions=rec,
+                #         targets=tar,
+                #         gaussian_likelihood=gauss_likelihood,
+                #         nm_likelihood=nm_likelihood,
+                #         nm_weight=model.loss_parameters.denoisplit_weight,
+                #         gaussian_weight=model.loss_parameters.musplit_weight,
+                #     )
+                #     rec_loss = {"loss": rec_loss}  # hacky, but ok for now
+
+                # # store rec loss values for first pred
+                # if mmse_idx == 0:
+                #     try:
+                #         losses.append(rec_loss["loss"].cpu().numpy())
+                #     except:
+                #         losses.append(rec_loss["loss"])
 
                 # update running PSNR
-                for i in range(num_channels):
-                    patch_psnr_channels[i].update(rec_img[:, i], tar[:, i])
+                # for i in range(num_channels):
+                #     patch_psnr_channels[i].update(rec_img[:, i], tar[:, i])
 
             # aggregate results
             samples = torch.cat(rec_img_list, dim=0)
             mmse_imgs = torch.mean(samples, dim=0)  # avg over MMSE dim
-            mmse_std = torch.std(samples, dim=0)
+            # mmse_std = torch.std(samples, dim=0)
             predictions.append(mmse_imgs.cpu().numpy())
-            predictions_std.append(mmse_std.cpu().numpy())
-
-    psnr = [x.get() for x in patch_psnr_channels]
-    return (
-        np.concatenate(predictions, axis=0),
-        np.concatenate(predictions_std, axis=0),
-        np.concatenate(logvar_arr),
-        np.array(losses),
-        psnr,
-    )
+            # predictions_std.append(mmse_std.cpu().numpy())
+
+    # psnr = [x.get() for x in patch_psnr_channels]
+    return np.concatenate(predictions, axis=0)
+    # np.concatenate(predictions_std, axis=0),
+    # np.concatenate(logvar_arr),
+    # np.array(losses),
+    # psnr, # TODO revisit !
 
 
 # ------------------------------------------------------------------------------------------
@@ -773,178 +778,85 @@ def stitch_predictions(predictions, dset, smoothening_pixelcount=0):
     return output
 
 
-# ------------------------------------------------------------------------------------------
-
-
-# ------------------------------------------------------------------------------------------
-### Classes and Functions used for Calibration
-class Calibration:
-
-    def __init__(
-        self, num_bins: int = 15, mode: Literal["pixelwise", "patchwise"] = "pixelwise"
-    ):
-        self._bins = num_bins
-        self._bin_boundaries = None
-        self._mode = mode
-        assert mode in ["pixelwise", "patchwise"]
-        self._boundary_mode = "uniform"
-        assert self._boundary_mode in ["quantile", "uniform"]
-        # self._bin_boundaries = {}
-
-    def logvar_to_std(self, logvar: np.ndarray) -> np.ndarray:
-        return np.exp(logvar / 2)
-
-    def compute_bin_boundaries(self, predict_logvar: np.ndarray) -> np.ndarray:
-        """
-        Compute the bin boundaries for `num_bins` bins and the given logvar values.
-        """
-        if self._boundary_mode == "quantile":
-            boundaries = np.quantile(
-                self.logvar_to_std(predict_logvar), np.linspace(0, 1, self._bins + 1)
-            )
-            return boundaries
-        else:
-            min_logvar = np.min(predict_logvar)
-            max_logvar = np.max(predict_logvar)
-            min_std = self.logvar_to_std(min_logvar)
-            max_std = self.logvar_to_std(max_logvar)
-        return np.linspace(min_std, max_std, self._bins + 1)
-
-    def compute_stats(
-        self, pred: np.ndarray, pred_logvar: np.ndarray, target: np.ndarray
-    ) -> Dict[int, Dict[str, Union[np.ndarray, List]]]:
-        """
-        It computes the bin-wise RMSE and RMV for each channel of the predicted image.
-
-        Recall that:
-            - RMSE = np.sqrt((pred - target)**2 / num_pixels)
-            - RMV = np.sqrt(np.mean(pred_std**2))
-
-        ALGORITHM
-        - For each channel:
-            - Given the bin boundaries, assign pixels of `std_ch` array to a specific bin index.
-            - For each bin index:
-                - Compute the RMSE, RMV, and number of pixels for that bin.
-
-        NOTE: each channel of the predicted image/logvar has its own stats.
-
-        Args:
-            pred: np.ndarray, shape (n, h, w, c)
-            pred_logvar: np.ndarray, shape (n, h, w, c)
-            target: np.ndarray, shape (n, h, w, c)
-        """
-        self._bin_boundaries = {}
-        stats = {}
-        for ch_idx in range(pred.shape[-1]):
-            stats[ch_idx] = {
-                "bin_count": [],
-                "rmv": [],
-                "rmse": [],
-                "bin_boundaries": None,
-                "bin_matrix": [],
-            }
-            pred_ch = pred[..., ch_idx]
-            logvar_ch = pred_logvar[..., ch_idx]
-            std_ch = self.logvar_to_std(logvar_ch)
-            target_ch = target[..., ch_idx]
-            if self._mode == "pixelwise":
-                boundaries = self.compute_bin_boundaries(logvar_ch)
-                stats[ch_idx]["bin_boundaries"] = boundaries
-                bin_matrix = np.digitize(std_ch.reshape(-1), boundaries)
-                bin_matrix = bin_matrix.reshape(std_ch.shape)
-                stats[ch_idx]["bin_matrix"] = bin_matrix
-                error = (pred_ch - target_ch) ** 2
-                for bin_idx in range(self._bins):
-                    bin_mask = bin_matrix == bin_idx
-                    bin_error = error[bin_mask]
-                    bin_size = np.sum(bin_mask)
-                    bin_error = (
-                        np.sqrt(np.sum(bin_error) / bin_size) if bin_size > 0 else None
-                    )  # RMSE
-                    bin_var = np.sqrt(np.mean(std_ch[bin_mask] ** 2))  # RMV
-                    stats[ch_idx]["rmse"].append(bin_error)
-                    stats[ch_idx]["rmv"].append(bin_var)
-                    stats[ch_idx]["bin_count"].append(bin_size)
-            else:
-                raise NotImplementedError("Patchwise mode is not implemented yet.")
-        return stats
-
-
-def nll(x, mean, logvar):
+# from disentangle.analysis.stitch_prediction import *
+def stitch_predictions_new(predictions, dset):
     """
-    Log of the probability density of the values x under the Normal
-    distribution with parameters mean and logvar.
-
-    :param x: tensor of points, with shape (batch, channels, dim1, dim2)
-    :param mean: tensor with mean of distribution, shape
-                 (batch, channels, dim1, dim2)
-    :param logvar: tensor with log-variance of distribution, shape has to be
-                   either scalar or broadcastable
-    """
-    var = torch.exp(logvar)
-    log_prob = -0.5 * (
-        ((x - mean) ** 2) / var + logvar + torch.tensor(2 * math.pi).log()
-    )
-    nll = -log_prob
-    return nll
-
-
-def get_calibrated_factor_for_stdev(
-    pred: Union[np.ndarray, torch.Tensor],
-    pred_logvar: Union[np.ndarray, torch.Tensor],
-    target: Union[np.ndarray, torch.Tensor],
-    batch_size: int = 32,
-    epochs: int = 500,
-    lr: float = 0.01,
-):
-    """
-    Here, we calibrate the uncertainty by multiplying the predicted std (mmse estimate or predicted logvar) with a scalar.
-    We return the calibrated scalar. This needs to be multiplied with the std.
-
-    NOTE: Why is the input logvar and not std? because the model typically predicts logvar and not std.
+    Args:
+        smoothening_pixelcount: number of pixels which can be interpolated
     """
-    # create a learnable scalar
-    scalar = torch.nn.Parameter(torch.tensor(2.0))
-    optimizer = torch.optim.Adam([scalar], lr=lr)
+    # Commented out since it is not used as of now
+    # if isinstance(dset, MultiFileDset):
+    #     cum_count = 0
+    #     output = []
+    #     for dset in dset.dsets:
+    #         cnt = dset.idx_manager.total_grid_count()
+    #         output.append(
+    #             stitch_predictions(predictions[cum_count:cum_count + cnt], dset))
+    #         cum_count += cnt
+    #     return output
 
-    bar = tqdm(range(epochs))
-    for _ in bar:
-        optimizer.zero_grad()
-        # Select a random batch of predictions
-        mask = np.random.randint(0, pred.shape[0], batch_size)
-        pred_batch = torch.Tensor(pred[mask]).cuda()
-        pred_logvar_batch = torch.Tensor(pred_logvar[mask]).cuda()
-        target_batch = torch.Tensor(target[mask]).cuda()
+    # else:
+    mng = dset.idx_manager
 
-        loss = torch.mean(
-            nll(target_batch, pred_batch, pred_logvar_batch + torch.log(scalar))
-        )
-        loss.backward()
-        optimizer.step()
-        bar.set_description(f"nll: {loss.item()} scalar: {scalar.item()}")
-
-    return np.sqrt(scalar.item())
-
-
-def plot_calibration(ax, calibration_stats):
-    first_idx = get_first_index(calibration_stats[0]["bin_count"], 0.001)
-    last_idx = get_last_index(calibration_stats[0]["bin_count"], 0.999)
-    ax.plot(
-        calibration_stats[0]["rmv"][first_idx:-last_idx],
-        calibration_stats[0]["rmse"][first_idx:-last_idx],
-        "o",
-        label=r"$\hat{C}_0$: Ch1",
-    )
+    # if there are more channels, use all of them.
+    shape = list(dset.get_data_shape())
+    shape[-1] = max(shape[-1], predictions.shape[1])
 
-    first_idx = get_first_index(calibration_stats[1]["bin_count"], 0.001)
-    last_idx = get_last_index(calibration_stats[1]["bin_count"], 0.999)
-    ax.plot(
-        calibration_stats[1]["rmv"][first_idx:-last_idx],
-        calibration_stats[1]["rmse"][first_idx:-last_idx],
-        "o",
-        label=r"$\hat{C}_1: : Ch2$",
-    )
+    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
+        # print('PS')
+        # print(ps)
+        # print(pe)
+
+        # 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)], 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 !
+        # print('VGS')
+        # print(gs)
+        # print(ge)
+
+        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)
+        # print('RS')
+        # print(rs)
+        # print(re)
+
+        # print(output.shape)
+        # print(predictions.shape)
+        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}")
 
-    ax.set_xlabel("RMV")
-    ax.set_ylabel("RMSE")
-    ax.legend()
+    return output

careamics/model_io/bioimage/_readme_factory.py

@@ -1,7 +1,6 @@
 """Functions used to create a README.md file for BMZ export."""
 
 from pathlib import Path
-from typing import Optional
 
 import yaml
 
@@ -28,7 +27,7 @@ def _yaml_block(yaml_str: str) -> str:
 def readme_factory(
     config: Configuration,
     careamics_version: str,
-    data_description: Optional[str] = None,
+    data_description: str,
 ) -> Path:
     """Create a README file for the model.
 
@@ -41,18 +40,14 @@ def readme_factory(
         CAREamics configuration.
     careamics_version : str
         CAREamics version.
-    data_description : Optional[str], optional
-        Description of the data, by default None.
+    data_description : str
+        Description of the data.
 
     Returns
     -------
     Path
         Path to the README file.
     """
-    algorithm = config.algorithm_config
-    training = config.training_config
-    data = config.data_config
-
     # create file
     # TODO use tempfile as in the bmz_io module
     with cwd(get_careamics_home()):
@@ -65,42 +60,39 @@ def readme_factory(
 
         description = [f"# {algorithm_pretty_name}\n\n"]
 
+        # data description
+        description.append("## Data description\n\n")
+        description.append(data_description)
+        description.append("\n\n")
+
         # algorithm description
-        description.append("Algorithm description:\n\n")
+        description.append("## Algorithm description:\n\n")
         description.append(config.get_algorithm_description())
         description.append("\n\n")
 
-        # algorithm details
+        # configuration description
+        description.append("## Configuration\n\n")
+
         description.append(
             f"{algorithm_flavour} was trained using CAREamics (version "
-            f"{careamics_version}) with the following algorithm "
-            f"parameters:\n\n"
-        )
-        description.append(
-            _yaml_block(yaml.dump(algorithm.model_dump(exclude_none=True)))
+            f"{careamics_version}) using the following configuration:\n\n"
         )
-        description.append("\n\n")
-
-        # data description
-        description.append("## Data description\n\n")
-        if data_description is not None:
-            description.append(data_description)
-            description.append("\n\n")
-
-        description.append("The data was processed using the following parameters:\n\n")
 
-        description.append(_yaml_block(yaml.dump(data.model_dump(exclude_none=True))))
+        description.append(_yaml_block(yaml.dump(config.model_dump(exclude_none=True))))
         description.append("\n\n")
 
-        # training description
-        description.append("## Training description\n\n")
-
-        description.append("The model was trained using the following parameters:\n\n")
+        # validation
+        description.append("# Validation\n\n")
 
         description.append(
-            _yaml_block(yaml.dump(training.model_dump(exclude_none=True)))
+            "In order to validate the model, we encourage users to acquire a "
+            "test dataset with ground-truth data. Comparing the ground-truth data "
+            "with the prediction allows unbiased evaluation of the model performances. "
+            "This can be done for instance by using metrics such as PSNR, SSIM, or"
+            "MicroSSIM. In the absence of ground-truth, inspecting the residual image "
+            "(difference between input and predicted image) can be helpful to identify "
+            "whether real signal is removed from the input image.\n\n"
         )
-        description.append("\n\n")
 
         # references
         reference = config.get_algorithm_references()
@@ -111,9 +103,9 @@ def readme_factory(
 
         # links
         description.append(
-            "## Links\n\n"
+            "# Links\n\n"
             "- [CAREamics repository](https://github.com/CAREamics/careamics)\n"
-            "- [CAREamics documentation](https://careamics.github.io/latest/)\n"
+            "- [CAREamics documentation](https://careamics.github.io/)\n"
         )
 
         readme.write_text("".join(description))

careamics/model_io/bioimage/cover_factory.py

@@ -0,0 +1,171 @@
+"""Convenience function to create covers for the BMZ."""
+
+from pathlib import Path
+
+import numpy as np
+from numpy.typing import NDArray
+from PIL import Image
+
+color_palette = np.array(
+    [
+        np.array([255, 195, 0]),  # grey
+        np.array([189, 226, 240]),
+        np.array([96, 60, 76]),
+        np.array([193, 225, 193]),
+    ]
+)
+
+
+def _get_norm_slice(array: NDArray) -> NDArray:
+    """Get the normalized middle slice of a 4D or 5D array (SC(Z)YX).
+
+    Parameters
+    ----------
+    array : NDArray
+        Array from which to get the middle slice.
+
+    Returns
+    -------
+    NDArray
+        Normalized middle slice of the input array.
+    """
+    if array.ndim not in (4, 5):
+        raise ValueError("Array must be 4D or 5D.")
+
+    channels = array.shape[1] > 1
+    z_stack = array.ndim == 5
+
+    # get slice
+    if z_stack:
+        array_slice = array[0, :, array.shape[2] // 2, ...]
+    else:
+        array_slice = array[0, ...]
+
+    # channels
+    if channels:
+        array_slice = np.moveaxis(array_slice, 0, -1)
+    else:
+        array_slice = array_slice[0, ...]
+
+    # normalize
+    array_slice = (
+        255
+        * (array_slice - array_slice.min())
+        / (array_slice.max() - array_slice.min())
+    )
+
+    return array_slice.astype(np.uint8)
+
+
+def _four_channel_image(array: NDArray) -> Image:
+    """Convert 4-channel array to Image.
+
+    Parameters
+    ----------
+    array : NDArray
+        Normalized array to convert.
+
+    Returns
+    -------
+    Image
+        Converted array.
+    """
+    colors = color_palette[np.newaxis, np.newaxis, :, :]
+    four_c_array = np.sum(array[..., :4, np.newaxis] * colors, axis=-2).astype(np.uint8)
+
+    return Image.fromarray(four_c_array).convert("RGB")
+
+
+def _convert_to_image(original_shape: tuple[int, ...], array: NDArray) -> Image:
+    """Convert to Image.
+
+    Parameters
+    ----------
+    original_shape : tuple
+        Original shape of the array.
+    array : NDArray
+        Normalized array to convert.
+
+    Returns
+    -------
+    Image
+        Converted array.
+    """
+    n_channels = original_shape[1]
+
+    if n_channels > 1:
+        if n_channels == 3:
+            return Image.fromarray(array).convert("RGB")
+        elif n_channels == 2:
+            # add an empty channel to the numpy array
+            array = np.concatenate([np.zeros_like(array[..., 0:1]), array], axis=-1)
+
+            return Image.fromarray(array).convert("RGB")
+        else:  # more than 4
+            return _four_channel_image(array[..., :4])
+    else:
+        return Image.fromarray(array).convert("L").convert("RGB")
+
+
+def create_cover(directory: Path, array_in: NDArray, array_out: NDArray) -> Path:
+    """Create a cover image from input and output arrays.
+
+    Input and output arrays are expected to be SC(Z)YX. For images with a Z
+    dimension, the middle slice is taken.
+
+    Parameters
+    ----------
+    directory : Path
+        Directory in which to save the cover.
+    array_in : numpy.ndarray
+        Array from which to create the cover image.
+    array_out : numpy.ndarray
+        Array from which to create the cover image.
+
+    Returns
+    -------
+    Path
+        Path to the saved cover image.
+    """
+    # extract slice and normalize arrays
+    slice_in = _get_norm_slice(array_in)
+    slice_out = _get_norm_slice(array_out)
+
+    horizontal_split = slice_in.shape[-1] == slice_out.shape[-1]
+    if not horizontal_split:
+        if slice_in.shape[-2] != slice_out.shape[-2]:
+            raise ValueError("Input and output arrays have different shapes.")
+
+    # convert to Image
+    image_in = _convert_to_image(array_in.shape, slice_in)
+    image_out = _convert_to_image(array_out.shape, slice_out)
+
+    # split horizontally or vertically
+    if horizontal_split:
+        width = image_in.width // 2
+
+        cover = Image.new("RGB", (image_in.width, image_in.height))
+        cover.paste(image_in.crop((0, 0, width, image_in.height)), (0, 0))
+        cover.paste(
+            image_out.crop(
+                (image_in.width - width, 0, image_in.width, image_in.height)
+            ),
+            (width, 0),
+        )
+    else:
+        height = image_in.height // 2
+
+        cover = Image.new("RGB", (image_in.width, image_in.height))
+        cover.paste(image_in.crop((0, 0, image_in.width, height)), (0, 0))
+        cover.paste(
+            image_out.crop(
+                (0, image_in.height - height, image_in.width, image_in.height)
+            ),
+            (0, height),
+        )
+
+    # save
+    cover_path = directory / "cover.png"
+    cover.save(cover_path)
+
+    return cover_path