torch_em.data.raw_dataset

  1import os
  2import warnings
  3import numpy as np
  4from typing import List, Union, Tuple, Optional, Any, Callable
  5
  6import torch
  7
  8from elf.wrapper import RoiWrapper
  9
 10from ..util import ensure_tensor_with_channels, ensure_patch_shape, load_data
 11
 12
 13class RawDataset(torch.utils.data.Dataset):
 14    """Dataset that provides raw data stored in a container data format for unsupervised training.
 15
 16    The dataset loads a patch from the raw data and returns a sample for a batch.
 17    The dataset supports all file formats that can be opened with `elf.io.open_file`, such as hdf5, zarr or n5.
 18    Use `raw_path` to specify the path to the file and `raw_key` to specify the internal dataset.
 19    It also supports regular image formats, such as .tif. For these cases set `raw_key=None`.
 20
 21    The dataset can also be used for contrastive learning that relies on two different views of the same data.
 22    You can use the `augmentations` argument for this.
 23
 24    Args:
 25        raw_path: The file path to the raw image data. May also be a list of file paths.
 26        raw_key: The key to the internal dataset containing the raw data.
 27        patch_shape: The patch shape for a training sample.
 28        raw_transform: Transformation applied to the raw data of a sample.
 29        transform: Transformation to the raw data. This can be used to implement data augmentations.
 30        roi: Region of interest in the raw data.
 31            If given, the raw data will only be loaded from the corresponding area.
 32        dtype: The return data type of the raw data.
 33        n_samples: The length of this dataset. If None, the length will be set to `len(raw_image_paths)`.
 34        sampler: Sampler for rejecting samples according to a defined criterion.
 35            The sampler must be a callable that accepts the raw data (as numpy arrays) as input.
 36        ndim: The spatial dimensionality of the data. If None, will be derived from the raw data.
 37        with_channels: Whether the raw data has channels.
 38        augmentations: Augmentations for contrastive learning. If given, these need to be two different callables.
 39            They will be applied to the sampled raw data to return two independent views of the raw data.
 40    """
 41    max_sampling_attempts = 500
 42    """The maximal number of sampling attempts, for loading a sample via `__getitem__`.
 43    This is used when `sampler` rejects a sample, to avoid an infinite loop if no valid sample can be found.
 44    """
 45
 46    @staticmethod
 47    def compute_len(shape, patch_shape):
 48        n_samples = int(np.prod([float(sh / csh) for sh, csh in zip(shape, patch_shape)]))
 49        return n_samples
 50
 51    def __init__(
 52        self,
 53        raw_path: Union[List[Any], str, os.PathLike],
 54        raw_key: Optional[str],
 55        patch_shape: Tuple[int, ...],
 56        raw_transform: Optional[Callable] = None,
 57        transform: Optional[Callable] = None,
 58        roi: Optional[Union[slice, Tuple[slice, ...]]] = None,
 59        dtype: torch.dtype = torch.float32,
 60        n_samples: Optional[int] = None,
 61        sampler: Optional[Callable] = None,
 62        ndim: Optional[int] = None,
 63        with_channels: bool = False,
 64        augmentations: Optional[Tuple[Callable, Callable]] = None,
 65    ):
 66        self.raw_path = raw_path
 67        self.raw_key = raw_key
 68        self.raw = load_data(raw_path, raw_key)
 69
 70        self._with_channels = with_channels
 71
 72        if roi is not None:
 73            if isinstance(roi, slice):
 74                roi = (roi,)
 75
 76            self.raw = RoiWrapper(self.raw, (slice(None),) + roi) if self._with_channels else RoiWrapper(self.raw, roi)
 77
 78        self.shape = self.raw.shape[1:] if self._with_channels else self.raw.shape
 79        self.roi = roi
 80
 81        self._ndim = len(self.shape) if ndim is None else ndim
 82        assert self._ndim in (2, 3, 4), f"Invalid data dimensions: {self._ndim}. Only 2d, 3d or 4d data is supported"
 83
 84        assert len(patch_shape) in (self._ndim, self._ndim + 1), f"{patch_shape}, {self._ndim}"
 85        self.patch_shape = patch_shape
 86
 87        self.raw_transform = raw_transform
 88        self.transform = transform
 89        self.sampler = sampler
 90        self.dtype = dtype
 91
 92        if augmentations is not None:
 93            assert len(augmentations) == 2
 94        self.augmentations = augmentations
 95
 96        self._len = self.compute_len(self.shape, self.patch_shape) if n_samples is None else n_samples
 97
 98        self.sample_shape = patch_shape
 99        self.trafo_halo = None
100        # TODO add support for trafo halo: asking for a bigger bounding box before applying the trafo,
101        # which is then cut. See code below; but this ne needs to be properly tested
102
103        # self.trafo_halo = None if self.transform is None else self.transform.halo(self.patch_shape)
104        # if self.trafo_halo is not None:
105        #     if len(self.trafo_halo) == 2 and self._ndim == 3:
106        #         self.trafo_halo = (0,) + self.trafo_halo
107        #     assert len(self.trafo_halo) == self._ndim
108        #     self.sample_shape = tuple(sh + ha for sh, ha in zip(self.patch_shape, self.trafo_halo))
109        #     self.inner_bb = tuple(slice(ha, sh - ha) for sh, ha in zip(self.patch_shape, self.trafo_halo))
110
111    def __len__(self):
112        return self._len
113
114    @property
115    def ndim(self):
116        return self._ndim
117
118    def _sample_bounding_box(self):
119        bb_start = [
120            np.random.randint(0, sh - psh) if sh - psh > 0 else 0
121            for sh, psh in zip(self.shape, self.sample_shape)
122        ]
123        return tuple(slice(start, start + psh) for start, psh in zip(bb_start, self.sample_shape))
124
125    def _get_sample(self, index):
126        if self.raw is None:
127            raise RuntimeError("RawDataset has not been properly deserialized.")
128        bb = self._sample_bounding_box()
129        raw = self.raw[(slice(None),) + bb] if self._with_channels else self.raw[bb]
130
131        if self.sampler is not None:
132            sample_id = 0
133            while not self.sampler(raw):
134                bb = self._sample_bounding_box()
135                raw = self.raw[(slice(None),) + bb] if self._with_channels else self.raw[bb]
136                sample_id += 1
137                if sample_id > self.max_sampling_attempts:
138                    raise RuntimeError(f"Could not sample a valid batch in {self.max_sampling_attempts} attempts")
139
140        if self.patch_shape is not None:
141            raw = ensure_patch_shape(
142                raw=raw, labels=None, patch_shape=self.patch_shape, have_raw_channels=self._with_channels
143            )
144
145        # squeeze the singleton spatial axis if we have a spatial shape that is larger by one than self._ndim
146        if len(self.patch_shape) == self._ndim + 1:
147            raw = raw.squeeze(1 if self._with_channels else 0)
148
149        return raw
150
151    def crop(self, tensor):
152        bb = self.inner_bb
153        if tensor.ndim > len(bb):
154            bb = (tensor.ndim - len(bb)) * (slice(None),) + bb
155        return tensor[bb]
156
157    def __getitem__(self, index):
158        raw = self._get_sample(index)
159
160        if self.raw_transform is not None:
161            raw = self.raw_transform(raw)
162
163        if self.transform is not None:
164            raw = self.transform(raw)
165            if isinstance(raw, list):
166                assert len(raw) == 1
167                raw = raw[0]
168
169            if self.trafo_halo is not None:
170                raw = self.crop(raw)
171
172        raw = ensure_tensor_with_channels(raw, ndim=self._ndim, dtype=self.dtype)
173        if self.augmentations is not None:
174            aug1, aug2 = self.augmentations
175            raw1, raw2 = aug1(raw), aug2(raw)
176            return raw1, raw2
177
178        return raw
179
180    # need to overwrite pickle to support h5py
181    def __getstate__(self):
182        state = self.__dict__.copy()
183        del state["raw"]
184        return state
185
186    def __setstate__(self, state):
187        raw_path, raw_key = state["raw_path"], state["raw_key"]
188        roi = state["roi"]
189        try:
190            raw = load_data(raw_path, raw_key)
191            if roi is not None:
192                raw = RoiWrapper(raw, (slice(None),) + roi) if state["_with_channels"] else RoiWrapper(raw, roi)
193            state["raw"] = raw
194        except Exception:
195            msg = f"RawDataset could not be deserialized because of missing {raw_path}, {raw_key}.\n"
196            msg += "The dataset is deserialized in order to allow loading trained models from a checkpoint.\n"
197            msg += "But it cannot be used for further training and wil throw an error."
198            warnings.warn(msg)
199            state["raw"] = None
200
201        self.__dict__.update(state)