torch_em.data.datasets.medical.wmh

  1"""The WMH dataset contains annotations for white matter hyperintensity segmentation
  2in brain MRI (FLAIR and T1) from the MICCAI 2017 WMH Segmentation Challenge.
  3
  4The data comprises 60 training and 110 test subjects, acquired at three sites (Utrecht, Singapore and Amsterdam)
  5on five different scanners. Each subject has a FLAIR and a T1 scan (the T1 has been registered to the FLAIR space
  6and both have been bias field corrected in the provided 'pre' folder) and a manual annotation in FLAIR space.
  7
  8NOTE: The label legend is as follows:
  9- background: 0, white matter hyperintensity: 1, other pathology: 2
 10
 11The dataset is located at https://dataverse.nl/dataset.xhtml?persistentId=doi:10.34894/AECRSD.
 12The challenge website is https://wmh.isi.uu.nl/.
 13
 14This dataset is from the publication https://doi.org/10.1109/TMI.2019.2905770.
 15Please cite it if you use this dataset in your research.
 16"""
 17
 18import os
 19from glob import glob
 20from natsort import natsorted
 21from typing import Union, Tuple, Literal, List, Optional
 22
 23import requests
 24
 25from torch.utils.data import Dataset, DataLoader
 26
 27import torch_em
 28
 29from .. import util
 30
 31
 32DOI = "doi:10.34894/AECRSD"
 33API_URL = "https://dataverse.nl/api"
 34URL = f"{API_URL}/datasets/:persistentId/?persistentId={DOI}"
 35
 36# NOTE: The dataset is downloaded file-by-file via the dataverse API (the full-dataset zip times out),
 37# so there is no single archive checksum.
 38CHECKSUM = None
 39
 40SPLIT_DIRS = {"train": "training", "test": "test"}
 41N_SUBJECTS = {"train": 60, "test": 110}
 42FILENAMES = ["FLAIR.nii.gz", "T1.nii.gz", "wmh.nii.gz"]
 43
 44
 45def _get_wmh_file_list(split):
 46    response = requests.get(URL)
 47    response.raise_for_status()
 48    files = response.json()["data"]["latestVersion"]["files"]
 49
 50    to_download = []
 51    for f in files:
 52        parts = f.get("directoryLabel", "").split("/")
 53        fname = f["dataFile"]["filename"]
 54        if parts[0] != SPLIT_DIRS[split] or fname not in FILENAMES:
 55            continue
 56
 57        # The FLAIR and T1 scans are used from the preprocessed 'pre' folder (bias field corrected, T1 in FLAIR space).
 58        # The annotation 'wmh.nii.gz' is located directly in the subject folder.
 59        if parts[-1] == "pre":
 60            parts = parts[:-1]
 61        elif fname != "wmh.nii.gz":
 62            continue
 63
 64        # The Amsterdam Philips scanner folder is named 'Philips_VU .PETMR_01.', we shorten it to 'Philips_VU'.
 65        parts = [p.split(" ")[0] for p in parts]
 66        to_download.append((os.path.join(*parts, fname), f["dataFile"]["id"]))
 67
 68    return to_download
 69
 70
 71def get_wmh_data(path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False) -> str:
 72    """Download the WMH dataset.
 73
 74    Args:
 75        path: Filepath to a folder where the data is downloaded for further processing.
 76        split: The choice of data split. Either 'train' or 'test'.
 77        download: Whether to download the data if it is not present.
 78
 79    Returns:
 80        Filepath where the data for the split is stored.
 81    """
 82    if split not in SPLIT_DIRS:
 83        raise ValueError(f"'{split}' is not a valid split. Please choose from {list(SPLIT_DIRS.keys())}.")
 84
 85    data_dir = os.path.join(path, SPLIT_DIRS[split])
 86    label_paths = glob(os.path.join(data_dir, "**", "wmh.nii.gz"), recursive=True)
 87    if len(label_paths) == N_SUBJECTS[split]:
 88        return data_dir
 89
 90    if not download:
 91        raise RuntimeError(f"Cannot find the data at {data_dir}, but download was set to False.")
 92
 93    for rel_path, file_id in _get_wmh_file_list(split):
 94        fpath = os.path.join(path, rel_path)
 95        os.makedirs(os.path.split(fpath)[0], exist_ok=True)
 96        util.download_source(path=fpath, url=f"{API_URL}/access/datafile/{file_id}", download=download, checksum=None)
 97
 98    return data_dir
 99
100
101def get_wmh_paths(
102    path: Union[os.PathLike, str],
103    split: Literal["train", "test"],
104    modality: Optional[Literal["FLAIR", "T1"]] = None,
105    site: Optional[Literal["Utrecht", "Singapore", "Amsterdam"]] = None,
106    download: bool = False,
107) -> Tuple[List[Union[str, Tuple[str, str]]], List[str]]:
108    """Get paths to the WMH data.
109
110    Args:
111        path: Filepath to a folder where the data is downloaded for further processing.
112        split: The choice of data split. Either 'train' or 'test'.
113        modality: The choice of modality. Either 'FLAIR' or 'T1'. By default, both are returned as channels.
114        site: The acquisition site. One of 'Utrecht', 'Singapore' or 'Amsterdam'. By default, all sites are returned.
115        download: Whether to download the data if it is not present.
116
117    Returns:
118        List of filepaths for the image data.
119        List of filepaths for the label data.
120    """
121    data_dir = get_wmh_data(path, split, download)
122
123    site_dir = "*" if site is None else site
124    label_paths = natsorted(glob(os.path.join(data_dir, site_dir, "**", "wmh.nii.gz"), recursive=True))
125    if len(label_paths) == 0:
126        raise ValueError(f"Could not find any data for split '{split}' and site '{site}'.")
127
128    flair_paths = [p.replace("wmh.nii.gz", "FLAIR.nii.gz") for p in label_paths]
129    t1_paths = [p.replace("wmh.nii.gz", "T1.nii.gz") for p in label_paths]
130    assert all(os.path.exists(p) for p in flair_paths + t1_paths)
131
132    if modality is None:
133        raw_paths = [(fp, tp) for fp, tp in zip(flair_paths, t1_paths)]
134    elif modality == "FLAIR":
135        raw_paths = flair_paths
136    elif modality == "T1":
137        raw_paths = t1_paths
138    else:
139        raise ValueError(f"'{modality}' is not a valid modality. Please choose from 'FLAIR' or 'T1'.")
140
141    return raw_paths, label_paths
142
143
144def get_wmh_dataset(
145    path: Union[os.PathLike, str],
146    patch_shape: Tuple[int, ...],
147    split: Literal["train", "test"],
148    modality: Optional[Literal["FLAIR", "T1"]] = None,
149    site: Optional[Literal["Utrecht", "Singapore", "Amsterdam"]] = None,
150    resize_inputs: bool = False,
151    download: bool = False,
152    **kwargs
153) -> Dataset:
154    """Get the WMH dataset for white matter hyperintensity segmentation.
155
156    Args:
157        path: Filepath to a folder where the data is downloaded for further processing.
158        patch_shape: The patch shape to use for training.
159        split: The choice of data split. Either 'train' or 'test'.
160        modality: The choice of modality. Either 'FLAIR' or 'T1'. By default, both are used as channels.
161        site: The acquisition site. One of 'Utrecht', 'Singapore' or 'Amsterdam'. By default, all sites are used.
162        resize_inputs: Whether to resize inputs to the desired patch shape.
163        download: Whether to download the data if it is not present.
164        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
165
166    Returns:
167        The segmentation dataset.
168    """
169    raw_paths, label_paths = get_wmh_paths(path, split, modality, site, download)
170
171    if resize_inputs:
172        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
173        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
174            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
175        )
176
177    return torch_em.default_segmentation_dataset(
178        raw_paths=raw_paths,
179        raw_key="data",
180        label_paths=label_paths,
181        label_key="data",
182        patch_shape=patch_shape,
183        is_seg_dataset=True,
184        with_channels=modality is None,
185        **kwargs
186    )
187
188
189def get_wmh_loader(
190    path: Union[os.PathLike, str],
191    batch_size: int,
192    patch_shape: Tuple[int, ...],
193    split: Literal["train", "test"],
194    modality: Optional[Literal["FLAIR", "T1"]] = None,
195    site: Optional[Literal["Utrecht", "Singapore", "Amsterdam"]] = None,
196    resize_inputs: bool = False,
197    download: bool = False,
198    **kwargs
199) -> DataLoader:
200    """Get the WMH dataloader for white matter hyperintensity segmentation.
201
202    Args:
203        path: Filepath to a folder where the data is downloaded for further processing.
204        batch_size: The batch size for training.
205        patch_shape: The patch shape to use for training.
206        split: The choice of data split. Either 'train' or 'test'.
207        modality: The choice of modality. Either 'FLAIR' or 'T1'. By default, both are used as channels.
208        site: The acquisition site. One of 'Utrecht', 'Singapore' or 'Amsterdam'. By default, all sites are used.
209        resize_inputs: Whether to resize inputs to the desired patch shape.
210        download: Whether to download the data if it is not present.
211        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
212
213    Returns:
214        The DataLoader.
215    """
216    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
217    dataset = get_wmh_dataset(path, patch_shape, split, modality, site, resize_inputs, download, **ds_kwargs)
218    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)