torch_em.data.datasets.electron_microscopy.lewy_mito

  1"""The Lewy Mito dataset contains annotations for mitochondria segmentation in transmission electron microscopy (TEM)
  2images of post-mortem human brain tissue.
  3
  4The tissue is from the substantia nigra (SN), entorhinal cortex (ENT) and hippocampus (HC) of donors with Parkinson's
  5disease (PD) and dementia with Lewy bodies (DLB), and of healthy controls. The images were acquired as part of a
  6correlative light and electron microscopy workflow. The data is organized into the following groups:
  7- 'ENT-DLB', 'ENT-PD', 'HC-ENT', 'HC-SN', 'SN-DLB', 'SN-PD': Images from the regions and conditions in their names.
  8- 'HN-ENT-DLB', 'HN-ENT-PD', 'HN-SN-DLB', 'HN-SN-PD': The same regions and conditions, from a second image set
  9  (called 'HN' in the original data, its meaning is not explained in the dataset description).
 10
 11There are roughly 16,500 small image crops (a few hundred pixels per side) with manual (and proofread) mitochondria
 12masks. In total the download size is ~7 GB. The files are only available individually (not as an archive),
 13so downloading the full dataset takes some time. Use `groups` to download only a subset.
 14
 15The data is available at https://www.ebi.ac.uk/biostudies/studies/S-BIAD3729 (doi: 10.6019/S-BIAD3729).
 16The licence of the BioStudies entry could not be confirmed (BioImage Archive entries are CC0 by default).
 17The dataset is from the publication "Correlative Ultrastructural Mapping of Lewy Pathology Reveals Regional Diversity
 18in Parkinson's and Dementia with Lewy bodies". Please cite it if you use this dataset in your research.
 19"""
 20
 21import os
 22import json
 23import time
 24import warnings
 25from concurrent.futures import ThreadPoolExecutor
 26from typing import List, Optional, Tuple, Union
 27
 28import numpy as np
 29from natsort import natsorted
 30from torch.utils.data import Dataset, DataLoader
 31
 32import torch_em
 33
 34from .. import util
 35
 36
 37BASE_URL = "https://ftp.ebi.ac.uk/biostudies/fire/S-BIAD/729/S-BIAD3729/Files"
 38FILE_LIST_URL = "https://www.ebi.ac.uk/biostudies/files/S-BIAD3729/Shafiei_Mito-Segmentation.json"
 39ROOT_FOLDER = "Shafiei_Mito-Segmentation"
 40GROUPS = [
 41    "ENT-DLB", "ENT-PD", "HC-ENT", "HC-SN", "HN-ENT-DLB", "HN-ENT-PD", "HN-SN-DLB", "HN-SN-PD", "SN-DLB", "SN-PD",
 42]
 43# NOTE: The file list of the study does not contain checksums, so the downloaded files are not verified.
 44
 45
 46def _get_file_list(path):
 47    list_path = os.path.join(path, "file_list.json")
 48    util.download_source(list_path, FILE_LIST_URL, download=True)
 49    with open(list_path) as f:
 50        return json.load(f)
 51
 52
 53def _download_group(path, group, file_list, num_workers):
 54    # Find all image-mask pairs of the group. Empty (broken) files and their partners are skipped.
 55    sizes = {x["path"]: x["size"] for x in file_list if x["path"].endswith(".tif")}
 56    pairs = []
 57    for rel_path, size in sizes.items():
 58        parts = rel_path.split("/")
 59        if parts[1] != group or parts[3] != "crops":
 60            continue
 61        mask_rel_path = "/".join(parts[:3] + ["masks", parts[-1]])
 62        if size == 0 or sizes.get(mask_rel_path, 0) == 0:
 63            continue
 64        pairs.append((rel_path, mask_rel_path))
 65
 66    def _is_complete(rel_path):
 67        file_path = os.path.join(path, rel_path)
 68        return os.path.exists(file_path) and os.path.getsize(file_path) == sizes[rel_path]
 69
 70    def _download(rel_path):
 71        # The server sometimes returns an html page (instead of the file) under load.
 72        # We check the downloaded file against the size in the file list and retry if necessary.
 73        file_path = os.path.join(path, rel_path)
 74        for attempt in range(5):
 75            try:
 76                if os.path.exists(file_path) and not _is_complete(rel_path):
 77                    os.remove(file_path)
 78                util.download_source(file_path, f"{BASE_URL}/{rel_path}", download=True)
 79                if _is_complete(rel_path):
 80                    return
 81            except Exception:
 82                pass
 83            time.sleep(2 ** attempt)
 84        return rel_path
 85
 86    todo = [p for pair in pairs for p in pair if not _is_complete(p)]
 87    for rel_path in todo:
 88        os.makedirs(os.path.dirname(os.path.join(path, rel_path)), exist_ok=True)
 89    with ThreadPoolExecutor(num_workers) as pool:
 90        failed = {rel_path for rel_path in pool.map(_download, todo) if rel_path is not None}
 91
 92    # Some files that are listed in the study cannot be retrieved from the server (it returns an index page).
 93    # We skip the image-mask pairs that are affected by this.
 94    if failed:
 95        warnings.warn(f"Skipping {len(failed)} files of group '{group}' that cannot be downloaded: {sorted(failed)}")
 96        for rel_path in failed:
 97            if os.path.exists(os.path.join(path, rel_path)):
 98                os.remove(os.path.join(path, rel_path))
 99        pairs = [pair for pair in pairs if not failed.intersection(pair)]
100
101    return pairs
102
103
104def get_lewy_mito_data(
105    path: Union[os.PathLike, str],
106    groups: Optional[List[str]] = None,
107    download: bool = False,
108    num_workers: int = 8,
109) -> List[Tuple[str, str]]:
110    """Download the Lewy Mito dataset.
111
112    Args:
113        path: Filepath to a folder where the downloaded data will be saved.
114        groups: The groups (region and condition) to download. By default all groups are used.
115            Choose from 'ENT-DLB', 'ENT-PD', 'HC-ENT', 'HC-SN', 'HN-ENT-DLB', 'HN-ENT-PD', 'HN-SN-DLB',
116            'HN-SN-PD', 'SN-DLB' and 'SN-PD'.
117        download: Whether to download the data if it is not present.
118        num_workers: The number of threads used for downloading the individual files.
119
120    Returns:
121        The relative filepaths (image, mask) of all image-mask pairs in the selected groups.
122    """
123    groups = GROUPS if groups is None else groups
124    for group in groups:
125        if group not in GROUPS:
126            raise ValueError(f"'{group}' is not a valid group. Please choose from {GROUPS}.")
127
128    os.makedirs(path, exist_ok=True)
129
130    # Cache the pair list per group, so that the data can be used offline once it was downloaded.
131    pairs = []
132    file_list = None
133    for group in groups:
134        pair_path = os.path.join(path, f"pairs_{group}.json")
135        if os.path.exists(pair_path):
136            with open(pair_path) as f:
137                group_pairs = json.load(f)
138        else:
139            if not download:
140                raise RuntimeError(f"Cannot find the data for group '{group}' at {path}. Set download=True.")
141            file_list = _get_file_list(path) if file_list is None else file_list
142            group_pairs = _download_group(path, group, file_list, num_workers)
143            with open(pair_path, "w") as f:
144                json.dump(group_pairs, f)
145        pairs.extend(group_pairs)
146
147    return pairs
148
149
150def get_lewy_mito_paths(
151    path: Union[os.PathLike, str],
152    groups: Optional[List[str]] = None,
153    download: bool = False,
154    num_workers: int = 8,
155) -> Tuple[List[str], List[str]]:
156    """Get paths to the Lewy Mito data.
157
158    Args:
159        path: Filepath to a folder where the downloaded data will be saved.
160        groups: The groups (region and condition) to use. By default all groups are used.
161        download: Whether to download the data if it is not present.
162        num_workers: The number of threads used for downloading the individual files.
163
164    Returns:
165        List of filepaths for the image data.
166        List of filepaths for the label data.
167    """
168    pairs = get_lewy_mito_data(path, groups, download, num_workers)
169    pairs = natsorted(pairs)
170    raw_paths = [os.path.join(path, raw) for raw, _ in pairs]
171    label_paths = [os.path.join(path, label) for _, label in pairs]
172    return raw_paths, label_paths
173
174
175def _binary_labels(labels):
176    return (labels > 0).astype("uint8")
177
178
179def get_lewy_mito_dataset(
180    path: Union[os.PathLike, str],
181    patch_shape: Tuple[int, int],
182    groups: Optional[List[str]] = None,
183    binary: bool = True,
184    download: bool = False,
185    **kwargs
186) -> Dataset:
187    """Get the Lewy Mito dataset for mitochondria segmentation in TEM.
188
189    Args:
190        path: Filepath to a folder where the downloaded data will be saved.
191        patch_shape: The patch shape to use for training.
192        groups: The groups (region and condition) to use. By default all groups are used.
193        binary: Whether to return a binary foreground mask. Otherwise the mask values are returned as they are.
194        download: Whether to download the data if it is not present.
195        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
196
197    Returns:
198        The segmentation dataset.
199    """
200    raw_paths, label_paths = get_lewy_mito_paths(path, groups, download)
201
202    if binary:
203        kwargs = util.update_kwargs(kwargs, "label_transform", _binary_labels)
204
205    return torch_em.default_segmentation_dataset(
206        raw_paths=raw_paths,
207        raw_key=None,
208        label_paths=label_paths,
209        label_key=None,
210        patch_shape=patch_shape,
211        is_seg_dataset=False,
212        **kwargs
213    )
214
215
216def get_lewy_mito_loader(
217    path: Union[os.PathLike, str],
218    batch_size: int,
219    patch_shape: Tuple[int, int],
220    groups: Optional[List[str]] = None,
221    binary: bool = True,
222    download: bool = False,
223    **kwargs
224) -> DataLoader:
225    """Get the Lewy Mito dataloader for mitochondria segmentation in TEM.
226
227    Args:
228        path: Filepath to a folder where the downloaded data will be saved.
229        batch_size: The batch size for training.
230        patch_shape: The patch shape to use for training.
231        groups: The groups (region and condition) to use. By default all groups are used.
232        binary: Whether to return a binary foreground mask. Otherwise the mask values are returned as they are.
233        download: Whether to download the data if it is not present.
234        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
235
236    Returns:
237        The DataLoader.
238    """
239    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
240    dataset = get_lewy_mito_dataset(path, patch_shape, groups, binary, download, **ds_kwargs)
241    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)