torch_em.data.datasets.medical.segthy

  1"""The SegThy dataset contains annotations for thyroid segmentation in MRI and US scans,
  2and additional annotations for vein and artery segmentation in MRI.
  3
  4NOTE: The label legends are described as following:
  51: For thyroid-only labels: (at 'MRI_thyroid' or 'US_thyroid')
  6- background: 0 and thyroid: 1
  72: For thyroid, jugular veins and carotid arteries (at 'MRI_thyroid+jugular+carotid_label')
  8- background: 0, thyroid: 1, jugular vein: 3 and 5, carotid artery: 2 and 4.
  9
 10The dataset is located at https://www.cs.cit.tum.de/camp/publications/segthy-dataset/.
 11
 12This dataset is from the publication https://doi.org/10.1371/journal.pone.0268550.
 13Please cite it if you use this dataset in your research.
 14"""
 15
 16import os
 17from glob import glob
 18from natsort import natsorted
 19from typing import Union, Tuple, Literal, List
 20
 21import numpy as np
 22
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30URLS = {
 31    "MRI": "https://www.campar.in.tum.de/public_datasets/2022_plosone_eilers/MRI_data.zip",
 32    "US": "https://www.campar.in.tum.de/public_datasets/2022_plosone_eilers/US_data.zip",
 33}
 34
 35CHECKSUMS = {
 36    "MRI": "e9d0599b305dfe36795c45282a8495d3bfb4a872851c221b321d59ed0b11e7eb",
 37    "US": "52c59ef4db08adfa0e6ea562c7fe747c612f2064e01f907a78b170b02fb459bb",
 38}
 39
 40
 41def get_segthy_data(path: Union[os.PathLike, str], source: Literal['MRI', 'US'], download: bool = False):
 42    """Download the SegThy dataset.
 43
 44    Args:
 45        path: Filepath to a folder where the data is downloaded for further processing.
 46        download: Whether to download the data if it is not present.
 47    """
 48    data_dir = os.path.join(path, f"{source}_volunteer_dataset")
 49    if os.path.exists(data_dir):
 50        return
 51
 52    os.makedirs(path, exist_ok=True)
 53
 54    zip_path = os.path.join(path, f"{source}_data.zip")
 55    util.download_source(path=zip_path, url=URLS[source], download=download, checksum=CHECKSUMS[source])
 56    util.unzip(zip_path=zip_path, dst=path)
 57
 58    # NOTE: There is one label with an empty channel.
 59    if source == "MRI":
 60        lpath = os.path.join(data_dir, "MRI_thyroid_label", "005_MRI_thyroid_label.nii.gz")
 61
 62        import nibabel as nib
 63        # Load the label volume and remove the empty channel.
 64        label = nib.load(lpath).get_fdata()
 65        label = label[..., 0]
 66
 67        # Store the updated label.
 68        label_nifti = nib.Nifti2Image(label, np.eye(4))
 69        nib.save(label_nifti, lpath)
 70
 71
 72def get_segthy_paths(
 73    path: Union[os.PathLike, str],
 74    split: Literal['train', 'val', 'test'],
 75    source: Literal['MRI', 'US'],
 76    region: Literal['thyroid', 'thyroid_and_vessels'] = "thyroid",
 77    download: bool = False
 78) -> Tuple[List[str], List[str]]:
 79    """Get paths to the SegThy data.
 80
 81    Args:
 82        path: Filepath to a folder where the data is downloaded for further processing.
 83        split: The choice of data split.
 84        source: The source of dataset. Either 'MRI' or 'US.
 85        region: The labeled regions for the corresponding volumes. Either 'thyroid' or 'thyroid_and_vessels'.
 86        download: Whether to download the data if it is not present.
 87
 88    Returns:
 89        List of filepaths for the image data.
 90        List of filepaths for the label data.
 91    """
 92    get_segthy_data(path, source, download)
 93
 94    if source == "MRI":
 95        ldir = "MRI_thyroid_label" if region == "thyroid" else "MRI_thyroid+jugular+carotid_label"
 96        label_paths = natsorted(glob(os.path.join(path, f"{source}_volunteer_dataset", ldir, "*.nii.gz")))
 97        raw_paths = [p.replace(ldir, "MRI") for p in label_paths]
 98
 99        if split == "train":
100            raw_paths = raw_paths[:15] if region == "thyroid" else raw_paths[:8]
101            label_paths = label_paths[:15] if region == "thyroid" else label_paths[:8]
102        elif split == "val":
103            raw_paths = raw_paths[15:20] if region == "thyroid" else raw_paths[8:10]
104            label_paths = label_paths[15:20] if region == "thyroid" else label_paths[8:10]
105        elif split == "test":
106            raw_paths = raw_paths[20:] if region == "thyroid" else raw_paths[10:]
107            label_paths = label_paths[20:] if region == "thyroid" else label_paths[10:]
108        else:
109            raise ValueError(f"'{split}' is not a valid split.")
110
111    else:  # US data
112        assert region != "thyroid_and_vessels", "US source does not have labels for both thyroid and vessels."
113        ldir = "ground_truth_data/US_thyroid_label"
114        label_paths = natsorted(glob(os.path.join(path, f"{source}_volunteer_dataset", ldir, "*.nii")))
115
116        raw_paths = [p.replace(ldir, "ground_truth_data/US") for p in label_paths]
117        raw_paths = [p.replace(".nii", "_US.nii") for p in raw_paths]
118
119        if split == "train":
120            raw_paths, label_paths = raw_paths[:20], label_paths[:20]
121        elif split == "val":
122            raw_paths, label_paths = raw_paths[20:25], label_paths[20:25]
123        elif split == "test":
124            raw_paths, label_paths = raw_paths[25:], label_paths[25:]
125        else:
126            raise ValueError(f"'{split}' is not a valid split.")
127
128    return raw_paths, label_paths
129
130
131def get_segthy_dataset(
132    path: Union[os.PathLike, str],
133    patch_shape: Tuple[int, ...],
134    split: Literal['train', 'val', 'test'],
135    source: Literal['MRI', 'US'],
136    region: Literal['thyroid', 'thyroid_and_vessels'] = "thyroid",
137    resize_inputs: bool = False,
138    download: bool = False,
139    **kwargs
140) -> Dataset:
141    """Get the SegThy dataset for thyroid (and vessel) segmentation.
142
143    Args:
144        path: Filepath to a folder where the data is downloaded for further processing.
145        patch_shape: The patch shape to use for training.
146        split: The choice of data split.
147        source: The source of dataset. Either 'MRI' or 'US.
148        region: The labeled regions for the corresponding volumes. Either 'thyroid' or 'thyroid_and_vessels'.
149        resize_inputs: Whether to resize inputs to the desired patch shape.
150        download: Whether to download the data if it is not present.
151        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
152
153    Returns:
154        The segmentation dataset.
155    """
156    raw_paths, label_paths = get_segthy_paths(path, split, source, region, download)
157
158    if resize_inputs:
159        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
160        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
161            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
162        )
163
164    return torch_em.default_segmentation_dataset(
165        raw_paths=raw_paths,
166        raw_key="data",
167        label_paths=label_paths,
168        label_key="data",
169        patch_shape=patch_shape,
170        is_seg_dataset=True,
171        **kwargs
172    )
173
174
175def get_segthy_loader(
176    path: Union[os.PathLike, str],
177    batch_size: int,
178    patch_shape: Tuple[int, ...],
179    split: Literal['train', 'val', 'test'],
180    source: Literal['MRI', 'US'],
181    region: Literal['thyroid', 'thyroid_and_vessels'] = "thyroid",
182    resize_inputs: bool = False,
183    download: bool = False,
184    **kwargs
185) -> DataLoader:
186    """Get the SegThy dataloader for thyroid (and vessel) segmentation.
187
188    Args:
189        path: Filepath to a folder where the data is downloaded for further processing.
190        batch_size: The batch size for training.
191        patch_shape: The patch shape to use for training.
192        split: The choice of data split.
193        source: The source of dataset. Either 'MRI' or 'US.
194        region: The labeled regions for the corresponding volumes. Either 'thyroid' or 'thyroid_and_vessels'.
195        resize_inputs: Whether to resize inputs to the desired patch shape.
196        download: Whether to download the data if it is not present.
197        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
198
199    Args:
200        The DataLoader.
201    """
202    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
203    dataset = get_segthy_dataset(path, patch_shape, split, source, region, resize_inputs, download, **ds_kwargs)
204    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)