torch_em.data.datasets.medical.ribseg

  1"""The RibSeg dataset contains annotations for rib segmentation and labeling in chest-abdomen CT scans.
  2
  3The annotations of RibSeg v2 cover all 660 CT scans of the RibFrac dataset. Each rib carries its own label id
  4from 1 to 24 (1-12 are the right ribs from top to bottom, 13-24 the left ribs), i.e. the segmentations are both
  5an instance segmentation of the individual ribs and a semantic labeling of the rib positions. Use the 'binary'
  6keyword argument of `torch_em.default_segmentation_dataset` to train on the binary rib mask instead.
  7
  8The rib labels are downloaded from the RibSeg v2 release on google drive and the CT scans from the RibFrac
  9records on zenodo. The 'split' argument selects the official RibFrac split: 'train' (420 scans, ca. 51 GB),
 10'val' (80 scans, ca. 8.7 GB) or 'test' (160 scans, ca. 18 GB). The centerline annotations that are part of the
 11RibSeg v2 release (a 24 x 500 x 3 point array per scan) are not exposed by this module.
 12
 13The dataset is located at https://github.com/M3DV/RibSeg and the CT scans at https://ribfrac.grand-challenge.org.
 14
 15This dataset is from the publication https://doi.org/10.1109/TMI.2023.3313627.
 16The CT scans are from the publication https://doi.org/10.1016/j.ebiom.2020.103106.
 17Please cite them if you use this dataset in your research.
 18"""
 19
 20import os
 21import shutil
 22from glob import glob
 23from natsort import natsorted
 24from typing import Union, Tuple, List, Literal
 25
 26from torch.utils.data import Dataset, DataLoader
 27
 28import torch_em
 29
 30from .. import util
 31
 32
 33URL_LABELS = "https://drive.google.com/uc?id=1ZZGGrhd0y1fLyOZGo_Y-wlVUP4lkHVgm"
 34CHECKSUM_LABELS = "6bf8a327f4a7f540a318caf66885e09d76fd8832a853806de4ad3d3ff4369714"
 35
 36URLS = {
 37    "train": [
 38        "https://zenodo.org/records/3893508/files/ribfrac-train-images-1.zip",
 39        "https://zenodo.org/records/3893498/files/ribfrac-train-images-2.zip",
 40    ],
 41    "val": ["https://zenodo.org/records/3893496/files/ribfrac-val-images.zip"],
 42    "test": ["https://zenodo.org/records/3993380/files/ribfrac-test-images.zip"],
 43}
 44
 45# NOTE: Only the checksum of the validation archive is known, the other archives are not verified.
 46CHECKSUMS = {
 47    "train": [None, None],
 48    "val": ["786cc14bf4ea55e93325d657d0e6490f4a1a7c8d073522404751b961dab95e24"],
 49    "test": [None],
 50}
 51
 52N_VOLUMES = {"train": 420, "val": 80, "test": 160}
 53
 54# The label ids are the individual ribs: 1-12 are the right ribs and 13-24 the left ribs, each counted
 55# from the first (top) to the twelfth (bottom) rib.
 56LABEL_IDS = {f"rib_{i}": i for i in range(1, 25)}
 57
 58
 59def get_ribseg_data(
 60    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False
 61) -> Tuple[str, str]:
 62    """Download the RibSeg dataset.
 63
 64    Args:
 65        path: Filepath to a folder where the data is downloaded for further processing.
 66        split: The choice of data split. Either 'train', 'val' or 'test'.
 67        download: Whether to download the data if it is not present.
 68
 69    Returns:
 70        Filepath to the folder with the CT scans.
 71        Filepath to the folder with the rib labels.
 72    """
 73    if split not in URLS:
 74        raise ValueError(f"'{split}' is not a valid split. Please choose one of {list(URLS.keys())}.")
 75
 76    os.makedirs(path, exist_ok=True)
 77
 78    label_dir = os.path.join(path, "ribseg_v2", "seg")
 79    if not os.path.exists(label_dir):
 80        zip_path = os.path.join(path, "ribseg_v2.zip")
 81        util.download_source_gdrive(path=zip_path, url=URL_LABELS, download=download, checksum=CHECKSUM_LABELS)
 82        util.unzip(zip_path=zip_path, dst=path)
 83
 84    image_dir = os.path.join(path, "images", split)
 85    if not os.path.exists(image_dir):
 86        if not download:
 87            raise RuntimeError(f"Cannot find the data at {image_dir}, but download was set to False.")
 88
 89        # The archives are extracted to a temporary folder, which is renamed once all parts are complete,
 90        # so that an interrupted download is not mistaken for a complete one.
 91        tmp_dir = f"{image_dir}.tmp"
 92        for url, checksum in zip(URLS[split], CHECKSUMS[split]):
 93            zip_path = os.path.join(path, os.path.basename(url))
 94            util.download_source(path=zip_path, url=url, download=download, checksum=checksum)
 95            util.unzip(zip_path=zip_path, dst=tmp_dir)
 96
 97        # The archives contain the volumes in a sub-folder, the name of which differs between the splits.
 98        for volume_path in glob(os.path.join(tmp_dir, "*", "*.nii.gz")):
 99            shutil.move(volume_path, os.path.join(tmp_dir, os.path.basename(volume_path)))
100        for sub_dir in glob(os.path.join(tmp_dir, "*", "")):
101            shutil.rmtree(sub_dir)
102        os.rename(tmp_dir, image_dir)
103
104    return image_dir, label_dir
105
106
107def get_ribseg_paths(
108    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False
109) -> Tuple[List[str], List[str]]:
110    """Get paths to the RibSeg data.
111
112    Args:
113        path: Filepath to a folder where the data is downloaded for further processing.
114        split: The choice of data split. Either 'train', 'val' or 'test'.
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    image_dir, label_dir = get_ribseg_data(path, split, download)
122
123    raw_paths = natsorted(glob(os.path.join(image_dir, "*-image.nii.gz")))
124    label_paths = [
125        os.path.join(label_dir, os.path.basename(p).replace("-image.nii.gz", "-rib-seg.nii.gz")) for p in raw_paths
126    ]
127    assert len(raw_paths) == N_VOLUMES[split] and all(os.path.exists(p) for p in label_paths)
128
129    return raw_paths, label_paths
130
131
132def get_ribseg_dataset(
133    path: Union[os.PathLike, str],
134    patch_shape: Tuple[int, ...],
135    split: Literal["train", "val", "test"],
136    resize_inputs: bool = False,
137    download: bool = False,
138    **kwargs
139) -> Dataset:
140    """Get the RibSeg dataset for rib segmentation and labeling.
141
142    Args:
143        path: Filepath to a folder where the data is downloaded for further processing.
144        patch_shape: The patch shape to use for training.
145        split: The choice of data split. Either 'train', 'val' or 'test'.
146        resize_inputs: Whether to resize inputs to the desired patch shape.
147        download: Whether to download the data if it is not present.
148        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
149
150    Returns:
151        The segmentation dataset.
152    """
153    raw_paths, label_paths = get_ribseg_paths(path, split, download)
154
155    if resize_inputs:
156        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
157        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
158            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
159        )
160
161    return torch_em.default_segmentation_dataset(
162        raw_paths=raw_paths,
163        raw_key="data",
164        label_paths=label_paths,
165        label_key="data",
166        patch_shape=patch_shape,
167        is_seg_dataset=True,
168        **kwargs
169    )
170
171
172def get_ribseg_loader(
173    path: Union[os.PathLike, str],
174    batch_size: int,
175    patch_shape: Tuple[int, ...],
176    split: Literal["train", "val", "test"],
177    resize_inputs: bool = False,
178    download: bool = False,
179    **kwargs
180) -> DataLoader:
181    """Get the RibSeg dataloader for rib segmentation and labeling.
182
183    Args:
184        path: Filepath to a folder where the data is downloaded for further processing.
185        batch_size: The batch size for training.
186        patch_shape: The patch shape to use for training.
187        split: The choice of data split. Either 'train', 'val' or 'test'.
188        resize_inputs: Whether to resize inputs to the desired patch shape.
189        download: Whether to download the data if it is not present.
190        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
191
192    Returns:
193        The DataLoader.
194    """
195    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
196    dataset = get_ribseg_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
197    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)