torch_em.data.datasets.electron_microscopy.cremi

  1"""CREMI is a dataset for neuron segmentation in EM.
  2
  3It contains three annotated volumes from the adult fruit-fly brain.
  4It was held as a challenge at MICCAI 2016. For details on the dataset check out https://cremi.org/.
  5Please cite the challenge if you use the dataset in your research.
  6"""
  7# TODO add support for realigned volumes
  8
  9import os
 10from typing import Any, Dict, List, Optional, Tuple, Union
 11
 12import numpy as np
 13
 14from torch.utils.data import Dataset, DataLoader
 15
 16import torch_em
 17
 18from .. import util
 19
 20
 21CREMI_URLS = {
 22    "original": {
 23        "A": "https://cremi.org/static/data/sample_A_20160501.hdf",
 24        "B": "https://cremi.org/static/data/sample_B_20160501.hdf",
 25        "C": "https://cremi.org/static/data/sample_C_20160501.hdf",
 26    },
 27    "realigned": {},
 28    "defects": "https://zenodo.org/record/5767036/files/sample_ABC_padded_defects.h5"
 29}
 30CHECKSUMS = {
 31    "original": {
 32        "A": "4c563d1b78acb2bcfb3ea958b6fe1533422f7f4a19f3e05b600bfa11430b510d",
 33        "B": "887e85521e00deead18c94a21ad71f278d88a5214c7edeed943130a1f4bb48b8",
 34        "C": "2874496f224d222ebc29d0e4753e8c458093e1d37bc53acd1b69b19ed1ae7052",
 35    },
 36    "realigned": {},
 37    "defects": "7b06ffa34733b2c32956ea5005e0cf345e7d3a27477f42f7c905701cdc947bd0"
 38}
 39
 40
 41def get_cremi_data(path: Union[os.PathLike, str], samples: Tuple[str], download: bool, use_realigned: bool = False):
 42    """Download the CREMI training data.
 43
 44    Args:
 45        path: Filepath to a folder where the downloaded data will be saved.
 46        samples: The CREMI samples to use. The available samples are 'A', 'B', 'C'.
 47        download: Whether to download the data if it is not present.
 48        use_realigned: Use the realigned instead of the original training data.
 49    """
 50    if use_realigned:
 51        # we need to sample batches in this case
 52        # sampler = torch_em.data.MinForegroundSampler(min_fraction=0.05, p_reject=.75)
 53        raise NotImplementedError
 54    else:
 55        urls = CREMI_URLS["original"]
 56        checksums = CHECKSUMS["original"]
 57
 58    os.makedirs(path, exist_ok=True)
 59    for name in samples:
 60        url = urls[name]
 61        checksum = checksums[name]
 62        data_path = os.path.join(path, f"sample{name}.h5")
 63        # CREMI SSL certificates expired, so we need to disable verification
 64        util.download_source(data_path, url, download, checksum, verify=False)
 65
 66
 67def get_cremi_paths(
 68    path: Union[os.PathLike, str],
 69    samples: Tuple[str, ...] = ("A", "B", "C"),
 70    use_realigned: bool = False,
 71    download: bool = False
 72) -> List[str]:
 73    """Get paths to the CREMI data.
 74
 75    Args:
 76        path: Filepath to a folder where the downloaded data will be saved.
 77        samples: The CREMI samples to use. The available samples are 'A', 'B', 'C'.
 78        use_realigned: Use the realigned instead of the original training data.
 79        download: Whether to download the data if it is not present.
 80
 81    Returns:
 82        The filepaths to the training data.
 83    """
 84    get_cremi_data(path, samples, download, use_realigned)
 85    data_paths = [os.path.join(path, f"sample{name}.h5") for name in samples]
 86    return data_paths
 87
 88
 89def get_cremi_dataset(
 90    path: Union[os.PathLike, str],
 91    patch_shape: Tuple[int, int, int],
 92    samples: Tuple[str, ...] = ("A", "B", "C"),
 93    use_realigned: bool = False,
 94    download: bool = False,
 95    offsets: Optional[List[List[int]]] = None,
 96    boundaries: bool = False,
 97    rois: Dict[str, Any] = {},
 98    defect_augmentation_kwargs: Dict[str, Any] = {
 99        "p_drop_slice": 0.025,
100        "p_low_contrast": 0.025,
101        "p_deform_slice": 0.0,
102        "deformation_mode": "compress",
103    },
104    **kwargs,
105) -> Dataset:
106    """Get the CREMI dataset for the segmentation of neurons in EM.
107
108    Args:
109        path: Filepath to a folder where the downloaded data will be saved.
110        patch_shape: The patch shape to use for training.
111        samples: The CREMI samples to use. The available samples are 'A', 'B', 'C'.
112        use_realigned: Use the realigned instead of the original training data.
113        download: Whether to download the data if it is not present.
114        offsets: Offset values for affinity computation used as target.
115        boundaries: Whether to compute boundaries as the target.
116        rois: The region of interests to use for the samples.
117        defect_augmentation_kwargs: Keyword arguments for defect augmentations.
118        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
119
120    Returns:
121       The segmentation dataset.
122    """
123    assert len(patch_shape) == 3
124    if rois is not None:
125        assert isinstance(rois, dict)
126
127    data_paths = get_cremi_paths(path, samples, use_realigned, download)
128    data_rois = [rois.get(name, np.s_[:, :, :]) for name in samples]
129
130    if defect_augmentation_kwargs is not None and "artifact_source" not in defect_augmentation_kwargs:
131        # download the defect volume
132        url = CREMI_URLS["defects"]
133        checksum = CHECKSUMS["defects"]
134        defect_path = os.path.join(path, "cremi_defects.h5")
135        util.download_source(defect_path, url, download, checksum)
136        defect_patch_shape = (1,) + tuple(patch_shape[1:])
137        artifact_source = torch_em.transform.get_artifact_source(
138            defect_path, defect_patch_shape,
139            min_mask_fraction=0.75,
140            raw_key="defect_sections/raw",
141            mask_key="defect_sections/mask"
142        )
143        defect_augmentation_kwargs.update({"artifact_source": artifact_source})
144
145    # defect augmentations
146    if defect_augmentation_kwargs is not None:
147        raw_transform = torch_em.transform.get_raw_transform(
148            augmentation1=torch_em.transform.EMDefectAugmentation(**defect_augmentation_kwargs)
149        )
150        kwargs = util.update_kwargs(kwargs, "raw_transform", raw_transform)
151
152    kwargs, _ = util.add_instance_label_transform(
153        kwargs, add_binary_target=False, boundaries=boundaries, offsets=offsets
154    )
155
156    return torch_em.default_segmentation_dataset(
157        raw_paths=data_paths,
158        raw_key="volumes/raw",
159        label_paths=data_paths,
160        label_key="volumes/labels/neuron_ids",
161        patch_shape=patch_shape,
162        rois=data_rois,
163        **kwargs
164    )
165
166
167def get_cremi_loader(
168    path: Union[os.PathLike, str],
169    patch_shape: Tuple[int, int, int],
170    batch_size: int,
171    samples: Tuple[str, ...] = ("A", "B", "C"),
172    use_realigned: bool = False,
173    download: bool = False,
174    offsets: Optional[List[List[int]]] = None,
175    boundaries: bool = False,
176    rois: Dict[str, Any] = {},
177    defect_augmentation_kwargs: Dict[str, Any] = {
178        "p_drop_slice": 0.025,
179        "p_low_contrast": 0.025,
180        "p_deform_slice": 0.0,
181        "deformation_mode": "compress",
182    },
183    **kwargs,
184) -> DataLoader:
185    """Get the DataLoader for EM neuron segmentation in the CREMI dataset.
186
187    Args:
188        path: Filepath to a folder where the downloaded data will be saved.
189        patch_shape: The patch shape to use for training.
190        batch_size: The batch size for training.
191        samples: The CREMI samples to use. The available samples are 'A', 'B', 'C'.
192        use_realigned: Use the realigned instead of the original training data.
193        download: Whether to download the data if it is not present.
194        offsets: Offset values for affinity computation used as target.
195        boundaries: Whether to compute boundaries as the target.
196        rois: The region of interests to use for the samples.
197        defect_augmentation_kwargs: Keyword arguments for defect augmentations.
198        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
199
200    Returns:
201        The DataLoader.
202    """
203    dataset_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
204    ds = get_cremi_dataset(
205        path=path,
206        patch_shape=patch_shape,
207        samples=samples,
208        use_realigned=use_realigned,
209        download=download,
210        offsets=offsets,
211        boundaries=boundaries,
212        rois=rois,
213        defect_augmentation_kwargs=defect_augmentation_kwargs,
214        **dataset_kwargs,
215    )
216    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)