torch_em.data.datasets.medical.flare

  1"""The FLARE dataset contains annotations for abdominal organ segmentation in CT scans.
  2
  3It comprises the labeled training set of the FLARE22 challenge (https://flare22.grand-challenge.org):
  450 abdominal CT volumes with a dense annotation of 13 organs. The 2000 unlabeled training volumes, the
  550 validation volumes and the 200 test volumes of the challenge are not included here.
  6
  7NOTE: The label legend is as follows:
  8- background: 0, liver: 1, right kidney: 2, spleen: 3, pancreas: 4, aorta: 5, inferior vena cava: 6,
  9  right adrenal gland: 7, left adrenal gland: 8, gallbladder: 9, esophagus: 10, stomach: 11,
 10  duodenum: 12, left kidney: 13
 11The ids were verified on the data: all 13 ids are present in every volume, the two adrenal glands (7, 8)
 12are by far the smallest structures, the liver (1) and stomach (11) the largest, and the gallbladder (9)
 13is the structure with the largest relative volume variation across cases.
 14
 15NOTE: The later editions of the challenge (FLARE23, FLARE24, ...) use different data with partial
 16annotations of additional structures. They are not covered here and would need separate modules.
 17
 18The dataset is located at https://doi.org/10.5281/zenodo.7860267 (CC BY 4.0).
 19See https://flare22.grand-challenge.org for the challenge.
 20
 21This dataset is from the publication https://doi.org/10.1109/TMI.2022.3230667.
 22Please cite it if you use this dataset in your research.
 23"""
 24
 25import os
 26from glob import glob
 27from natsort import natsorted
 28from typing import Union, Tuple, List
 29
 30from torch.utils.data import Dataset, DataLoader
 31
 32import torch_em
 33
 34from .. import util
 35
 36
 37URL = "https://zenodo.org/records/7860267/files/FLARE22Train.zip"
 38CHECKSUM = "d57e201dc9002bd4a6fc1f5e0d4525285f46bcc5be1a3c8f12543151e6675154"
 39
 40ORGAN_NAMES = [
 41    "liver", "right_kidney", "spleen", "pancreas", "aorta", "inferior_vena_cava", "right_adrenal_gland",
 42    "left_adrenal_gland", "gallbladder", "esophagus", "stomach", "duodenum", "left_kidney",
 43]
 44
 45LABEL_IDS = {"background": 0, **{name: i + 1 for i, name in enumerate(ORGAN_NAMES)}}
 46
 47
 48def get_flare_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 49    """Download the FLARE22 dataset.
 50
 51    Args:
 52        path: Filepath to a folder where the data is downloaded for further processing.
 53        download: Whether to download the data if it is not present.
 54
 55    Returns:
 56        Filepath where the data is stored.
 57    """
 58    data_dir = os.path.join(path, "FLARE22Train")
 59    if os.path.exists(data_dir):
 60        return data_dir
 61
 62    os.makedirs(path, exist_ok=True)
 63
 64    zip_path = os.path.join(path, "FLARE22Train.zip")
 65    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 66    util.unzip(zip_path=zip_path, dst=path)
 67
 68    return data_dir
 69
 70
 71def get_flare_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 72    """Get paths to the FLARE22 data.
 73
 74    Args:
 75        path: Filepath to a folder where the data is downloaded for further processing.
 76        download: Whether to download the data if it is not present.
 77
 78    Returns:
 79        List of filepaths for the image data.
 80        List of filepaths for the label data.
 81    """
 82    data_dir = get_flare_data(path, download)
 83
 84    raw_paths = natsorted(glob(os.path.join(data_dir, "images", "*.nii.gz")))
 85    label_paths = [
 86        os.path.join(data_dir, "labels", os.path.basename(p).replace("_0000.nii.gz", ".nii.gz")) for p in raw_paths
 87    ]
 88    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in label_paths)
 89
 90    return raw_paths, label_paths
 91
 92
 93def get_flare_dataset(
 94    path: Union[os.PathLike, str],
 95    patch_shape: Tuple[int, ...],
 96    resize_inputs: bool = False,
 97    download: bool = False,
 98    **kwargs
 99) -> Dataset:
100    """Get the FLARE22 dataset for abdominal organ segmentation.
101
102    Args:
103        path: Filepath to a folder where the data is downloaded for further processing.
104        patch_shape: The patch shape to use for training.
105        resize_inputs: Whether to resize inputs to the desired patch shape.
106        download: Whether to download the data if it is not present.
107        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
108
109    Returns:
110        The segmentation dataset.
111    """
112    raw_paths, label_paths = get_flare_paths(path, download)
113
114    if resize_inputs:
115        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
116        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
117            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
118        )
119
120    return torch_em.default_segmentation_dataset(
121        raw_paths=raw_paths,
122        raw_key="data",
123        label_paths=label_paths,
124        label_key="data",
125        patch_shape=patch_shape,
126        is_seg_dataset=True,
127        **kwargs
128    )
129
130
131def get_flare_loader(
132    path: Union[os.PathLike, str],
133    batch_size: int,
134    patch_shape: Tuple[int, ...],
135    resize_inputs: bool = False,
136    download: bool = False,
137    **kwargs
138) -> DataLoader:
139    """Get the FLARE22 dataloader for abdominal organ segmentation.
140
141    Args:
142        path: Filepath to a folder where the data is downloaded for further processing.
143        batch_size: The batch size for training.
144        patch_shape: The patch shape to use for training.
145        resize_inputs: Whether to resize inputs to the desired patch shape.
146        download: Whether to download the data if it is not present.
147        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
148
149    Returns:
150        The DataLoader.
151    """
152    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
153    dataset = get_flare_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
154    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)