careamics 0.0.19__py3-none-any.whl

careamics/utils/logging.py

@@ -0,0 +1,323 @@
+"""
+Logging submodule.
+
+The methods are responsible for the in-console logger.
+"""
+
+import logging
+import sys
+import time
+from collections.abc import Generator
+from pathlib import Path
+from typing import Any, Union
+
+LOGGERS: dict = {}
+
+
+def get_logger(
+    name: str,
+    log_level: int = logging.INFO,
+    log_path: Union[str, Path] | None = None,
+) -> logging.Logger:
+    """
+    Create a python logger instance with configured handlers.
+
+    Parameters
+    ----------
+    name : str
+        Name of the logger.
+    log_level : int, optional
+        Log level (info, error etc.), by default logging.INFO.
+    log_path : Optional[Union[str, Path]], optional
+        Path in which to save the log, by default None.
+
+    Returns
+    -------
+    logging.Logger
+        Logger.
+    """
+    logger = logging.getLogger(name)
+    logger.propagate = False
+
+    if name in LOGGERS:
+        return logger
+
+    for logger_name in LOGGERS:
+        if name.startswith(logger_name):
+            return logger
+
+    logger.propagate = False
+
+    if log_path:
+        handlers = [
+            logging.StreamHandler(),
+            logging.FileHandler(log_path),
+        ]
+    else:
+        handlers = [logging.StreamHandler()]
+
+    formatter = logging.Formatter("%(message)s")
+
+    for handler in handlers:
+        handler.setFormatter(formatter)  # type: ignore
+        handler.setLevel(log_level)  # type: ignore
+        logger.addHandler(handler)  # type: ignore
+
+    logger.setLevel(log_level)
+    LOGGERS[name] = True
+
+    logger.propagate = False
+
+    return logger
+
+
+class ProgressBar:
+    """
+    Keras style progress bar.
+
+    Adapted from https://github.com/yueyericardo/pkbar.
+
+    Parameters
+    ----------
+    max_value : Optional[int], optional
+        Maximum progress bar value, by default None.
+    epoch : Optional[int], optional
+        Zero-indexed current epoch, by default None.
+    num_epochs : Optional[int], optional
+        Total number of epochs, by default None.
+    stateful_metrics : Optional[list], optional
+        Iterable of string names of metrics that should *not* be averaged over time.
+        Metrics in this list will be displayed as-is. All others will be averaged by
+        the progress bar before display, by default None.
+    always_stateful : bool, optional
+            Whether to set all metrics to be stateful, by default False.
+    mode : str, optional
+        Mode, one of "train", "val", or "predict", by default "train".
+    """
+
+    def __init__(
+        self,
+        max_value: int | None = None,
+        epoch: int | None = None,
+        num_epochs: int | None = None,
+        stateful_metrics: list | None = None,
+        always_stateful: bool = False,
+        mode: str = "train",
+    ) -> None:
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        max_value : Optional[int], optional
+            Maximum progress bar value, by default None.
+        epoch : Optional[int], optional
+            Zero-indexed current epoch, by default None.
+        num_epochs : Optional[int], optional
+            Total number of epochs, by default None.
+        stateful_metrics : Optional[list], optional
+            Iterable of string names of metrics that should *not* be averaged over time.
+            Metrics in this list will be displayed as-is. All others will be averaged by
+            the progress bar before display, by default None.
+        always_stateful : bool, optional
+             Whether to set all metrics to be stateful, by default False.
+        mode : str, optional
+            Mode, one of "train", "val", or "predict", by default "train".
+        """
+        self.max_value = max_value
+        # Width of the progress bar
+        self.width = 30
+        self.always_stateful = always_stateful
+
+        if (epoch is not None) and (num_epochs is not None):
+            print(f"Epoch: {epoch + 1}/{num_epochs}")
+
+        if stateful_metrics:
+            self.stateful_metrics = set(stateful_metrics)
+        else:
+            self.stateful_metrics = set()
+
+        self._dynamic_display = (
+            (hasattr(sys.stdout, "isatty") and sys.stdout.isatty())
+            or "ipykernel" in sys.modules
+            or "posix" in sys.modules
+        )
+        self._total_width = 0
+        self._seen_so_far = 0
+        # We use a dict + list to avoid garbage collection
+        # issues found in OrderedDict
+        self._values: dict[Any, Any] = {}
+        self._values_order: list[Any] = []
+        self._start = time.time()
+        self._last_update = 0.0
+        self.spin = self.spinning_cursor() if self.max_value is None else None
+        if mode == "train" and self.max_value is None:
+            self.message = "Estimating dataset size"
+        elif mode == "val":
+            self.message = "Validating"
+        elif mode == "predict":
+            self.message = "Denoising"
+
+    def update(
+        self, current_step: int, batch_size: int = 1, values: list | None = None
+    ) -> None:
+        """
+        Update the progress bar.
+
+        Parameters
+        ----------
+        current_step : int
+            Index of the current step.
+        batch_size : int, optional
+            Batch size, by default 1.
+        values : Optional[list], optional
+            Updated metrics values, by default None.
+        """
+        values = values or []
+        for k, v in values:
+            # if torch tensor, convert it to numpy
+            if str(type(v)) == "<class 'torch.Tensor'>":
+                v = v.detach().cpu().numpy()
+
+            if k not in self._values_order:
+                self._values_order.append(k)
+            if k not in self.stateful_metrics and not self.always_stateful:
+                if k not in self._values:
+                    self._values[k] = [
+                        v * (current_step - self._seen_so_far),
+                        current_step - self._seen_so_far,
+                    ]
+                else:
+                    self._values[k][0] += v * (current_step - self._seen_so_far)
+                    self._values[k][1] += current_step - self._seen_so_far
+            else:
+                # Stateful metrics output a numeric value. This representation
+                # means "take an average from a single value" but keeps the
+                # numeric formatting.
+                self._values[k] = [v, 1]
+
+        self._seen_so_far = current_step
+
+        now = time.time()
+        info = f" - {(now - self._start):.0f}s"
+
+        prev_total_width = self._total_width
+        if self._dynamic_display:
+            sys.stdout.write("\b" * prev_total_width)
+            sys.stdout.write("\r")
+        else:
+            sys.stdout.write("\n")
+
+        if self.max_value is not None:
+            bar = f"{current_step}/{self.max_value} ["
+            progress = float(current_step) / self.max_value
+            progress_width = int(self.width * progress)
+            if progress_width > 0:
+                bar += "=" * (progress_width - 1)
+                if current_step < self.max_value:
+                    bar += ">"
+                else:
+                    bar += "="
+            bar += "." * (self.width - progress_width)
+            bar += "]"
+        else:
+            bar = (
+                f"{self.message} {next(self.spin)}, tile "  # type: ignore
+                f"No. {current_step * batch_size}"
+            )
+
+        self._total_width = len(bar)
+        sys.stdout.write(bar)
+
+        if current_step > 0:
+            time_per_unit = (now - self._start) / current_step
+        else:
+            time_per_unit = 0
+
+        if time_per_unit >= 1 or time_per_unit == 0:
+            info += f" {time_per_unit:.0f}s/step"
+        elif time_per_unit >= 1e-3:
+            info += f" {time_per_unit * 1e3:.0f}ms/step"
+        else:
+            info += f" {time_per_unit * 1e6:.0f}us/step"
+
+        for k in self._values_order:
+            info += f" - {k}:"
+            if isinstance(self._values[k], list):
+                avg = self._values[k][0] / max(1, self._values[k][1])
+                if abs(avg) > 1e-3:
+                    info += f" {avg:.4f}"
+                else:
+                    info += f" {avg:.4e}"
+            else:
+                info += f" {self._values[k]}s"
+
+        self._total_width += len(info)
+        if prev_total_width > self._total_width:
+            info += " " * (prev_total_width - self._total_width)
+
+        if self.max_value is not None and current_step >= self.max_value:
+            info += "\n"
+
+        sys.stdout.write(info)
+        sys.stdout.flush()
+
+        self._last_update = now
+
+    def add(self, n: int, values: list | None = None) -> None:
+        """
+        Update the progress bar by n steps.
+
+        Parameters
+        ----------
+        n : int
+            Number of steps to increase the progress bar with.
+        values : Optional[list], optional
+            Updated metrics values, by default None.
+        """
+        self.update(self._seen_so_far + n, 1, values=values)
+
+    def spinning_cursor(self) -> Generator:
+        """
+        Generate a spinning cursor animation.
+
+        Taken from https://github.com/manrajgrover/py-spinners/tree/master.
+
+        Returns
+        -------
+        Generator
+            Generator of animation frames.
+        """
+        while True:
+            yield from [
+                "▓ ----- ▒",
+                "▓ ----- ▒",
+                "▓ ----- ▒",
+                "▓ ->--- ▒",
+                "▓ ->--- ▒",
+                "▓ ->--- ▒",
+                "▓ -->-- ▒",
+                "▓ -->-- ▒",
+                "▓ -->-- ▒",
+                "▓ --->- ▒",
+                "▓ --->- ▒",
+                "▓ --->- ▒",
+                "▓ ----> ▒",
+                "▓ ----> ▒",
+                "▓ ----> ▒",
+                "▒ ----- ░",
+                "▒ ----- ░",
+                "▒ ----- ░",
+                "▒ ->--- ░",
+                "▒ ->--- ░",
+                "▒ ->--- ░",
+                "▒ -->-- ░",
+                "▒ -->-- ░",
+                "▒ -->-- ░",
+                "▒ --->- ░",
+                "▒ --->- ░",
+                "▒ --->- ░",
+                "▒ ----> ░",
+                "▒ ----> ░",
+                "▒ ----> ░",
+            ]

careamics/utils/metrics.py

@@ -0,0 +1,394 @@
+"""
+Metrics submodule.
+
+This module contains various metrics and a metrics tracking class.
+"""
+
+from collections.abc import Callable
+from typing import Union
+
+import numpy as np
+import torch
+from skimage.metrics import peak_signal_noise_ratio, structural_similarity
+from torchmetrics.image import MultiScaleStructuralSimilarityIndexMeasure
+
+# TODO: does this add additional dependency?
+
+
+# TODO revisit metric for notebook
+def avg_range_invariant_psnr(
+    pred: np.ndarray,
+    target: np.ndarray,
+) -> float:
+    """Compute the average range-invariant PSNR.
+
+    Parameters
+    ----------
+    pred : np.ndarray
+        Predicted images.
+    target : np.ndarray
+        Target images.
+
+    Returns
+    -------
+    float
+        Average range-invariant PSNR value.
+    """
+    psnr_arr = []
+    for i in range(pred.shape[0]):
+        psnr_arr.append(scale_invariant_psnr(pred[i], target[i]))
+    return np.mean(psnr_arr)
+
+
+def psnr(gt: np.ndarray, pred: np.ndarray, data_range: float) -> float:
+    """
+    Peak Signal to Noise Ratio.
+
+    This method calls skimage.metrics.peak_signal_noise_ratio. See:
+    https://scikit-image.org/docs/dev/api/skimage.metrics.html.
+
+    NOTE: to avoid unwanted behaviors (e.g., data_range inferred from array dtype),
+    the data_range parameter is mandatory.
+
+    Parameters
+    ----------
+    gt : np.ndarray
+        Ground truth array.
+    pred : np.ndarray
+        Predicted array.
+    data_range : float
+        The images pixel range.
+
+    Returns
+    -------
+    float
+        PSNR value.
+    """
+    return peak_signal_noise_ratio(gt, pred, data_range=data_range)
+
+
+def _zero_mean(x: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]:
+    """
+    Zero the mean of an array.
+
+    Parameters
+    ----------
+    x : numpy.ndarray or torch.Tensor
+        Input array.
+
+    Returns
+    -------
+    numpy.ndarray or torch.Tensor
+        Zero-mean array.
+    """
+    return x - x.mean()
+
+
+def _fix_range(
+    gt: Union[np.ndarray, torch.Tensor], x: Union[np.ndarray, torch.Tensor]
+) -> Union[np.ndarray, torch.Tensor]:
+    """
+    Adjust the range of an array based on a reference ground-truth array.
+
+    Parameters
+    ----------
+    gt : Union[np.ndarray, torch.Tensor]
+        Ground truth array.
+    x : Union[np.ndarray, torch.Tensor]
+        Input array.
+
+    Returns
+    -------
+    Union[np.ndarray, torch.Tensor]
+        Range-adjusted array.
+    """
+    a = (gt * x).sum() / (x * x).sum()
+    return x * a
+
+
+def _fix(
+    gt: Union[np.ndarray, torch.Tensor], x: Union[np.ndarray, torch.Tensor]
+) -> Union[np.ndarray, torch.Tensor]:
+    """
+    Zero mean a groud truth array and adjust the range of the array.
+
+    Parameters
+    ----------
+    gt : Union[np.ndarray, torch.Tensor]
+        Ground truth image.
+    x : Union[np.ndarray, torch.Tensor]
+        Input array.
+
+    Returns
+    -------
+    Union[np.ndarray, torch.Tensor]
+        Zero-mean and range-adjusted array.
+    """
+    gt_ = _zero_mean(gt)
+    return _fix_range(gt_, _zero_mean(x))
+
+
+def scale_invariant_psnr(
+    gt: np.ndarray, pred: np.ndarray
+) -> Union[float, torch.tensor]:
+    """
+    Scale invariant PSNR.
+
+    Parameters
+    ----------
+    gt : np.ndarray
+        Ground truth image.
+    pred : np.ndarray
+        Predicted image.
+
+    Returns
+    -------
+    Union[float, torch.tensor]
+        Scale invariant PSNR value.
+    """
+    range_parameter = (np.max(gt) - np.min(gt)) / np.std(gt)
+    gt_ = _zero_mean(gt) / np.std(gt)
+    return psnr(_zero_mean(gt_), _fix(gt_, pred), range_parameter)
+
+
+class RunningPSNR:
+    """Compute the running PSNR during validation step in training.
+
+    This class allows to compute the PSNR on the entire validation set
+    one batch at the time.
+
+    Attributes
+    ----------
+    N : int
+        Number of elements seen so far during the epoch.
+    mse_sum : float
+        Running sum of the MSE over the N elements seen so far.
+    max : float
+        Running max value of the N target images seen so far.
+    min : float
+        Running min value of the N target images seen so far.
+    """
+
+    def __init__(self):
+        """Constructor."""
+        self.N = None
+        self.mse_sum = None
+        self.max = self.min = None
+        self.reset()
+
+    def reset(self):
+        """Reset the running PSNR computation.
+
+        Usually called at the end of each epoch.
+        """
+        self.mse_sum = 0
+        self.N = 0
+        self.max = self.min = None
+
+    def update(self, rec: torch.Tensor, tar: torch.Tensor) -> None:
+        """Update the running PSNR statistics given a new batch.
+
+        Parameters
+        ----------
+        rec : torch.Tensor
+            Reconstructed batch.
+        tar : torch.Tensor
+            Target batch.
+        """
+        ins_max = torch.max(tar).item()
+        ins_min = torch.min(tar).item()
+        if self.max is None:
+            assert self.min is None
+            self.max = ins_max
+            self.min = ins_min
+        else:
+            self.max = max(self.max, ins_max)
+            self.min = min(self.min, ins_min)
+
+        mse = (rec - tar) ** 2
+        elementwise_mse = torch.mean(mse.view(len(mse), -1), dim=1)
+        self.mse_sum += torch.nansum(elementwise_mse)
+        self.N += len(elementwise_mse) - torch.sum(torch.isnan(elementwise_mse))
+
+    def get(self) -> torch.Tensor | None:
+        """Get the actual PSNR value given the running statistics.
+
+        Returns
+        -------
+        Optional[torch.Tensor]
+            PSNR value.
+        """
+        if self.N == 0 or self.N is None:
+            return None
+        rmse = torch.sqrt(self.mse_sum / self.N)
+        return 20 * torch.log10((self.max - self.min) / rmse)
+
+
+def _range_invariant_multiscale_ssim(
+    gt_: Union[np.ndarray, torch.Tensor], pred_: Union[np.ndarray, torch.Tensor]
+) -> float:
+    """Compute range invariant multiscale SSIM for a single channel.
+
+    The advantage of this metric in comparison to commonly used SSIM is that
+    it is invariant to scalar multiplications in the prediction.
+    # TODO: Add reference to the paper.
+
+    NOTE: images fed to this function should have channels dimension as the last one.
+
+    Parameters
+    ----------
+    gt_ : Union[np.ndarray, torch.Tensor]
+        Ground truth image with shape (N, H, W).
+    pred_ : Union[np.ndarray, torch.Tensor]
+        Predicted image with shape (N, H, W).
+
+    Returns
+    -------
+    float
+        Range invariant multiscale SSIM value.
+    """
+    shape = gt_.shape
+    gt_ = torch.Tensor(gt_.reshape((shape[0], -1)))
+    pred_ = torch.Tensor(pred_.reshape((shape[0], -1)))
+    gt_ = _zero_mean(gt_)
+    pred_ = _zero_mean(pred_)
+    pred_ = _fix(gt_, pred_)
+    pred_ = pred_.reshape(shape)
+    gt_ = gt_.reshape(shape)
+
+    ms_ssim = MultiScaleStructuralSimilarityIndexMeasure(
+        data_range=gt_.max() - gt_.min()
+    )
+    return ms_ssim(torch.Tensor(pred_[:, None]), torch.Tensor(gt_[:, None])).item()
+
+
+def multiscale_ssim(
+    gt_: Union[np.ndarray, torch.Tensor],
+    pred_: Union[np.ndarray, torch.Tensor],
+    range_invariant: bool = True,
+) -> list[Union[float, None]]:
+    """Compute channel-wise multiscale SSIM for each channel.
+
+    It allows to use either standard multiscale SSIM or its range-invariant version.
+
+    NOTE: images fed to this function should have channels dimension as the last one.
+    # TODO: do we want to allow this behavior? or we want the usual (N, C, H, W)?
+
+    Parameters
+    ----------
+    gt_ : Union[np.ndarray, torch.Tensor]
+        Ground truth image with shape (N, H, W, C).
+    pred_ : Union[np.ndarray, torch.Tensor]
+        Predicted image with shape (N, H, W, C).
+    range_invariant : bool
+        Whether to use standard or range invariant multiscale SSIM.
+
+    Returns
+    -------
+    list[float]
+        List of SSIM values for each channel.
+    """
+    ms_ssim_values = {}
+    for ch_idx in range(gt_.shape[-1]):
+        tar_tmp = gt_[..., ch_idx]
+        pred_tmp = pred_[..., ch_idx]
+        if range_invariant:
+            ms_ssim_values[ch_idx] = _range_invariant_multiscale_ssim(
+                gt_=tar_tmp, pred_=pred_tmp
+            )
+        else:
+            ms_ssim = MultiScaleStructuralSimilarityIndexMeasure(
+                data_range=tar_tmp.max() - tar_tmp.min()
+            )
+            ms_ssim_values[ch_idx] = ms_ssim(
+                torch.Tensor(pred_tmp[:, None]), torch.Tensor(tar_tmp[:, None])
+            ).item()
+
+    return [ms_ssim_values[i] for i in range(gt_.shape[-1])]  # type: ignore
+
+
+def _avg_psnr(target: np.ndarray, prediction: np.ndarray, psnr_fn: Callable) -> float:
+    """Compute the average PSNR over a batch of images.
+
+    Parameters
+    ----------
+    target : np.ndarray
+        Array of ground truth images, shape is (N, C, H, W).
+    prediction : np.ndarray
+        Array of predicted images, shape is (N, C, H, W).
+    psnr_fn : Callable
+        PSNR function to use.
+
+    Returns
+    -------
+    float
+        Average PSNR value over the batch.
+    """
+    return np.mean(
+        [
+            psnr_fn(target[i : i + 1], prediction[i : i + 1]).item()
+            for i in range(len(prediction))
+        ]
+    )
+
+
+def avg_range_inv_psnr(target: np.ndarray, prediction: np.ndarray) -> float:
+    """Compute the average range-invariant PSNR over a batch of images.
+
+    Parameters
+    ----------
+    target : np.ndarray
+        Array of ground truth images, shape is (N, C, H, W).
+    prediction : np.ndarray
+        Array of predicted images, shape is (N, C, H, W).
+
+    Returns
+    -------
+    float
+        Average range-invariant PSNR value over the batch.
+    """
+    return _avg_psnr(target, prediction, scale_invariant_psnr)
+
+
+def avg_psnr(target: np.ndarray, prediction: np.ndarray) -> float:
+    """Compute the average PSNR over a batch of images.
+
+    Parameters
+    ----------
+    target : np.ndarray
+        Array of ground truth images, shape is (N, C, H, W).
+    prediction : np.ndarray
+        Array of predicted images, shape is (N, C, H, W).
+
+    Returns
+    -------
+    float
+        Average PSNR value over the batch.
+    """
+    return _avg_psnr(target, prediction, psnr)
+
+
+def avg_ssim(
+    target: Union[np.ndarray, torch.Tensor], prediction: Union[np.ndarray, torch.Tensor]
+) -> tuple[float, float]:
+    """Compute the average Structural Similarity (SSIM) over a batch of images.
+
+    Parameters
+    ----------
+    target : np.ndarray
+        Array of ground truth images, shape is (N, C, H, W).
+    prediction : np.ndarray
+        Array of predicted images, shape is (N, C, H, W).
+
+    Returns
+    -------
+    tuple[float, float]
+        Mean and standard deviation of SSIM values over the batch.
+    """
+    ssim = [
+        structural_similarity(
+            target[i], prediction[i], data_range=(target[i].max() - target[i].min())
+        )
+        for i in range(len(target))
+    ]
+    return np.mean(ssim), np.std(ssim)