torch_em.data.datasets.electron_microscopy.isbi2012

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