torch_em.data.datasets.medical.tongue_mri

  1"""The Tongue MRI dataset contains annotations for tongue musculature segmentation in MRI.
  2
  3The dataset consists of 47 healthy subjects collated from three studies: BMC (24 subjects,
  4T1-weighted MRI), EATT4MND (19 subjects, T2-weighted MRI) and BeLong (4 subjects, T2-weighted
  5MRI). Each subject has a manually corrected semantic segmentation of four tongue muscles,
  6following `LABEL_IDS`: the genioglossus, the superior longitudinal, the inferior longitudinal,
  7and the combined transverse / vertical muscle.
  8
  9The dataset also ships study templates for each site (in the 'template' folders next to the
 10'images' and 'labels' folders that this module uses), which are not returned by this module.
 11
 12The dataset is located at https://osf.io/wt9fc/ and is distributed under the CC0 1.0 license.
 13
 14This dataset is from the publication https://doi.org/10.1038/s41597-025-05092-8.
 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
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30# The site sub-folders are generated on-the-fly by OSF, hence the checksums of the archives are not reliable.
 31URLS = {
 32    "BMC": {
 33        "images": "https://files.osf.io/v1/resources/wt9fc/providers/osfstorage/66e224fb34c37b7bbdcb074e/?zip=",
 34        "labels": "https://files.osf.io/v1/resources/wt9fc/providers/osfstorage/66e224ff87c914991b64b5cb/?zip=",
 35    },
 36    "EATT4MND": {
 37        "images": "https://files.osf.io/v1/resources/wt9fc/providers/osfstorage/66e2259e2d9d10b9f9cb0ab2/?zip=",
 38        "labels": "https://files.osf.io/v1/resources/wt9fc/providers/osfstorage/66e225a28cec7392871d9456/?zip=",
 39    },
 40    "BeLong": {
 41        "images": "https://files.osf.io/v1/resources/wt9fc/providers/osfstorage/66e2248cf57074722a64c057/?zip=",
 42        "labels": "https://files.osf.io/v1/resources/wt9fc/providers/osfstorage/66e224931bc7424b3dc05a46/?zip=",
 43    },
 44}
 45
 46SITES = list(URLS.keys())
 47
 48LABEL_IDS = {
 49    "background": 0,
 50    "genioglossus": 1,
 51    "superior_longitudinal": 2,
 52    "inferior_longitudinal": 3,
 53    "transverse_vertical": 4,
 54}
 55"""The semantic label ids of the tongue muscle classes."""
 56
 57
 58def get_tongue_mri_data(
 59    path: Union[os.PathLike, str], site: Literal["BMC", "EATT4MND", "BeLong"] = "BMC", download: bool = False
 60) -> str:
 61    """Download the Tongue MRI dataset.
 62
 63    Args:
 64        path: Filepath to a folder where the data is downloaded for further processing.
 65        site: The study site to download. One of 'BMC', 'EATT4MND', 'BeLong'.
 66        download: Whether to download the data if it is not present.
 67
 68    Returns:
 69        Filepath to the site folder where the data is downloaded.
 70    """
 71    if site not in SITES:
 72        raise ValueError(f"'{site}' is not a valid site. Choose one of {SITES}.")
 73
 74    site_dir = os.path.join(path, site)
 75    os.makedirs(site_dir, exist_ok=True)
 76
 77    for name in ["images", "labels"]:
 78        data_dir = os.path.join(site_dir, name)
 79        if os.path.exists(data_dir):
 80            continue
 81
 82        zip_path = os.path.join(site_dir, f"{name}.zip")
 83        util.download_source(path=zip_path, url=URLS[site][name], download=download, checksum=None)
 84        util.unzip(zip_path=zip_path, dst=data_dir)
 85
 86    return site_dir
 87
 88
 89def get_tongue_mri_paths(
 90    path: Union[os.PathLike, str],
 91    site: Optional[Literal["BMC", "EATT4MND", "BeLong"]] = None,
 92    download: bool = False,
 93) -> Tuple[List[str], List[str]]:
 94    """Get paths to the Tongue MRI data.
 95
 96    Args:
 97        path: Filepath to a folder where the data is downloaded for further processing.
 98        site: The study site to use. One of 'BMC', 'EATT4MND', 'BeLong'. By default, all sites are used.
 99        download: Whether to download the data if it is not present.
100
101    Returns:
102        List of filepaths for the image data.
103        List of filepaths for the label data.
104    """
105    sites = SITES if site is None else [site]
106
107    raw_paths, label_paths = [], []
108    for s in sites:
109        site_dir = get_tongue_mri_data(path, s, download)
110        cur_label_paths = natsorted(glob(os.path.join(site_dir, "labels", "*_labels.nii.gz")))
111        cur_raw_paths = [
112            os.path.join(site_dir, "images", os.path.basename(p).replace("_labels.nii.gz", ".nii.gz"))
113            for p in cur_label_paths
114        ]
115        assert all(os.path.exists(p) for p in cur_raw_paths), f"Some image volumes are missing for site '{s}'."
116        raw_paths.extend(cur_raw_paths)
117        label_paths.extend(cur_label_paths)
118
119    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
120
121    return raw_paths, label_paths
122
123
124def get_tongue_mri_dataset(
125    path: Union[os.PathLike, str],
126    patch_shape: Tuple[int, ...],
127    site: Optional[Literal["BMC", "EATT4MND", "BeLong"]] = None,
128    resize_inputs: bool = False,
129    download: bool = False,
130    **kwargs
131) -> Dataset:
132    """Get the Tongue MRI dataset for tongue musculature segmentation.
133
134    Args:
135        path: Filepath to a folder where the data is downloaded for further processing.
136        patch_shape: The patch shape to use for training.
137        site: The study site to use. One of 'BMC', 'EATT4MND', 'BeLong'. By default, all sites are used.
138        resize_inputs: Whether to resize inputs to the desired patch shape.
139        download: Whether to download the data if it is not present.
140        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
141
142    Returns:
143        The segmentation dataset.
144    """
145    raw_paths, label_paths = get_tongue_mri_paths(path, site, download)
146
147    if resize_inputs:
148        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
149        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
150            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
151        )
152
153    return torch_em.default_segmentation_dataset(
154        raw_paths=raw_paths,
155        raw_key="data",
156        label_paths=label_paths,
157        label_key="data",
158        patch_shape=patch_shape,
159        is_seg_dataset=True,
160        **kwargs
161    )
162
163
164def get_tongue_mri_loader(
165    path: Union[os.PathLike, str],
166    batch_size: int,
167    patch_shape: Tuple[int, ...],
168    site: Optional[Literal["BMC", "EATT4MND", "BeLong"]] = None,
169    resize_inputs: bool = False,
170    download: bool = False,
171    **kwargs
172) -> DataLoader:
173    """Get the Tongue MRI dataloader for tongue musculature segmentation.
174
175    Args:
176        path: Filepath to a folder where the data is downloaded for further processing.
177        batch_size: The batch size for training.
178        patch_shape: The patch shape to use for training.
179        site: The study site to use. One of 'BMC', 'EATT4MND', 'BeLong'. By default, all sites are used.
180        resize_inputs: Whether to resize inputs to the desired patch shape.
181        download: Whether to download the data if it is not present.
182        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
183
184    Returns:
185        The DataLoader.
186    """
187    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
188    dataset = get_tongue_mri_dataset(path, patch_shape, site, resize_inputs, download, **ds_kwargs)
189    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)