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

careamics/lvae_training/dataset/multifile_dataset.py

@@ -0,0 +1,334 @@
+from typing import Union, Callable, Sequence
+
+import numpy as np
+from numpy.typing import NDArray
+
+from .config import DatasetConfig
+from .multich_dataset import MultiChDloader
+from .types import DataSplitType
+from .lc_dataset import LCMultiChDloader
+
+
+class TwoChannelData(Sequence):
+    """
+    each element in data_arr should be a N*H*W array
+    """
+
+    def __init__(self, data_arr1, data_arr2, paths_data1=None, paths_data2=None):
+        assert len(data_arr1) == len(data_arr2)
+        self.paths1 = paths_data1
+        self.paths2 = paths_data2
+
+        self._data = []
+        for i in range(len(data_arr1)):
+            assert data_arr1[i].shape == data_arr2[i].shape
+            assert (
+                len(data_arr1[i].shape) == 3
+            ), f"Each element in data arrays should be a N*H*W, but {data_arr1[i].shape}"
+            self._data.append(
+                np.concatenate(
+                    [data_arr1[i][..., None], data_arr2[i][..., None]], axis=-1
+                )
+            )
+
+    def __len__(self):
+        n = 0
+        for x in self._data:
+            n += x.shape[0]
+        return n
+
+    def __getitem__(self, idx):
+        n = 0
+        for dataidx, x in enumerate(self._data):
+            if idx < n + x.shape[0]:
+                if self.paths1 is None:
+                    return x[idx - n], None
+                else:
+                    return x[idx - n], (self.paths1[dataidx], self.paths2[dataidx])
+            n += x.shape[0]
+        raise IndexError("Index out of range")
+
+
+class MultiChannelData(Sequence):
+    """
+    each element in data_arr should be a N*H*W array
+    """
+
+    def __init__(self, data_arr, paths=None):
+        self.paths = paths
+
+        self._data = data_arr
+
+    def __len__(self):
+        n = 0
+        for x in self._data:
+            n += x.shape[0]
+        return n
+
+    def __getitem__(self, idx):
+        n = 0
+        for dataidx, x in enumerate(self._data):
+            if idx < n + x.shape[0]:
+                if self.paths is None:
+                    return x[idx - n], None
+                else:
+                    return x[idx - n], (self.paths[dataidx])
+            n += x.shape[0]
+        raise IndexError("Index out of range")
+
+
+class SingleFileLCDset(LCMultiChDloader):
+    def __init__(
+        self,
+        preloaded_data: NDArray,
+        data_config: DatasetConfig,
+        fpath: str,
+        load_data_fn: Callable,
+        val_fraction=None,
+        test_fraction=None,
+    ):
+        self._preloaded_data = preloaded_data
+        super().__init__(
+            data_config,
+            fpath,
+            load_data_fn=load_data_fn,
+            val_fraction=val_fraction,
+            test_fraction=test_fraction,
+        )
+
+    @property
+    def data_path(self):
+        return self._fpath
+
+    def rm_bkground_set_max_val_and_upperclip_data(self, max_val, datasplit_type):
+        pass
+
+    def load_data(
+        self,
+        data_config: DatasetConfig,
+        datasplit_type: DataSplitType,
+        load_data_fn: Callable,
+        val_fraction=None,
+        test_fraction=None,
+        allow_generation=None,
+    ):
+        self._data = self._preloaded_data
+        assert "channel_1" not in data_config or isinstance(data_config.channel_1, str)
+        assert "channel_2" not in data_config or isinstance(data_config.channel_2, str)
+        assert "channel_3" not in data_config or isinstance(data_config.channel_3, str)
+        self._loaded_data_preprocessing(data_config)
+
+
+class SingleFileDset(MultiChDloader):
+    def __init__(
+        self,
+        preloaded_data: NDArray,
+        data_config: DatasetConfig,
+        fpath: str,
+        load_data_fn: Callable,
+        val_fraction=None,
+        test_fraction=None,
+    ):
+        self._preloaded_data = preloaded_data
+        super().__init__(
+            data_config,
+            fpath,
+            load_data_fn=load_data_fn,
+            val_fraction=val_fraction,
+            test_fraction=test_fraction,
+        )
+
+    def rm_bkground_set_max_val_and_upperclip_data(self, max_val, datasplit_type):
+        pass
+
+    @property
+    def data_path(self):
+        return self._fpath
+
+    def load_data(
+        self,
+        data_config: DatasetConfig,
+        datasplit_type: DataSplitType,
+        load_data_fn: Callable[..., NDArray],
+        val_fraction=None,
+        test_fraction=None,
+        allow_generation=None,
+    ):
+        self._data = self._preloaded_data
+        assert (
+            "channel_1" not in data_config
+        ), "Outdated config file. Please remove channel_1, channel_2, channel_3 from the config file."
+        assert (
+            "channel_2" not in data_config
+        ), "Outdated config file. Please remove channel_1, channel_2, channel_3 from the config file."
+        assert (
+            "channel_3" not in data_config
+        ), "Outdated config file. Please remove channel_1, channel_2, channel_3 from the config file."
+        self._loaded_data_preprocessing(data_config)
+
+
+class MultiFileDset:
+    """
+    Here, we handle dataset having multiple files. Each file can have a different spatial dimension and number of frames (Z stack).
+    """
+
+    def __init__(
+        self,
+        data_config: DatasetConfig,
+        fpath: str,
+        load_data_fn: Callable[..., Union[TwoChannelData, MultiChannelData]],
+        val_fraction=None,
+        test_fraction=None,
+    ):
+        self._fpath = fpath
+        data: Union[TwoChannelData, MultiChannelData] = load_data_fn(
+            data_config,
+            self._fpath,
+            data_config.datasplit_type,
+            val_fraction=val_fraction,
+            test_fraction=test_fraction,
+        )
+        self.dsets = []
+
+        for i in range(len(data)):
+            prefetched_data, fpath_tuple = data[i]
+            if (
+                data_config.multiscale_lowres_count is not None
+                and data_config.multiscale_lowres_count > 1
+            ):
+
+                self.dsets.append(
+                    SingleFileLCDset(
+                        prefetched_data[None],
+                        data_config,
+                        fpath_tuple,
+                        load_data_fn,
+                        val_fraction=val_fraction,
+                        test_fraction=test_fraction,
+                    )
+                )
+
+            else:
+                self.dsets.append(
+                    SingleFileDset(
+                        prefetched_data[None],
+                        data_config,
+                        fpath_tuple,
+                        load_data_fn,
+                        val_fraction=val_fraction,
+                        test_fraction=test_fraction,
+                    )
+                )
+
+        self.rm_bkground_set_max_val_and_upperclip_data(
+            data_config.max_val, data_config.datasplit_type
+        )
+        count = 0
+        avg_height = 0
+        avg_width = 0
+        for dset in self.dsets:
+            shape = dset.get_data_shape()
+            avg_height += shape[1]
+            avg_width += shape[2]
+            count += shape[0]
+
+        avg_height = int(avg_height / len(self.dsets))
+        avg_width = int(avg_width / len(self.dsets))
+        print(
+            f"{self.__class__.__name__} avg height: {avg_height}, avg width: {avg_width}, count: {count}"
+        )
+
+    def rm_bkground_set_max_val_and_upperclip_data(self, max_val, datasplit_type):
+        self.set_max_val(max_val, datasplit_type)
+        self.upperclip_data()
+
+    def set_mean_std(self, mean_val, std_val):
+        for dset in self.dsets:
+            dset.set_mean_std(mean_val, std_val)
+
+    def get_mean_std(self):
+        return self.dsets[0].get_mean_std()
+
+    def compute_max_val(self):
+        max_val_arr = []
+        for dset in self.dsets:
+            max_val_arr.append(dset.compute_max_val())
+        return np.max(max_val_arr)
+
+    def set_max_val(self, max_val, datasplit_type):
+        if datasplit_type == DataSplitType.Train:
+            assert max_val is None
+            max_val = self.compute_max_val()
+        for dset in self.dsets:
+            dset.set_max_val(max_val, datasplit_type)
+
+    def upperclip_data(self):
+        for dset in self.dsets:
+            dset.upperclip_data()
+
+    def get_max_val(self):
+        return self.dsets[0].get_max_val()
+
+    def get_img_sz(self):
+        return self.dsets[0].get_img_sz()
+
+    def set_img_sz(self, image_size, grid_size):
+        for dset in self.dsets:
+            dset.set_img_sz(image_size, grid_size)
+
+    def compute_mean_std(self):
+        cur_mean = {"target": 0, "input": 0}
+        cur_std = {"target": 0, "input": 0}
+        for dset in self.dsets:
+            mean, std = dset.compute_mean_std()
+            cur_mean["target"] += mean["target"]
+            cur_mean["input"] += mean["input"]
+
+            cur_std["target"] += std["target"]
+            cur_std["input"] += std["input"]
+
+        cur_mean["target"] /= len(self.dsets)
+        cur_mean["input"] /= len(self.dsets)
+        cur_std["target"] /= len(self.dsets)
+        cur_std["input"] /= len(self.dsets)
+        return cur_mean, cur_std
+
+    def compute_individual_mean_std(self):
+        cum_mean = 0
+        cum_std = 0
+        for dset in self.dsets:
+            mean, std = dset.compute_individual_mean_std()
+            cum_mean += mean
+            cum_std += std
+        return cum_mean / len(self.dsets), cum_std / len(self.dsets)
+
+    def get_num_frames(self):
+        return len(self.dsets)
+
+    def reduce_data(
+        self, t_list=None, h_start=None, h_end=None, w_start=None, w_end=None
+    ):
+        assert h_start is None
+        assert h_end is None
+        assert w_start is None
+        assert w_end is None
+        self.dsets = [self.dsets[t] for t in t_list]
+        print(
+            f"[{self.__class__.__name__}] Data reduced. New data count: {len(self.dsets)}"
+        )
+
+    def __len__(self):
+        out = 0
+        for dset in self.dsets:
+            out += len(dset)
+        return out
+
+    def __getitem__(self, idx):
+        cum_len = 0
+        for dset in self.dsets:
+            cum_len += len(dset)
+            if idx < cum_len:
+                rel_idx = idx - (cum_len - len(dset))
+                return dset[rel_idx]
+
+        raise IndexError("Index out of range")

careamics/lvae_training/dataset/types.py

@@ -0,0 +1,43 @@
+from enum import Enum
+
+
+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
+    NicolaData = 25
+
+
+class DataSplitType(Enum):
+    All = 0
+    Train = 1
+    Val = 2
+    Test = 3
+
+
+class TilingMode(Enum):
+    TrimBoundary = 0
+    PadBoundary = 1
+    ShiftBoundary = 2

careamics/lvae_training/dataset/utils/data_utils.py

@@ -0,0 +1,114 @@
+"""
+Utility functions needed by dataloader & co.
+"""
+
+from typing import List
+
+import numpy as np
+from skimage.io import imread, imsave
+
+
+def load_tiff(path):
+    """
+    Returns a 4d numpy array: num_imgs*h*w*num_channels
+    """
+    data = imread(path, plugin="tifffile")
+    return data
+
+
+def save_tiff(path, data):
+    imsave(path, data, plugin="tifffile")
+
+
+def load_tiffs(paths):
+    data = [load_tiff(path) for path in paths]
+    return np.concatenate(data, axis=0)
+
+
+def split_in_half(s, e):
+    n = e - s
+    s1 = list(np.arange(n // 2))
+    s2 = list(np.arange(n // 2, n))
+    return [x + s for x in s1], [x + s for x in s2]
+
+
+def adjust_for_imbalance_in_fraction_value(
+    val: List[int],
+    test: List[int],
+    val_fraction: float,
+    test_fraction: float,
+    total_size: int,
+):
+    """
+    here, val and test are divided almost equally. Here, we need to take into account their respective fractions
+    and pick elements rendomly from one array and put in the other array.
+    """
+    if val_fraction == 0:
+        test += val
+        val = []
+    elif test_fraction == 0:
+        val += test
+        test = []
+    else:
+        diff_fraction = test_fraction - val_fraction
+        if diff_fraction > 0:
+            imb_count = int(diff_fraction * total_size / 2)
+            val = list(np.random.RandomState(seed=955).permutation(val))
+            test += val[:imb_count]
+            val = val[imb_count:]
+        elif diff_fraction < 0:
+            imb_count = int(-1 * diff_fraction * total_size / 2)
+            test = list(np.random.RandomState(seed=955).permutation(test))
+            val += test[:imb_count]
+            test = test[imb_count:]
+    return val, test
+
+
+def get_datasplit_tuples(
+    val_fraction: float,
+    test_fraction: float,
+    total_size: int,
+    starting_test: bool = False,
+):
+    if starting_test:
+        # test => val => train
+        test = list(range(0, int(total_size * test_fraction)))
+        val = list(range(test[-1] + 1, test[-1] + 1 + int(total_size * val_fraction)))
+        train = list(range(val[-1] + 1, total_size))
+    else:
+        # {test,val}=> train
+        test_val_size = int((val_fraction + test_fraction) * total_size)
+        train = list(range(test_val_size, total_size))
+
+        if test_val_size == 0:
+            test = []
+            val = []
+            return train, val, test
+
+        # Split the test and validation in chunks.
+        chunksize = max(1, min(3, test_val_size // 2))
+
+        nchunks = test_val_size // chunksize
+
+        test = []
+        val = []
+        s = 0
+        for i in range(nchunks):
+            if i % 2 == 0:
+                val += list(np.arange(s, s + chunksize))
+            else:
+                test += list(np.arange(s, s + chunksize))
+            s += chunksize
+
+        if i % 2 == 0:
+            test += list(np.arange(s, test_val_size))
+        else:
+            p1, p2 = split_in_half(s, test_val_size)
+            test += p1
+            val += p2
+
+    val, test = adjust_for_imbalance_in_fraction_value(
+        val, test, val_fraction, test_fraction, total_size
+    )
+
+    return train, val, test

careamics/lvae_training/dataset/utils/empty_patch_fetcher.py

@@ -0,0 +1,65 @@
+import numpy as np
+from tqdm import tqdm
+
+
+class EmptyPatchFetcher:
+    """
+    The idea is to fetch empty patches so that real content can be replaced with this.
+    """
+
+    def __init__(self, idx_manager, patch_size, data_frames, max_val_threshold=None):
+        self._frames = data_frames
+        self._idx_manager = idx_manager
+        self._max_val_threshold = max_val_threshold
+        self._idx_list = []
+        self._patch_size = patch_size
+        self._grid_size = 1
+        self.set_empty_idx()
+
+        print(f"[{self.__class__.__name__}] MaxVal:{self._max_val_threshold}")
+
+    def compute_max(self, window):
+        """
+        Rolling compute.
+        """
+        N, H, W = self._frames.shape
+        randnum = -954321
+        assert self._grid_size == 1
+        max_data = np.zeros((N, H - window, W - window)) * randnum
+
+        for h in tqdm(range(H - window)):
+            for w in range(W - window):
+                max_data[:, h, w] = self._frames[:, h : h + window, w : w + window].max(
+                    axis=(1, 2)
+                )
+
+        assert (max_data != 954321).any()
+        return max_data
+
+    def set_empty_idx(self):
+        max_data = self.compute_max(self._patch_size)
+        empty_loc = np.where(
+            np.logical_and(max_data >= 0, max_data < self._max_val_threshold)
+        )
+        # print(max_data.shape, len(empty_loc))
+        self._idx_list = []
+        for idx in range(len(empty_loc[0])):
+            n_idx = empty_loc[0][idx]
+            h_start = empty_loc[1][idx]
+            w_start = empty_loc[2][idx]
+            # print(n_idx,h_start,w_start)
+            # channel_idx = self._idx_manager.get_location_from_dataset_idx(0)[-1]
+            loc = (n_idx, h_start, w_start, 0)
+            idx = self._idx_manager.get_dataset_idx_from_location(loc)
+            t, h, w, _ = self._idx_manager.get_location_from_dataset_idx(idx)
+            assert h == h_start, f"{h} != {h_start}"
+            assert w == w_start, f"{w} != {w_start}"
+            assert t == n_idx, f"{t} != {n_idx}"
+            self._idx_list.append(idx)
+
+        self._idx_list = np.array(self._idx_list)
+
+        assert len(self._idx_list) > 0
+
+    def sample(self):
+        return (np.random.choice(self._idx_list), self._grid_size)