careamics 0.1.0rc6__py3-none-any.whl → 0.1.0rc8__py3-none-any.whl

careamics/lvae_training/data_modules.py

@@ -0,0 +1,1220 @@
+"""
+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,
+)
+
+
+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)
+
+
+class MultiChDloader:
+
+    def __init__(
+        self,
+        data_config: ml_collections.ConfigDict,
+        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
+
+        # 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._train_index_switcher = None
+        # 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)
+
+        if 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:
+            self._datausage_fraction = 1.0
+        else:
+            self._datausage_fraction = 1.0
+
+        self.load_data(
+            data_config,
+            datasplit_type,
+            val_fraction=val_fraction,
+            test_fraction=test_fraction,
+            allow_generation=allow_generation,
+        )
+        self._normalized_input = normalized_input
+        self._quantile = 1.0
+        self._channelwise_quantile = False
+        self._background_quantile = 0.0
+        self._clip_background_noise_to_zero = False
+        self._skip_normalization_using_mean = False
+        self._empty_patch_replacement_enabled = False
+
+        self._background_values = None
+
+        self._grid_alignment = grid_alignment
+        self._overlapping_padding_kwargs = overlapping_padding_kwargs
+        if self._grid_alignment == GridAlignement.LeftTop:
+            assert (
+                self._overlapping_padding_kwargs is None
+                or data_config.multiscale_lowres_count is not None
+            ), "Padding is not used with this alignement style"
+        elif self._grid_alignment == GridAlignement.Center:
+            assert (
+                self._overlapping_padding_kwargs is not None
+            ), "With Center grid alignment, padding is needed."
+
+        self._is_train = 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._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.set_img_sz(
+                data_config.image_size,
+                (
+                    data_config.grid_size
+                    if "grid_size" in data_config
+                    else data_config.image_size
+                ),
+            )
+
+            # 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
+
+        else:
+            self.set_img_sz(
+                data_config.image_size,
+                (
+                    data_config.grid_size
+                    if "grid_size" in data_config
+                    else data_config.image_size
+                ),
+            )
+
+        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.rm_bkground_set_max_val_and_upperclip_data(max_val, 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
+        # Hardcoded
+        self._target_separate_normalization = True
+
+        self._enable_rotation = enable_rotation_aug
+        self._enable_random_cropping = enable_random_cropping
+        self._uncorrelated_channels = (
+            data_config.get("uncorrelated_channels", False) and self._is_train
+        )
+        assert self._is_train or self._uncorrelated_channels is False
+        assert (
+            self._enable_random_cropping is True or self._uncorrelated_channels is False
+        )
+        # Randomly rotate [-90,90]
+
+        self._rotation_transform = None
+        if self._enable_rotation:
+            raise NotImplementedError(
+                "Augmentation by means of rotation is not supported yet."
+            )
+            self._rotation_transform = A.Compose([A.Flip(), A.RandomRotate90()])
+
+        if print_vars:
+            msg = self._init_msg()
+            print(msg)
+
+    def disable_noise(self):
+        assert (
+            self._poisson_noise_factor is None
+        ), "This is not supported. Poisson noise is added to the data itself and so the noise cannot be disabled."
+        self._disable_noise = True
+
+    def enable_noise(self):
+        self._disable_noise = False
+
+    def get_data_shape(self):
+        return self._data.shape
+
+    def load_data(
+        self,
+        data_config,
+        datasplit_type,
+        val_fraction=None,
+        test_fraction=None,
+        allow_generation=None,
+    ):
+        self._data = get_train_val_data(
+            data_config,
+            self._fpath,
+            datasplit_type,
+            val_fraction=val_fraction,
+            test_fraction=test_fraction,
+            allow_generation=allow_generation,
+        )
+
+        old_shape = self._data.shape
+        if self._datausage_fraction < 1.0:
+            framepixelcount = np.prod(self._data.shape[1:3])
+            pixelcount = int(
+                len(self._data) * framepixelcount * self._datausage_fraction
+            )
+            frame_count = int(np.ceil(pixelcount / framepixelcount))
+            last_frame_reduced_size, _ = IndexSwitcher.get_reduced_frame_size(
+                self._data.shape[:3], self._datausage_fraction
+            )
+            self._data = self._data[:frame_count].copy()
+            if frame_count == 1:
+                self._data = self._data[
+                    :, :last_frame_reduced_size, :last_frame_reduced_size
+                ].copy()
+            print(
+                f"[{self.__class__.__name__}] New data shape: {self._data.shape} Old: {old_shape}"
+            )
+
+        msg = ""
+        if data_config.get("poisson_noise_factor", -1) > 0:
+            self._poisson_noise_factor = data_config.poisson_noise_factor
+            msg += f"Adding Poisson noise with factor {self._poisson_noise_factor}.\t"
+            self._data = (
+                np.random.poisson(self._data / self._poisson_noise_factor)
+                * self._poisson_noise_factor
+            )
+
+        if data_config.get("enable_gaussian_noise", False):
+            synthetic_scale = data_config.get("synthetic_gaussian_scale", 0.1)
+            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):
+                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)
+        assert (
+            self._data.shape[-1] == self._num_channels
+        ), "Number of channels in data and config do not match."
+
+    def save_background(self, channel_idx, frame_idx, background_value):
+        self._background_values[frame_idx, channel_idx] = background_value
+
+    def get_background(self, channel_idx, frame_idx):
+        return self._background_values[frame_idx, channel_idx]
+
+    def remove_background(self):
+
+        self._background_values = np.zeros((self._data.shape[0], self._data.shape[-1]))
+
+        if self._background_quantile == 0.0:
+            assert (
+                self._clip_background_noise_to_zero is False
+            ), "This operation currently happens later in this function."
+            return
+
+        if self._data.dtype in [np.uint16]:
+            # unsigned integer creates havoc
+            self._data = self._data.astype(np.int32)
+
+        for ch in range(self._data.shape[-1]):
+            for idx in range(self._data.shape[0]):
+                qval = np.quantile(self._data[idx, ..., ch], self._background_quantile)
+                assert (
+                    np.abs(qval) > 20
+                ), "We are truncating the qval to an integer which will only make sense if it is large enough"
+                # NOTE: Here, there can be an issue if you work with normalized data
+                qval = int(qval)
+                self.save_background(ch, idx, qval)
+                self._data[idx, ..., ch] -= qval
+
+        if self._clip_background_noise_to_zero:
+            self._data[self._data < 0] = 0
+
+    def rm_bkground_set_max_val_and_upperclip_data(self, max_val, datasplit_type):
+        self.remove_background()
+        self.set_max_val(max_val, datasplit_type)
+        self.upperclip_data()
+
+    def upperclip_data(self):
+        if isinstance(self.max_val, list):
+            chN = self._data.shape[-1]
+            assert chN == len(self.max_val)
+            for ch in range(chN):
+                ch_data = self._data[..., ch]
+                ch_q = self.max_val[ch]
+                ch_data[ch_data > ch_q] = ch_q
+                self._data[..., ch] = ch_data
+        else:
+            self._data[self._data > self.max_val] = self.max_val
+
+    def compute_max_val(self):
+        if self._channelwise_quantile:
+            max_val_arr = [
+                np.quantile(self._data[..., i], self._quantile)
+                for i in range(self._data.shape[-1])
+            ]
+            return max_val_arr
+        else:
+            return np.quantile(self._data, self._quantile)
+
+    def set_max_val(self, max_val, datasplit_type):
+
+        if max_val is None:
+            assert datasplit_type == DataSplitType.Train
+            self.max_val = self.compute_max_val()
+        else:
+            assert max_val is not None
+            self.max_val = max_val
+
+    def get_max_val(self):
+        return self.max_val
+
+    def get_img_sz(self):
+        return self._img_sz
+
+    def reduce_data(
+        self, t_list=None, h_start=None, h_end=None, w_start=None, w_end=None
+    ):
+        if t_list is None:
+            t_list = list(range(self._data.shape[0]))
+        if h_start is None:
+            h_start = 0
+        if h_end is None:
+            h_end = self._data.shape[1]
+        if w_start is None:
+            w_start = 0
+        if w_end is None:
+            w_end = self._data.shape[2]
+
+        self._data = self._data[t_list, h_start:h_end, w_start:w_end, :].copy()
+        if self._noise_data is not None:
+            self._noise_data = self._noise_data[
+                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 set_img_sz(self, image_size, grid_size):
+        """
+        If one wants to change the image size on the go, then this can be used.
+        Args:
+            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
+        self.idx_manager = GridIndexManager(
+            self._data.shape, self._grid_sz, self._img_sz, self._grid_alignment
+        )
+        self.set_repeat_factor()
+
+    def set_repeat_factor(self):
+        if self._grid_sz > 1:
+            self._repeat_factor = self.idx_manager.grid_rows(
+                self._grid_sz
+            ) * self.idx_manager.grid_cols(self._grid_sz)
+        else:
+            self._repeat_factor = self.idx_manager.grid_rows(
+                self._img_sz
+            ) * self.idx_manager.grid_cols(self._img_sz)
+
+    def _init_msg(
+        self,
+    ):
+        msg = (
+            f"[{self.__class__.__name__}] Train:{int(self._is_train)} Sz:{self._img_sz}"
+        )
+        msg += f" N:{self.N} NumPatchPerN:{self._repeat_factor}"
+        # msg += f' NormInp:{self._normalized_input}'
+        # msg += f' SingleNorm:{self._use_one_mu_std}'
+        msg += f" Rot:{self._enable_rotation}"
+        msg += f" RandCrop:{self._enable_random_cropping}"
+        msg += f" Channel:{self._num_channels}"
+        # msg += f' Q:{self._quantile}'
+        if self._input_is_sum:
+            msg += f" SummedInput:{self._input_is_sum}"
+
+        if self._empty_patch_replacement_enabled:
+            msg += f" ReplaceWithRandSample:{self._empty_patch_replacement_enabled}"
+        if self._uncorrelated_channels:
+            msg += f" Uncorr:{self._uncorrelated_channels}"
+        if self._empty_patch_replacement_enabled:
+            msg += f"-{self._empty_patch_replacement_channel_idx}-{self._empty_patch_replacement_probab}"
+        if self._background_quantile > 0.0:
+            msg += f" BckQ:{self._background_quantile}"
+
+        if self._start_alpha_arr is not None:
+            msg += f" Alpha:[{self._start_alpha_arr},{self._end_alpha_arr}]"
+        return msg
+
+    def _crop_imgs(self, index, *img_tuples: np.ndarray):
+        h, w = img_tuples[0].shape[-2:]
+        if self._img_sz is None:
+            return (
+                *img_tuples,
+                {"h": [0, h], "w": [0, w], "hflip": False, "wflip": False},
+            )
+
+        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_imgs = []
+        for img in img_tuples:
+            img = self._crop_flip_img(img, h_start, w_start, 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:
+            # 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
+            ]
+            return new_img
+        elif self._grid_alignment == GridAlignement.Center:
+            # 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)
+
+    def get_begin_end_padding(self, start_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.
+        """
+        pad_start = 0
+        pad_end = 0
+        if start_pos < 0:
+            pad_start = -1 * start_pos
+
+        pad_end = max(0, start_pos + self._img_sz - 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
+        # 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 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
+    ):
+        new_img = self._crop_img(img, h_start, w_start)
+        if h_flip:
+            new_img = new_img[..., ::-1, :]
+        if w_flip:
+            new_img = new_img[..., :, ::-1]
+
+        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]:
+        """
+        Returns the channels and also the respective noise channels.
+        """
+        if isinstance(index, int) or isinstance(index, np.int64):
+            idx = index
+        else:
+            idx = index[0]
+
+        imgs = self._data[self.idx_manager.get_t(idx)]
+        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]
+                for i in range(self._noise_data.shape[-1])
+            ]
+        return tuple(loaded_imgs), tuple(noise)
+
+    def get_mean_std(self):
+        return self._mean, self._std
+
+    def set_mean_std(self, mean_val, std_val):
+        self._mean = mean_val
+        self._std = std_val
+
+    def normalize_img(self, *img_tuples):
+        mean, std = self.get_mean_std()
+        mean = mean["target"]
+        std = std["target"]
+        mean = mean.squeeze()
+        std = std.squeeze()
+        normalized_imgs = []
+        for i, img in enumerate(img_tuples):
+            img = (img - mean[i]) / std[i]
+            normalized_imgs.append(img)
+        return tuple(normalized_imgs)
+
+    def get_grid_size(self):
+        return self._grid_sz
+
+    def get_idx_manager(self):
+        return self.idx_manager
+
+    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 _get_deterministic_hw(self, index: Union[int, Tuple[int, 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
+
+    def compute_individual_mean_std(self):
+        # numpy 1.19.2 has issues in computing for large arrays. https://github.com/numpy/numpy/issues/8869
+        # mean = np.mean(self._data, axis=(0, 1, 2))
+        # std = np.std(self._data, axis=(0, 1, 2))
+        mean_arr = []
+        std_arr = []
+        for ch_idx in range(self._data.shape[-1]):
+            mean_ = (
+                0.0
+                if self._skip_normalization_using_mean
+                else self._data[..., ch_idx].mean()
+            )
+            if self._noise_data is not None:
+                std_ = (
+                    self._data[..., ch_idx] + self._noise_data[..., ch_idx + 1]
+                ).std()
+            else:
+                std_ = self._data[..., ch_idx].std()
+
+            mean_arr.append(mean_)
+            std_arr.append(std_)
+
+        mean = np.array(mean_arr)
+        std = np.array(std_arr)
+
+        return mean[None, :, None, None], std[None, :, None, None]
+
+    def compute_mean_std(self, allow_for_validation_data=False):
+        """
+        Note that we must compute this only for training data.
+        """
+        assert (
+            self._is_train is True or allow_for_validation_data
+        ), "This is just allowed for training data"
+        assert self._use_one_mu_std is True, "This is the only supported case"
+
+        if self._input_idx is not None:
+            assert (
+                self._tar_idx_list is not None
+            ), "tar_idx_list must be set if input_idx is set."
+            assert self._noise_data is None, "This is not supported with noise"
+            assert (
+                self._target_separate_normalization is True
+            ), "This is not supported with target_separate_normalization=False"
+
+            mean, std = self.compute_individual_mean_std()
+            mean_dict = {
+                "input": mean[:, self._input_idx : self._input_idx + 1],
+                "target": mean[:, self._tar_idx_list],
+            }
+            std_dict = {
+                "input": std[:, self._input_idx : self._input_idx + 1],
+                "target": std[:, self._tar_idx_list],
+            }
+            return mean_dict, std_dict
+
+        if self._input_is_sum:
+            assert self._noise_data is None, "This is not supported with noise"
+            mean = [
+                np.mean(self._data[..., k : k + 1], keepdims=True)
+                for k in range(self._num_channels)
+            ]
+            mean = np.sum(mean, keepdims=True)[0]
+            std = np.linalg.norm(
+                [
+                    np.std(self._data[..., k : k + 1], keepdims=True)
+                    for k in range(self._num_channels)
+                ],
+                keepdims=True,
+            )[0]
+        else:
+            mean = np.mean(self._data, keepdims=True).reshape(1, 1, 1, 1)
+            if self._noise_data is not None:
+                std = np.std(
+                    self._data + self._noise_data[..., 1:], keepdims=True
+                ).reshape(1, 1, 1, 1)
+            else:
+                std = np.std(self._data, keepdims=True).reshape(1, 1, 1, 1)
+
+        mean = np.repeat(mean, self._num_channels, axis=1)
+        std = np.repeat(std, self._num_channels, axis=1)
+
+        if self._skip_normalization_using_mean:
+            mean = np.zeros_like(mean)
+
+        mean_dict = {"input": mean}  # , 'target':mean}
+        std_dict = {"input": std}  # , 'target':std}
+
+        if self._target_separate_normalization:
+            mean, std = self.compute_individual_mean_std()
+
+        mean_dict["target"] = mean
+        std_dict["target"] = std
+        return mean_dict, std_dict
+
+    def _get_random_hw(self, h: int, w: int):
+        """
+        Random starting position for the crop for the img with index `index`.
+        """
+        if h != self._img_sz:
+            h_start = np.random.choice(h - self._img_sz)
+            w_start = np.random.choice(w - self._img_sz)
+        else:
+            h_start = 0
+            w_start = 0
+        return h_start, w_start
+
+    def _get_img(self, index: Union[int, Tuple[int, int]]):
+        """
+        Loads an image.
+        Crops the image such that cropped image has content.
+        """
+        img_tuples, noise_tuples = self._load_img(index)
+        cropped_img_tuples = self._crop_imgs(index, *img_tuples, *noise_tuples)[:-1]
+        cropped_noise_tuples = cropped_img_tuples[len(img_tuples) :]
+        cropped_img_tuples = cropped_img_tuples[: len(img_tuples)]
+        return cropped_img_tuples, cropped_noise_tuples
+
+    def replace_with_empty_patch(self, img_tuples):
+        empty_index = self._empty_patch_fetcher.sample()
+        empty_img_tuples = self._get_img(empty_index)
+        final_img_tuples = []
+        for tuple_idx in range(len(img_tuples)):
+            if tuple_idx == self._empty_patch_replacement_channel_idx:
+                final_img_tuples.append(empty_img_tuples[tuple_idx])
+            else:
+                final_img_tuples.append(img_tuples[tuple_idx])
+        return tuple(final_img_tuples)
+
+    def get_mean_std_for_input(self):
+        mean, std = self.get_mean_std()
+        return mean["input"], std["input"]
+
+    def _compute_target(self, img_tuples, alpha):
+        if self._tar_idx_list is not None and isinstance(self._tar_idx_list, int):
+            target = img_tuples[self._tar_idx_list]
+        else:
+            if self._tar_idx_list is not None:
+                assert isinstance(self._tar_idx_list, list) or isinstance(
+                    self._tar_idx_list, tuple
+                )
+                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)
+        return target
+
+    def _compute_input_with_alpha(self, img_tuples, alpha_list):
+        # assert self._normalized_input is True, "normalization should happen here"
+        if self._input_idx is not None:
+            inp = img_tuples[self._input_idx]
+        else:
+            inp = 0
+            for alpha, img in zip(alpha_list, img_tuples):
+                inp += img * alpha
+
+            if self._normalized_input is False:
+                return inp.astype(np.float32)
+
+        mean, std = self.get_mean_std_for_input()
+        mean = mean.squeeze()
+        std = std.squeeze()
+        if mean.size == 1:
+            mean = mean.reshape(
+                1,
+            )
+            std = std.reshape(
+                1,
+            )
+
+        for i in range(len(mean)):
+            assert mean[0] == mean[i]
+            assert std[0] == std[i]
+
+        inp = (inp - mean[0]) / std[0]
+        return inp.astype(np.float32)
+
+    def _sample_alpha(self):
+        alpha_arr = []
+        for i in range(self._num_channels):
+            alpha_pos = np.random.rand()
+            alpha = self._start_alpha_arr[i] + alpha_pos * (
+                self._end_alpha_arr[i] - self._start_alpha_arr[i]
+            )
+            alpha_arr.append(alpha)
+        return alpha_arr
+
+    def _compute_input(self, img_tuples):
+        alpha = [1 / len(img_tuples) for _ in range(len(img_tuples))]
+        if self._start_alpha_arr is not None:
+            alpha = self._sample_alpha()
+
+        inp = self._compute_input_with_alpha(img_tuples, alpha)
+        if self._input_is_sum:
+            inp = len(img_tuples) * inp
+        return inp, alpha
+
+    def _get_index_from_valid_target_logic(self, index):
+        if self._validtarget_rand_fract is not None:
+            if np.random.rand() < self._validtarget_rand_fract:
+                index = self._train_index_switcher.get_valid_target_index()
+            else:
+                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):
+            for k in range(len(img)):
+                img_kwargs[f"img{i}_{k}"] = img[k]
+
+        noise_kwargs = {}
+        for i, nimg in enumerate(noise_tuples):
+            for k in range(len(nimg)):
+                noise_kwargs[f"noise{i}_{k}"] = nimg[k]
+
+        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(rot_dic[f"img{i}_0"][None])
+            else:
+                rotated_img_tuples.append(
+                    np.concatenate(
+                        [rot_dic[f"img{i}_{k}"][None] for k in range(len(img))], axis=0
+                    )
+                )
+
+        rotated_noise_tuples = []
+        for i, nimg in enumerate(noise_tuples):
+            if len(nimg) == 1:
+                rotated_noise_tuples.append(rot_dic[f"noise{i}_0"][None])
+            else:
+                rotated_noise_tuples.append(
+                    np.concatenate(
+                        [rot_dic[f"noise{i}_{k}"][None] for k in range(len(nimg))],
+                        axis=0,
+                    )
+                )
+
+        return rotated_img_tuples, rotated_noise_tuples
+
+    def get_uncorrelated_img_tuples(self, index):
+        img_tuples, noise_tuples = self._get_img(index)
+        assert len(noise_tuples) == 0
+        img_tuples = [img_tuples[0]]
+        for ch_idx in range(1, len(img_tuples)):
+            new_index = np.random.randint(len(self))
+            other_img_tuples, _ = self._get_img(new_index)
+            img_tuples.append(other_img_tuples[ch_idx])
+        return img_tuples, noise_tuples
+
+    def __getitem__(
+        self, index: Union[int, Tuple[int, int]]
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        if self._train_index_switcher is not None:
+            index = self._get_index_from_valid_target_logic(index)
+
+        if self._uncorrelated_channels:
+            img_tuples, noise_tuples = self.get_uncorrelated_img_tuples(index)
+        else:
+            img_tuples, noise_tuples = self._get_img(index)
+
+        assert (
+            self._empty_patch_replacement_enabled != True
+        ), "This is not supported with noise"
+
+        if self._empty_patch_replacement_enabled:
+            if np.random.rand() < self._empty_patch_replacement_probab:
+                img_tuples = self.replace_with_empty_patch(img_tuples)
+
+        if self._enable_rotation:
+            img_tuples, noise_tuples = self._rotate(img_tuples, noise_tuples)
+
+        # add noise to input
+        if 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
+        inp, alpha = self._compute_input(input_tuples)
+
+        # add noise to target.
+        if len(noise_tuples) >= 1:
+            img_tuples = [x + noise for x, noise in zip(img_tuples, noise_tuples[1:])]
+
+        target = self._compute_target(img_tuples, alpha)
+
+        output = [inp, target]
+
+        if self._return_alpha:
+            output.append(alpha)
+
+        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)