torch_em.data.datasets.electron_microscopy.vem_nuclei

  1"""This dataset contains nucleus instance segmentations in volume electron microscopy (vEM) data.
  2
  3It has 9 2D array tomography (AT) SEM images and 2 3D FIB-SEM volumes of cancer cells and tumor tissue
  4(human breast cancer cells, mouse lung tumor, mouse 4T1 cells and human MB-231 cells treated with ONC-201),
  5with paired nucleus instance labels. The 2D labels were predicted with the NucleoNet model and manually proofread,
  6the 3D labels (a single nucleus per volume) were obtained by 2D inference in orthogonal planes, voting and proofreading.
  7
  8Nuclei touching the image border are removed from the labels of most 2D images and kept in others,
  9see `EDGE_NUCLEI_INCLUDED`. The 2D images are very large (up to 49530 x 7448 pixels).
 10
 11The record lists a 10th 2D image ('id_9'), but its file in the archive is truncated (4096 bytes instead of 100 MB),
 12so it is not available here.
 13
 14The data is available at https://www.ebi.ac.uk/biostudies/studies/S-BIAD2822 (CC-BY-4.0).
 15The record does not list a publication. Please cite the BioImage Archive entry S-BIAD2822 if you use this dataset.
 16"""
 17
 18import os
 19from typing import List, Literal, Optional, Sequence, Tuple, Union
 20
 21from torch.utils.data import Dataset, DataLoader
 22
 23import torch_em
 24
 25from .. import util
 26
 27
 28BASE_URL = "https://www.ebi.ac.uk/biostudies/files/S-BIAD2822"
 29
 30SAMPLES = {
 31    1: ("2d", "img_2bf732c1.tif", "lbl_2bf732c1.tif"),
 32    2: ("2d", "img_97dde2ee.tiff", "lbl_97dde2ee.tif"),
 33    3: ("2d", "img_626dde70.tif", "lbl_626dde70.tif"),
 34    4: ("2d", "img_57f48663.tif", "lbl_57f48663.tif"),
 35    5: ("2d", "img_22668234.tiff", "lbl_22668234.tif"),
 36    6: ("2d", "img_c69fbe8a.tif", "lbl_c69fbe8a.tif"),
 37    7: ("2d", "img_4ff7ad90.tif", "lbl_4ff7ad90.tif"),
 38    8: ("2d", "img_ee855e08.tif", "lbl_ee855e08.tif"),
 39    10: ("2d", "img_87356f94.tif", "lbl_87356f94.tif"),
 40    11: ("3d", "img_4808a925.tif", "lbl_4808a925.tif"),
 41    12: ("3d", "img_596d5ecb.tif", "lbl_596d5ecb.tif"),
 42}
 43
 44# The 2D images in which the nuclei touching the image border are kept in the labels.
 45EDGE_NUCLEI_INCLUDED = (8, 10)
 46
 47CHECKSUMS = {
 48    "id_1/img_2bf732c1.tif": "d7705c8a7768738a4a685f15fe998c0d1832090bcd316a7c4564cc8392b05328",
 49    "id_1/lbl_2bf732c1.tif": "55660f4a1093877da8c330937498efcfea542e0c7b4cf5690ae80ac84182bcaf",
 50    "id_2/img_97dde2ee.tiff": "9753295d17fe82e6d91d87a2f49625b48fabceea35798263a8cb918c840f8120",
 51    "id_2/lbl_97dde2ee.tif": "f093ab210d0e64aa27040178de3a3a47373dc3d4df342efab8143920bafb2dc7",
 52    "id_3/img_626dde70.tif": "fa0ce1ef7cf3623a5f0bf8e74d0b40d482359244abb1168c329350d81af5c10e",
 53    "id_3/lbl_626dde70.tif": "8046479b75eb3a854eff34a4b8ef9906219a8aa7f15b18be7d3834111244bcda",
 54    "id_4/img_57f48663.tif": "0e02b8349e05b407afb25c643d16b21d7ead570b2389e6ebd0628abb3b2ccaa8",
 55    "id_4/lbl_57f48663.tif": "a63c7efbb12aa514c81aa6f2c00cfb5d1d7092eacc868ba9b496d8b46b8e0e30",
 56    "id_5/img_22668234.tiff": "fa7d6f4b75017317c42418e021581060dbc8bf7e6c4b14374adf0bf673ebb6a7",
 57    "id_5/lbl_22668234.tif": "861ded58e610ca45d9162d5a9f99147611a8deb728bdfa611cde8477b9dc23af",
 58    "id_6/img_c69fbe8a.tif": "abf005036147a765a2212ac5c00055b32aa1591db41db871bd5880d00926dacc",
 59    "id_6/lbl_c69fbe8a.tif": "edc06af9061744c0993297a616258a000c32e6727fa16dbef1b1195ae883e88d",
 60    "id_7/img_4ff7ad90.tif": "e5032ffe2bd3f6f12638a3f25409b682cda410df780dc44224d203c68106d6c7",
 61    "id_7/lbl_4ff7ad90.tif": "b9cc9e2868c0775df3fef0b587aa05ca4d6904a85015dd5ed6f94f044fa30248",
 62    "id_8/img_ee855e08.tif": "2f85a64b60f47530d24c43baa0b7f727516fb87ee7d6ebcf84d338c685bc9c06",
 63    "id_8/lbl_ee855e08.tif": "2dc89dc2a051c74cd366c96861f6d1dda07d0774eec5db4dfcae5075a147bf59",
 64    "id_10/img_87356f94.tif": "5e39af57acc6b69b979f587c896b7e7617d7d790cd5d598c22ca553565104127",
 65    "id_10/lbl_87356f94.tif": "788179894654f4488cf404856752269b4e7402d15353e83d01c259de251f2cc6",
 66    "id_11/img_4808a925.tif": "2978908ebf14933e906c826dcf836831b6f9707dee49abc250876556e209613f",
 67    "id_11/lbl_4808a925.tif": "c253fea05e70cb6229d65ea80f04101e3be8b7f1f07ea235d1d914b4d7dfb87a",
 68    "id_12/img_596d5ecb.tif": "17505809d0317e3278ce40cd11a787af089f3727fa23bc4b6b793693653031ea",
 69    "id_12/lbl_596d5ecb.tif": "6245dc72027876f1f00380840ea176313a6a87386ccb27d69e18288a63e187f4",
 70}
 71
 72
 73def _select_sample_ids(dimensionality, sample_ids):
 74    if dimensionality not in ("2d", "3d"):
 75        raise ValueError(f"'{dimensionality}' is not a valid dimensionality. Choose from '2d' or '3d'.")
 76
 77    available = [sample_id for sample_id, (dim, _, _) in SAMPLES.items() if dim == dimensionality]
 78    if sample_ids is None:
 79        return available
 80
 81    sample_ids = [sample_ids] if isinstance(sample_ids, int) else list(sample_ids)
 82    invalid = [sample_id for sample_id in sample_ids if sample_id not in available]
 83    if invalid:
 84        raise ValueError(f"The sample ids {invalid} are not available for '{dimensionality}'. Choose from {available}.")
 85    return sorted(sample_ids)
 86
 87
 88def get_vem_nuclei_data(
 89    path: Union[os.PathLike, str], sample_ids: Optional[Sequence[int]] = None, download: bool = False
 90) -> str:
 91    """Download the vEM nuclei data.
 92
 93    Args:
 94        path: Filepath to a folder where the downloaded data will be saved.
 95        sample_ids: The ids of the samples to download, see `SAMPLES`. By default all samples are downloaded
 96            (about 1.8 GB).
 97        download: Whether to download the data if it is not present.
 98
 99    Returns:
100        The filepath to the folder with the data.
101    """
102    sample_ids = list(SAMPLES) if sample_ids is None else list(sample_ids)
103    for sample_id in sample_ids:
104        _, img_name, lbl_name = SAMPLES[sample_id]
105        for name in (img_name, lbl_name):
106            rel_path = f"id_{sample_id}/{name}"
107            file_path = os.path.join(path, rel_path)
108            if not os.path.exists(file_path):
109                os.makedirs(os.path.dirname(file_path), exist_ok=True)
110            util.download_source(file_path, f"{BASE_URL}/{rel_path}", download, checksum=CHECKSUMS[rel_path])
111    return str(path)
112
113
114def get_vem_nuclei_paths(
115    path: Union[os.PathLike, str],
116    dimensionality: Literal["2d", "3d"] = "2d",
117    sample_ids: Optional[Union[int, Sequence[int]]] = None,
118    download: bool = False,
119) -> Tuple[List[str], List[str]]:
120    """Get paths to the vEM nuclei data.
121
122    Args:
123        path: Filepath to a folder where the downloaded data will be saved.
124        dimensionality: The data to use, either '2d' (array tomography images with ids 1-8 and 10)
125            or '3d' (FIB-SEM volumes with ids 11 and 12).
126        sample_ids: The samples to use. By default all samples of the dimensionality are used.
127        download: Whether to download the data if it is not present.
128
129    Returns:
130        The filepaths to the images.
131        The filepaths to the labels.
132    """
133    sample_ids = _select_sample_ids(dimensionality, sample_ids)
134    get_vem_nuclei_data(path, sample_ids, download)
135
136    raw_paths = [os.path.join(path, f"id_{sample_id}", SAMPLES[sample_id][1]) for sample_id in sample_ids]
137    label_paths = [os.path.join(path, f"id_{sample_id}", SAMPLES[sample_id][2]) for sample_id in sample_ids]
138    return raw_paths, label_paths
139
140
141def get_vem_nuclei_dataset(
142    path: Union[os.PathLike, str],
143    patch_shape: Tuple[int, ...],
144    dimensionality: Literal["2d", "3d"] = "2d",
145    sample_ids: Optional[Union[int, Sequence[int]]] = None,
146    download: bool = False,
147    offsets: Optional[List[List[int]]] = None,
148    boundaries: bool = False,
149    binary: bool = False,
150    **kwargs
151) -> Dataset:
152    """Get the vEM nuclei dataset for nucleus segmentation in EM.
153
154    Args:
155        path: Filepath to a folder where the downloaded data will be saved.
156        patch_shape: The patch shape to use for training. Use (H, W) for '2d' and (D, H, W) for '3d'.
157        dimensionality: The data to use, either '2d' (array tomography images) or '3d' (FIB-SEM volumes).
158        sample_ids: The samples to use. By default all samples of the dimensionality are used.
159        download: Whether to download the data if it is not present.
160        offsets: Offset values for affinity computation used as target.
161        boundaries: Whether to compute boundaries as the target.
162        binary: Whether to return a binary segmentation target.
163        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
164
165    Returns:
166        The segmentation dataset.
167    """
168    assert len(patch_shape) == (2 if dimensionality == "2d" else 3), f"{dimensionality}, {patch_shape}"
169
170    raw_paths, label_paths = get_vem_nuclei_paths(path, dimensionality, sample_ids, download)
171
172    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
173    kwargs, _ = util.add_instance_label_transform(
174        kwargs, add_binary_target=True, binary=binary, boundaries=boundaries, offsets=offsets
175    )
176
177    return torch_em.default_segmentation_dataset(
178        raw_paths=raw_paths, raw_key=None, label_paths=label_paths, label_key=None, patch_shape=patch_shape, **kwargs
179    )
180
181
182def get_vem_nuclei_loader(
183    path: Union[os.PathLike, str],
184    batch_size: int,
185    patch_shape: Tuple[int, ...],
186    dimensionality: Literal["2d", "3d"] = "2d",
187    sample_ids: Optional[Union[int, Sequence[int]]] = None,
188    download: bool = False,
189    offsets: Optional[List[List[int]]] = None,
190    boundaries: bool = False,
191    binary: bool = False,
192    **kwargs
193) -> DataLoader:
194    """Get the vEM nuclei dataloader for nucleus segmentation in EM.
195
196    Args:
197        path: Filepath to a folder where the downloaded data will be saved.
198        batch_size: The batch size for training.
199        patch_shape: The patch shape to use for training. Use (H, W) for '2d' and (D, H, W) for '3d'.
200        dimensionality: The data to use, either '2d' (array tomography images) or '3d' (FIB-SEM volumes).
201        sample_ids: The samples to use. By default all samples of the dimensionality are used.
202        download: Whether to download the data if it is not present.
203        offsets: Offset values for affinity computation used as target.
204        boundaries: Whether to compute boundaries as the target.
205        binary: Whether to return a binary segmentation target.
206        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
207
208    Returns:
209        The DataLoader.
210    """
211    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
212    dataset = get_vem_nuclei_dataset(
213        path=path, patch_shape=patch_shape, dimensionality=dimensionality, sample_ids=sample_ids, download=download,
214        offsets=offsets, boundaries=boundaries, binary=binary, **ds_kwargs,
215    )
216    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)