torch_em.data.datasets.electron_microscopy.elmi_mito

  1"""The ELMI Mito dataset contains annotations for mitochondria instance segmentation in transmission electron
  2microscopy (TEM) images.
  3
  4It consists of five large images (about 2000 x 2000 pixels each) with manually annotated mitochondria instances.
  5The images were shared by the authors as part of the training data for a MitoNet (empanada) model that was fine-tuned
  6on mitochondria in these ELMI images. In total the download size is ~40 MB.
  7
  8The data is available at https://owncloud.cesnet.cz/index.php/s/GRxauIeDbnSLiZV.
  9The share does not state a licence or a publication to cite, and does not explain the abbreviation 'ELMI'.
 10Please check with the data providers before using this dataset in your research.
 11The share also contains the fine-tuned model and its predictions on the images, which are not downloaded here.
 12"""
 13
 14import os
 15from typing import List, Tuple, Union
 16
 17from natsort import natsorted
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25BASE_URL = "https://owncloud.cesnet.cz/index.php/s/GRxauIeDbnSLiZV/download"
 26IMAGE_NAMES = [f"ELMI_0{i}.tif" for i in range(1, 6)]
 27CHECKSUMS = {
 28    "images": {
 29        "ELMI_01.tif": "17dd6971fb03566cd32dfbf3395640459d9f982195f7ade92a8169f8b0033090",
 30        "ELMI_02.tif": "dbe01aa9773aaf25898282513f35fc352e8f2e6eb11fa917e7c31c4a4bd8b982",
 31        "ELMI_03.tif": "e9ea36b449abf9126e5dfc057e9f11b3aa5f114f571e8cec5bc5225dccfb79f8",
 32        "ELMI_04.tif": "43ac3a04f26c845872ade63efdd6d8d63488d589546b4240576cef2d80aa0a40",
 33        "ELMI_05.tif": "e6005e1ee2a96a2205bd7c2883cfb25111c9c7434abb563d3e2e3e1b002c2d78",
 34    },
 35    "labels": {
 36        "ELMI_01.tif": "6844b7f078446c351ab0cb3149a479c9183eb5ca1381d9a5b574753fd8719035",
 37        "ELMI_02.tif": "17d7f6a85ca536770fd65e696e06d6898c6fe5dc7ca45b945952218f90334021",
 38        "ELMI_03.tif": "74c82d0d5c735dcc0d8955dd462420d4760dd8cb23991bc4fba343727763a597",
 39        "ELMI_04.tif": "144f3fd0ddd4f19780295c48454d88fd9362871cabec78b45d40c874b30bf213",
 40        "ELMI_05.tif": "65de9d6839ae70703cc2c2cb72b6ff864db60632b0457fa899825f0df799c6e5",
 41    },
 42}
 43# The folders of the images and the labels on the share.
 44FOLDERS = {"images": "ELMI_orig_images", "labels": "ELMI_manual_annotation"}
 45
 46
 47def get_elmi_mito_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 48    """Download the ELMI Mito dataset.
 49
 50    Args:
 51        path: Filepath to a folder where the downloaded data will be saved.
 52        download: Whether to download the data if it is not present.
 53
 54    Returns:
 55        The filepath to the folder with the data.
 56    """
 57    for kind, folder in FOLDERS.items():
 58        os.makedirs(os.path.join(path, folder), exist_ok=True)
 59        for name in IMAGE_NAMES:
 60            url = f"{BASE_URL}?path=%2FELMI_images%2F{folder}&files={name}"
 61            util.download_source(os.path.join(path, folder, name), url, download, checksum=CHECKSUMS[kind][name])
 62
 63    return path
 64
 65
 66def get_elmi_mito_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 67    """Get paths to the ELMI Mito data.
 68
 69    Args:
 70        path: Filepath to a folder where the downloaded data will be saved.
 71        download: Whether to download the data if it is not present.
 72
 73    Returns:
 74        List of filepaths for the image data.
 75        List of filepaths for the label data.
 76    """
 77    get_elmi_mito_data(path, download)
 78    raw_paths = natsorted(os.path.join(path, FOLDERS["images"], name) for name in IMAGE_NAMES)
 79    label_paths = natsorted(os.path.join(path, FOLDERS["labels"], name) for name in IMAGE_NAMES)
 80    return raw_paths, label_paths
 81
 82
 83def _binary_labels(labels):
 84    return (labels > 0).astype("uint8")
 85
 86
 87def get_elmi_mito_dataset(
 88    path: Union[os.PathLike, str],
 89    patch_shape: Tuple[int, int],
 90    binary: bool = False,
 91    download: bool = False,
 92    **kwargs
 93) -> Dataset:
 94    """Get the ELMI Mito dataset for mitochondria segmentation in TEM.
 95
 96    Args:
 97        path: Filepath to a folder where the downloaded data will be saved.
 98        patch_shape: The patch shape to use for training.
 99        binary: Whether to return a binary foreground mask. Otherwise the mitochondria instances are returned.
100        download: Whether to download the data if it is not present.
101        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
102
103    Returns:
104        The segmentation dataset.
105    """
106    raw_paths, label_paths = get_elmi_mito_paths(path, download)
107
108    if binary:
109        kwargs = util.update_kwargs(kwargs, "label_transform", _binary_labels)
110
111    return torch_em.default_segmentation_dataset(
112        raw_paths=raw_paths,
113        raw_key=None,
114        label_paths=label_paths,
115        label_key=None,
116        patch_shape=patch_shape,
117        is_seg_dataset=False,
118        **kwargs
119    )
120
121
122def get_elmi_mito_loader(
123    path: Union[os.PathLike, str],
124    batch_size: int,
125    patch_shape: Tuple[int, int],
126    binary: bool = False,
127    download: bool = False,
128    **kwargs
129) -> DataLoader:
130    """Get the ELMI Mito dataloader for mitochondria segmentation in TEM.
131
132    Args:
133        path: Filepath to a folder where the downloaded data will be saved.
134        batch_size: The batch size for training.
135        patch_shape: The patch shape to use for training.
136        binary: Whether to return a binary foreground mask. Otherwise the mitochondria instances are returned.
137        download: Whether to download the data if it is not present.
138        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
139
140    Returns:
141        The DataLoader.
142    """
143    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
144    dataset = get_elmi_mito_dataset(path, patch_shape, binary, download, **ds_kwargs)
145    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)