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

careamics/lvae_training/data_utils.py

@@ -0,0 +1,618 @@
+"""
+Utility functions needed by dataloader & co.
+"""
+
+from typing import List
+
+import numpy as np
+from skimage.io import imread, imsave
+
+from careamics.models.lvae.utils 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
+    # 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
+
+
+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
+
+
+def get_mrc_data(fpath):
+    # HXWXN
+    _, data = read_mrc(fpath)
+    data = data[None]
+    data = np.swapaxes(data, 0, 3)
+    return data[..., 0]
+
+
+class GridIndexManager:
+
+    def __init__(self, data_shape, grid_size, patch_size, grid_alignement) -> None:
+        self._data_shape = data_shape
+        self._default_grid_size = grid_size
+        self.patch_size = patch_size
+        self.N = self._data_shape[0]
+        self._align = grid_alignement
+
+    def get_data_shape(self):
+        return self._data_shape
+
+    def use_default_grid(self, grid_size):
+        return grid_size is None or grid_size < 0
+
+    def grid_rows(self, grid_size):
+        if self._align == GridAlignement.LeftTop:
+            extra_pixels = self.patch_size - grid_size
+        elif self._align == GridAlignement.Center:
+            # Center is exclusively used during evaluation. In this case, we use the padding to handle edge cases.
+            # So, here, we will ideally like to cover all pixels and so extra_pixels is set to 0.
+            # If there was no padding, then it should be set to (self.patch_size - grid_size) // 2
+            extra_pixels = 0
+
+        return (self._data_shape[-3] - extra_pixels) // grid_size
+
+    def grid_cols(self, grid_size):
+        if self._align == GridAlignement.LeftTop:
+            extra_pixels = self.patch_size - grid_size
+        elif self._align == GridAlignement.Center:
+            extra_pixels = 0
+
+        return (self._data_shape[-2] - extra_pixels) // grid_size
+
+    def grid_count(self, grid_size=None):
+        if self.use_default_grid(grid_size):
+            grid_size = self._default_grid_size
+
+        return self.N * self.grid_rows(grid_size) * self.grid_cols(grid_size)
+
+    def hwt_from_idx(self, index, grid_size=None):
+        t = self.get_t(index)
+        return (*self.get_deterministic_hw(index, grid_size=grid_size), t)
+
+    def idx_from_hwt(self, h_start, w_start, t, grid_size=None):
+        """
+        Given h,w,t (where h,w constitutes the top left corner of the patch), it returns the corresponding index.
+        """
+        if grid_size is None:
+            grid_size = self._default_grid_size
+
+        nth_row = h_start // grid_size
+        nth_col = w_start // grid_size
+
+        index = self.grid_cols(grid_size) * nth_row + nth_col
+        return index * self._data_shape[0] + t
+
+    def get_t(self, index):
+        return index % self.N
+
+    def get_top_nbr_idx(self, index, grid_size=None):
+        if self.use_default_grid(grid_size):
+            grid_size = self._default_grid_size
+
+        ncols = self.grid_cols(grid_size)
+        index -= ncols * self.N
+        if index < 0:
+            return None
+
+        return index
+
+    def get_bottom_nbr_idx(self, index, grid_size=None):
+        if self.use_default_grid(grid_size):
+            grid_size = self._default_grid_size
+
+        ncols = self.grid_cols(grid_size)
+        index += ncols * self.N
+        if index > self.grid_count(grid_size=grid_size):
+            return None
+
+        return index
+
+    def get_left_nbr_idx(self, index, grid_size=None):
+        if self.on_left_boundary(index, grid_size=grid_size):
+            return None
+
+        index -= self.N
+        return index
+
+    def get_right_nbr_idx(self, index, grid_size=None):
+        if self.on_right_boundary(index, grid_size=grid_size):
+            return None
+        index += self.N
+        return index
+
+    def on_left_boundary(self, index, grid_size=None):
+        if self.use_default_grid(grid_size):
+            grid_size = self._default_grid_size
+
+        factor = index // self.N
+        ncols = self.grid_cols(grid_size)
+
+        left_boundary = (factor // ncols) != (factor - 1) // ncols
+        return left_boundary
+
+    def on_right_boundary(self, index, grid_size=None):
+        if self.use_default_grid(grid_size):
+            grid_size = self._default_grid_size
+
+        factor = index // self.N
+        ncols = self.grid_cols(grid_size)
+
+        right_boundary = (factor // ncols) != (factor + 1) // ncols
+        return right_boundary
+
+    def on_top_boundary(self, index, grid_size=None):
+        if self.use_default_grid(grid_size):
+            grid_size = self._default_grid_size
+
+        ncols = self.grid_cols(grid_size)
+        return index < self.N * ncols
+
+    def on_bottom_boundary(self, index, grid_size=None):
+        if self.use_default_grid(grid_size):
+            grid_size = self._default_grid_size
+
+        ncols = self.grid_cols(grid_size)
+        return index + self.N * ncols > self.grid_count(grid_size=grid_size)
+
+    def on_boundary(self, idx, grid_size=None):
+        if self.on_left_boundary(idx, grid_size=grid_size):
+            return True
+
+        if self.on_right_boundary(idx, grid_size=grid_size):
+            return True
+
+        if self.on_top_boundary(idx, grid_size=grid_size):
+            return True
+
+        if self.on_bottom_boundary(idx, grid_size=grid_size):
+            return True
+        return False
+
+    def get_deterministic_hw(self, index: int, grid_size=None):
+        """
+        Fixed starting position for the crop for the img with index `index`.
+        """
+        if self.use_default_grid(grid_size):
+            grid_size = self._default_grid_size
+
+        # _, h, w, _ = self._data_shape
+        # assert h == w
+        factor = index // self.N
+        ncols = self.grid_cols(grid_size)
+
+        ith_row = factor // ncols
+        jth_col = factor % ncols
+        h_start = ith_row * grid_size
+        w_start = jth_col * grid_size
+        return h_start, w_start
+
+
+class IndexSwitcher:
+    """
+    The idea is to switch from valid indices for target to invalid indices for target.
+    If index in invalid for the target, then we return all zero vector as target.
+    This combines both logic:
+    1. Using less amount of total data.
+    2. Using less amount of target data but using full data.
+    """
+
+    def __init__(self, idx_manager, data_config, patch_size) -> None:
+        self.idx_manager = idx_manager
+        self._data_shape = self.idx_manager.get_data_shape()
+        self._training_validtarget_fraction = data_config.get(
+            "training_validtarget_fraction", 1.0
+        )
+        self._validtarget_ceilT = int(
+            np.ceil(self._data_shape[0] * self._training_validtarget_fraction)
+        )
+        self._patch_size = patch_size
+        assert (
+            data_config.deterministic_grid is True
+        ), "This only works when the dataset has deterministic grid. Needed randomness comes from this class."
+        assert (
+            "grid_size" in data_config and data_config.grid_size == 1
+        ), "We need a one to one mapping between index and h, w, t"
+
+        self._h_validmax, self._w_validmax = self.get_reduced_frame_size(
+            self._data_shape[:3], self._training_validtarget_fraction
+        )
+        if self._h_validmax < self._patch_size or self._w_validmax < self._patch_size:
+            print(
+                "WARNING: The valid target size is smaller than the patch size. This will result in all zero target. so, we are ignoring this frame for target."
+            )
+            self._h_validmax = 0
+            self._w_validmax = 0
+
+        print(
+            f"[{self.__class__.__name__}] Target Indices: [0,{self._validtarget_ceilT-1}]. Index={self._validtarget_ceilT-1} has shape [:{self._h_validmax},:{self._w_validmax}].  Available data: {self._data_shape[0]}"
+        )
+
+    def get_valid_target_index(self):
+        """
+        Returns an index which corresponds to a frame which is expected to have a target.
+        """
+        _, h, w, _ = self._data_shape
+        framepixelcount = h * w
+        targetpixels = np.array(
+            [framepixelcount] * (self._validtarget_ceilT - 1)
+            + [self._h_validmax * self._w_validmax]
+        )
+        targetpixels = targetpixels / np.sum(targetpixels)
+        t = np.random.choice(self._validtarget_ceilT, p=targetpixels)
+        # t = np.random.randint(0, self._validtarget_ceilT) if self._validtarget_ceilT >= 1 else 0
+        h, w = self.get_valid_target_hw(t)
+        index = self.idx_manager.idx_from_hwt(h, w, t)
+        # print('Valid', index, h,w,t)
+        return index
+
+    def get_invalid_target_index(self):
+        # if self._validtarget_ceilT == 0:
+        # TODO: There may not be enough data for this to work. The better way is to skip using 0 for invalid target.
+        # t = np.random.randint(1, self._data_shape[0])
+        # elif self._validtarget_ceilT < self._data_shape[0]:
+        #     t = np.random.randint(self._validtarget_ceilT, self._data_shape[0])
+        # else:
+        #     t = self._validtarget_ceilT - 1
+        # 5
+        # 1.2 => 2
+        total_t, h, w, _ = self._data_shape
+        framepixelcount = h * w
+        available_h = h - self._h_validmax
+        if available_h < self._patch_size:
+            available_h = 0
+        available_w = w - self._w_validmax
+        if available_w < self._patch_size:
+            available_w = 0
+
+        targetpixels = np.array(
+            [available_h * available_w]
+            + [framepixelcount] * (total_t - self._validtarget_ceilT)
+        )
+        t_probab = targetpixels / np.sum(targetpixels)
+        t = np.random.choice(
+            np.arange(self._validtarget_ceilT - 1, total_t), p=t_probab
+        )
+
+        h, w = self.get_invalid_target_hw(t)
+        index = self.idx_manager.idx_from_hwt(h, w, t)
+        # print('Invalid', index, h,w,t)
+        return index
+
+    def get_valid_target_hw(self, t):
+        """
+        This is the opposite of get_invalid_target_hw. It returns a h,w which is valid for target.
+        This is only valid for single frame setup.
+        """
+        if t == self._validtarget_ceilT - 1:
+            h = np.random.randint(0, self._h_validmax - self._patch_size)
+            w = np.random.randint(0, self._w_validmax - self._patch_size)
+        else:
+            h = np.random.randint(0, self._data_shape[1] - self._patch_size)
+            w = np.random.randint(0, self._data_shape[2] - self._patch_size)
+        return h, w
+
+    def get_invalid_target_hw(self, t):
+        """
+        This is the opposite of get_valid_target_hw. It returns a h,w which is not valid for target.
+        This is only valid for single frame setup.
+        """
+        if t == self._validtarget_ceilT - 1:
+            h = np.random.randint(
+                self._h_validmax, self._data_shape[1] - self._patch_size
+            )
+            w = np.random.randint(
+                self._w_validmax, self._data_shape[2] - self._patch_size
+            )
+        else:
+            h = np.random.randint(0, self._data_shape[1] - self._patch_size)
+            w = np.random.randint(0, self._data_shape[2] - self._patch_size)
+        return h, w
+
+    def _get_tidx(self, index):
+        if isinstance(index, int) or isinstance(index, np.int64):
+            idx = index
+        else:
+            idx = index[0]
+        return self.idx_manager.get_t(idx)
+
+    def index_should_have_target(self, index):
+        tidx = self._get_tidx(index)
+        if tidx < self._validtarget_ceilT - 1:
+            return True
+        elif tidx > self._validtarget_ceilT - 1:
+            return False
+        else:
+            h, w, _ = self.idx_manager.hwt_from_idx(index)
+            return (
+                h + self._patch_size < self._h_validmax
+                and w + self._patch_size < self._w_validmax
+            )
+
+    @staticmethod
+    def get_reduced_frame_size(data_shape_nhw, fraction):
+        n, h, w = data_shape_nhw
+
+        framepixelcount = h * w
+        targetpixelcount = int(n * framepixelcount * fraction)
+
+        # We are currently supporting this only when there is just one frame.
+        # if np.ceil(pixelcount / framepixelcount) > 1:
+        #     return None, None
+
+        lastframepixelcount = targetpixelcount % framepixelcount
+        assert data_shape_nhw[1] == data_shape_nhw[2]
+        if lastframepixelcount > 0:
+            new_size = int(np.sqrt(lastframepixelcount))
+            return new_size, new_size
+        else:
+            assert (
+                targetpixelcount / framepixelcount >= 1
+            ), "This is not possible in euclidean space :D (so this is a bug)"
+            return h, w
+
+
+rec_header_dtd = [
+    ("nx", "i4"),  # Number of columns
+    ("ny", "i4"),  # Number of rows
+    ("nz", "i4"),  # Number of sections
+    ("mode", "i4"),  # Types of pixels in the image. Values used by IMOD:
+    #  0 = unsigned or signed bytes depending on flag in imodFlags
+    #  1 = signed short integers (16 bits)
+    #  2 = float (32 bits)
+    #  3 = short * 2, (used for complex data)
+    #  4 = float * 2, (used for complex data)
+    #  6 = unsigned 16-bit integers (non-standard)
+    # 16 = unsigned char * 3 (for rgb data, non-standard)
+    ("nxstart", "i4"),  # Starting point of sub-image (not used in IMOD)
+    ("nystart", "i4"),
+    ("nzstart", "i4"),
+    ("mx", "i4"),  # Grid size in X, Y and Z
+    ("my", "i4"),
+    ("mz", "i4"),
+    ("xlen", "f4"),  # Cell size; pixel spacing = xlen/mx, ylen/my, zlen/mz
+    ("ylen", "f4"),
+    ("zlen", "f4"),
+    ("alpha", "f4"),  # Cell angles - ignored by IMOD
+    ("beta", "f4"),
+    ("gamma", "f4"),
+    # These need to be set to 1, 2, and 3 for pixel spacing to be interpreted correctly
+    ("mapc", "i4"),  # map column  1=x,2=y,3=z.
+    ("mapr", "i4"),  # map row     1=x,2=y,3=z.
+    ("maps", "i4"),  # map section 1=x,2=y,3=z.
+    # These need to be set for proper scaling of data
+    ("amin", "f4"),  # Minimum pixel value
+    ("amax", "f4"),  # Maximum pixel value
+    ("amean", "f4"),  # Mean pixel value
+    ("ispg", "i4"),  # space group number (ignored by IMOD)
+    (
+        "next",
+        "i4",
+    ),  # number of bytes in extended header (called nsymbt in MRC standard)
+    ("creatid", "i2"),  # used to be an ID number, is 0 as of IMOD 4.2.23
+    ("extra_data", "V30"),  # (not used, first two bytes should be 0)
+    # These two values specify the structure of data in the extended header; their meaning depend on whether the
+    # extended header has the Agard format, a series of 4-byte integers then real numbers, or has data
+    # produced by SerialEM, a series of short integers. SerialEM stores a float as two shorts, s1 and s2, by:
+    # value = (sign of s1)*(|s1|*256 + (|s2| modulo 256)) * 2**((sign of s2) * (|s2|/256))
+    ("nint", "i2"),
+    # Number of integers per section (Agard format) or number of bytes per section (SerialEM format)
+    ("nreal", "i2"),  # Number of reals per section (Agard format) or bit
+    # Number of reals per section (Agard format) or bit
+    # flags for which types of short data (SerialEM format):
+    # 1 = tilt angle * 100  (2 bytes)
+    # 2 = piece coordinates for montage  (6 bytes)
+    # 4 = Stage position * 25    (4 bytes)
+    # 8 = Magnification / 100 (2 bytes)
+    # 16 = Intensity * 25000  (2 bytes)
+    # 32 = Exposure dose in e-/A2, a float in 4 bytes
+    # 128, 512: Reserved for 4-byte items
+    # 64, 256, 1024: Reserved for 2-byte items
+    # If the number of bytes implied by these flags does
+    # not add up to the value in nint, then nint and nreal
+    # are interpreted as ints and reals per section
+    ("extra_data2", "V20"),  # extra data (not used)
+    ("imodStamp", "i4"),  # 1146047817 indicates that file was created by IMOD
+    ("imodFlags", "i4"),  # Bit flags: 1 = bytes are stored as signed
+    # Explanation of type of data
+    ("idtype", "i2"),  # ( 0 = mono, 1 = tilt, 2 = tilts, 3 = lina, 4 = lins)
+    ("lens", "i2"),
+    # ("nd1", "i2"),  # for idtype = 1, nd1 = axis (1, 2, or 3)
+    # ("nd2", "i2"),
+    ("nphase", "i4"),
+    ("vd1", "i2"),  # vd1 = 100. * tilt increment
+    ("vd2", "i2"),  # vd2 = 100. * starting angle
+    # Current angles are used to rotate a model to match a new rotated image.  The three values in each set are
+    # rotations about X, Y, and Z axes, applied in the order Z, Y, X.
+    ("triangles", "f4", 6),  # 0,1,2 = original:  3,4,5 = current
+    ("xorg", "f4"),  # Origin of image
+    ("yorg", "f4"),
+    ("zorg", "f4"),
+    ("cmap", "S4"),  # Contains "MAP "
+    (
+        "stamp",
+        "u1",
+        4,
+    ),  # First two bytes have 17 and 17 for big-endian or 68 and 65 for little-endian
+    ("rms", "f4"),  # RMS deviation of densities from mean density
+    ("nlabl", "i4"),  # Number of labels with useful data
+    ("labels", "S80", 10),  # 10 labels of 80 charactors
+]
+
+
+def read_mrc(filename, filetype="image"):
+
+    fd = open(filename, "rb")
+    header = np.fromfile(fd, dtype=rec_header_dtd, count=1)
+
+    nx, ny, nz = header["nx"][0], header["ny"][0], header["nz"][0]
+
+    if header[0][3] == 1:
+        data_type = "int16"
+    elif header[0][3] == 2:
+        data_type = "float32"
+    elif header[0][3] == 4:
+        data_type = "single"
+        nx = nx * 2
+    elif header[0][3] == 6:
+        data_type = "uint16"
+
+    data = np.ndarray(shape=(nx, ny, nz))
+    imgrawdata = np.fromfile(fd, data_type)
+    fd.close()
+
+    if filetype == "image":
+        for iz in range(nz):
+            data_2d = imgrawdata[nx * ny * iz : nx * ny * (iz + 1)]
+            data[:, :, iz] = data_2d.reshape(nx, ny, order="F")
+    else:
+        data = imgrawdata
+
+    return header, data