torch_em.data.datasets.medical.word

  1"""The WORD dataset contains annotations for 16 abdominal organs in CT scans.
  2
  3The dataset consists of 150 abdominal CT volumes with an official split into 100 training, 20 validation
  4and 30 test volumes. The label ids are: 1: liver, 2: spleen, 3: kidney (left), 4: kidney (right), 5: stomach,
  56: gallbladder, 7: esophagus, 8: pancreas, 9: duodenum, 10: colon, 11: intestine, 12: adrenal, 13: rectum,
  614: bladder, 15: head of femur (left), 16: head of femur (right). See also `CLASS_IDS`.
  7
  8The dataset is located at https://github.com/HiLab-git/WORD. It is distributed as a password protected zip archive
  9via Google Drive (the password 'word@uestc' is given in the repository).
 10
 11NOTE: The Google Drive download is quota limited, so it intermittently fails with
 12'Too many users have viewed or downloaded this file recently'. Retrying later usually works. Otherwise download
 13'WORD-V0.1.0.zip' manually from https://drive.google.com/drive/folders/16qwlCxH7XtJD9MyPnAbmY4ATxu2mKu67
 14and place it in `path`.
 15
 16This dataset is from the publication https://doi.org/10.1016/j.media.2022.102642.
 17Please cite it if you use this dataset in your research.
 18"""
 19
 20import os
 21import zipfile
 22from glob import glob
 23from shutil import which
 24from subprocess import run
 25from natsort import natsorted
 26from typing import Union, Tuple, Literal, List
 27
 28from torch.utils.data import Dataset, DataLoader
 29
 30import torch_em
 31
 32from .. import util
 33
 34
 35URLS = {
 36    "data": "https://drive.google.com/uc?id=19OWCXZGrimafREhXm8O8w2HBHZTfxEgU",
 37    "test_labels": "https://github.com/HiLab-git/WORD/raw/main/WORD_V0.1.0_labelsTs.zip",
 38}
 39
 40CHECKSUMS = {
 41    "data": "1ef1e48f8b41d72d733c5cdd91e1238d9b345d88c0de0c902e356310901f3115",
 42    "test_labels": "0552b1e208a8a5345a1c12003d4f1d010621c6ef65bdd62bff85a32b2976a46d",
 43}
 44
 45PASSWORD = "word@uestc"
 46
 47CLASS_NAMES = [
 48    "liver", "spleen", "kidney_left", "kidney_right", "stomach", "gallbladder", "esophagus", "pancreas", "duodenum",
 49    "colon", "intestine", "adrenal", "rectum", "bladder", "head_of_femur_left", "head_of_femur_right",
 50]
 51"""The organs of the WORD dataset. The label id of an organ is its 1-based index."""
 52
 53CLASS_IDS = {name: i + 1 for i, name in enumerate(CLASS_NAMES)}
 54"""Mapping from the organ name to its label id."""
 55
 56SPLIT_DIRS = {"train": "Tr", "val": "Val", "test": "Ts"}
 57
 58
 59def _unzip_with_password(zip_path, dst, password):
 60    try:
 61        with zipfile.ZipFile(zip_path) as f:
 62            f.extractall(dst, pwd=password.encode())
 63    except (NotImplementedError, RuntimeError) as e:
 64        # The python zipfile module does not support AES encryption, so we fall back to the 7z CLI.
 65        if which("7z") is None:
 66            raise RuntimeError(
 67                f"Could not extract '{zip_path}' with the zipfile module ({e}). Please install the '7z' CLI "
 68                "('conda install -c conda-forge p7zip') or extract the archive manually."
 69            )
 70        run(["7z", "x", f"-o{dst}", f"-p{password}", "-y", zip_path], check=True)
 71
 72
 73def _find_data_dir(path):
 74    candidates = glob(os.path.join(path, "imagesTr")) + glob(os.path.join(path, "*", "imagesTr"))
 75    return os.path.dirname(candidates[0]) if candidates else None
 76
 77
 78def get_word_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 79    """Download the WORD dataset.
 80
 81    Args:
 82        path: Filepath to a folder where the data is downloaded for further processing.
 83        download: Whether to download the data if it is not present.
 84
 85    Returns:
 86        Filepath where the data is downloaded.
 87    """
 88    data_dir = _find_data_dir(path)
 89    if data_dir is None:
 90        os.makedirs(path, exist_ok=True)
 91        zip_path = os.path.join(path, "WORD-V0.1.0.zip")
 92        util.download_source_gdrive(path=zip_path, url=URLS["data"], download=download, checksum=CHECKSUMS["data"])
 93        if not os.path.exists(zip_path):
 94            raise RuntimeError(
 95                "The automatic download of the WORD dataset failed, most likely because the Google Drive download "
 96                "quota of the archive is exceeded. Please download 'WORD-V0.1.0.zip' manually from "
 97                f"https://drive.google.com/drive/folders/16qwlCxH7XtJD9MyPnAbmY4ATxu2mKu67 and place it at "
 98                f"'{zip_path}'."
 99            )
100        _unzip_with_password(zip_path, path, PASSWORD)
101        data_dir = _find_data_dir(path)
102        if data_dir is None:
103            raise RuntimeError(f"Could not find the 'imagesTr' folder of the WORD dataset in '{path}'.")
104
105    # The test labels were released separately via the GitHub repository.
106    if not os.path.exists(os.path.join(data_dir, "labelsTs")):
107        zip_path = os.path.join(path, "WORD_V0.1.0_labelsTs.zip")
108        util.download_source(
109            path=zip_path, url=URLS["test_labels"], download=download, checksum=CHECKSUMS["test_labels"]
110        )
111        util.unzip(zip_path=zip_path, dst=data_dir, remove=False)
112
113    return data_dir
114
115
116def get_word_paths(
117    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False
118) -> Tuple[List[str], List[str]]:
119    """Get paths to the WORD data.
120
121    Args:
122        path: Filepath to a folder where the data is downloaded for further processing.
123        split: The choice of data split. Either 'train', 'val' or 'test'.
124        download: Whether to download the data if it is not present.
125
126    Returns:
127        List of filepaths for the image data.
128        List of filepaths for the label data.
129    """
130    if split not in SPLIT_DIRS:
131        raise ValueError(f"'{split}' is not a valid split. Choose from {list(SPLIT_DIRS.keys())}.")
132
133    data_dir = get_word_data(path, download)
134
135    raw_paths = natsorted(glob(os.path.join(data_dir, f"images{SPLIT_DIRS[split]}", "*.nii.gz")))
136    label_paths = [p.replace(f"images{SPLIT_DIRS[split]}", f"labels{SPLIT_DIRS[split]}") for p in raw_paths]
137
138    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
139    assert all(os.path.exists(p) for p in label_paths)
140
141    return raw_paths, label_paths
142
143
144def get_word_dataset(
145    path: Union[os.PathLike, str],
146    patch_shape: Tuple[int, ...],
147    split: Literal["train", "val", "test"],
148    resize_inputs: bool = False,
149    download: bool = False,
150    **kwargs
151) -> Dataset:
152    """Get the WORD dataset for abdominal organ segmentation.
153
154    Args:
155        path: Filepath to a folder where the data is downloaded for further processing.
156        patch_shape: The patch shape to use for training.
157        split: The choice of data split. Either 'train', 'val' or 'test'.
158        resize_inputs: Whether to resize inputs to the desired patch shape.
159        download: Whether to download the data if it is not present.
160        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
161
162    Returns:
163        The segmentation dataset.
164    """
165    raw_paths, label_paths = get_word_paths(path, split, download)
166
167    if resize_inputs:
168        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
169        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
170            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
171        )
172
173    return torch_em.default_segmentation_dataset(
174        raw_paths=raw_paths,
175        raw_key="data",
176        label_paths=label_paths,
177        label_key="data",
178        patch_shape=patch_shape,
179        is_seg_dataset=True,
180        **kwargs
181    )
182
183
184def get_word_loader(
185    path: Union[os.PathLike, str],
186    batch_size: int,
187    patch_shape: Tuple[int, ...],
188    split: Literal["train", "val", "test"],
189    resize_inputs: bool = False,
190    download: bool = False,
191    **kwargs
192) -> DataLoader:
193    """Get the WORD dataloader for abdominal organ segmentation.
194
195    Args:
196        path: Filepath to a folder where the data is downloaded for further processing.
197        batch_size: The batch size for training.
198        patch_shape: The patch shape to use for training.
199        split: The choice of data split. Either 'train', 'val' or 'test'.
200        resize_inputs: Whether to resize inputs to the desired patch shape.
201        download: Whether to download the data if it is not present.
202        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
203
204    Returns:
205        The DataLoader.
206    """
207    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
208    dataset = get_word_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
209    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)