torch_em.data.datasets.medical.multi_organ_abdominal_ct

  1"""The Multi-organ Abdominal CT dataset contains reference standard annotations for multiple abdominal organs
  2in CT, released with the paper 'Automatic Multi-organ Segmentation on Abdominal CT with Dense V-networks'.
  3
  4The dataset comprises annotations for 90 abdominal CT volumes: 43 volumes from the TCIA Pancreas-CT collection
  5(one of them, PANCREAS_0025, has since been removed from TCIA, so 42 are available) and 47 volumes from the
  6Beyond the Cranial Vault (BTCV) abdomen challenge. The Zenodo record contains only the annotations,
  7the CT volumes have to be obtained from the original sources:
  8- The TCIA Pancreas-CT volumes are public and downloaded automatically from the TCIA NBIA API.
  9- The BTCV volumes require registration at Synapse. Please download 'RawData.zip' from
 10  https://www.synapse.org/#!Synapse:syn3193805 (Abdomen) and place it in the folder passed as 'path'.
 11
 12NOTE: The label legend is as follows (labels marked with * are only annotated in the BTCV volumes):
 13- background: 0, spleen: 1, right kidney*: 2, left kidney: 3, gallbladder: 4, esophagus: 5, liver: 6, stomach: 7,
 14  aorta*: 8, inferior vena cava*: 9, portal vein and splenic vein*: 10, pancreas: 11, right adrenal gland*: 12,
 15  left adrenal gland*: 13, duodenum: 14
 16The annotations may be incomplete outside of the cropping region specified in 'cropping.csv' (see the Zenodo record).
 17
 18The dataset is located at https://zenodo.org/records/1169361.
 19
 20This dataset is from the publication https://doi.org/10.1109/TMI.2018.2806309.
 21Please cite it if you use this dataset in your research.
 22"""
 23
 24import os
 25import shutil
 26from glob import glob
 27from warnings import warn
 28from tqdm import tqdm
 29from natsort import natsorted
 30from typing import Union, Tuple, Literal, List
 31
 32import numpy as np
 33import requests
 34
 35from torch.utils.data import Dataset, DataLoader
 36
 37import torch_em
 38
 39from .. import util
 40
 41
 42URLS = {
 43    "tcia": "https://zenodo.org/records/1169361/files/label_tciapancreasct_multiorgan.tar.gz?download=1",
 44    "btcv": "https://zenodo.org/records/1169361/files/label_btcv_multiorgan.tar.gz?download=1",
 45    "cropping": "https://zenodo.org/records/1169361/files/cropping.csv?download=1",
 46}
 47
 48CHECKSUMS = {
 49    "tcia": "1790e252ba0732cc06ec727f00245c8d9c1d5f0bb9829fd3e9915863c5100c61",
 50    "btcv": "bb080c7de1094cc0ee46a5e1bef5b66e635075f8e9a3080ca681086cd2723998",
 51    "cropping": "d6503bda3d776c10698523aae14446409bc9d3722fb405b0025b4a5f59094f8f",
 52}
 53
 54NBIA_API_URL = "https://services.cancerimagingarchive.net/nbia-api/services/v1"
 55TCIA_COLLECTION = "Pancreas-CT"
 56
 57# NOTE: The case PANCREAS_0025 has been removed from the TCIA Pancreas-CT collection, so its annotation is skipped.
 58MISSING_TCIA_CASES = ["0025"]
 59
 60LABEL_DIRS = {"tcia": "label_tcia_multiorgan", "btcv": "label_btcv_multiorgan"}
 61IMAGE_DIRS = {"tcia": "image_tcia_multiorgan", "btcv": "image_btcv_multiorgan"}
 62
 63ORGANS = {
 64    "spleen": 1, "right kidney": 2, "left kidney": 3, "gallbladder": 4, "esophagus": 5, "liver": 6, "stomach": 7,
 65    "aorta": 8, "inferior vena cava": 9, "portal vein and splenic vein": 10, "pancreas": 11,
 66    "right adrenal gland": 12, "left adrenal gland": 13, "duodenum": 14,
 67}
 68
 69
 70def _get_label_id(label_path):
 71    return os.path.basename(label_path).replace("label", "").replace(".nii.gz", "")
 72
 73
 74def _download_labels(path, source, download):
 75    label_dir = os.path.join(path, LABEL_DIRS[source])
 76    if not os.path.exists(label_dir):
 77        tar_path = os.path.join(path, f"{LABEL_DIRS[source]}.tar.gz")
 78        util.download_source(path=tar_path, url=URLS[source], download=download, checksum=CHECKSUMS[source])
 79        util.unzip_tarfile(tar_path=tar_path, dst=path)
 80
 81    csv_path = os.path.join(path, "cropping.csv")
 82    if not os.path.exists(csv_path):
 83        util.download_source(path=csv_path, url=URLS["cropping"], download=download, checksum=CHECKSUMS["cropping"])
 84
 85    return label_dir
 86
 87
 88def _get_tcia_series_uids():
 89    response = requests.get(f"{NBIA_API_URL}/getSeries", params={"Collection": TCIA_COLLECTION})
 90    response.raise_for_status()
 91    return {series["PatientID"]: series["SeriesInstanceUID"] for series in response.json()}
 92
 93
 94def _convert_tcia_dicom_to_nifti(dicom_dir, label_path, image_path):
 95    import imageio.v2 as imageio
 96    import nibabel as nib
 97
 98    # The DICOM slices are sorted by their position and stacked to a (Z, Y, X) volume,
 99    # which is transposed to match the (X, Y, Z) axis order of the nifti annotations.
100    volume = np.asarray(imageio.volread(dicom_dir, format="DICOM"))
101    volume = volume.transpose(2, 1, 0)
102
103    # The reference annotations share the geometry of the CT volumes, so we store the CT with the label's affine.
104    label = nib.load(label_path)
105    assert volume.shape == label.shape, f"Shape mismatch between CT {volume.shape} and label {label.shape}."
106    nib.save(nib.Nifti1Image(volume, label.affine), image_path)
107
108
109def _get_label_paths(label_dir, source):
110    label_paths = natsorted(glob(os.path.join(label_dir, "label*.nii.gz")))
111    if source == "tcia":
112        label_paths = [p for p in label_paths if _get_label_id(p) not in MISSING_TCIA_CASES]
113    return label_paths
114
115
116def _prepare_tcia_images(path, label_dir, download):
117    image_dir = os.path.join(path, IMAGE_DIRS["tcia"])
118    label_paths = _get_label_paths(label_dir, "tcia")
119    if len(glob(os.path.join(image_dir, "*.nii.gz"))) == len(label_paths):
120        return image_dir
121
122    if not download:
123        raise RuntimeError(f"Cannot find the CT volumes at {image_dir}, but download was set to False.")
124
125    os.makedirs(image_dir, exist_ok=True)
126    dicom_root = os.path.join(path, "tcia_dicom")
127    os.makedirs(dicom_root, exist_ok=True)
128
129    series_uids = _get_tcia_series_uids()
130    for label_path in tqdm(label_paths, desc="Download and convert the TCIA Pancreas-CT volumes"):
131        label_id = _get_label_id(label_path)
132        image_path = os.path.join(image_dir, f"img{label_id}.nii.gz")
133        if os.path.exists(image_path):
134            continue
135
136        patient_id = f"PANCREAS_{label_id}"
137        if patient_id not in series_uids:
138            warn(f"The case '{patient_id}' is not available in the TCIA Pancreas-CT collection and will be skipped.")
139            continue
140
141        zip_path = os.path.join(dicom_root, f"{patient_id}.zip")
142        dicom_dir = os.path.join(dicom_root, patient_id)
143        if not os.path.exists(dicom_dir):
144            url = f"{NBIA_API_URL}/getImage?SeriesInstanceUID={series_uids[patient_id]}"
145            util.download_source(path=zip_path, url=url, download=download, checksum=None)
146            util.unzip(zip_path=zip_path, dst=dicom_dir)
147
148        _convert_tcia_dicom_to_nifti(dicom_dir, label_path, image_path)
149        # The DICOM files are removed after the conversion to save disk space.
150        shutil.rmtree(dicom_dir)
151
152    return image_dir
153
154
155def _prepare_btcv_images(path, label_dir):
156    image_dir = os.path.join(path, IMAGE_DIRS["btcv"])
157    label_paths = _get_label_paths(label_dir, "btcv")
158    if len(glob(os.path.join(image_dir, "*.nii.gz"))) == len(label_paths):
159        return image_dir
160
161    zip_path = os.path.join(path, "RawData.zip")
162    if not os.path.exists(zip_path):
163        raise RuntimeError(
164            "The BTCV CT volumes cannot be downloaded automatically. Please register at Synapse, join the challenge "
165            "at https://www.synapse.org/#!Synapse:syn3193805 and download 'RawData.zip' from the 'Abdomen' folder "
166            f"at https://www.synapse.org/#!Synapse:syn3376386. Place the file at '{zip_path}' and try again."
167        )
168
169    raw_dir = os.path.join(path, "btcv_raw")
170    if not os.path.exists(raw_dir):
171        util.unzip(zip_path=zip_path, dst=raw_dir, remove=False)
172
173    os.makedirs(image_dir, exist_ok=True)
174    for label_path in label_paths:
175        label_id = _get_label_id(label_path)
176        image_path = os.path.join(image_dir, f"img{label_id}.nii.gz")
177        if os.path.exists(image_path):
178            continue
179
180        # The BTCV ids 0001-0040 are part of the challenge training set, 0061-0080 of the challenge test set.
181        candidates = glob(os.path.join(raw_dir, "**", f"img{label_id}.nii.gz"), recursive=True)
182        if len(candidates) != 1:
183            raise RuntimeError(f"Could not find the BTCV volume 'img{label_id}.nii.gz' in '{raw_dir}'.")
184        os.symlink(os.path.abspath(candidates[0]), image_path)
185
186    return image_dir
187
188
189def get_multi_organ_abdominal_ct_data(
190    path: Union[os.PathLike, str], source: Literal["tcia", "btcv"], download: bool = False
191) -> Tuple[str, str]:
192    """Download the Multi-organ Abdominal CT dataset.
193
194    Args:
195        path: Filepath to a folder where the data is downloaded for further processing.
196        source: The source of the CT volumes. Either 'tcia' (Pancreas-CT) or 'btcv' (Beyond the Cranial Vault).
197        download: Whether to download the data if it is not present.
198
199    Returns:
200        Filepath where the CT volumes are stored.
201        Filepath where the annotations are stored.
202    """
203    if source not in LABEL_DIRS:
204        raise ValueError(f"'{source}' is not a valid source. Please choose from {list(LABEL_DIRS.keys())}.")
205
206    os.makedirs(path, exist_ok=True)
207    label_dir = _download_labels(path, source, download)
208
209    if source == "tcia":
210        image_dir = _prepare_tcia_images(path, label_dir, download)
211    else:
212        image_dir = _prepare_btcv_images(path, label_dir)
213
214    return image_dir, label_dir
215
216
217def get_multi_organ_abdominal_ct_paths(
218    path: Union[os.PathLike, str], source: Literal["tcia", "btcv"], download: bool = False
219) -> Tuple[List[str], List[str]]:
220    """Get paths to the Multi-organ Abdominal CT data.
221
222    Args:
223        path: Filepath to a folder where the data is downloaded for further processing.
224        source: The source of the CT volumes. Either 'tcia' (Pancreas-CT) or 'btcv' (Beyond the Cranial Vault).
225        download: Whether to download the data if it is not present.
226
227    Returns:
228        List of filepaths for the image data.
229        List of filepaths for the label data.
230    """
231    image_dir, label_dir = get_multi_organ_abdominal_ct_data(path, source, download)
232
233    label_paths = _get_label_paths(label_dir, source)
234    raw_paths = [os.path.join(image_dir, f"img{_get_label_id(p)}.nii.gz") for p in label_paths]
235    assert all(os.path.exists(p) for p in raw_paths)
236
237    return raw_paths, label_paths
238
239
240def get_multi_organ_abdominal_ct_dataset(
241    path: Union[os.PathLike, str],
242    patch_shape: Tuple[int, ...],
243    source: Literal["tcia", "btcv"],
244    resize_inputs: bool = False,
245    download: bool = False,
246    **kwargs
247) -> Dataset:
248    """Get the Multi-organ Abdominal CT dataset for abdominal organ segmentation.
249
250    Args:
251        path: Filepath to a folder where the data is downloaded for further processing.
252        patch_shape: The patch shape to use for training.
253        source: The source of the CT volumes. Either 'tcia' (Pancreas-CT) or 'btcv' (Beyond the Cranial Vault).
254        resize_inputs: Whether to resize inputs to the desired patch shape.
255        download: Whether to download the data if it is not present.
256        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
257
258    Returns:
259        The segmentation dataset.
260    """
261    raw_paths, label_paths = get_multi_organ_abdominal_ct_paths(path, source, download)
262
263    if resize_inputs:
264        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
265        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
266            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
267        )
268
269    return torch_em.default_segmentation_dataset(
270        raw_paths=raw_paths,
271        raw_key="data",
272        label_paths=label_paths,
273        label_key="data",
274        patch_shape=patch_shape,
275        is_seg_dataset=True,
276        **kwargs
277    )
278
279
280def get_multi_organ_abdominal_ct_loader(
281    path: Union[os.PathLike, str],
282    batch_size: int,
283    patch_shape: Tuple[int, ...],
284    source: Literal["tcia", "btcv"],
285    resize_inputs: bool = False,
286    download: bool = False,
287    **kwargs
288) -> DataLoader:
289    """Get the Multi-organ Abdominal CT dataloader for abdominal organ segmentation.
290
291    Args:
292        path: Filepath to a folder where the data is downloaded for further processing.
293        batch_size: The batch size for training.
294        patch_shape: The patch shape to use for training.
295        source: The source of the CT volumes. Either 'tcia' (Pancreas-CT) or 'btcv' (Beyond the Cranial Vault).
296        resize_inputs: Whether to resize inputs to the desired patch shape.
297        download: Whether to download the data if it is not present.
298        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
299
300    Returns:
301        The DataLoader.
302    """
303    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
304    dataset = get_multi_organ_abdominal_ct_dataset(path, patch_shape, source, resize_inputs, download, **ds_kwargs)
305    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)