torch_em.data.datasets.electron_microscopy.platynereis

  1"""Dataset for the segmentation of different structures in EM volume of a
  2platynereis larve. Contains annotations for the segmentation of:
  3- Cuticle
  4- Cilia
  5- Cells
  6- Nuclei
  7
  8This dataset is from the publication https://doi.org/10.1016/j.cell.2021.07.017.
  9Please cite it if you use this dataset for a publication.
 10"""
 11
 12import os
 13from glob import glob
 14from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 15
 16import numpy as np
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25URLS = {
 26    "cells": "https://zenodo.org/record/3675220/files/membrane.zip",
 27    "nuclei": "https://zenodo.org/record/3675220/files/nuclei.zip",
 28    "cilia": "https://zenodo.org/record/3675220/files/cilia.zip",
 29    "cuticle": "https://zenodo.org/record/3675220/files/cuticle.zip"
 30}
 31
 32CHECKSUMS = {
 33    "cells": "30eb50c39e7e9883e1cd96e0df689fac37a56abb11e8ed088907c94a5980d6a3",
 34    "nuclei": "a05033c5fbc6a3069479ac6595b0a430070f83f5281f5b5c8913125743cf5510",
 35    "cilia": "6d2b47f63d39a671789c02d8b66cad5e4cf30eb14cdb073da1a52b7defcc5e24",
 36    "cuticle": "464f75d30133e8864958049647fe3c2216ddf2d4327569738ad72d299c991843"
 37}
 38
 39FILE_TEMPLATES = {
 40    "cells": "train_data_membrane_%02i.n5",
 41    "nuclei": "train_data_nuclei_%02i.h5",
 42    "cilia": "train_data_cilia_%02i.h5",
 43    "cuticle": "train_data_%02i.n5",
 44}
 45
 46
 47#
 48# TODO data-loader for more classes:
 49# - mitos
 50#
 51
 52
 53def _check_data(path, prefix, extension, n_files):
 54    if not os.path.exists(path):
 55        return False
 56    files = glob(os.path.join(path, f"{prefix}*{extension}"))
 57    return len(files) == n_files
 58
 59
 60def get_platynereis_data(path: Union[os.PathLike, str], name: str, download: bool) -> Tuple[str, int]:
 61    """Download the platynereis dataset.
 62
 63    Args:
 64        path: Filepath to a folder where the downloaded data will be saved.
 65        name: Name of the segmentation task. Available tasks: 'cuticle', 'cilia', 'cells' or 'nuclei'.
 66        download: Whether to download the data if it is not present.
 67
 68    Returns:
 69        The path to the folder where the data has been downloaded.
 70        The number of files downloaded.
 71    """
 72    data_root = os.path.join(path, name)
 73
 74    if name == "cuticle":
 75        ext, prefix, n_files = ".n5", "train_data_", 5
 76    elif name == "cilia":
 77        ext, prefix, n_files = ".h5", "train_data_cilia_", 3
 78    elif name == "cells":
 79        data_root = os.path.join(path, "membrane")
 80        ext, prefix, n_files = ".n5", "train_data_membrane_", 9
 81    elif name == "nuclei":
 82        ext, prefix, n_files = ".h5", "train_data_nuclei_", 12
 83    else:
 84        raise ValueError(f"Invalid name {name}. Expect one of 'cuticle', 'cilia', 'cell' or 'nuclei'.")
 85
 86    data_is_complete = _check_data(data_root, prefix, ext, n_files)
 87    if data_is_complete:
 88        return data_root, n_files
 89
 90    os.makedirs(path, exist_ok=True)
 91    url = URLS[name]
 92    checksum = CHECKSUMS[name]
 93
 94    zip_path = os.path.join(path, f"data-{name}.zip")
 95    util.download_source(zip_path, url, download=download, checksum=checksum)
 96    util.unzip(zip_path, path, remove=True)
 97
 98    return data_root, n_files
 99
100
101def get_platynereis_paths(path, sample_ids, name, rois={}, download=False, return_rois=False):
102    """Get paths to the platynereis data.
103
104    Args:
105        path: Filepath to a folder where the downloaded data will be saved.
106        sample_ids: The sample ids to use for the dataset
107        name: Name of the segmentation task. Available tasks: 'cuticle', 'cilia', 'cells' or 'nuclei'.
108        rois: The region of interests to use for the data blocks.
109        download: Whether to download the data if it is not present.
110        return_rois: Whether to return the extracted rois.
111
112    Returns:
113        The filepaths for the stored data.
114    """
115    root, n_files = get_platynereis_data(path, name, download)
116    template = os.path.join(root, FILE_TEMPLATES[name])
117
118    if sample_ids is None:
119        sample_ids = list(range(1, n_files + 1))
120    else:
121        assert min(sample_ids) >= 1 and max(sample_ids) <= n_files
122        sample_ids.sort()
123    paths = [template % sample for sample in sample_ids]
124    data_rois = [rois.get(sample, np.s_[:, :, :]) for sample in sample_ids]
125
126    if return_rois:
127        return paths, data_rois
128    else:
129        return paths
130
131
132def get_platynereis_cuticle_dataset(
133    path: Union[os.PathLike, str],
134    patch_shape: Tuple[int, int, int],
135    sample_ids: Optional[Sequence[int]] = None,
136    download: bool = False,
137    rois: Dict[int, Any] = {},
138    **kwargs
139) -> Dataset:
140    """Get the dataset for cuticle segmentation in platynereis.
141
142    Args:
143        path: Filepath to a folder where the downloaded data will be saved.
144        patch_shape: The patch shape to use for training.
145        sample_ids: The sample ids to use for the dataset
146        download: Whether to download the data if it is not present.
147        rois: The region of interests to use for the data blocks.
148        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
149
150    Returns:
151        The segmentation dataset.
152    """
153    paths, data_rois = get_platynereis_paths(
154        path=path, sample_ids=sample_ids, name="cuticle", rois=rois, download=download, return_rois=True,
155    )
156    return torch_em.default_segmentation_dataset(
157        raw_paths=paths,
158        raw_key="volumes/raw",
159        label_paths=paths,
160        label_key="volumes/labels/segmentation",
161        patch_shape=patch_shape,
162        rois=data_rois,
163        **kwargs
164    )
165
166
167def get_platynereis_cuticle_loader(
168    path: Union[os.PathLike, str],
169    patch_shape: Tuple[int, int, int],
170    batch_size: int,
171    sample_ids: Optional[Sequence[int]] = None,
172    download: bool = False,
173    rois: Dict[int, Any] = {},
174    **kwargs
175) -> DataLoader:
176    """Get the dataloader for cuticle segmentation in platynereis.
177
178    Args:
179        path: Filepath to a folder where the downloaded data will be saved.
180        patch_shape: The patch shape to use for training.
181        batch_size: The batch size for training.
182        sample_ids: The sample ids to use for the dataset
183        download: Whether to download the data if it is not present.
184        rois: The region of interests to use for the data blocks.
185        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
186
187    Returns:
188        The DataLoader.
189    """
190    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
191    ds = get_platynereis_cuticle_dataset(
192        path, patch_shape, sample_ids=sample_ids, download=download, rois=rois, **ds_kwargs,
193    )
194    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)
195
196
197def get_platynereis_cilia_dataset(
198    path: Union[os.PathLike, str],
199    patch_shape: Tuple[int, int, int],
200    sample_ids: Optional[Sequence[int]] = None,
201    offsets: Optional[List[List[int]]] = None,
202    boundaries: bool = False,
203    binary: bool = False,
204    rois: Dict[int, Any] = {},
205    download: bool = False,
206    **kwargs
207) -> Dataset:
208    """Get the dataset for cilia segmentation in platynereis.
209
210    Args:
211        path: Filepath to a folder where the downloaded data will be saved.
212        patch_shape: The patch shape to use for training.
213        sample_ids: The sample ids to use for the dataset
214        offsets: Offset values for affinity computation used as target.
215        boundaries: Whether to compute boundaries as the target.
216        binary: Whether to use a binary segmentation target.
217        rois: The region of interests to use for the data blocks.
218        download: Whether to download the data if it is not present.
219        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
220
221    Returns:
222        The segmentation dataset.
223    """
224    paths, rois = get_platynereis_paths(
225        path=path, sample_ids=sample_ids, name="cilia", rois=rois, download=download, return_rois=True,
226    )
227    kwargs = util.update_kwargs(kwargs, "rois", rois)
228    kwargs, _ = util.add_instance_label_transform(
229        kwargs, add_binary_target=True, boundaries=boundaries, offsets=offsets, binary=binary,
230    )
231    return torch_em.default_segmentation_dataset(
232        raw_paths=paths,
233        raw_key="volumes/raw",
234        label_paths=paths,
235        label_key="volumes/labels/segmentation",
236        patch_shape=patch_shape,
237        **kwargs
238    )
239
240
241def get_platynereis_cilia_loader(
242    path: Union[os.PathLike, str],
243    patch_shape: Tuple[int, int, int],
244    batch_size: int,
245    sample_ids: Optional[Sequence[int]] = None,
246    offsets: Optional[List[List[int]]] = None,
247    boundaries: bool = False,
248    binary: bool = False,
249    rois: Dict[int, Any] = {},
250    download: bool = False,
251    **kwargs
252) -> DataLoader:
253    """Get the dataloader for cilia segmentation in platynereis.
254
255    Args:
256        path: Filepath to a folder where the downloaded data will be saved.
257        patch_shape: The patch shape to use for training.
258        batch_size: The batch size for training.
259        sample_ids: The sample ids to use for the dataset
260        offsets: Offset values for affinity computation used as target.
261        boundaries: Whether to compute boundaries as the target.
262        binary: Whether to return a binary segmentation target.
263        rois: The region of interests to use for the data blocks.
264        download: Whether to download the data if it is not present.
265        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
266
267    Returns:
268        The DataLoader.
269    """
270    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
271    ds = get_platynereis_cilia_dataset(
272        path, patch_shape, sample_ids=sample_ids,
273        offsets=offsets, boundaries=boundaries, binary=binary,
274        rois=rois, download=download, **ds_kwargs,
275    )
276    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)
277
278
279def get_platynereis_cell_dataset(
280    path: Union[os.PathLike, str],
281    patch_shape: Tuple[int, int, int],
282    sample_ids: Optional[Sequence[int]] = None,
283    offsets: Optional[List[List[int]]] = None,
284    boundaries: bool = False,
285    rois: Dict[int, Any] = {},
286    download: bool = False,
287    **kwargs
288) -> Dataset:
289    """Get the dataset for cell segmentation in platynereis.
290
291    Args:
292        path: Filepath to a folder where the downloaded data will be saved.
293        patch_shape: The patch shape to use for training.
294        sample_ids: The sample ids to use for the dataset
295        offsets: Offset values for affinity computation used as target.
296        boundaries: Whether to compute boundaries as the target.
297        rois: The region of interests to use for the data blocks.
298        download: Whether to download the data if it is not present.
299        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
300
301    Returns:
302        The segmentation dataset.
303    """
304    data_paths, data_rois = get_platynereis_paths(
305        path=path, sample_ids=sample_ids, name="cells", rois=rois, download=download, return_rois=True,
306    )
307
308    kwargs = util.update_kwargs(kwargs, "rois", data_rois)
309    kwargs, _ = util.add_instance_label_transform(
310        kwargs, add_binary_target=False, boundaries=boundaries, offsets=offsets,
311    )
312
313    return torch_em.default_segmentation_dataset(
314        raw_paths=data_paths,
315        raw_key="volumes/raw/s1",
316        label_paths=data_paths,
317        label_key="volumes/labels/segmentation/s1",
318        patch_shape=patch_shape,
319        **kwargs
320    )
321
322
323def get_platynereis_cell_loader(
324    path: Union[os.PathLike, str],
325    patch_shape: Tuple[int, int, int],
326    batch_size: int,
327    sample_ids: Optional[Sequence[int]] = None,
328    offsets: Optional[List[List[int]]] = None,
329    boundaries: bool = False,
330    rois: Dict[int, Any] = {},
331    download: bool = False,
332    **kwargs
333) -> DataLoader:
334    """Get the dataloader for cell segmentation in platynereis.
335
336    Args:
337        path: Filepath to a folder where the downloaded data will be saved.
338        patch_shape: The patch shape to use for training.
339        batch_size: The batch size for training.
340        sample_ids: The sample ids to use for the dataset
341        offsets: Offset values for affinity computation used as target.
342        boundaries: Whether to compute boundaries as the target.
343        rois: The region of interests to use for the data blocks.
344        download: Whether to download the data if it is not present.
345        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
346
347    Returns:
348        The DataLoader.
349    """
350    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
351    ds = get_platynereis_cell_dataset(
352        path, patch_shape, sample_ids, rois=rois,
353        offsets=offsets, boundaries=boundaries, download=download,
354        **ds_kwargs,
355    )
356    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)
357
358
359def get_platynereis_nuclei_dataset(
360    path: Union[os.PathLike, str],
361    patch_shape: Tuple[int, int, int],
362    sample_ids: Optional[Sequence[int]] = None,
363    offsets: Optional[List[List[int]]] = None,
364    boundaries: bool = False,
365    binary: bool = False,
366    rois: Dict[int, Any] = {},
367    download: bool = False,
368    **kwargs
369) -> Dataset:
370    """Get the dataset for nucleus segmentation in platynereis.
371
372    Args:
373        path: Filepath to a folder where the downloaded data will be saved.
374        patch_shape: The patch shape to use for training.
375        sample_ids: The sample ids to use for the dataset
376        offsets: Offset values for affinity computation used as target.
377        boundaries: Whether to compute boundaries as the target.
378        binary: Whether to return a binary segmentation target.
379        rois: The region of interests to use for the data blocks.
380        download: Whether to download the data if it is not present.
381        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
382
383    Returns:
384        The segmentation dataset.
385    """
386    _, n_files = get_platynereis_data(path, "nuclei", download)
387
388    if sample_ids is None:
389        sample_ids = list(range(1, n_files + 1))
390    assert min(sample_ids) >= 1 and max(sample_ids) <= n_files
391    sample_ids.sort()
392
393    data_paths, data_rois = get_platynereis_paths(
394        path=path, sample_ids=sample_ids, name="nuclei", rois=rois, download=download, return_rois=True,
395    )
396
397    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
398    kwargs = util.update_kwargs(kwargs, "rois", data_rois)
399    kwargs, _ = util.add_instance_label_transform(
400        kwargs, add_binary_target=True, boundaries=boundaries, offsets=offsets, binary=binary,
401    )
402
403    return torch_em.default_segmentation_dataset(
404        raw_paths=data_paths,
405        raw_key="volumes/raw",
406        label_paths=data_paths,
407        label_key="volumes/labels/nucleus_instance_labels",
408        patch_shape=patch_shape,
409        **kwargs
410    )
411
412
413def get_platynereis_nuclei_loader(
414    path: Union[os.PathLike, str],
415    patch_shape: Tuple[int, int, int],
416    batch_size: int,
417    sample_ids: Optional[Sequence[int]] = None,
418    offsets: Optional[List[List[int]]] = None,
419    boundaries: bool = False,
420    binary: bool = False,
421    rois: Dict[int, Any] = {},
422    download: bool = False,
423    **kwargs
424) -> DataLoader:
425    """Get the dataloader for nucleus segmentation in platynereis.
426
427    Args:
428        path: Filepath to a folder where the downloaded data will be saved.
429        patch_shape: The patch shape to use for training.
430        batch_size: The batch size for training.
431        sample_ids: The sample ids to use for the dataset
432        offsets: Offset values for affinity computation used as target.
433        boundaries: Whether to compute boundaries as the target.
434        binary: Whether to return a binary segmentation target.
435        rois: The region of interests to use for the data blocks.
436        download: Whether to download the data if it is not present.
437        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
438
439    Returns:
440        The DataLoader.
441    """
442    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
443    ds = get_platynereis_nuclei_dataset(
444        path, patch_shape, sample_ids=sample_ids, rois=rois,
445        offsets=offsets, boundaries=boundaries, binary=binary, download=download,
446        **ds_kwargs,
447    )
448    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)