torch_em.data.datasets.medical.atlas_stroke

  1"""The ATLAS Stroke dataset contains annotations for stroke lesion segmentation in T1-weighted brain MRI.
  2
  3It is version 2.0 of ATLAS (Anatomical Tracings of Lesions After Stroke, https://atlas.grand-challenge.org),
  4a collection of 1271 T1-weighted MRI scans of chronic stroke patients pooled from 44 research cohorts, in
  5which the lesions were traced manually. The release is split into a training set of 655 scans with public
  6lesion masks, a test set of 300 scans whose masks are withheld for the challenge evaluation, and a
  7generalizability set of 316 scans whose scans and masks are both withheld. This module therefore provides
  8the 655 annotated scans of the training set.
  9
 10The scans are defaced, intensity normalized and registered to the MNI-152 template
 11('space-MNI152NLin2009aSym'), so all of them have a shape of (197, 233, 189) at 1 mm isotropic resolution.
 12
 13The annotations are binary, see `LABEL_IDS`: 1 = stroke lesion.
 14
 15NOTE: The official release at https://fcon_1000.projects.nitrc.org/indi/retro/atlas.html is only handed out
 16after agreeing to the terms of use in a request form, and is distributed as an encrypted archive, so it
 17cannot be downloaded automatically. This module downloads a public mirror of the training set at
 18https://huggingface.co/datasets/jayzzzzz0134/atlas-stroke instead. If the official BIDS release is extracted
 19into the folder passed as 'path', so that files such as
 20'<path>/**/sub-r001s001/ses-1/anat/sub-r001s001_ses-1_space-MNI152NLin2009aSym_T1w.nii.gz' and the matching
 21'..._label-L_desc-T1lesion_mask.nii.gz' exist, it is used instead of the mirror.
 22
 23The scans are used as nifti volumes directly (the key is 'data'). They are loaded with the axis order
 24reversed with respect to the nifti file, i.e. (Z, Y, X), so that a 2d patch shape selects axial slices.
 25
 26This dataset is from the publication https://doi.org/10.1038/s41597-022-01401-7.
 27Please cite it if you use this dataset in your research.
 28"""
 29
 30import os
 31import json
 32from glob import glob
 33from tqdm import tqdm
 34from natsort import natsorted
 35from typing import Union, Tuple, List
 36
 37from torch.utils.data import Dataset, DataLoader
 38
 39import torch_em
 40
 41from .. import util
 42
 43
 44URL_BASE = "https://huggingface.co/datasets/jayzzzzz0134/atlas-stroke/resolve/main"
 45
 46API_URL = "https://huggingface.co/api/datasets/jayzzzzz0134/atlas-stroke/tree/main/masks"
 47
 48LABEL_IDS = {"background": 0, "stroke_lesion": 1}
 49
 50N_SUBJECTS = 655
 51
 52N_RETRIES = 5
 53
 54
 55def _get_subject_ids(path, download):
 56    """List the subjects of the mirror via the huggingface API and cache the listing next to the data."""
 57    listing_path = os.path.join(path, "subject_ids.json")
 58    if os.path.exists(listing_path):
 59        with open(listing_path, "r") as f:
 60            return json.load(f)
 61
 62    if not download:
 63        raise RuntimeError(f"Cannot find the data at '{path}', but download was set to False.")
 64
 65    import requests
 66
 67    subject_ids, cursor = [], None
 68    while True:
 69        params = {"limit": 1000}
 70        if cursor is not None:
 71            params["cursor"] = cursor
 72
 73        response = requests.get(API_URL, params=params)
 74        response.raise_for_status()
 75        subject_ids.extend(
 76            os.path.basename(entry["path"]).replace("_lesion_mask.nii.gz", "") for entry in response.json()
 77        )
 78
 79        link = response.headers.get("Link", "")
 80        if 'rel="next"' not in link:
 81            break
 82        cursor = link.split("cursor=")[1].split("&")[0].split(">")[0]
 83
 84    subject_ids = natsorted(subject_ids)
 85    assert len(subject_ids) == N_SUBJECTS, f"Expected {N_SUBJECTS} subjects in the mirror, got {len(subject_ids)}."
 86
 87    with open(listing_path, "w") as f:
 88        json.dump(subject_ids, f)
 89
 90    return subject_ids
 91
 92
 93def _find_official_data(path):
 94    """Find the scans of the official BIDS release, in case it was downloaded manually."""
 95    pattern = os.path.join(path, "**", "sub-*_space-MNI152NLin2009aSym_T1w.nii.gz")
 96    raw_paths = natsorted(glob(pattern, recursive=True))
 97    label_paths = [p.replace("_T1w.nii.gz", "_label-L_desc-T1lesion_mask.nii.gz") for p in raw_paths]
 98
 99    keep = [i for i, p in enumerate(label_paths) if os.path.exists(p)]
100    return [raw_paths[i] for i in keep], [label_paths[i] for i in keep]
101
102
103def _download_volumes(path, download):
104    for folder in ["images", "masks"]:
105        os.makedirs(os.path.join(path, folder), exist_ok=True)
106
107    raw_paths, label_paths = [], []
108    for subject_id in tqdm(_get_subject_ids(path, download), desc="Downloading the ATLAS Stroke scans"):
109        for folder, fname in [("images", f"{subject_id}_T1w.nii.gz"), ("masks", f"{subject_id}_lesion_mask.nii.gz")]:
110            fpath = os.path.join(path, folder, fname)
111            # The mirror is fetched file by file, so a transient error is retried instead of failing the download.
112            for attempt in range(N_RETRIES):
113                try:
114                    util.download_source(path=fpath, url=f"{URL_BASE}/{folder}/{fname}", download=download)
115                    break
116                except Exception:
117                    if attempt == N_RETRIES - 1:
118                        raise
119
120            (raw_paths if folder == "images" else label_paths).append(fpath)
121
122    return raw_paths, label_paths
123
124
125def get_atlas_stroke_data(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
126    """Download the ATLAS v2.0 dataset.
127
128    Args:
129        path: Filepath to a folder where the data is downloaded for further processing.
130        download: Whether to download the data if it is not present.
131
132    Returns:
133        List of filepaths for the image data.
134        List of filepaths for the label data.
135    """
136    os.makedirs(path, exist_ok=True)
137
138    raw_paths, label_paths = _find_official_data(path)
139    if len(raw_paths) > 0:
140        return raw_paths, label_paths
141
142    return _download_volumes(path, download)
143
144
145def get_atlas_stroke_paths(
146    path: Union[os.PathLike, str], download: bool = False
147) -> Tuple[List[str], List[str]]:
148    """Get paths to the ATLAS v2.0 data.
149
150    Args:
151        path: Filepath to a folder where the data is downloaded for further processing.
152        download: Whether to download the data if it is not present.
153
154    Returns:
155        List of filepaths for the image data.
156        List of filepaths for the label data.
157    """
158    raw_paths, label_paths = get_atlas_stroke_data(path, download)
159    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0, f"Could not find the scans in '{path}'."
160    return raw_paths, label_paths
161
162
163def get_atlas_stroke_dataset(
164    path: Union[os.PathLike, str],
165    patch_shape: Tuple[int, ...],
166    resize_inputs: bool = False,
167    download: bool = False,
168    **kwargs
169) -> Dataset:
170    """Get the ATLAS v2.0 dataset for stroke lesion segmentation.
171
172    Args:
173        path: Filepath to a folder where the data is downloaded for further processing.
174        patch_shape: The patch shape to use for training.
175        resize_inputs: Whether to resize inputs to the desired patch shape.
176        download: Whether to download the data if it is not present.
177        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
178
179    Returns:
180        The segmentation dataset.
181    """
182    raw_paths, label_paths = get_atlas_stroke_paths(path, download)
183
184    if resize_inputs:
185        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
186        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
187            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
188        )
189
190    return torch_em.default_segmentation_dataset(
191        raw_paths=raw_paths,
192        raw_key="data",
193        label_paths=label_paths,
194        label_key="data",
195        patch_shape=patch_shape,
196        is_seg_dataset=True,
197        **kwargs
198    )
199
200
201def get_atlas_stroke_loader(
202    path: Union[os.PathLike, str],
203    batch_size: int,
204    patch_shape: Tuple[int, ...],
205    resize_inputs: bool = False,
206    download: bool = False,
207    **kwargs
208) -> DataLoader:
209    """Get the ATLAS v2.0 dataloader for stroke lesion segmentation.
210
211    Args:
212        path: Filepath to a folder where the data is downloaded for further processing.
213        batch_size: The batch size for training.
214        patch_shape: The patch shape to use for training.
215        resize_inputs: Whether to resize inputs to the desired patch shape.
216        download: Whether to download the data if it is not present.
217        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
218
219    Returns:
220        The DataLoader.
221    """
222    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
223    dataset = get_atlas_stroke_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
224    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)