careamics 0.0.19__py3-none-any.whl

careamics/lvae_training/calibration.py

@@ -0,0 +1,191 @@
+from typing import Union, Optional
+
+import numpy as np
+import torch
+from scipy import stats
+
+
+def get_last_index(bin_count, quantile):
+    cumsum = np.cumsum(bin_count)
+    normalized_cumsum = cumsum / cumsum[-1]
+    for i in range(1, len(normalized_cumsum)):
+        if normalized_cumsum[-i] < quantile:
+            return i - 1
+    return None
+
+
+def get_first_index(bin_count, quantile):
+    cumsum = np.cumsum(bin_count)
+    normalized_cumsum = cumsum / cumsum[-1]
+    for i in range(len(normalized_cumsum)):
+        if normalized_cumsum[i] > quantile:
+            return i
+    return None
+
+
+class Calibration:
+    """Calibrate the uncertainty computed over samples from LVAE model.
+
+    Calibration is done by learning a scalar that maps the pixel-wise standard
+    deviation of the the predicted samples into the actual prediction error.
+    """
+
+    def __init__(self, num_bins: int = 15):
+        self._bins = num_bins
+        self._bin_boundaries = None
+
+    def compute_bin_boundaries(self, predict_std: np.ndarray) -> np.ndarray:
+        """Compute the bin boundaries for `num_bins` bins and predicted std values."""
+        min_std = np.min(predict_std)
+        max_std = np.max(predict_std)
+        return np.linspace(min_std, max_std, self._bins + 1)
+
+    def compute_stats(
+        self, pred: np.ndarray, pred_std: 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.
+
+        Parameters
+        ----------
+        pred: np.ndarray
+            Predicted patches, shape (n, h, w, c).
+        pred_std: np.ndarray
+            Std computed over the predicted patches, shape (n, h, w, c).
+        target: np.ndarray
+            Target GT image, shape (n, h, w, c).
+        """
+        self._bin_boundaries = {}
+        stats_dict = {}
+        for ch_idx in range(pred.shape[-1]):
+            stats_dict[ch_idx] = {
+                "bin_count": [],
+                "rmv": [],
+                "rmse": [],
+                "bin_boundaries": None,
+                "bin_matrix": [],
+                "rmse_err": [],
+            }
+            pred_ch = pred[..., ch_idx]
+            std_ch = pred_std[..., ch_idx]
+            target_ch = target[..., ch_idx]
+            boundaries = self.compute_bin_boundaries(std_ch)
+            stats_dict[ch_idx]["bin_boundaries"] = boundaries
+            bin_matrix = np.digitize(std_ch.reshape(-1), boundaries)
+            bin_matrix = bin_matrix.reshape(std_ch.shape)
+            stats_dict[ch_idx]["bin_matrix"] = bin_matrix
+            error = (pred_ch - target_ch) ** 2
+            for bin_idx in range(1, 1 + 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
+                )
+                stderr = (
+                    np.std(error[bin_mask]) / np.sqrt(bin_size)
+                    if bin_size > 0
+                    else None
+                )
+                rmse_stderr = np.sqrt(stderr) if stderr is not None else None
+
+                bin_var = np.mean(std_ch[bin_mask] ** 2)
+                stats_dict[ch_idx]["rmse"].append(bin_error)
+                stats_dict[ch_idx]["rmse_err"].append(rmse_stderr)
+                stats_dict[ch_idx]["rmv"].append(np.sqrt(bin_var))
+                stats_dict[ch_idx]["bin_count"].append(bin_size)
+        self.stats_dict = stats_dict
+        return stats_dict
+
+    def get_calibrated_factor_for_stdev(
+        self,
+        pred: Optional[np.ndarray] = None,
+        pred_std: Optional[np.ndarray] = None,
+        target: Optional[np.ndarray] = None,
+        q_s: float = 0.00001,
+        q_e: float = 0.99999,
+    ) -> dict[str, float]:
+        """Calibrate the uncertainty by multiplying the predicted std with a scalar.
+
+        Parameters
+        ----------
+        stats_dict : dict[int, dict[str, Union[np.ndarray, list]]]
+            Dictionary containing the stats for each channel.
+        q_s : float, optional
+            Start quantile, by default 0.00001.
+        q_e : float, optional
+            End quantile, by default 0.99999.
+
+        Returns
+        -------
+        dict[str, float]
+            Calibrated factor for each channel (slope + intercept).
+        """
+        if not hasattr(self, "stats_dict"):
+            print("No stats found. Computing stats...")
+            if any(v is None for v in [pred, pred_std, target]):
+                raise ValueError("pred, pred_std, and target must be provided.")
+            self.stats_dict = self.compute_stats(
+                pred=pred, pred_std=pred_std, target=target
+            )
+        outputs = {}
+        for ch_idx in self.stats_dict.keys():
+            y = self.stats_dict[ch_idx]["rmse"]
+            x = self.stats_dict[ch_idx]["rmv"]
+            count = self.stats_dict[ch_idx]["bin_count"]
+
+            first_idx = get_first_index(count, q_s)
+            last_idx = get_last_index(count, q_e)
+            x = x[first_idx:-last_idx]
+            y = y[first_idx:-last_idx]
+            slope, intercept, *_ = stats.linregress(x, y)
+            output = {"scalar": slope, "offset": intercept}
+            outputs[ch_idx] = output
+        factors = self.get_factors_array(factors_dict=outputs)
+        return outputs, factors
+
+    def get_factors_array(self, factors_dict: list[dict]):
+        """Get the calibration factors as a numpy array."""
+        calib_scalar = [factors_dict[i]["scalar"] for i in range(len(factors_dict))]
+        calib_scalar = np.array(calib_scalar).reshape(1, 1, 1, -1)
+        calib_offset = [
+            factors_dict[i].get("offset", 0.0) for i in range(len(factors_dict))
+        ]
+        calib_offset = np.array(calib_offset).reshape(1, 1, 1, -1)
+        return {"scalar": calib_scalar, "offset": calib_offset}
+
+
+def plot_calibration(ax, calibration_stats):
+    first_idx = get_first_index(calibration_stats[0]["bin_count"], 0.0001)
+    last_idx = get_last_index(calibration_stats[0]["bin_count"], 0.9999)
+    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",
+    )
+
+    first_idx = get_first_index(calibration_stats[1]["bin_count"], 0.0001)
+    last_idx = get_last_index(calibration_stats[1]["bin_count"], 0.9999)
+    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",
+    )
+    # TODO add multichannel
+    ax.set_xlabel("RMV")
+    ax.set_ylabel("RMSE")
+    ax.legend()

careamics/lvae_training/dataset/__init__.py

@@ -0,0 +1,20 @@
+from .config import MicroSplitDataConfig
+from .lc_dataset import LCMultiChDloader
+from .ms_dataset_ref import MultiChDloaderRef
+from .multich_dataset import MultiChDloader
+from .multicrop_dset import MultiCropDset
+from .multifile_dataset import MultiFileDset
+from .types import DataSplitType, DataType, TilingMode
+
+__all__ = [
+    "DataSplitType",
+    "DataType",
+    "LCMultiChDloader",
+    "LCMultiChDloaderRef",
+    "MicroSplitDataConfig",
+    "MultiChDloader",
+    "MultiChDloaderRef",
+    "MultiCropDset",
+    "MultiFileDset",
+    "TilingMode",
+]

careamics/lvae_training/dataset/config.py

@@ -0,0 +1,135 @@
+from typing import Any, Union
+
+from pydantic import BaseModel, ConfigDict
+
+from .types import DataSplitType, DataType, TilingMode
+
+
+# TODO: check if any bool logic can be removed
+class MicroSplitDataConfig(BaseModel):
+    model_config = ConfigDict(validate_assignment=True, extra="allow")
+
+    data_type: Union[DataType, str] | None  # TODO remove or refactor!!
+    """Type of the dataset, should be one of DataType"""
+
+    depth3D: int | None = 1
+    """Number of slices in 3D. If data is 2D depth3D is equal to 1"""
+
+    datasplit_type: DataSplitType | None = None
+    """Whether to return training, validation or test split, should be one of
+    DataSplitType"""
+
+    num_channels: int | None = 2
+    """Number of channels in the input"""
+
+    # TODO: remove ch*_fname parameters, should be parsed automatically from a name list
+    ch1_fname: str | None = None
+    ch2_fname: str | None = None
+    ch_input_fname: str | None = None
+
+    input_is_sum: bool | None = False
+    """Whether the input is the sum or average of channels"""
+
+    input_idx: int | None = None
+    """Index of the channel where the input is stored in the data"""
+
+    target_idx_list: list[int] | None = None
+    """Indices of the channels where the targets are stored in the data"""
+
+    # TODO: where are there used?
+    start_alpha: Any | None = None
+    end_alpha: Any | None = None
+
+    image_size: tuple  # TODO: revisit, new model_config uses tuple
+    """Size of one patch of data"""
+
+    grid_size: Union[int, tuple[int, int, int]] | None = None
+    """Frame is divided into square grids of this size. A patch centered on a grid 
+    having size `image_size` is returned. Grid size not used in training,
+    used only during val / test, grid size controls the overlap of the patches"""
+
+    empty_patch_replacement_enabled: bool | None = False
+    """Whether to replace the content of one of the channels
+    with background with given probability"""
+    empty_patch_replacement_channel_idx: Any | None = None
+    empty_patch_replacement_probab: Any | None = None
+    empty_patch_max_val_threshold: Any | None = None
+
+    uncorrelated_channels: bool | None = False
+    """Replace the content in one of the channels with given probability to make
+    channel content 'uncorrelated'"""
+    uncorrelated_channel_probab: float | None = 0.5
+
+    poisson_noise_factor: float | None = -1
+    """The added poisson noise factor"""
+
+    synthetic_gaussian_scale: float | None = 0.1
+
+    # TODO: set to True in training code, recheck
+    input_has_dependant_noise: bool | None = False
+
+    # TODO: sometimes max_val differs between runs with fixed seeds with noise enabled
+    enable_gaussian_noise: bool | None = False
+    """Whether to enable gaussian noise"""
+
+    # TODO: is this parameter used?
+    allow_generation: bool = False
+
+    # TODO: both used in IndexSwitcher, insure correct passing
+    training_validtarget_fraction: Any = None
+    deterministic_grid: Any = None
+
+    # TODO: why is this not used?
+    enable_rotation_aug: bool | None = False
+
+    max_val: Union[float, tuple] | None = None
+    """Maximum data in the dataset. Is calculated for train split, and should be
+    externally set for val and test splits."""
+
+    overlapping_padding_kwargs: Any = None
+    """Parameters for np.pad method"""
+
+    # TODO: remove this parameter, controls debug print
+    print_vars: bool | None = False
+
+    # Hard-coded parameters (used to be in the config file)
+    normalized_input: bool = True
+    """If this is set to true, then one mean and stdev is used
+                for both channels. Otherwise, two different mean and stdev are used."""
+    use_one_mu_std: bool | None = True
+
+    # TODO: is this parameter used?
+    train_aug_rotate: bool | None = False
+    enable_random_cropping: bool | None = True
+
+    multiscale_lowres_count: int | None = None
+    """Number of LC scales"""
+
+    tiling_mode: TilingMode | None = TilingMode.ShiftBoundary
+
+    target_separate_normalization: bool | None = True
+
+    mode_3D: bool | None = False
+    """If training in 3D mode or not"""
+
+    trainig_datausage_fraction: float | None = 1.0
+
+    validtarget_random_fraction: float | None = None
+
+    validation_datausage_fraction: float | None = 1.0
+
+    random_flip_z_3D: bool | None = False
+
+    padding_kwargs: dict = {"mode": "reflect"}  # TODO remove !!
+
+    def __init__(self, **data):
+        # Convert string data_type to enum if needed
+        if "data_type" in data and isinstance(data["data_type"], str):
+            try:
+                data["data_type"] = DataType[data["data_type"]]
+            except KeyError:
+                # Keep original value to let validation handle the error
+                pass
+        super().__init__(**data)
+
+    # TODO add validators !

careamics/lvae_training/dataset/lc_dataset.py

@@ -0,0 +1,274 @@
+"""
+A place for Datasets and Dataloaders.
+"""
+
+import logging
+import math
+from pathlib import Path
+from typing import Any, Callable, Optional, Union
+
+import numpy as np
+from skimage.transform import resize
+
+from .config import MicroSplitDataConfig
+from .multich_dataset import MultiChDloader
+
+
+class LCMultiChDloader(MultiChDloader):
+    """Multi-channel dataset loader for LC-style datasets."""
+
+    def __init__(
+        self,
+        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._padding_kwargs = (
+            data_config.padding_kwargs  # mode=padding_mode, constant_values=constant_value
+        )
+        self._uncorrelated_channel_probab = data_config.uncorrelated_channel_probab
+
+        super().__init__(
+            data_config,
+            datapath,
+            load_data_fn=load_data_fn,
+            val_fraction=val_fraction,
+            test_fraction=test_fraction,
+        )
+
+        if data_config.overlapping_padding_kwargs is not None:
+            assert (
+                self._padding_kwargs == data_config.overlapping_padding_kwargs
+            ), "During evaluation, overlapping_padding_kwargs should be same as padding_args. \
+                It should be so since we just use overlapping_padding_kwargs when it is not None"
+
+        else:
+            self._overlapping_padding_kwargs = data_config.padding_kwargs
+
+        self.multiscale_lowres_count = data_config.multiscale_lowres_count
+        assert self.multiscale_lowres_count is not None
+        self._scaled_data = [self._data]
+        self._scaled_noise_data = [self._noise_data]
+
+        assert (
+            isinstance(self.multiscale_lowres_count, int)
+            and self.multiscale_lowres_count >= 1
+        )
+        assert isinstance(self._padding_kwargs, dict)
+        assert "mode" in self._padding_kwargs
+
+        for _ in range(1, self.multiscale_lowres_count):
+            shape = self._scaled_data[-1].shape
+            assert len(shape) == 4
+            new_shape = (shape[0], shape[1] // 2, shape[2] // 2, shape[3])
+            ds_data = resize(
+                self._scaled_data[-1].astype(np.float32), new_shape
+            ).astype(self._scaled_data[-1].dtype)
+            # NOTE: These asserts are important. the resize method expects np.float32. otherwise, one gets weird results.
+            assert (
+                ds_data.max() / self._scaled_data[-1].max() < 5
+            ), "Downsampled image should not have very different values"
+            assert (
+                ds_data.max() / self._scaled_data[-1].max() > 0.2
+            ), "Downsampled image should not have very different values"
+
+            self._scaled_data.append(ds_data)
+            # do the same for noise
+            if self._noise_data is not None:
+                noise_data = resize(self._scaled_noise_data[-1], new_shape)
+                self._scaled_noise_data.append(noise_data)
+
+    def reduce_data(
+        self, t_list=None, h_start=None, h_end=None, w_start=None, w_end=None
+    ):
+        assert t_list is not None
+        assert h_start is None
+        assert h_end is None
+        assert w_start is None
+        assert w_end is None
+
+        self._data = self._data[t_list].copy()
+        self._scaled_data = [
+            self._scaled_data[i][t_list].copy() for i in range(len(self._scaled_data))
+        ]
+
+        if self._noise_data is not None:
+            self._noise_data = self._noise_data[t_list].copy()
+            self._scaled_noise_data = [
+                self._scaled_noise_data[i][t_list].copy()
+                for i in range(len(self._scaled_noise_data))
+            ]
+
+        self.N = len(t_list)
+        # TODO where tf is self._img_sz defined?
+        self.set_img_sz([self._img_sz, self._img_sz], self._grid_sz)
+        print(
+            f"[{self.__class__.__name__}] Data reduced. New data shape: {self._data.shape}"
+        )
+
+    def _init_msg(self):
+        msg = super()._init_msg()
+        msg += f" Pad:{self._padding_kwargs}"
+        if self._uncorrelated_channels:
+            msg += f" UncorrChProbab:{self._uncorrelated_channel_probab}"
+        return msg
+
+    def _load_scaled_img(
+        self, scaled_index, index: Union[int, tuple[int, int]]
+    ) -> tuple[np.ndarray, np.ndarray]:
+        if isinstance(index, int):
+            idx = index
+        else:
+            idx, _ = index
+
+        # tidx = self.idx_manager.get_t(idx)
+        patch_loc_list = self.idx_manager.get_patch_location_from_dataset_idx(idx)
+        nidx = patch_loc_list[0]
+
+        imgs = self._scaled_data[scaled_index][nidx]
+        imgs = tuple([imgs[None, ..., i] for i in range(imgs.shape[-1])])
+        if self._noise_data is not None:
+            noisedata = self._scaled_noise_data[scaled_index][nidx]
+            noise = tuple([noisedata[None, ..., i] for i in range(noisedata.shape[-1])])
+            factor = np.sqrt(2) if self._input_is_sum else 1.0
+            imgs = tuple([img + noise[0] * factor for img in imgs])
+        return imgs
+
+    def _crop_img(self, img: np.ndarray, patch_start_loc: tuple):
+        """
+        Here, h_start, w_start could be negative. That simply means we need to pick the content from 0. So,
+        the cropped image will be smaller than self._img_sz * self._img_sz
+        """
+        max_len_vals = list(self.idx_manager.data_shape[1:-1])
+        max_len_vals[-2:] = img.shape[-2:]
+        return self._crop_img_with_padding(
+            img, patch_start_loc, max_len_vals=max_len_vals
+        )
+
+    def _get_img(self, index: int):
+        """
+        Returns the primary patch along with low resolution patches centered on the primary patch.
+        """
+        # Noise_tuples is populated when there is synthetic noise in training
+        # Should have similar type of noise with the noise model
+        # Starting with microsplit, dump the noise, use it instead as an augmentation if nessesary
+        img_tuples, noise_tuples = self._load_img(index)
+        assert self._img_sz is not None
+        h, w = img_tuples[0].shape[-2:]
+        if self._enable_random_cropping:
+            patch_start_loc = self._get_random_hw(h, w)
+            if self._5Ddata:
+                patch_start_loc = (
+                    np.random.choice(img_tuples[0].shape[-3] - self._depth3D),
+                ) + patch_start_loc
+        else:
+            patch_start_loc = self._get_deterministic_loc(index)
+
+        # LC logic is located here, the function crops the image of the highest resolution
+        cropped_img_tuples = [
+            self._crop_flip_img(img, patch_start_loc, False, False)
+            for img in img_tuples
+        ]
+        cropped_noise_tuples = [
+            self._crop_flip_img(noise, patch_start_loc, False, False)
+            for noise in noise_tuples
+        ]
+        patch_start_loc = list(patch_start_loc)
+        h_start, w_start = patch_start_loc[-2], patch_start_loc[-1]
+        h_center = h_start + self._img_sz // 2
+        w_center = w_start + self._img_sz // 2
+        allres_versions = {
+            i: [cropped_img_tuples[i]] for i in range(len(cropped_img_tuples))
+        }
+        for scale_idx in range(1, self.multiscale_lowres_count):
+            # Returning the image of the lower resolution
+            scaled_img_tuples = self._load_scaled_img(scale_idx, index)
+
+            h_center = h_center // 2
+            w_center = w_center // 2
+
+            h_start = h_center - self._img_sz // 2
+            w_start = w_center - self._img_sz // 2
+            patch_start_loc[-2:] = [h_start, w_start]
+            scaled_cropped_img_tuples = [
+                self._crop_flip_img(img, patch_start_loc, False, False)
+                for img in scaled_img_tuples
+            ]
+            for ch_idx in range(len(img_tuples)):
+                allres_versions[ch_idx].append(scaled_cropped_img_tuples[ch_idx])
+
+        output_img_tuples = tuple(
+            [
+                np.concatenate(allres_versions[ch_idx])
+                for ch_idx in range(len(img_tuples))
+            ]
+        )
+        return output_img_tuples, cropped_noise_tuples
+
+    def __getitem__(self, index: Union[int, tuple[int, int]]):
+        img_tuples, noise_tuples = self._get_img(index)
+        if self._uncorrelated_channels:
+            assert (
+                self._input_idx is None
+            ), "Uncorrelated channels is not implemented when there is a separate input channel."
+            if np.random.rand() < self._uncorrelated_channel_probab:
+                img_tuples_new = [None] * len(img_tuples)
+                img_tuples_new[0] = img_tuples[0]
+                for i in range(1, len(img_tuples)):
+                    new_index = np.random.randint(len(self))
+                    img_tuples_tmp, _ = self._get_img(new_index)
+                    img_tuples_new[i] = img_tuples_tmp[i]
+                img_tuples = img_tuples_new
+
+        if self._is_train:
+            if self._empty_patch_replacement_enabled:
+                if np.random.rand() < self._empty_patch_replacement_probab:
+                    img_tuples = self.replace_with_empty_patch(img_tuples)
+
+        if self._enable_rotation:
+            img_tuples, noise_tuples = self._rotate(img_tuples, noise_tuples)
+
+        # add noise to input, if noise is present combine it with the image
+        # factor is for the compute input not to have too much noise because the average of two gaussians
+        if len(noise_tuples) > 0:
+            factor = np.sqrt(2) if self._input_is_sum else 1.0
+            input_tuples = []
+            for x in img_tuples:
+                x = (
+                    x.copy()
+                )  # to avoid changing the original image since it is later used for target
+                # NOTE: other LC levels already have noise added. So, we just need to add noise to the highest resolution.
+                x[0] = x[0] + noise_tuples[0] * factor
+                input_tuples.append(x)
+        else:
+            input_tuples = img_tuples
+
+        # Compute the input by sum / average the channels
+        # Alpha is an amount of weight which is applied to the channels when combining them
+        # How to sample alpha is still under research
+        inp, alpha = self._compute_input(input_tuples)
+        target_tuples = [img[:1] for img in img_tuples]
+        # add noise to target.
+        if len(noise_tuples) >= 1:
+            target_tuples = [
+                x + noise for x, noise in zip(target_tuples, noise_tuples[1:])
+            ]
+
+        target = self._compute_target(target_tuples, alpha)
+
+        norm_target = self.normalize_target(target)
+
+        output = [inp, norm_target]
+
+        if self._return_alpha:
+            output.append(alpha)
+
+        if isinstance(index, int):
+            return tuple(output)
+
+        _, grid_size = index
+        output.append(grid_size)
+        return tuple(output)