careamics 0.0.3__py3-none-any.whl → 0.0.4.1__py3-none-any.whl

careamics/lvae_training/dataset/utils/index_manager.py

@@ -0,0 +1,232 @@
+from dataclasses import dataclass
+
+import numpy as np
+
+from careamics.lvae_training.dataset.types import TilingMode
+
+
+@dataclass
+class GridIndexManager:
+    data_shape: tuple
+    grid_shape: tuple
+    patch_shape: tuple
+    tiling_mode: TilingMode
+
+    # Patch is centered on index in the grid, grid size not used in training,
+    # used only during val / test, grid size controls the overlap of the patches
+    # in training you only get random patches every time
+    # For borders - just cropped the data, so it perfectly divisible
+
+    def __post_init__(self):
+        assert len(self.data_shape) == len(
+            self.grid_shape
+        ), f"Data shape:{self.data_shape} and grid size:{self.grid_shape} must have the same dimension"
+        assert len(self.data_shape) == len(
+            self.patch_shape
+        ), f"Data shape:{self.data_shape} and patch shape:{self.patch_shape} must have the same dimension"
+        innerpad = np.array(self.patch_shape) - np.array(self.grid_shape)
+        for dim, pad in enumerate(innerpad):
+            if pad < 0:
+                raise ValueError(
+                    f"Patch shape:{self.patch_shape} must be greater than or equal to grid shape:{self.grid_shape} in dimension {dim}"
+                )
+            if pad % 2 != 0:
+                raise ValueError(
+                    f"Patch shape:{self.patch_shape} must have even padding in dimension {dim}"
+                )
+
+    def patch_offset(self):
+        return (np.array(self.patch_shape) - np.array(self.grid_shape)) // 2
+
+    def get_individual_dim_grid_count(self, dim: int):
+        """
+        Returns the number of the grid in the specified dimension, ignoring all other dimensions.
+        """
+        assert dim < len(
+            self.data_shape
+        ), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+        assert dim >= 0, "Dimension must be greater than or equal to 0"
+
+        if self.grid_shape[dim] == 1 and self.patch_shape[dim] == 1:
+            return self.data_shape[dim]
+        elif self.tiling_mode == TilingMode.PadBoundary:
+            return int(np.ceil(self.data_shape[dim] / self.grid_shape[dim]))
+        elif self.tiling_mode == TilingMode.ShiftBoundary:
+            excess_size = self.patch_shape[dim] - self.grid_shape[dim]
+            return int(
+                np.ceil((self.data_shape[dim] - excess_size) / self.grid_shape[dim])
+            )
+        else:
+            excess_size = self.patch_shape[dim] - self.grid_shape[dim]
+            return int(
+                np.floor((self.data_shape[dim] - excess_size) / self.grid_shape[dim])
+            )
+
+    def total_grid_count(self):
+        """
+        Returns the total number of grids in the dataset.
+        """
+        return self.grid_count(0) * self.get_individual_dim_grid_count(0)
+
+    def grid_count(self, dim: int):
+        """
+        Returns the total number of grids for one value in the specified dimension.
+        """
+        assert dim < len(
+            self.data_shape
+        ), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+        assert dim >= 0, "Dimension must be greater than or equal to 0"
+        if dim == len(self.data_shape) - 1:
+            return 1
+
+        return self.get_individual_dim_grid_count(dim + 1) * self.grid_count(dim + 1)
+
+    def get_grid_index(self, dim: int, coordinate: int):
+        """
+        Returns the index of the grid in the specified dimension.
+        """
+        assert dim < len(
+            self.data_shape
+        ), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+        assert dim >= 0, "Dimension must be greater than or equal to 0"
+        assert (
+            coordinate < self.data_shape[dim]
+        ), f"Coordinate {coordinate} is out of bounds for data shape {self.data_shape}"
+
+        if self.grid_shape[dim] == 1 and self.patch_shape[dim] == 1:
+            return coordinate
+        elif self.tiling_mode == TilingMode.PadBoundary:  # self.trim_boundary is False:
+            return np.floor(coordinate / self.grid_shape[dim])
+        elif self.tiling_mode == TilingMode.TrimBoundary:
+            excess_size = (self.patch_shape[dim] - self.grid_shape[dim]) // 2
+            # can be <0 if coordinate is in [0,grid_shape[dim]]
+            return max(0, np.floor((coordinate - excess_size) / self.grid_shape[dim]))
+        elif self.tiling_mode == TilingMode.ShiftBoundary:
+            excess_size = (self.patch_shape[dim] - self.grid_shape[dim]) // 2
+            if coordinate + self.grid_shape[dim] + excess_size == self.data_shape[dim]:
+                return self.get_individual_dim_grid_count(dim) - 1
+            else:
+                # can be <0 if coordinate is in [0,grid_shape[dim]]
+                return max(
+                    0, np.floor((coordinate - excess_size) / self.grid_shape[dim])
+                )
+
+        else:
+            raise ValueError(f"Unsupported tiling mode {self.tiling_mode}")
+
+    def dataset_idx_from_grid_idx(self, grid_idx: tuple):
+        """
+        Returns the index of the grid in the dataset.
+        """
+        assert len(grid_idx) == len(
+            self.data_shape
+        ), f"Dimension indices {grid_idx} must have the same dimension as data shape {self.data_shape}"
+        index = 0
+        for dim in range(len(grid_idx)):
+            index += grid_idx[dim] * self.grid_count(dim)
+        return index
+
+    def get_patch_location_from_dataset_idx(self, dataset_idx: int):
+        """
+        Returns the patch location of the grid in the dataset.
+        """
+        grid_location = self.get_location_from_dataset_idx(dataset_idx)
+        offset = self.patch_offset()
+        return tuple(np.array(grid_location) - np.array(offset))
+
+    def get_dataset_idx_from_grid_location(self, location: tuple):
+        assert len(location) == len(
+            self.data_shape
+        ), f"Location {location} must have the same dimension as data shape {self.data_shape}"
+        grid_idx = [
+            self.get_grid_index(dim, location[dim]) for dim in range(len(location))
+        ]
+        return self.dataset_idx_from_grid_idx(tuple(grid_idx))
+
+    def get_gridstart_location_from_dim_index(self, dim: int, dim_index: int):
+        """
+        Returns the grid-start coordinate of the grid in the specified dimension.
+        """
+        assert dim < len(
+            self.data_shape
+        ), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+        assert dim >= 0, "Dimension must be greater than or equal to 0"
+        assert dim_index < self.get_individual_dim_grid_count(
+            dim
+        ), f"Dimension index {dim_index} is out of bounds for data shape {self.data_shape}"
+
+        if self.grid_shape[dim] == 1 and self.patch_shape[dim] == 1:
+            return dim_index
+        elif self.tiling_mode == TilingMode.PadBoundary:
+            return dim_index * self.grid_shape[dim]
+        elif self.tiling_mode == TilingMode.TrimBoundary:
+            excess_size = (self.patch_shape[dim] - self.grid_shape[dim]) // 2
+            return dim_index * self.grid_shape[dim] + excess_size
+        elif self.tiling_mode == TilingMode.ShiftBoundary:
+            excess_size = (self.patch_shape[dim] - self.grid_shape[dim]) // 2
+            if dim_index < self.get_individual_dim_grid_count(dim) - 1:
+                return dim_index * self.grid_shape[dim] + excess_size
+            else:
+                # on boundary. grid should be placed such that the patch covers the entire data.
+                return self.data_shape[dim] - self.grid_shape[dim] - excess_size
+        else:
+            raise ValueError(f"Unsupported tiling mode {self.tiling_mode}")
+
+    def get_location_from_dataset_idx(self, dataset_idx: int):
+        """
+        Returns the start location of the grid in the dataset.
+        """
+        grid_idx = []
+        for dim in range(len(self.data_shape)):
+            grid_idx.append(dataset_idx // self.grid_count(dim))
+            dataset_idx = dataset_idx % self.grid_count(dim)
+        location = [
+            self.get_gridstart_location_from_dim_index(dim, grid_idx[dim])
+            for dim in range(len(self.data_shape))
+        ]
+        return tuple(location)
+
+    def on_boundary(self, dataset_idx: int, dim: int, only_end: bool = False):
+        """
+        Returns True if the grid is on the boundary in the specified dimension.
+        """
+        assert dim < len(
+            self.data_shape
+        ), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+        assert dim >= 0, "Dimension must be greater than or equal to 0"
+
+        if dim > 0:
+            dataset_idx = dataset_idx % self.grid_count(dim - 1)
+
+        dim_index = dataset_idx // self.grid_count(dim)
+        if only_end:
+            return dim_index == self.get_individual_dim_grid_count(dim) - 1
+
+        return (
+            dim_index == 0 or dim_index == self.get_individual_dim_grid_count(dim) - 1
+        )
+
+    def next_grid_along_dim(self, dataset_idx: int, dim: int):
+        """
+        Returns the index of the grid in the specified dimension in the specified direction.
+        """
+        assert dim < len(
+            self.data_shape
+        ), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+        assert dim >= 0, "Dimension must be greater than or equal to 0"
+        new_idx = dataset_idx + self.grid_count(dim)
+        if new_idx >= self.total_grid_count():
+            return None
+        return new_idx
+
+    def prev_grid_along_dim(self, dataset_idx: int, dim: int):
+        """
+        Returns the index of the grid in the specified dimension in the specified direction.
+        """
+        assert dim < len(
+            self.data_shape
+        ), f"Dimension {dim} is out of bounds for data shape {self.data_shape}"
+        assert dim >= 0, "Dimension must be greater than or equal to 0"
+        new_idx = dataset_idx - self.grid_count(dim)
+        if new_idx < 0:
+            return None

careamics/lvae_training/dataset/utils/index_switcher.py

@@ -0,0 +1,165 @@
+import numpy as np
+
+
+class IndexSwitcher:
+    """
+    The idea is to switch from valid indices for target to invalid indices for target.
+    If index in invalid for the target, then we return all zero vector as target.
+    This combines both logic:
+    1. Using less amount of total data.
+    2. Using less amount of target data but using full data.
+    """
+
+    def __init__(self, idx_manager, data_config, patch_size) -> None:
+        self.idx_manager = idx_manager
+        self._data_shape = self.idx_manager.get_data_shape()
+        self._training_validtarget_fraction = data_config.get(
+            "training_validtarget_fraction", 1.0
+        )
+        self._validtarget_ceilT = int(
+            np.ceil(self._data_shape[0] * self._training_validtarget_fraction)
+        )
+        self._patch_size = patch_size
+        assert (
+            data_config.deterministic_grid is True
+        ), "This only works when the dataset has deterministic grid. Needed randomness comes from this class."
+        assert (
+            "grid_size" in data_config and data_config.grid_size == 1
+        ), "We need a one to one mapping between index and h, w, t"
+
+        self._h_validmax, self._w_validmax = self.get_reduced_frame_size(
+            self._data_shape[:3], self._training_validtarget_fraction
+        )
+        if self._h_validmax < self._patch_size or self._w_validmax < self._patch_size:
+            print(
+                "WARNING: The valid target size is smaller than the patch size. This will result in all zero target. so, we are ignoring this frame for target."
+            )
+            self._h_validmax = 0
+            self._w_validmax = 0
+
+        print(
+            f"[{self.__class__.__name__}] Target Indices: [0,{self._validtarget_ceilT - 1}]. Index={self._validtarget_ceilT - 1} has shape [:{self._h_validmax},:{self._w_validmax}].  Available data: {self._data_shape[0]}"
+        )
+
+    def get_valid_target_index(self):
+        """
+        Returns an index which corresponds to a frame which is expected to have a target.
+        """
+        _, h, w, _ = self._data_shape
+        framepixelcount = h * w
+        targetpixels = np.array(
+            [framepixelcount] * (self._validtarget_ceilT - 1)
+            + [self._h_validmax * self._w_validmax]
+        )
+        targetpixels = targetpixels / np.sum(targetpixels)
+        t = np.random.choice(self._validtarget_ceilT, p=targetpixels)
+        # t = np.random.randint(0, self._validtarget_ceilT) if self._validtarget_ceilT >= 1 else 0
+        h, w = self.get_valid_target_hw(t)
+        index = self.idx_manager.idx_from_hwt(h, w, t)
+        # print('Valid', index, h,w,t)
+        return index
+
+    def get_invalid_target_index(self):
+        # if self._validtarget_ceilT == 0:
+        # TODO: There may not be enough data for this to work. The better way is to skip using 0 for invalid target.
+        # t = np.random.randint(1, self._data_shape[0])
+        # elif self._validtarget_ceilT < self._data_shape[0]:
+        #     t = np.random.randint(self._validtarget_ceilT, self._data_shape[0])
+        # else:
+        #     t = self._validtarget_ceilT - 1
+        # 5
+        # 1.2 => 2
+        total_t, h, w, _ = self._data_shape
+        framepixelcount = h * w
+        available_h = h - self._h_validmax
+        if available_h < self._patch_size:
+            available_h = 0
+        available_w = w - self._w_validmax
+        if available_w < self._patch_size:
+            available_w = 0
+
+        targetpixels = np.array(
+            [available_h * available_w]
+            + [framepixelcount] * (total_t - self._validtarget_ceilT)
+        )
+        t_probab = targetpixels / np.sum(targetpixels)
+        t = np.random.choice(
+            np.arange(self._validtarget_ceilT - 1, total_t), p=t_probab
+        )
+
+        h, w = self.get_invalid_target_hw(t)
+        index = self.idx_manager.idx_from_hwt(h, w, t)
+        # print('Invalid', index, h,w,t)
+        return index
+
+    def get_valid_target_hw(self, t):
+        """
+        This is the opposite of get_invalid_target_hw. It returns a h,w which is valid for target.
+        This is only valid for single frame setup.
+        """
+        if t == self._validtarget_ceilT - 1:
+            h = np.random.randint(0, self._h_validmax - self._patch_size)
+            w = np.random.randint(0, self._w_validmax - self._patch_size)
+        else:
+            h = np.random.randint(0, self._data_shape[1] - self._patch_size)
+            w = np.random.randint(0, self._data_shape[2] - self._patch_size)
+        return h, w
+
+    def get_invalid_target_hw(self, t):
+        """
+        This is the opposite of get_valid_target_hw. It returns a h,w which is not valid for target.
+        This is only valid for single frame setup.
+        """
+        if t == self._validtarget_ceilT - 1:
+            h = np.random.randint(
+                self._h_validmax, self._data_shape[1] - self._patch_size
+            )
+            w = np.random.randint(
+                self._w_validmax, self._data_shape[2] - self._patch_size
+            )
+        else:
+            h = np.random.randint(0, self._data_shape[1] - self._patch_size)
+            w = np.random.randint(0, self._data_shape[2] - self._patch_size)
+        return h, w
+
+    def _get_tidx(self, index):
+        if isinstance(index, int) or isinstance(index, np.int64):
+            idx = index
+        else:
+            idx = index[0]
+        return self.idx_manager.get_t(idx)
+
+    def index_should_have_target(self, index):
+        tidx = self._get_tidx(index)
+        if tidx < self._validtarget_ceilT - 1:
+            return True
+        elif tidx > self._validtarget_ceilT - 1:
+            return False
+        else:
+            h, w, _ = self.idx_manager.hwt_from_idx(index)
+            return (
+                h + self._patch_size < self._h_validmax
+                and w + self._patch_size < self._w_validmax
+            )
+
+    @staticmethod
+    def get_reduced_frame_size(data_shape_nhw, fraction):
+        n, h, w = data_shape_nhw
+
+        framepixelcount = h * w
+        targetpixelcount = int(n * framepixelcount * fraction)
+
+        # We are currently supporting this only when there is just one frame.
+        # if np.ceil(pixelcount / framepixelcount) > 1:
+        #     return None, None
+
+        lastframepixelcount = targetpixelcount % framepixelcount
+        assert data_shape_nhw[1] == data_shape_nhw[2]
+        if lastframepixelcount > 0:
+            new_size = int(np.sqrt(lastframepixelcount))
+            return new_size, new_size
+        else:
+            assert (
+                targetpixelcount / framepixelcount >= 1
+            ), "This is not possible in euclidean space :D (so this is a bug)"
+            return h, w

careamics/lvae_training/eval_utils.py

@@ -14,13 +14,19 @@ import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
 import torch
+from torch import nn
+from torch.utils.data import Dataset
 from matplotlib.gridspec import GridSpec
 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 .metrics import RangeInvariantPsnr, RunningPSNR
+from careamics.utils.metrics import scale_invariant_psnr, RunningPSNR
 
 
 # ------------------------------------------------------------------------------------------------
@@ -40,28 +46,26 @@ def clean_ax(ax):
     ax.tick_params(left=False, right=False, top=False, bottom=False)
 
 
-def get_plots_output_dir(
+def get_eval_output_dir(
     saveplotsdir: str, patch_size: int, mmse_count: int = 50
 ) -> str:
     """
     Given the path to a root directory to save plots, patch size, and mmse count,
     it returns the specific directory to save the plots.
     """
-    plotsrootdir = os.path.join(
-        saveplotsdir, f"plots/patch_{patch_size}_mmse_{mmse_count}"
+    eval_out_dir = os.path.join(
+        saveplotsdir, f"eval_outputs/patch_{patch_size}_mmse_{mmse_count}"
     )
-    os.makedirs(plotsrootdir, exist_ok=True)
-    print(plotsrootdir)
-    return plotsrootdir
+    os.makedirs(eval_out_dir, exist_ok=True)
+    print(eval_out_dir)
+    return eval_out_dir
 
 
 def get_psnr_str(tar_hsnr, pred, col_idx):
     """
     Compute PSNR between the ground truth (`tar_hsnr`) and the predicted image (`pred`).
     """
-    return (
-        f"{RangeInvariantPsnr(tar_hsnr[col_idx][None], pred[col_idx][None]).item():.1f}"
-    )
+    return f"{scale_invariant_psnr(tar_hsnr[col_idx][None], pred[col_idx][None]).item():.1f}"
 
 
 def add_psnr_str(ax_, psnr):
@@ -499,20 +503,40 @@ def get_predictions(idx, val_dset, model, mmse_count=50, patch_size=256):
 
 
 def get_dset_predictions(
-    model,
-    dset,
+    model: VAEModule,
+    dset: Dataset,
     batch_size: int,
-    model_type: ModelType = None,
+    loss_type: Literal["musplit", "denoisplit", "denoisplit_musplit"],
     mmse_count: int = 1,
     num_workers: int = 4,
-):
-    """
-    Get predictions from a model for the entire dataset.
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, List[float]]:
+    """Get patch-wise predictions from a model for the entire dataset.
 
     Parameters
     ----------
-    mmse_count : int
-        Number of samples to generate for each input and then to average over for MMSE estimation.
+    model : VAEModule
+        Lightning model used for prediction.
+    dset : Dataset
+        Dataset to predict on.
+    batch_size : int
+        Batch size to use for prediction.
+    loss_type :
+        Type of reconstruction loss used by the model, by default `None`.
+    mmse_count : int, optional
+        Number of samples to generate for each input and then to average over for
+        MMSE estimation, by default 1.
+    num_workers : int, optional
+        Number of workers to use for DataLoader, by default 4.
+
+    Returns
+    -------
+    tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, List[float]]
+        Tuple containing:
+            - predictions: Predicted images for the dataset.
+            - predictions_std: Standard deviation of the predicted images.
+            - logvar_arr: Log variance of the predicted images.
+            - losses: Reconstruction losses for the predictions.
+            - psnr: PSNR values for the predictions.
     """
     dloader = DataLoader(
         dset,
@@ -521,69 +545,90 @@ def get_dset_predictions(
         shuffle=False,
         batch_size=batch_size,
     )
-    likelihood = model.model.likelihood
+
+    gauss_likelihood = model.gaussian_likelihood
+    nm_likelihood = model.noise_model_likelihood
+
     predictions = []
     predictions_std = []
     losses = []
     logvar_arr = []
-    patch_psnr_channels = [RunningPSNR() for _ in range(dset[0][1].shape[0])]
+    num_channels = dset[0][1].shape[0]
+    patch_psnr_channels = [RunningPSNR() for _ in range(num_channels)]
     with torch.no_grad():
-        for batch in tqdm(dloader):
-            inp, tar = batch[:2]
+        for batch in tqdm(dloader, desc="Predicting patches"):
+            inp, tar = batch
             inp = inp.cuda()
             tar = tar.cuda()
 
-            recon_img_list = []
+            rec_img_list = []
             for mmse_idx in range(mmse_count):
-                if model_type == ModelType.Denoiser:
-                    assert model.denoise_channel in [
-                        "Ch1",
-                        "Ch2",
-                        "input",
-                    ], '"all" denoise channel not supported for evaluation. Pick one of "Ch1", "Ch2", "input"'
-
-                    x_normalized_new, tar_new = model.get_new_input_target(
-                        (inp, tar, *batch[2:])
+
+                # TODO: case of HDN left for future refactoring
+                # if model_type == ModelType.Denoiser:
+                #     assert model.denoise_channel in [
+                #         "Ch1",
+                #         "Ch2",
+                #         "input",
+                #     ], '"all" denoise channel not supported for evaluation. Pick one of "Ch1", "Ch2", "input"'
+
+                #     x_normalized_new, tar_new = model.get_new_input_target(
+                #         (inp, tar, *batch[2:])
+                #     )
+                #     rec, _ = model(x_normalized_new)
+                #     rec_loss, imgs = model.get_reconstruction_loss(
+                #         rec,
+                #         tar,
+                #         x_normalized_new,
+                #         return_predicted_img=True,
+                #     )
+
+                # get model output
+                rec, _ = model(inp)
+
+                # get reconstructed img
+                if model.model.predict_logvar is None:
+                    rec_img = rec
+                    logvar = torch.tensor([-1])
+                else:
+                    rec_img, logvar = torch.chunk(rec, chunks=2, dim=1)
+                rec_img_list.append(rec_img.cpu().unsqueeze(0))  # add MMSE dim
+                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
                     )
-                    tar_normalized = model.normalize_target(tar_new)
-                    recon_normalized, _ = model(x_normalized_new)
-                    rec_loss, imgs = model.get_reconstruction_loss(
-                        recon_normalized,
-                        tar_normalized,
-                        x_normalized_new,
-                        return_predicted_img=True,
+                elif loss_type == "denoisplit":
+                    rec_loss = get_reconstruction_loss(
+                        reconstruction=rec, target=tar, likelihood_obj=nm_likelihood
                     )
-                else:
-                    x_normalized = model.normalize_input(inp)
-                    tar_normalized = model.normalize_target(tar)
-                    recon_normalized, _ = model(x_normalized)
-                    rec_loss, imgs = model.get_reconstruction_loss(
-                        recon_normalized, tar_normalized, inp, return_predicted_img=True
+                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:
-                    q_dic = (
-                        likelihood.distr_params(recon_normalized)
-                        if likelihood is not None
-                        else {"logvar": None}
-                    )
-                    if q_dic["logvar"] is not None:
-                        logvar_arr.append(q_dic["logvar"].cpu().numpy())
-                    else:
-                        logvar_arr.append(np.array([-1]))
-
                     try:
                         losses.append(rec_loss["loss"].cpu().numpy())
                     except:
                         losses.append(rec_loss["loss"])
 
-                for i in range(imgs.shape[1]):
-                    patch_psnr_channels[i].update(imgs[:, i], tar_normalized[:, i])
+                # update running PSNR
+                for i in range(num_channels):
+                    patch_psnr_channels[i].update(rec_img[:, i], tar[:, i])
 
-                recon_img_list.append(imgs.cpu()[None])
-
-            samples = torch.cat(recon_img_list, dim=0)
-            mmse_imgs = torch.mean(samples, dim=0)
+            # 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)
             predictions.append(mmse_imgs.cpu().numpy())
             predictions_std.append(mmse_std.cpu().numpy())
@@ -591,10 +636,10 @@ def get_dset_predictions(
     psnr = [x.get() for x in patch_psnr_channels]
     return (
         np.concatenate(predictions, axis=0),
-        np.array(losses),
+        np.concatenate(predictions_std, axis=0),
         np.concatenate(logvar_arr),
+        np.array(losses),
         psnr,
-        np.concatenate(predictions_std, axis=0),
     )
 
 

careamics/lvae_training/get_config.py

@@ -6,7 +6,7 @@ import os
 
 import ml_collections
 
-from careamics.lvae_training.dataset.data_utils import DataType
+from careamics.lvae_training.dataset.utils.data_utils import DataType
 from careamics.models.lvae.utils import LossType
 
 

careamics/lvae_training/train_lvae.py

@@ -24,7 +24,7 @@ from careamics.lvae_training.dataset.data_modules import (
     LCMultiChDloader,
     MultiChDloader,
 )
-from careamics.lvae_training.dataset.data_utils import DataSplitType
+from careamics.lvae_training.dataset.utils.data_utils import DataSplitType
 from careamics.lvae_training.lightning_module import LadderVAELight
 from careamics.lvae_training.train_utils import *