careamics 0.0.19__py3-none-any.whl

careamics/lvae_training/dataset/multicrop_dset.py

@@ -0,0 +1,196 @@
+"""
+Here, we have multiple folders, each containing images of a single channel.
+"""
+
+from collections import defaultdict
+from functools import cache
+
+import numpy as np
+
+from .types import DataSplitType
+
+
+def l2(x):
+    return np.sqrt(np.mean(np.array(x) ** 2))
+
+
+class MultiCropDset:
+    def __init__(
+        self,
+        data_config,
+        fpath: str,
+        load_data_fn=None,
+        val_fraction=None,
+        test_fraction=None,
+    ):
+
+        assert (
+            data_config.input_is_sum == True
+        ), "This dataset is designed for sum of images"
+
+        self._img_sz = data_config.image_size
+        self._enable_rotation = data_config.enable_rotation_aug
+
+        self._background_values = data_config.background_values
+        self._data = load_data_fn(
+            data_config, fpath, data_config.datasplit_type, val_fraction, test_fraction
+        )
+
+        # remove upper quantiles, crucial for removing puncta
+        self.max_val = data_config.max_val
+        if self.max_val is not None:
+            for ch_idx, data in enumerate(self._data):
+                if self.max_val[ch_idx] is not None:
+                    for idx in range(len(data)):
+                        data[idx][data[idx] > self.max_val[ch_idx]] = self.max_val[
+                            ch_idx
+                        ]
+
+        # remove background values
+        if self._background_values is not None:
+            final_data_arr = []
+            for ch_idx, data in enumerate(self._data):
+                data_float = [x.astype(np.float32) for x in data]
+                final_data_arr.append(
+                    [x - self._background_values[ch_idx] for x in data_float]
+                )
+            self._data = final_data_arr
+
+        print(
+            f"{self.__class__.__name__} N:{len(self)} Rot:{self._enable_rotation} Ch:{len(self._data)} MaxVal:{self.max_val} Bg:{self._background_values}"
+        )
+
+    def get_max_val(self):
+        return self.max_val
+
+    def compute_mean_std(self):
+        mean_tar_dict = defaultdict(list)
+        std_tar_dict = defaultdict(list)
+        mean_inp = []
+        std_inp = []
+        for _ in range(30000):
+            crops = []
+            for ch_idx in range(len(self._data)):
+                crop = self.sample_crop(ch_idx)
+                mean_tar_dict[ch_idx].append(np.mean(crop))
+                std_tar_dict[ch_idx].append(np.std(crop))
+                crops.append(crop)
+
+            inp = 0
+            for img in crops:
+                inp += img
+
+            mean_inp.append(np.mean(inp))
+            std_inp.append(np.std(inp))
+
+        output_mean = defaultdict(list)
+        output_std = defaultdict(list)
+        NC = len(self._data)
+        for ch_idx in range(NC):
+            output_mean["target"].append(np.mean(mean_tar_dict[ch_idx]))
+            output_std["target"].append(l2(std_tar_dict[ch_idx]))
+
+        output_mean["target"] = np.array(output_mean["target"]).reshape(NC, 1, 1)
+        output_std["target"] = np.array(output_std["target"]).reshape(NC, 1, 1)
+
+        output_mean["input"] = np.array([np.mean(mean_inp)]).reshape(1, 1, 1)
+        output_std["input"] = np.array([l2(std_inp)]).reshape(1, 1, 1)
+        return dict(output_mean), dict(output_std)
+
+    def set_mean_std(self, mean_dict, std_dict):
+        self._data_mean = mean_dict
+        self._data_std = std_dict
+
+    def get_mean_std(self):
+        return self._data_mean, self._data_std
+
+    def get_num_frames(self):
+        return len(self._data)
+
+    @cache
+    def crop_probablities(self, ch_idx):
+        sizes = np.array([np.prod(x.shape) for x in self._data[ch_idx]])
+        return sizes / sizes.sum()
+
+    def sample_crop(self, ch_idx):
+        idx = None
+        count = 0
+        while idx is None:
+            count += 1
+            idx = np.random.choice(
+                len(self._data[ch_idx]), p=self.crop_probablities(ch_idx)
+            )
+            data = self._data[ch_idx][idx]
+            if data.shape[0] >= self._img_sz[0] and data.shape[1] >= self._img_sz[1]:
+                h = np.random.randint(0, data.shape[0] - self._img_sz[0])
+                w = np.random.randint(0, data.shape[1] - self._img_sz[1])
+                return data[h : h + self._img_sz[0], w : w + self._img_sz[1]]
+            elif count > 100:
+                raise ValueError("Cannot find a valid crop")
+            else:
+                idx = None
+
+        return None
+
+    def len_per_channel(self, ch_idx):
+        return np.sum([np.prod(x.shape) for x in self._data[ch_idx]]) / np.prod(
+            self._img_sz
+        )
+
+    def imgs_for_patch(self):
+        return [self.sample_crop(ch_idx) for ch_idx in range(len(self._data))]
+
+    def __len__(self):
+        len_per_channel = [
+            self.len_per_channel(ch_idx) for ch_idx in range(len(self._data))
+        ]
+        return int(np.max(len_per_channel))
+
+    def _rotate(self, img_tuples):
+        return self._rotate2D(img_tuples)
+
+    def _rotate2D(self, img_tuples):
+        img_kwargs = {}
+        for i, img in enumerate(img_tuples):
+            for k in range(len(img)):
+                img_kwargs[f"img{i}_{k}"] = img[k]
+
+        keys = list(img_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)
+
+        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
+                    )
+                )
+
+        return rotated_img_tuples
+
+    def _compute_input(self, imgs):
+        inp = 0
+        for img in imgs:
+            inp += img
+
+        inp = inp[None]
+        inp = (inp - self._data_mean["input"]) / (self._data_std["input"])
+        return inp
+
+    def _compute_target(self, imgs):
+        imgs = np.stack(imgs)
+        target = (imgs - self._data_mean["target"]) / (self._data_std["target"])
+        return target
+
+    def __getitem__(self, idx):
+        imgs = self.imgs_for_patch()
+        if self._enable_rotation:
+            imgs = self._rotate(imgs)
+
+        inp = self._compute_input(imgs)
+        target = self._compute_target(imgs)
+        return inp, target

careamics/lvae_training/dataset/multifile_dataset.py

@@ -0,0 +1,335 @@
+from collections.abc import Sequence
+from typing import Callable, Union
+
+import numpy as np
+from numpy.typing import NDArray
+
+from .config import MicroSplitDataConfig
+from .lc_dataset import LCMultiChDloader
+from .multich_dataset import MultiChDloader
+from .types import DataSplitType
+
+
+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: MicroSplitDataConfig,
+        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: MicroSplitDataConfig,
+        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: MicroSplitDataConfig,
+        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: MicroSplitDataConfig,
+        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: MicroSplitDataConfig,
+        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,32 @@
+from enum import Enum
+
+
+class DataType(Enum):
+    HTH24Data = 0
+    HTLIF24Data = 1
+    PaviaP24Data = 2
+    TavernaSox2GolgiV2 = 3
+    Dao3ChannelWithInput = 4
+    ExpMicroscopyV1 = 5
+    ExpMicroscopyV2 = 6
+    Dao3Channel = 7
+    TavernaSox2Golgi = 8
+    HTIba1Ki67 = 9
+    OptiMEM100_014 = 10
+    SeparateTiffData = 11
+    BioSR_MRC = 12
+    HTH23BData = 13  # puncta, in case we have differently sized crops for each channel.
+    Care3D = 14
+
+
+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