careamics 0.0.4.2__py3-none-any.whl → 0.0.6__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/__init__.py

@@ -1,6 +1,6 @@
 """Model I/O utilities."""
 
-__all__ = ["load_pretrained", "export_to_bmz"]
+__all__ = ["export_to_bmz", "load_pretrained"]
 
 
 from .bmz_io import export_to_bmz

careamics/model_io/bioimage/__init__.py

@@ -1,10 +1,10 @@
 """Bioimage Model Zoo format functions."""
 
 __all__ = [
+    "create_env_text",
     "create_model_description",
     "extract_model_path",
     "get_unzip_path",
-    "create_env_text",
 ]
 
 from .bioimage_utils import create_env_text, get_unzip_path

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()):
@@ -60,47 +55,44 @@ def readme_factory(
         readme.touch()
 
         # algorithm pretty name
-        algorithm_flavour = config.get_algorithm_flavour()
+        algorithm_flavour = config.get_algorithm_friendly_name()
         algorithm_pretty_name = algorithm_flavour + " - CAREamics"
 
         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))