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

careamics/lvae_training/{data_modules.py → dataset/vae_dataset.py}

@@ -2,189 +2,65 @@
 A place for Datasets and Dataloaders.
 """
 
-import os
 from typing import Tuple, Union
 
-# import albumentations as A
-import ml_collections
 import numpy as np
-from skimage.transform import resize
 
 from .data_utils import (
-    DataSplitType,
-    DataType,
-    GridAlignement,
     GridIndexManager,
     IndexSwitcher,
-    get_datasplit_tuples,
-    get_mrc_data,
-    load_tiff,
+    get_train_val_data,
 )
-
-
-def get_train_val_data(
-    data_config,
-    fpath,
-    datasplit_type: DataSplitType,
-    val_fraction=None,
-    test_fraction=None,
-    allow_generation=None,
-    ignore_specific_datapoints=None,
-):
-    """
-    Load the data from the given path and split them in training, validation and test sets.
-
-    Ensure that the shape of data should be N*H*W*C: N is number of data points. H,W are the image dimensions.
-    C is the number of channels.
-    """
-    if data_config.data_type == DataType.SeparateTiffData:
-        fpath1 = os.path.join(fpath, data_config.ch1_fname)
-        fpath2 = os.path.join(fpath, data_config.ch2_fname)
-        fpaths = [fpath1, fpath2]
-        fpath0 = ""
-        if "ch_input_fname" in data_config:
-            fpath0 = os.path.join(fpath, data_config.ch_input_fname)
-            fpaths = [fpath0] + fpaths
-
-        print(
-            f"Loading from {fpath} Channels: "
-            f"{fpath1},{fpath2}, inp:{fpath0} Mode:{DataSplitType.name(datasplit_type)}"
-        )
-
-        data = np.concatenate([load_tiff(fpath)[..., None] for fpath in fpaths], axis=3)
-        if data_config.data_type == DataType.PredictedTiffData:
-            assert len(data.shape) == 5 and data.shape[-1] == 1
-            data = data[..., 0].copy()
-        # data = data[::3].copy()
-        # NOTE: This was not the correct way to do it. It is so because the noise present in the input was directly related
-        # to the noise present in the channels and so this is not the way we would get the data.
-        # We need to add the noise independently to the input and the target.
-
-        # if data_config.get('poisson_noise_factor', False):
-        #     data = np.random.poisson(data)
-        # if data_config.get('enable_gaussian_noise', False):
-        #     synthetic_scale = data_config.get('synthetic_gaussian_scale', 0.1)
-        #     print('Adding Gaussian noise with scale', synthetic_scale)
-        #     noise = np.random.normal(0, synthetic_scale, data.shape)
-        #     data = data + noise
-
-        if datasplit_type == DataSplitType.All:
-            return data.astype(np.float32)
-
-        train_idx, val_idx, test_idx = get_datasplit_tuples(
-            val_fraction, test_fraction, len(data), starting_test=True
-        )
-        if datasplit_type == DataSplitType.Train:
-            return data[train_idx].astype(np.float32)
-        elif datasplit_type == DataSplitType.Val:
-            return data[val_idx].astype(np.float32)
-        elif datasplit_type == DataSplitType.Test:
-            return data[test_idx].astype(np.float32)
-
-    elif data_config.data_type == DataType.BioSR_MRC:
-        num_channels = data_config.get("num_channels", 2)
-        fpaths = []
-        data_list = []
-        for i in range(num_channels):
-            fpath1 = os.path.join(fpath, data_config.get(f"ch{i + 1}_fname"))
-            fpaths.append(fpath1)
-            data = get_mrc_data(fpath1)[..., None]
-            data_list.append(data)
-
-        dirname = os.path.dirname(os.path.dirname(fpaths[0])) + "/"
-
-        msg = ",".join([x[len(dirname) :] for x in fpaths])
-        print(
-            f"Loaded from {dirname} Channels:{len(fpaths)} {msg} Mode:{DataSplitType.name(datasplit_type)}"
-        )
-        N = data_list[0].shape[0]
-        for data in data_list:
-            N = min(N, data.shape[0])
-
-        cropped_data = []
-        for data in data_list:
-            cropped_data.append(data[:N])
-
-        data = np.concatenate(cropped_data, axis=3)
-
-        if datasplit_type == DataSplitType.All:
-            return data.astype(np.float32)
-
-        train_idx, val_idx, test_idx = get_datasplit_tuples(
-            val_fraction, test_fraction, len(data), starting_test=True
-        )
-        if datasplit_type == DataSplitType.Train:
-            return data[train_idx].astype(np.float32)
-        elif datasplit_type == DataSplitType.Val:
-            return data[val_idx].astype(np.float32)
-        elif datasplit_type == DataSplitType.Test:
-            return data[test_idx].astype(np.float32)
+from .vae_data_config import VaeDatasetConfig, DataSplitType, GridAlignement
 
 
 class MultiChDloader:
-
     def __init__(
         self,
-        data_config: ml_collections.ConfigDict,
+        data_config: VaeDatasetConfig,
         fpath: str,
-        datasplit_type: DataSplitType = None,
         val_fraction: float = None,
         test_fraction: float = None,
-        normalized_input=None,
-        enable_rotation_aug: bool = False,
-        enable_random_cropping: bool = False,
-        use_one_mu_std=None,
-        allow_generation: bool = False,
-        max_val: float = None,
-        grid_alignment=GridAlignement.LeftTop,
-        overlapping_padding_kwargs=None,
-        print_vars: bool = True,
     ):
-        """
-        Here, an image is split into grids of size img_sz.
-        Args:
-            repeat_factor: Since we are doing a random crop, repeat_factor is
-                given which can repeatedly sample from the same image. If self.N=12
-                and repeat_factor is 5, then index upto 12*5 = 60 is allowed.
-            use_one_mu_std: If this is set to true, then one mean and stdev is used
-                for both channels. Otherwise, two different meean and stdev are used.
-
-        """
+        """ """
         self._data_type = data_config.data_type
         self._fpath = fpath
         self._data = self.N = self._noise_data = None
-
+        self.Z = 1
+        self._trim_boundary = data_config.trim_boundary
         # Hardcoded params, not included in the config file.
 
         # by default, if the noise is present, add it to the input and target.
         self._disable_noise = False  # to add synthetic noise
+        self._poisson_noise_factor = None
         self._train_index_switcher = None
+        self._depth3D = data_config.depth3D
         # NOTE: Input is the sum of the different channels. It is not the average of the different channels.
-        self._input_is_sum = data_config.get("input_is_sum", False)
-        self._num_channels = data_config.get("num_channels", 2)
-        self._input_idx = data_config.get("input_idx", None)
-        self._tar_idx_list = data_config.get("target_idx_list", None)
+        self._input_is_sum = data_config.input_is_sum
+        self._num_channels = data_config.num_channels
+        self._input_idx = data_config.input_idx
+        self._tar_idx_list = data_config.target_idx_list
 
-        if datasplit_type == DataSplitType.Train:
+        if data_config.datasplit_type == DataSplitType.Train:
             self._datausage_fraction = 1.0
             # assert self._datausage_fraction == 1.0, 'Not supported. Use validtarget_random_fraction and training_validtarget_fraction to get the same effect'
             self._validtarget_rand_fract = None
             # self._validtarget_random_fraction_final = data_config.get('validtarget_random_fraction_final', None)
             # self._validtarget_random_fraction_stepepoch = data_config.get('validtarget_random_fraction_stepepoch', None)
             # self._idx_count = 0
-        elif datasplit_type == DataSplitType.Val:
+        elif data_config.datasplit_type == DataSplitType.Val:
             self._datausage_fraction = 1.0
         else:
             self._datausage_fraction = 1.0
 
         self.load_data(
             data_config,
-            datasplit_type,
+            data_config.datasplit_type,
             val_fraction=val_fraction,
             test_fraction=test_fraction,
-            allow_generation=allow_generation,
+            allow_generation=data_config.allow_generation,
         )
-        self._normalized_input = normalized_input
+        self._normalized_input = data_config.normalized_input
         self._quantile = 1.0
         self._channelwise_quantile = False
         self._background_quantile = 0.0
@@ -194,8 +70,8 @@ class MultiChDloader:
 
         self._background_values = None
 
-        self._grid_alignment = grid_alignment
-        self._overlapping_padding_kwargs = overlapping_padding_kwargs
+        self._grid_alignment = data_config.grid_alignment
+        self._overlapping_padding_kwargs = data_config.overlapping_padding_kwargs
         if self._grid_alignment == GridAlignement.LeftTop:
             assert (
                 self._overlapping_padding_kwargs is None
@@ -205,20 +81,28 @@ class MultiChDloader:
             assert (
                 self._overlapping_padding_kwargs is not None
             ), "With Center grid alignment, padding is needed."
+        if self._trim_boundary:
+            if (
+                self._overlapping_padding_kwargs is None
+                or data_config.multiscale_lowres_count is not None
+            ):
+                # raise warning
+                print("Padding is not used with this alignement style")
+        else:
+            assert (
+                self._overlapping_padding_kwargs is not None
+            ), "When not trimming boudnary, padding is needed."
 
-        self._is_train = datasplit_type == DataSplitType.Train
+        self._is_train = data_config.datasplit_type == DataSplitType.Train
 
         # input = alpha * ch1 + (1-alpha)*ch2.
         # alpha is sampled randomly between these two extremes
-        self._start_alpha_arr = self._end_alpha_arr = self._return_alpha = (
-            self._alpha_weighted_target
-        ) = None
+        self._start_alpha_arr = self._end_alpha_arr = self._return_alpha = None
 
         self._img_sz = self._grid_sz = self._repeat_factor = self.idx_manager = None
         if self._is_train:
-            self._start_alpha_arr = None
-            self._end_alpha_arr = None
-            self._alpha_weighted_target = False
+            self._start_alpha_arr = data_config.start_alpha
+            self._end_alpha_arr = data_config.end_alpha
 
             self.set_img_sz(
                 data_config.image_size,
@@ -229,11 +113,13 @@ class MultiChDloader:
                 ),
             )
 
-            # if self._validtarget_rand_fract is not None:
-            #     self._train_index_switcher = IndexSwitcher(self.idx_manager, data_config, self._img_sz)
-            #     self._std_background_arr = None
+            if self._validtarget_rand_fract is not None:
+                self._train_index_switcher = IndexSwitcher(
+                    self.idx_manager, data_config, self._img_sz
+                )
 
         else:
+
             self.set_img_sz(
                 data_config.image_size,
                 (
@@ -246,32 +132,42 @@ class MultiChDloader:
         self._return_alpha = False
         self._return_index = False
 
-        # self._empty_patch_replacement_enabled = data_config.get("empty_patch_replacement_enabled",
-        #                                                         False) and self._is_train
-        # if self._empty_patch_replacement_enabled:
-        #     self._empty_patch_replacement_channel_idx = data_config.empty_patch_replacement_channel_idx
-        #     self._empty_patch_replacement_probab = data_config.empty_patch_replacement_probab
-        #     data_frames = self._data[..., self._empty_patch_replacement_channel_idx]
-        #     # NOTE: This is on the raw data. So, it must be called before removing the background.
-        #     self._empty_patch_fetcher = EmptyPatchFetcher(self.idx_manager,
-        #                                                   self._img_sz,
-        #                                                   data_frames,
-        #                                                   max_val_threshold=data_config.empty_patch_max_val_threshold)
+        self._empty_patch_replacement_enabled = (
+            data_config.empty_patch_replacement_enabled and self._is_train
+        )
+        if self._empty_patch_replacement_enabled:
+            self._empty_patch_replacement_channel_idx = (
+                data_config.empty_patch_replacement_channel_idx
+            )
+            self._empty_patch_replacement_probab = (
+                data_config.empty_patch_replacement_probab
+            )
+            data_frames = self._data[..., self._empty_patch_replacement_channel_idx]
+            # NOTE: This is on the raw data. So, it must be called before removing the background.
+            # TODO: missing import, needs fixing asap!
+            self._empty_patch_fetcher = EmptyPatchFetcher(
+                self.idx_manager,
+                self._img_sz,
+                data_frames,
+                max_val_threshold=data_config.empty_patch_max_val_threshold,
+            )
 
-        self.rm_bkground_set_max_val_and_upperclip_data(max_val, datasplit_type)
+        self.rm_bkground_set_max_val_and_upperclip_data(
+            data_config.max_val, data_config.datasplit_type
+        )
 
         # For overlapping dloader, image_size and repeat_factors are not related. hence a different function.
 
         self._mean = None
         self._std = None
-        self._use_one_mu_std = use_one_mu_std
+        self._use_one_mu_std = data_config.use_one_mu_std
         # Hardcoded
         self._target_separate_normalization = True
 
-        self._enable_rotation = enable_rotation_aug
-        self._enable_random_cropping = enable_random_cropping
+        self._enable_rotation = data_config.enable_rotation_aug
+        self._enable_random_cropping = data_config.enable_random_cropping
         self._uncorrelated_channels = (
-            data_config.get("uncorrelated_channels", False) and self._is_train
+            data_config.uncorrelated_channels and self._is_train
         )
         assert self._is_train or self._uncorrelated_channels is False
         assert (
@@ -286,9 +182,10 @@ class MultiChDloader:
             )
             self._rotation_transform = A.Compose([A.Flip(), A.RandomRotate90()])
 
-        if print_vars:
-            msg = self._init_msg()
-            print(msg)
+        # TODO: remove print log messages
+        # if print_vars:
+        #     msg = self._init_msg()
+        #     print(msg)
 
     def disable_noise(self):
         assert (
@@ -339,7 +236,7 @@ class MultiChDloader:
             )
 
         msg = ""
-        if data_config.get("poisson_noise_factor", -1) > 0:
+        if data_config.poisson_noise_factor > 0:
             self._poisson_noise_factor = data_config.poisson_noise_factor
             msg += f"Adding Poisson noise with factor {self._poisson_noise_factor}.\t"
             self._data = (
@@ -347,20 +244,26 @@ class MultiChDloader:
                 * self._poisson_noise_factor
             )
 
-        if data_config.get("enable_gaussian_noise", False):
-            synthetic_scale = data_config.get("synthetic_gaussian_scale", 0.1)
+        if data_config.enable_gaussian_noise:
+            synthetic_scale = data_config.synthetic_gaussian_scale
             msg += f"Adding Gaussian noise with scale {synthetic_scale}"
             # 0 => noise for input. 1: => noise for all targets.
             shape = self._data.shape
             self._noise_data = np.random.normal(
                 0, synthetic_scale, (*shape[:-1], shape[-1] + 1)
             )
-            if data_config.get("input_has_dependant_noise", False):
+            if data_config.input_has_dependant_noise:
                 msg += ". Moreover, input has dependent noise"
                 self._noise_data[..., 0] = np.mean(self._noise_data[..., 1:], axis=-1)
         print(msg)
 
-        self.N = len(self._data)
+        self._5Ddata = len(self._data.shape) == 5
+        if self._5Ddata:
+            self.Z = self._data.shape[1]
+
+        if self._depth3D > 1:
+            assert self._5Ddata, "Data must be 5D:NxZxHxWxC for 3D data"
+
         assert (
             self._data.shape[-1] == self._num_channels
         ), "Number of channels in data and config do not match."
@@ -441,9 +344,13 @@ class MultiChDloader:
     def get_img_sz(self):
         return self._img_sz
 
+    def get_num_frames(self):
+        return self._data.shape[0]
+
     def reduce_data(
         self, t_list=None, h_start=None, h_end=None, w_start=None, w_end=None
     ):
+        assert not self._5Ddata, "This function is not supported for 3D data."
         if t_list is None:
             t_list = list(range(self._data.shape[0]))
         if h_start is None:
@@ -461,12 +368,22 @@ class MultiChDloader:
                 t_list, h_start:h_end, w_start:w_end, :
             ].copy()
 
-        self.N = len(t_list)
         self.set_img_sz(self._img_sz, self._grid_sz)
         print(
             f"[{self.__class__.__name__}] Data reduced. New data shape: {self._data.shape}"
         )
 
+    def get_idx_manager_shapes(self, patch_size: int, grid_size: int):
+        numC = self._data.shape[-1]
+        if self._5Ddata:
+            grid_shape = (1, 1, grid_size, grid_size, numC)
+            patch_shape = (1, self._depth3D, patch_size, patch_size, numC)
+        else:
+            grid_shape = (1, grid_size, grid_size, numC)
+            patch_shape = (1, patch_size, patch_size, numC)
+
+        return patch_shape, grid_shape
+
     def set_img_sz(self, image_size, grid_size):
         """
         If one wants to change the image size on the go, then this can be used.
@@ -474,12 +391,23 @@ class MultiChDloader:
             image_size: size of one patch
             grid_size: frame is divided into square grids of this size. A patch centered on a grid having size `image_size` is returned.
         """
+
         self._img_sz = image_size
         self._grid_sz = grid_size
+        shape = self._data.shape
+
+        patch_shape, grid_shape = self.get_idx_manager_shapes(
+            self._img_sz, self._grid_sz
+        )
         self.idx_manager = GridIndexManager(
-            self._data.shape, self._grid_sz, self._img_sz, self._grid_alignment
+            shape, grid_shape, patch_shape, self._trim_boundary
         )
-        self.set_repeat_factor()
+        # self.set_repeat_factor()
+
+    def __len__(self):
+        # Vera: N is the number of frames in Z stack
+        # Repeat factor is n_rows * n_cols
+        return self.idx_manager.total_grid_count()
 
     def set_repeat_factor(self):
         if self._grid_sz > 1:
@@ -497,7 +425,14 @@ class MultiChDloader:
         msg = (
             f"[{self.__class__.__name__}] Train:{int(self._is_train)} Sz:{self._img_sz}"
         )
+        dim_sizes = [
+            self.idx_manager.get_individual_dim_grid_count(dim)
+            for dim in range(len(self._data.shape))
+        ]
+        dim_sizes = ",".join([str(x) for x in dim_sizes])
         msg += f" N:{self.N} NumPatchPerN:{self._repeat_factor}"
+        msg += f"{self.idx_manager.total_grid_count()} DimSz:({dim_sizes})"
+        msg += f" TrimB:{self._trim_boundary}"
         # msg += f' NormInp:{self._normalized_input}'
         # msg += f' SingleNorm:{self._use_one_mu_std}'
         msg += f" Rot:{self._enable_rotation}"
@@ -529,40 +464,52 @@ class MultiChDloader:
             )
 
         if self._enable_random_cropping:
-            h_start, w_start = self._get_random_hw(h, w)
+            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:
-            h_start, w_start = self._get_deterministic_hw(index)
+            patch_start_loc = self._get_deterministic_loc(index)
 
         cropped_imgs = []
         for img in img_tuples:
-            img = self._crop_flip_img(img, h_start, w_start, False, False)
+            img = self._crop_flip_img(img, patch_start_loc, False, False)
             cropped_imgs.append(img)
 
         return (
             *tuple(cropped_imgs),
             {
-                "h": [h_start, h_start + self._img_sz],
-                "w": [w_start, w_start + self._img_sz],
                 "hflip": False,
                 "wflip": False,
             },
         )
 
-    def _crop_img(self, img: np.ndarray, h_start: int, w_start: int):
-        if self._grid_alignment == GridAlignement.LeftTop:
+    def _crop_img(self, img: np.ndarray, patch_start_loc: Tuple):
+        if self._trim_boundary:
             # In training, this is used.
             # NOTE: It is my opinion that if I just use self._crop_img_with_padding, it will work perfectly fine.
             # The only benefit this if else loop provides is that it makes it easier to see what happens during training.
-            new_img = img[
-                ..., h_start : h_start + self._img_sz, w_start : w_start + self._img_sz
-            ]
+            patch_end_loc = (
+                np.array(patch_start_loc, dtype=np.int32)
+                + self.idx_manager.patch_shape[1:-1]
+            )
+            if self._5Ddata:
+                z_start, h_start, w_start = patch_start_loc
+                z_end, h_end, w_end = patch_end_loc
+                new_img = img[..., z_start:z_end, h_start:h_end, w_start:w_end]
+            else:
+                h_start, w_start = patch_start_loc
+                h_end, w_end = patch_end_loc
+                new_img = img[..., h_start:h_end, w_start:w_end]
+
             return new_img
-        elif self._grid_alignment == GridAlignement.Center:
+        else:
             # During evaluation, this is used. In this situation, we can have negative h_start, w_start. Or h_start +self._img_sz can be larger than frame
             # In these situations, we need some sort of padding. This is not needed  in the LeftTop alignement.
-            return self._crop_img_with_padding(img, h_start, w_start)
+            return self._crop_img_with_padding(img, patch_start_loc)
 
-    def get_begin_end_padding(self, start_pos, max_len):
+    def get_begin_end_padding(self, start_pos, end_pos, max_len):
         """
         The effect is that the image with size self._grid_sz is in the center of the patch with sufficient
         padding on all four sides so that the final patch size is self._img_sz.
@@ -572,44 +519,56 @@ class MultiChDloader:
         if start_pos < 0:
             pad_start = -1 * start_pos
 
-        pad_end = max(0, start_pos + self._img_sz - max_len)
+        pad_end = max(0, end_pos - max_len)
 
         return pad_start, pad_end
 
-    def _crop_img_with_padding(self, img: np.ndarray, h_start: int, w_start: int):
-        _, H, W = img.shape
-        h_on_boundary = self.on_boundary(h_start, H)
-        w_on_boundary = self.on_boundary(w_start, W)
-
-        assert h_start < H
-        assert w_start < W
-
-        assert h_start + self._img_sz <= H or h_on_boundary
-        assert w_start + self._img_sz <= W or w_on_boundary
+    def _crop_img_with_padding(
+        self, img: np.ndarray, patch_start_loc, max_len_vals=None
+    ):
+        if max_len_vals is None:
+            max_len_vals = self.idx_manager.data_shape[1:-1]
+        patch_end_loc = np.array(patch_start_loc, dtype=int) + np.array(
+            self.idx_manager.patch_shape[1:-1], dtype=int
+        )
+        boundary_crossed = []
+        valid_slice = []
+        padding = [[0, 0]]
+        for start_idx, end_idx, max_len in zip(
+            patch_start_loc, patch_end_loc, max_len_vals
+        ):
+            boundary_crossed.append(end_idx > max_len or start_idx < 0)
+            valid_slice.append((max(0, start_idx), min(max_len, end_idx)))
+            pad = [0, 0]
+            if boundary_crossed[-1]:
+                pad = self.get_begin_end_padding(start_idx, end_idx, max_len)
+            padding.append(pad)
         # max() is needed since h_start could be negative.
-        new_img = img[
-            ...,
-            max(0, h_start) : h_start + self._img_sz,
-            max(0, w_start) : w_start + self._img_sz,
-        ]
-        padding = np.array([[0, 0], [0, 0], [0, 0]])
-
-        if h_on_boundary:
-            pad = self.get_begin_end_padding(h_start, H)
-            padding[1] = pad
-        if w_on_boundary:
-            pad = self.get_begin_end_padding(w_start, W)
-            padding[2] = pad
+        if self._5Ddata:
+            new_img = img[
+                ...,
+                valid_slice[0][0] : valid_slice[0][1],
+                valid_slice[1][0] : valid_slice[1][1],
+                valid_slice[2][0] : valid_slice[2][1],
+            ]
+        else:
+            new_img = img[
+                ...,
+                valid_slice[0][0] : valid_slice[0][1],
+                valid_slice[1][0] : valid_slice[1][1],
+            ]
 
+        # print(np.array(padding).shape, img.shape, new_img.shape)
+        # print(padding)
         if not np.all(padding == 0):
             new_img = np.pad(new_img, padding, **self._overlapping_padding_kwargs)
 
         return new_img
 
     def _crop_flip_img(
-        self, img: np.ndarray, h_start: int, w_start: int, h_flip: bool, w_flip: bool
+        self, img: np.ndarray, patch_start_loc: Tuple, h_flip: bool, w_flip: bool
     ):
-        new_img = self._crop_img(img, h_start, w_start)
+        new_img = self._crop_img(img, patch_start_loc)
         if h_flip:
             new_img = new_img[..., ::-1, :]
         if w_flip:
@@ -617,9 +576,6 @@ class MultiChDloader:
 
         return new_img.astype(np.float32)
 
-    def __len__(self):
-        return self.N * self._repeat_factor
-
     def _load_img(
         self, index: Union[int, Tuple[int, int]]
     ) -> Tuple[np.ndarray, np.ndarray]:
@@ -631,12 +587,21 @@ class MultiChDloader:
         else:
             idx = index[0]
 
-        imgs = self._data[self.idx_manager.get_t(idx)]
+        patch_loc_list = self.idx_manager.get_patch_location_from_dataset_idx(idx)
+        imgs = self._data[patch_loc_list[0]]
+        # if self._5Ddata:
+        #     assert self._noise_data is None, 'Noise is not supported for 5D data'
+        #     n_loc, z_loc = patch_loc_list[:2]
+        #     z_loc_interval = range(z_loc, z_loc + self._depth3D)
+        #     imgs = self._data[n_loc, z_loc_interval]
+        # else:
+        #     imgs = self._data[patch_loc_list[0]]
+
         loaded_imgs = [imgs[None, ..., i] for i in range(imgs.shape[-1])]
         noise = []
         if self._noise_data is not None and not self._disable_noise:
             noise = [
-                self._noise_data[self.idx_manager.get_t(idx)][None, ..., i]
+                self._noise_data[patch_loc_list[0]][None, ..., i]
                 for i in range(self._noise_data.shape[-1])
             ]
         return tuple(loaded_imgs), tuple(noise)
@@ -669,27 +634,16 @@ class MultiChDloader:
     def per_side_overlap_pixelcount(self):
         return (self._img_sz - self._grid_sz) // 2
 
-    def on_boundary(self, cur_loc, frame_size):
-        return cur_loc + self._img_sz > frame_size or cur_loc < 0
+    # def on_boundary(self, cur_loc, frame_size):
+    #     return cur_loc + self._img_sz > frame_size or cur_loc < 0
 
-    def _get_deterministic_hw(self, index: Union[int, Tuple[int, int]]):
+    def _get_deterministic_loc(self, index: int):
         """
         It returns the top-left corner of the patch corresponding to index.
         """
-        if isinstance(index, int) or isinstance(index, np.int64):
-            idx = index
-            grid_size = self._grid_sz
-        else:
-            idx, grid_size = index
-
-        h_start, w_start = self.idx_manager.get_deterministic_hw(
-            idx, grid_size=grid_size
-        )
-        if self._grid_alignment == GridAlignement.LeftTop:
-            return h_start, w_start
-        elif self._grid_alignment == GridAlignement.Center:
-            pad = self.per_side_overlap_pixelcount()
-            return h_start - pad, w_start - pad
+        loc_list = self.idx_manager.get_patch_location_from_dataset_idx(index)
+        # last dim is channel. we need to take the third and the second last element.
+        return loc_list[1:-1]
 
     def compute_individual_mean_std(self):
         # numpy 1.19.2 has issues in computing for large arrays. https://github.com/numpy/numpy/issues/8869
@@ -715,6 +669,10 @@ class MultiChDloader:
 
         mean = np.array(mean_arr)
         std = np.array(std_arr)
+        if (
+            self._5Ddata
+        ):  # NOTE: IDEALLY this should be only when the model expects 3D data.
+            return mean[None, :, None, None, None], std[None, :, None, None, None]
 
         return mean[None, :, None, None], std[None, :, None, None]
 
@@ -776,6 +734,10 @@ class MultiChDloader:
         if self._skip_normalization_using_mean:
             mean = np.zeros_like(mean)
 
+        if self._5Ddata:
+            mean = mean[:, :, None]
+            std = std[:, :, None]
+
         mean_dict = {"input": mean}  # , 'target':mean}
         std_dict = {"input": std}  # , 'target':std}
 
@@ -810,8 +772,14 @@ class MultiChDloader:
         return cropped_img_tuples, cropped_noise_tuples
 
     def replace_with_empty_patch(self, img_tuples):
+        """
+        Replaces the content of one of the channels with background
+        """
         empty_index = self._empty_patch_fetcher.sample()
-        empty_img_tuples = self._get_img(empty_index)
+        empty_img_tuples, empty_img_noise_tuples = self._get_img(empty_index)
+        assert (
+            len(empty_img_noise_tuples) == 0
+        ), "Noise is not supported with empty patch replacement"
         final_img_tuples = []
         for tuple_idx in range(len(img_tuples)):
             if tuple_idx == self._empty_patch_replacement_channel_idx:
@@ -834,14 +802,7 @@ class MultiChDloader:
                 )
                 img_tuples = [img_tuples[i] for i in self._tar_idx_list]
 
-            if self._alpha_weighted_target:
-                assert self._input_is_sum is False
-                target = []
-                for i in range(len(img_tuples)):
-                    target.append(img_tuples[i] * alpha[i])
-                target = np.concatenate(target, axis=0)
-            else:
-                target = np.concatenate(img_tuples, axis=0)
+            target = np.concatenate(img_tuples, axis=0)
         return target
 
     def _compute_input_with_alpha(self, img_tuples, alpha_list):
@@ -902,9 +863,6 @@ class MultiChDloader:
                 index = self._train_index_switcher.get_invalid_target_index()
         return index
 
-    def _rotate(self, img_tuples, noise_tuples):
-        return self._rotate2D(img_tuples, noise_tuples)
-
     def _rotate2D(self, img_tuples, noise_tuples):
         img_kwargs = {}
         for i, img in enumerate(img_tuples):
@@ -921,6 +879,7 @@ class MultiChDloader:
         rot_dic = self._rotation_transform(
             image=img_tuples[0][0], **img_kwargs, **noise_kwargs
         )
+
         rotated_img_tuples = []
         for i, img in enumerate(img_tuples):
             if len(img) == 1:
@@ -946,7 +905,90 @@ class MultiChDloader:
 
         return rotated_img_tuples, rotated_noise_tuples
 
+    def _rotate(self, img_tuples, noise_tuples):
+        if self._depth3D > 1:
+            return self._rotate3D(img_tuples, noise_tuples)
+        else:
+            return self._rotate2D(img_tuples, noise_tuples)
+
+    def _rotate3D(self, img_tuples, noise_tuples):
+        img_kwargs = {}
+        for i, img in enumerate(img_tuples):
+            for j in range(self._depth3D):
+                for k in range(len(img)):
+                    img_kwargs[f"img{i}_{j}_{k}"] = img[k, j]
+
+        noise_kwargs = {}
+        for i, nimg in enumerate(noise_tuples):
+            for j in range(self._depth3D):
+                for k in range(len(nimg)):
+                    noise_kwargs[f"noise{i}_{j}_{k}"] = nimg[k, j]
+
+        keys = list(img_kwargs.keys()) + list(noise_kwargs.keys())
+        self._rotation_transform.add_targets({k: "image" for k in keys})
+        rot_dic = self._rotation_transform(
+            image=img_tuples[0][0], **img_kwargs, **noise_kwargs
+        )
+        rotated_img_tuples = []
+        for i, img in enumerate(img_tuples):
+            if len(img) == 1:
+                rotated_img_tuples.append(
+                    np.concatenate(
+                        [
+                            rot_dic[f"img{i}_{j}_0"][None, None]
+                            for j in range(self._depth3D)
+                        ],
+                        axis=1,
+                    )
+                )
+            else:
+                temp_arr = []
+                for k in range(len(img)):
+                    temp_arr.append(
+                        np.concatenate(
+                            [
+                                rot_dic[f"img{i}_{j}_{k}"][None, None]
+                                for j in range(self._depth3D)
+                            ],
+                            axis=1,
+                        )
+                    )
+                rotated_img_tuples.append(np.concatenate(temp_arr, axis=0))
+
+        rotated_noise_tuples = []
+        for i, nimg in enumerate(noise_tuples):
+            if len(nimg) == 1:
+                rotated_noise_tuples.append(
+                    np.concatenate(
+                        [
+                            rot_dic[f"noise{i}_{j}_0"][None, None]
+                            for j in range(self._depth3D)
+                        ],
+                        axis=1,
+                    )
+                )
+            else:
+                temp_arr = []
+                for k in range(len(nimg)):
+                    temp_arr.append(
+                        np.concatenate(
+                            [
+                                rot_dic[f"noise{i}_{j}_{k}"][None, None]
+                                for j in range(self._depth3D)
+                            ],
+                            axis=1,
+                        )
+                    )
+                rotated_noise_tuples.append(np.concatenate(temp_arr, axis=0))
+
+        return rotated_img_tuples, rotated_noise_tuples
+
     def get_uncorrelated_img_tuples(self, index):
+        """
+        Content of channels like actin and nuclei is "correlated" in its
+        respective location, this function allows to pick channels' content
+        from different patches of the image to make it "uncorrelated".
+        """
         img_tuples, noise_tuples = self._get_img(index)
         assert len(noise_tuples) == 0
         img_tuples = [img_tuples[0]]
@@ -959,6 +1001,8 @@ class MultiChDloader:
     def __getitem__(
         self, index: Union[int, Tuple[int, int]]
     ) -> Tuple[np.ndarray, np.ndarray]:
+        # Vera: input can be both real microscopic image and two separate channels that are summed in the code
+
         if self._train_index_switcher is not None:
             index = self._get_index_from_valid_target_logic(index)
 
@@ -971,22 +1015,29 @@ class MultiChDloader:
             self._empty_patch_replacement_enabled != True
         ), "This is not supported with noise"
 
+        # Replace the content of one of the channels
+        # with background with given probability
         if self._empty_patch_replacement_enabled:
             if np.random.rand() < self._empty_patch_replacement_probab:
                 img_tuples = self.replace_with_empty_patch(img_tuples)
 
+        # Noise tuples are not needed for the paper
+        # the image tuples are noisy by default
+        # TODO: remove noise tuples completely?
         if self._enable_rotation:
             img_tuples, noise_tuples = self._rotate(img_tuples, noise_tuples)
 
-        # add noise to input
+        # Add noise tuples with image tuples to create the input
         if len(noise_tuples) > 0:
             factor = np.sqrt(2) if self._input_is_sum else 1.0
             input_tuples = [x + noise_tuples[0] * factor for x in img_tuples]
         else:
             input_tuples = img_tuples
+
+        # Weight the individual channels, typically alpha is fixed
         inp, alpha = self._compute_input(input_tuples)
 
-        # add noise to target.
+        # Add noise tuples to the image tuples to create the target
         if len(noise_tuples) >= 1:
             img_tuples = [x + noise for x, noise in zip(img_tuples, noise_tuples[1:])]
 
@@ -1000,221 +1051,4 @@ class MultiChDloader:
         if self._return_index:
             output.append(index)
 
-        if isinstance(index, int) or isinstance(index, np.int64):
-            return tuple(output)
-
-        _, grid_size = index
-        output.append(grid_size)
-        return tuple(output)
-
-
-class LCMultiChDloader(MultiChDloader):
-
-    def __init__(
-        self,
-        data_config,
-        fpath: str,
-        datasplit_type: DataSplitType = None,
-        val_fraction=None,
-        test_fraction=None,
-        normalized_input=None,
-        enable_rotation_aug: bool = False,
-        use_one_mu_std=None,
-        num_scales: int = None,
-        enable_random_cropping=False,
-        padding_kwargs: dict = None,
-        allow_generation: bool = False,
-        lowres_supervision=None,
-        max_val=None,
-        grid_alignment=GridAlignement.LeftTop,
-        overlapping_padding_kwargs=None,
-        print_vars=True,
-    ):
-        """
-        Args:
-            num_scales: The number of resolutions at which we want the input. Note that the target is formed at the
-                        highest resolution.
-        """
-        self._padding_kwargs = (
-            padding_kwargs  # mode=padding_mode, constant_values=constant_value
-        )
-        if overlapping_padding_kwargs is not None:
-            assert (
-                self._padding_kwargs == 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:
-            overlapping_padding_kwargs = padding_kwargs
-
-        super().__init__(
-            data_config,
-            fpath,
-            datasplit_type=datasplit_type,
-            val_fraction=val_fraction,
-            test_fraction=test_fraction,
-            normalized_input=normalized_input,
-            enable_rotation_aug=enable_rotation_aug,
-            enable_random_cropping=enable_random_cropping,
-            use_one_mu_std=use_one_mu_std,
-            allow_generation=allow_generation,
-            max_val=max_val,
-            grid_alignment=grid_alignment,
-            overlapping_padding_kwargs=overlapping_padding_kwargs,
-            print_vars=print_vars,
-        )
-        self.num_scales = num_scales
-        assert self.num_scales is not None
-        self._scaled_data = [self._data]
-        self._scaled_noise_data = [self._noise_data]
-
-        assert isinstance(self.num_scales, int) and self.num_scales >= 1
-        self._lowres_supervision = lowres_supervision
-        assert isinstance(self._padding_kwargs, dict)
-        assert "mode" in self._padding_kwargs
-
-        for _ in range(1, self.num_scales):
-            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 _init_msg(self):
-        msg = super()._init_msg()
-        msg += f" Pad:{self._padding_kwargs}"
-        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
-        imgs = self._scaled_data[scaled_index][idx % self.N]
-        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][idx % self.N]
-            noise = tuple(
-                [noisedata[None, :, :, i] for i in range(noisedata.shape[-1])]
-            )
-            factor = np.sqrt(2) if self._input_is_sum else 1.0
-            # since we are using this lowres images for just the input, we need to add the noise of the input.
-            assert self._lowres_supervision is None or self._lowres_supervision is False
-            imgs = tuple([img + noise[0] * factor for img in imgs])
-        return imgs
-
-    def _crop_img(self, img: np.ndarray, h_start: int, w_start: int):
-        """
-        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
-        """
-        return self._crop_img_with_padding(img, h_start, w_start)
-
-    def _get_img(self, index: int):
-        """
-        Returns the primary patch along with low resolution patches centered on the primary patch.
-        """
-        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:
-            h_start, w_start = self._get_random_hw(h, w)
-        else:
-            h_start, w_start = self._get_deterministic_hw(index)
-
-        cropped_img_tuples = [
-            self._crop_flip_img(img, h_start, w_start, False, False)
-            for img in img_tuples
-        ]
-        cropped_noise_tuples = [
-            self._crop_flip_img(noise, h_start, w_start, False, False)
-            for noise in noise_tuples
-        ]
-        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.num_scales):
-            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
-
-            scaled_cropped_img_tuples = [
-                self._crop_flip_img(img, h_start, w_start, 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]]):
-        if self._uncorrelated_channels:
-            img_tuples, noise_tuples = self.get_uncorrelated_img_tuples(index)
-        else:
-            img_tuples, noise_tuples = self._get_img(index)
-
-        if self._enable_rotation:
-            img_tuples, noise_tuples = self._rotate(img_tuples, noise_tuples)
-
-        assert self._lowres_supervision != True
-        # add noise to input
-        if len(noise_tuples) > 0:
-            factor = np.sqrt(2) if self._input_is_sum else 1.0
-            input_tuples = []
-            for x in img_tuples:
-                # 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
-
-        inp, alpha = self._compute_input(input_tuples)
-        # assert self._alpha_weighted_target in [False, None]
-        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)
-
-        output = [inp, 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)