torch_em.data.datasets.electron_microscopy.isbi2012

  1"""The ISBI2012 dataset was the first neuron segmentation challenge, held at the ISBI 2012 competition.
  2It contains a small annotated EM volume from the fruit-fly brain.
  3
  4If you use this dataset in your research please cite the following publication:
  5https://doi.org/10.3389/fnana.2015.00142.
  6"""
  7
  8import os
  9from typing import List, Optional, Tuple, Union
 10
 11from torch.utils.data import Dataset, DataLoader
 12
 13import torch_em
 14from .. import util
 15
 16
 17ISBI_URL = "https://oc.embl.de/index.php/s/h0TkwqxU0PJDdMd/download"
 18CHECKSUM = "0e10fe909a1243084d91773470856993b7d40126a12e85f0f1345a7a9e512f29"
 19
 20
 21def get_isbi_data(path: Union[os.PathLike, str], download: bool = False):
 22    """Download the ISBI2012 dataset.
 23
 24    Args:
 25        path: Filepath to a folder where the downloaded data will be saved.
 26        download: Whether to download the data if it is not present.
 27    """
 28    os.makedirs(path, exist_ok=True)
 29    util.download_source(os.path.join(path, "isbi.h5"), ISBI_URL, download, CHECKSUM)
 30
 31
 32def get_isbi_paths(path: Union[os.PathLike, str], download: bool = False) -> str:
 33    """Get path to ISBI data.
 34
 35    Args:
 36        path: Filepath to a folder where the downloaded data will be saved.
 37        download: Whether to download the data if it is not present.
 38
 39    Returns:
 40        The filepath for the stored data.
 41    """
 42    get_isbi_data(path, download)
 43    volume_path = os.path.join(path, "isbi.h5")
 44    return volume_path
 45
 46
 47def get_isbi_dataset(
 48    path: Union[os.PathLike, str],
 49    patch_shape: Tuple[int, int, int],
 50    download: bool = False,
 51    offsets: Optional[List[List[int]]] = None,
 52    boundaries: bool = False,
 53    use_original_labels: bool = False,
 54    **kwargs
 55) -> Dataset:
 56    """Get the dataset for EM neuron segmentation in ISBI 2012.
 57
 58    Args:
 59        path: Filepath to a folder where the downloaded data will be saved.
 60        patch_shape: The patch shape to use for training.
 61        download: Whether to download the data if it is not present.
 62        offsets: Offset values for affinity computation used as target.
 63        boundaries: Whether to compute boundaries as the target.
 64        use_original_labels: Whether to use the original annotations or postprocessed 3d annotations.
 65        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 66
 67    Returns:
 68       The segmentation dataset.
 69    """
 70    assert len(patch_shape) == 3
 71
 72    volume_path = get_isbi_paths(path, download)
 73
 74    ndim = 2 if patch_shape[0] == 1 else 3
 75    kwargs = util.update_kwargs(kwargs, "ndim", ndim)
 76
 77    kwargs, _ = util.add_instance_label_transform(
 78        kwargs, add_binary_target=False, boundaries=boundaries, offsets=offsets
 79    )
 80
 81    return torch_em.default_segmentation_dataset(
 82        raw_paths=volume_path,
 83        raw_key="raw",
 84        label_paths=volume_path,
 85        label_key="labels/membranes" if use_original_labels else "labels/gt_segmentation",
 86        patch_shape=patch_shape,
 87        **kwargs
 88    )
 89
 90
 91def get_isbi_loader(
 92    path: Union[os.PathLike, str],
 93    patch_shape: Tuple[int, int, int],
 94    batch_size: int,
 95    download: bool = False,
 96    offsets: Optional[List[List[int]]] = None,
 97    boundaries: bool = False,
 98    use_original_labels: bool = False,
 99    **kwargs
100) -> DataLoader:
101    """Get the DataLoader for EM neuron segmentation in ISBI 2012.
102
103    Args:
104        path: Filepath to a folder where the downloaded data will be saved.
105        patch_shape: The patch shape to use for training.
106        batch_size: The batch size for training.
107        download: Whether to download the data if it is not present.
108        offsets: Offset values for affinity computation used as target.
109        boundaries: Whether to compute boundaries as the target.
110        use_original_labels: Whether to use the original annotations or postprocessed 3d annotations.
111        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
112
113    Returns:
114        The DataLoader.
115    """
116    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
117    dataset = get_isbi_dataset(
118        path, patch_shape, download=download, offsets=offsets,
119        boundaries=boundaries, use_original_labels=use_original_labels, **ds_kwargs
120    )
121    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)