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

careamics/lvae_training/dataset/{vae_data_config.py → config.py}

@@ -1,63 +1,13 @@
 from typing import Any, Optional
-from enum import Enum
-
-from pydantic import BaseModel, ConfigDict, computed_field
-
-
-# TODO: get rid of unnecessary enums
-class DataType(Enum):
-    MNIST = 0
-    Places365 = 1
-    NotMNIST = 2
-    OptiMEM100_014 = 3
-    CustomSinosoid = 4
-    Prevedel_EMBL = 5
-    AllenCellMito = 6
-    SeparateTiffData = 7
-    CustomSinosoidThreeCurve = 8
-    SemiSupBloodVesselsEMBL = 9
-    Pavia2 = 10
-    Pavia2VanillaSplitting = 11
-    ExpansionMicroscopyMitoTub = 12
-    ShroffMitoEr = 13
-    HTIba1Ki67 = 14
-    BSD68 = 15
-    BioSR_MRC = 16
-    TavernaSox2Golgi = 17
-    Dao3Channel = 18
-    ExpMicroscopyV2 = 19
-    Dao3ChannelWithInput = 20
-    TavernaSox2GolgiV2 = 21
-    TwoDset = 22
-    PredictedTiffData = 23
-    Pavia3SeqData = 24
-    # Here, we have 16 splitting tasks.
-    NicolaData = 25
-
-
-class DataSplitType(Enum):
-    All = 0
-    Train = 1
-    Val = 2
-    Test = 3
-
-
-class GridAlignement(Enum):
-    """
-    A patch is formed by padding the grid with content. If the grids are 'Center' aligned, then padding is to done equally on all 4 sides.
-    On the other hand, if grids are 'LeftTop' aligned, padding is to be done on the right and bottom end of the grid.
-    In the former case, one needs (patch_size - grid_size)//2 amount of content on the right end of the frame.
-    In the latter case, one needs patch_size - grid_size amount of content on the right end of the frame.
-    """
-
-    LeftTop = 0
-    Center = 1
-
-
-# TODO: for all bool params check if they are taking different values in Disentangle repo
+
+from pydantic import BaseModel, ConfigDict
+
+from .types import DataType, DataSplitType, TilingMode
+
+
 # TODO: check if any bool logic can be removed
-class VaeDatasetConfig(BaseModel):
-    model_config = ConfigDict(validate_assignment=True)
+class DatasetConfig(BaseModel):
+    model_config = ConfigDict(validate_assignment=True, extra="forbid")
 
     data_type: Optional[DataType]
     """Type of the dataset, should be one of DataType"""
@@ -132,15 +82,10 @@ class VaeDatasetConfig(BaseModel):
     # TODO: why is this not used?
     enable_rotation_aug: Optional[bool] = False
 
-    grid_alignment: GridAlignement = GridAlignement.LeftTop
-
     max_val: Optional[float] = None
     """Maximum data in the dataset. Is calculated for train split, and should be 
     externally set for val and test splits."""
 
-    trim_boundary: Optional[bool] = True
-    """Whether to trim boundary of the image"""
-
     overlapping_padding_kwargs: Any = None
     """Parameters for np.pad method"""
 
@@ -157,23 +102,22 @@ class VaeDatasetConfig(BaseModel):
     train_aug_rotate: Optional[bool] = False
     enable_random_cropping: Optional[bool] = True
 
-    # TODO: not used?
     multiscale_lowres_count: Optional[int] = None
+    """Number of LC scales"""
+
+    tiling_mode: Optional[TilingMode] = TilingMode.ShiftBoundary
+
+    target_separate_normalization: Optional[bool] = True
+
+    mode_3D: Optional[bool] = False
+    """If training in 3D mode or not"""
+
+    trainig_datausage_fraction: Optional[float] = 1.0
+
+    validtarget_random_fraction: Optional[float] = None
+
+    validation_datausage_fraction: Optional[float] = 1.0
+
+    random_flip_z_3D: Optional[bool] = False
 
-    @computed_field
-    @property
-    def padding_kwargs(self) -> dict:
-        kwargs_dict = {}
-        padding_kwargs = {}
-        if (
-            self.multiscale_lowres_count is not None
-            and self.multiscale_lowres_count is not None
-        ):
-            # Get padding attributes
-            if "padding_kwargs" not in kwargs_dict:
-                padding_kwargs = {}
-                padding_kwargs["mode"] = "constant"
-                padding_kwargs["constant_values"] = 0
-            else:
-                padding_kwargs = kwargs_dict.pop("padding_kwargs")
-        return padding_kwargs
+    padding_kwargs: Optional[dict] = None

careamics/lvae_training/dataset/lc_dataset.py

@@ -2,34 +2,37 @@
 A place for Datasets and Dataloaders.
 """
 
-from typing import Tuple, Union
+from typing import Tuple, Union, Callable
 
 import numpy as np
 from skimage.transform import resize
 
-from .lc_dataset_config import LCVaeDatasetConfig
-from .vae_dataset import MultiChDloader
+from .config import DatasetConfig
+from .multich_dataset import MultiChDloader
 
 
 class LCMultiChDloader(MultiChDloader):
-
     def __init__(
         self,
-        data_config: LCVaeDatasetConfig,
+        data_config: DatasetConfig,
         fpath: str,
+        load_data_fn: Callable,
         val_fraction=None,
         test_fraction=None,
     ):
-        """
-        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 = (
             data_config.padding_kwargs  # mode=padding_mode, constant_values=constant_value
         )
         self._uncorrelated_channel_probab = data_config.uncorrelated_channel_probab
 
+        super().__init__(
+            data_config,
+            fpath,
+            load_data_fn=load_data_fn,
+            val_fraction=val_fraction,
+            test_fraction=test_fraction,
+        )
+
         if data_config.overlapping_padding_kwargs is not None:
             assert (
                 self._padding_kwargs == data_config.overlapping_padding_kwargs
@@ -37,21 +40,21 @@ class LCMultiChDloader(MultiChDloader):
                 It should be so since we just use overlapping_padding_kwargs when it is not None"
 
         else:
-            overlapping_padding_kwargs = data_config.padding_kwargs
+            self._overlapping_padding_kwargs = data_config.padding_kwargs
 
-        super().__init__(
-            data_config, fpath, val_fraction=val_fraction, test_fraction=test_fraction
-        )
-        self.num_scales = data_config.num_scales
-        assert self.num_scales is not None
+        self.multiscale_lowres_count = data_config.multiscale_lowres_count
+        assert self.multiscale_lowres_count is not None
         self._scaled_data = [self._data]
         self._scaled_noise_data = [self._noise_data]
 
-        assert isinstance(self.num_scales, int) and self.num_scales >= 1
+        assert (
+            isinstance(self.multiscale_lowres_count, int)
+            and self.multiscale_lowres_count >= 1
+        )
         assert isinstance(self._padding_kwargs, dict)
         assert "mode" in self._padding_kwargs
 
-        for _ in range(1, self.num_scales):
+        for _ in range(1, self.multiscale_lowres_count):
             shape = self._scaled_data[-1].shape
             assert len(shape) == 4
             new_shape = (shape[0], shape[1] // 2, shape[2] // 2, shape[3])
@@ -173,7 +176,7 @@ class LCMultiChDloader(MultiChDloader):
         allres_versions = {
             i: [cropped_img_tuples[i]] for i in range(len(cropped_img_tuples))
         }
-        for scale_idx in range(1, self.num_scales):
+        for scale_idx in range(1, self.multiscale_lowres_count):
             # Returning the image of the lower resolution
             scaled_img_tuples = self._load_scaled_img(scale_idx, index)
 
@@ -227,6 +230,9 @@ class LCMultiChDloader(MultiChDloader):
             factor = np.sqrt(2) if self._input_is_sum else 1.0
             input_tuples = []
             for x in img_tuples:
+                x = (
+                    x.copy()
+                )  # to avoid changing the original image since it is later used for target
                 # NOTE: other LC levels already have noise added. So, we just need to add noise to the highest resolution.
                 x[0] = x[0] + noise_tuples[0] * factor
                 input_tuples.append(x)
@@ -246,7 +252,9 @@ class LCMultiChDloader(MultiChDloader):
 
         target = self._compute_target(target_tuples, alpha)
 
-        output = [inp, target]
+        norm_target = self.normalize_target(target)
+
+        output = [inp, norm_target]
 
         if self._return_alpha:
             output.append(alpha)

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

@@ -2,39 +2,39 @@
 A place for Datasets and Dataloaders.
 """
 
-from typing import Tuple, Union
+from typing import Tuple, Union, Callable
 
 import numpy as np
 
-from .data_utils import (
-    GridIndexManager,
-    IndexSwitcher,
-    get_train_val_data,
-)
-from .vae_data_config import VaeDatasetConfig, DataSplitType, GridAlignement
+from .utils.empty_patch_fetcher import EmptyPatchFetcher
+from .utils.index_manager import GridIndexManager
+from .utils.index_switcher import IndexSwitcher
+from .config import DatasetConfig
+from .types import DataSplitType, TilingMode
 
 
 class MultiChDloader:
     def __init__(
         self,
-        data_config: VaeDatasetConfig,
+        data_config: DatasetConfig,
         fpath: str,
+        load_data_fn: Callable,
         val_fraction: float = None,
         test_fraction: float = None,
     ):
         """ """
         self._data_type = data_config.data_type
         self._fpath = fpath
-        self._data = self.N = self._noise_data = None
+        self._data = self._noise_data = None
         self.Z = 1
-        self._trim_boundary = data_config.trim_boundary
-        # Hardcoded params, not included in the config file.
-
+        self._5Ddata = False
+        self._tiling_mode = data_config.tiling_mode
         # 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
+        self._mode_3D = data_config.mode_3D
         # NOTE: Input is the sum of the different channels. It is not the average of the different channels.
         self._input_is_sum = data_config.input_is_sum
         self._num_channels = data_config.num_channels
@@ -42,20 +42,21 @@ class MultiChDloader:
         self._tar_idx_list = data_config.target_idx_list
 
         if data_config.datasplit_type == DataSplitType.Train:
-            self._datausage_fraction = 1.0
+            self._datausage_fraction = data_config.trainig_datausage_fraction
             # 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_rand_fract = data_config.validtarget_random_fraction
             # 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 data_config.datasplit_type == DataSplitType.Val:
-            self._datausage_fraction = 1.0
+            self._datausage_fraction = data_config.validation_datausage_fraction
         else:
             self._datausage_fraction = 1.0
 
         self.load_data(
             data_config,
             data_config.datasplit_type,
+            load_data_fn=load_data_fn,
             val_fraction=val_fraction,
             test_fraction=test_fraction,
             allow_generation=data_config.allow_generation,
@@ -70,18 +71,8 @@ class MultiChDloader:
 
         self._background_values = None
 
-        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
-                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."
-        if self._trim_boundary:
+        if self._tiling_mode in [TilingMode.TrimBoundary, TilingMode.ShiftBoundary]:
             if (
                 self._overlapping_padding_kwargs is None
                 or data_config.multiscale_lowres_count is not None
@@ -144,7 +135,6 @@ class MultiChDloader:
             )
             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,
@@ -161,14 +151,18 @@ class MultiChDloader:
         self._mean = None
         self._std = None
         self._use_one_mu_std = data_config.use_one_mu_std
-        # Hardcoded
-        self._target_separate_normalization = True
+
+        self._target_separate_normalization = data_config.target_separate_normalization
 
         self._enable_rotation = data_config.enable_rotation_aug
+        flipz_3D = data_config.random_flip_z_3D
+        self._flipz_3D = flipz_3D and self._enable_rotation
+
         self._enable_random_cropping = data_config.enable_random_cropping
         self._uncorrelated_channels = (
             data_config.uncorrelated_channels and self._is_train
         )
+        self._uncorrelated_channel_probab = data_config.uncorrelated_channel_probab
         assert self._is_train or self._uncorrelated_channels is False
         assert (
             self._enable_random_cropping is True or self._uncorrelated_channels is False
@@ -177,9 +171,9 @@ class MultiChDloader:
 
         self._rotation_transform = None
         if self._enable_rotation:
-            raise NotImplementedError(
-                "Augmentation by means of rotation is not supported yet."
-            )
+            # TODO: fix this import
+            import albumentations as A
+
             self._rotation_transform = A.Compose([A.Flip(), A.RandomRotate90()])
 
         # TODO: remove print log messages
@@ -203,11 +197,12 @@ class MultiChDloader:
         self,
         data_config,
         datasplit_type,
+        load_data_fn: Callable,
         val_fraction=None,
         test_fraction=None,
         allow_generation=None,
     ):
-        self._data = get_train_val_data(
+        self._data = load_data_fn(
             data_config,
             self._fpath,
             datasplit_type,
@@ -215,7 +210,9 @@ class MultiChDloader:
             test_fraction=test_fraction,
             allow_generation=allow_generation,
         )
+        self._loaded_data_preprocessing(data_config)
 
+    def _loaded_data_preprocessing(self, data_config):
         old_shape = self._data.shape
         if self._datausage_fraction < 1.0:
             framepixelcount = np.prod(self._data.shape[1:3])
@@ -239,10 +236,7 @@ class MultiChDloader:
         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 = (
-                np.random.poisson(self._data / self._poisson_noise_factor)
-                * self._poisson_noise_factor
-            )
+            self._data = np.random.poisson(self._data / self._poisson_noise_factor)
 
         if data_config.enable_gaussian_noise:
             synthetic_scale = data_config.synthetic_gaussian_scale
@@ -257,7 +251,13 @@ class MultiChDloader:
                 self._noise_data[..., 0] = np.mean(self._noise_data[..., 1:], axis=-1)
         print(msg)
 
-        self._5Ddata = len(self._data.shape) == 5
+        if len(self._data.shape) == 5:
+            if self._mode_3D:
+                self._5Ddata = True
+            else:
+                assert self._depth3D == 1, "Depth3D must be 1 for 2D training"
+                self._data = self._data.reshape(-1, *self._data.shape[2:])
+
         if self._5Ddata:
             self.Z = self._data.shape[1]
 
@@ -373,18 +373,28 @@ class MultiChDloader:
             f"[{self.__class__.__name__}] Data reduced. New data shape: {self._data.shape}"
         )
 
-    def get_idx_manager_shapes(self, patch_size: int, grid_size: int):
+    def get_idx_manager_shapes(
+        self, patch_size: int, grid_size: Union[int, Tuple[int, int, 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)
+            if isinstance(grid_size, int):
+                grid_shape = (1, 1, grid_size, grid_size, numC)
+            else:
+                assert len(grid_size) == 3
+                assert all(
+                    [g <= p for g, p in zip(grid_size, patch_shape[1:-1])]
+                ), f"Grid size {grid_size} must be less than patch size {patch_shape[1:-1]}"
+                grid_shape = (1, grid_size[0], grid_size[1], grid_size[2], numC)
         else:
+            assert isinstance(grid_size, int)
             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):
+    def set_img_sz(self, image_size, grid_size: Union[int, Tuple[int, int, int]]):
         """
         If one wants to change the image size on the go, then this can be used.
         Args:
@@ -400,7 +410,7 @@ class MultiChDloader:
             self._img_sz, self._grid_sz
         )
         self.idx_manager = GridIndexManager(
-            shape, grid_shape, patch_shape, self._trim_boundary
+            shape, grid_shape, patch_shape, self._tiling_mode
         )
         # self.set_repeat_factor()
 
@@ -432,10 +442,13 @@ class MultiChDloader:
         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" TrimB:{self._tiling_mode}"
         # msg += f' NormInp:{self._normalized_input}'
         # msg += f' SingleNorm:{self._use_one_mu_std}'
         msg += f" Rot:{self._enable_rotation}"
+        if self._flipz_3D:
+            msg += f" FlipZ:{self._flipz_3D}"
+
         msg += f" RandCrop:{self._enable_random_cropping}"
         msg += f" Channel:{self._num_channels}"
         # msg += f' Q:{self._quantile}'
@@ -467,7 +480,7 @@ class MultiChDloader:
             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),
+                    np.random.choice(1 + img_tuples[0].shape[-3] - self._depth3D),
                 ) + patch_start_loc
         else:
             patch_start_loc = self._get_deterministic_loc(index)
@@ -486,7 +499,7 @@ class MultiChDloader:
         )
 
     def _crop_img(self, img: np.ndarray, patch_start_loc: Tuple):
-        if self._trim_boundary:
+        if self._tiling_mode in [TilingMode.TrimBoundary, TilingMode.ShiftBoundary]:
             # 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.
@@ -625,6 +638,18 @@ class MultiChDloader:
             normalized_imgs.append(img)
         return tuple(normalized_imgs)
 
+    def normalize_input(self, x):
+        mean_dict, std_dict = self.get_mean_std()
+        mean_ = mean_dict["input"].mean()
+        std_ = std_dict["input"].mean()
+        return (x - mean_) / std_
+
+    def normalize_target(self, target):
+        mean_dict, std_dict = self.get_mean_std()
+        mean_ = mean_dict["target"].squeeze(0)
+        std_ = std_dict["target"].squeeze(0)
+        return (target - mean_) / std_
+
     def get_grid_size(self):
         return self._grid_sz
 
@@ -906,28 +931,39 @@ class MultiChDloader:
         return rotated_img_tuples, rotated_noise_tuples
 
     def _rotate(self, img_tuples, noise_tuples):
-        if self._depth3D > 1:
+
+        if self._5Ddata:
             return self._rotate3D(img_tuples, noise_tuples)
         else:
             return self._rotate2D(img_tuples, noise_tuples)
 
     def _rotate3D(self, img_tuples, noise_tuples):
         img_kwargs = {}
+        # random flip in z direction
+        flip_z = self._flipz_3D and np.random.rand() < 0.5
         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]
+                    if flip_z:
+                        z_idx = self._depth3D - 1 - j
+                    else:
+                        z_idx = j
+                    img_kwargs[f"img{i}_{z_idx}_{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]
+                    if flip_z:
+                        z_idx = self._depth3D - 1 - j
+                    else:
+                        z_idx = j
+                    noise_kwargs[f"noise{i}_{z_idx}_{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
+            image=img_tuples[0][0][0], **img_kwargs, **noise_kwargs
         )
         rotated_img_tuples = []
         for i, img in enumerate(img_tuples):
@@ -1006,7 +1042,10 @@ class MultiChDloader:
         if self._train_index_switcher is not None:
             index = self._get_index_from_valid_target_logic(index)
 
-        if self._uncorrelated_channels:
+        if (
+            self._uncorrelated_channels
+            and np.random.rand() < self._uncorrelated_channel_probab
+        ):
             img_tuples, noise_tuples = self.get_uncorrelated_img_tuples(index)
         else:
             img_tuples, noise_tuples = self._get_img(index)
@@ -1042,8 +1081,9 @@ class MultiChDloader:
             img_tuples = [x + noise for x, noise in zip(img_tuples, noise_tuples[1:])]
 
         target = self._compute_target(img_tuples, alpha)
+        norm_target = self.normalize_target(target)
 
-        output = [inp, target]
+        output = [inp, norm_target]
 
         if self._return_alpha:
             output.append(alpha)