torch_em.data.datasets.medical.rumc_kidney

  1"""The RUMC Kidney dataset contains annotations for kidney and kidney abnormality
  2segmentation in contrast-enhanced thorax-abdomen CT scans.
  3
  4The scans were collected at the Radboud University Medical Center (RUMC), Nijmegen.
  5
  6The label ids are - kidney: 1, abnormality: 2
  7
  8NOTE: The abnormality class merges the five types the authors annotated (tumors, cysts, masses,
  9lesions and metastases) into one id, i.e. tumors cannot be told apart from cysts. Use
 10`torch_em.data.datasets.medical.kits` if you need labels that separate tumors from cysts.
 11
 12The dataset is located at https://doi.org/10.5281/zenodo.8014290.
 13This dataset is from the publication https://doi.org/10.48550/arXiv.2309.03383.
 14Please cite it if you use this dataset for your research.
 15"""
 16
 17import os
 18import json
 19from glob import glob
 20from tqdm import tqdm
 21from natsort import natsorted
 22from typing import Union, Tuple, List, Literal
 23
 24import numpy as np
 25
 26from sklearn.model_selection import train_test_split
 27
 28from torch.utils.data import Dataset, DataLoader
 29
 30import torch_em
 31
 32from .. import util
 33
 34
 35URL = "https://zenodo.org/records/8014290/files/{}?download=1"
 36
 37CHECKSUMS = {
 38    "images1.zip": "5b7af2786d0b2281c73a6e0c69eb4982998786a964dc9256fa7fee931a035d7a",
 39    "images2.zip": "2f4e3ebfb1f0aeeffd85c84bccbcae33cd8623193e839cc12a5cd58abb6e8a9e",
 40    "images3.zip": "c7df9dc1cb017e2ca77f74c7f7d90f5f1e2dbb4a16eb7c7b49393bdc65107eed",
 41    "images4.zip": "73aa389b54791a5abd722051b9a0ea2337aa1098f6b87b59e2aa1486a89157f1",
 42    "segmentations.zip": "7d711e306e52fe37216c126c3b282935f916db594193a2d1e3fdf13a2757c163",
 43}
 44
 45IMAGE_DIRS = ("images1", "images2", "images3", "images4")
 46LABEL_IDS = {"kidney": 1, "abnormality": 2}
 47VALID_SPLITS = ("train", "val", "test")
 48SUFFIX = ".nii.gz"
 49METAIMAGE_EXTS = (".mha", ".mhd")
 50
 51# The NIfTI extension code for free-form text, which we use to store the original header as json.
 52NIFTI_COMMENT_ECODE = 6
 53
 54
 55def _read_metaimage_header(path):
 56    header, is_mha = {}, path.endswith(".mha")
 57    with open(path, "rb") as f:
 58        for line in f:
 59            if b"=" not in line:
 60                break
 61            key, value = line.decode("latin-1").split("=", 1)
 62            key = key.strip()
 63            header[key] = value.strip()
 64            # For '.mha' the binary data directly follows the header, so we must not read any further.
 65            if is_mha and key == "ElementDataFile":
 66                break
 67    return header
 68
 69
 70def _itk_geometry_to_ras_affine(spacing, origin, direction):
 71    affine = np.eye(4, dtype="float64")
 72    direction = np.asarray(direction, dtype="float64").reshape(3, 3)
 73    affine[:3, :3] = direction @ np.diag(np.asarray(spacing, dtype="float64"))
 74    affine[:3, 3] = np.asarray(origin, dtype="float64")
 75    # ITK stores the geometry in LPS, whereas NIfTI expects RAS, so the first two axes flip sign.
 76    return np.diag([-1.0, -1.0, 1.0, 1.0]) @ affine
 77
 78
 79def _convert_mha_to_nifti(path: str, output_path: str, keep_metadata: bool = True) -> str:
 80    """Convert a MetaImage file to the NIfTI format.
 81
 82    The voxel spacing, origin and direction are mapped to the NIfTI affine, so that the converted
 83    volume keeps its physical geometry. The original MetaImage header is stored verbatim in a NIfTI
 84    header extension, so that the metadata stays inside the converted file.
 85
 86    Requires the SimpleITK python library.
 87
 88    Args:
 89        path: Filepath to the MetaImage file ('.mha' or '.mhd').
 90        output_path: Filepath for the converted file. Use a '.nii' suffix to write an uncompressed
 91            volume, which can be memory-mapped for lazy loading, or '.nii.gz' to write a compressed one.
 92        keep_metadata: Whether to store the original header in a NIfTI header extension.
 93
 94    Returns:
 95        The filepath to the converted file.
 96    """
 97    import nibabel as nib
 98    import SimpleITK as sitk
 99
100    if not path.endswith(METAIMAGE_EXTS):
101        raise ValueError(f"The provided file ({path}) isn't in MetaImage format.")
102
103    image = sitk.ReadImage(path)
104    affine = _itk_geometry_to_ras_affine(image.GetSpacing(), image.GetOrigin(), image.GetDirection())
105
106    # SimpleITK returns the array in 'zyx' order, whereas the NIfTI affine refers to 'xyz'.
107    data = sitk.GetArrayFromImage(image).transpose(2, 1, 0)
108
109    nifti = nib.Nifti1Image(data, affine)
110    nifti.set_data_dtype(data.dtype)
111    nifti.header.set_xyzt_units(xyz="mm")
112
113    if keep_metadata:
114        metadata = {
115            "metaimage_header": _read_metaimage_header(path),
116            "itk_metadata": {k: image.GetMetaData(k) for k in image.GetMetaDataKeys()},
117        }
118        payload = json.dumps(metadata).encode("utf-8")
119        nifti.header.extensions.append(nib.nifti1.Nifti1Extension(NIFTI_COMMENT_ECODE, payload))
120
121    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
122    nib.save(nifti, output_path)
123    return output_path
124
125
126def _read_nifti_metadata(path: str) -> dict:
127    """Read the metadata that `_convert_mha_to_nifti` stored in a NIfTI header extension.
128
129    Args:
130        path: Filepath to the NIfTI file.
131
132    Returns:
133        The original metadata. Empty if the file does not carry any.
134    """
135    import nibabel as nib
136
137    for extension in nib.load(path).header.extensions:
138        if extension.get_code() != NIFTI_COMMENT_ECODE:
139            continue
140        try:
141            return json.loads(extension.get_content())
142        except (ValueError, UnicodeDecodeError):  # Some other tool wrote a plain text comment.
143            continue
144    return {}
145
146
147class _SelectLabel:
148    def __init__(self, label_id):
149        self.label_id = label_id
150
151    def __call__(self, labels):
152        return (labels == self.label_id).astype("uint8")
153
154
155def _find_image_path(path, case_id):
156    for image_dir in IMAGE_DIRS:
157        image_path = os.path.join(path, image_dir, f"{case_id}.mha")
158        if os.path.exists(image_path):
159            return image_path
160
161    raise FileNotFoundError(f"Could not find the image for case '{case_id}'.")
162
163
164def _preprocess_inputs(path):
165    # The volumes are converted to nifti, because elf has no file wrapper for MetaImage and would
166    # have to load each volume into memory as a whole.
167    image_dir = os.path.join(path, "preprocessed", "images")
168    label_dir = os.path.join(path, "preprocessed", "labels")
169    os.makedirs(image_dir, exist_ok=True)
170    os.makedirs(label_dir, exist_ok=True)
171
172    label_paths = natsorted(glob(os.path.join(path, "segmentations", "*.mha")))
173    if not label_paths:
174        raise RuntimeError(f"Could not find the segmentations in '{path}'.")
175
176    for label_path in tqdm(label_paths, desc="Preprocessing inputs"):
177        case_id = os.path.basename(label_path)[:-len("_segmentations.mha")]
178
179        target_image_path = os.path.join(image_dir, f"{case_id}{SUFFIX}")
180        target_label_path = os.path.join(label_dir, f"{case_id}{SUFFIX}")
181        if os.path.exists(target_image_path) and os.path.exists(target_label_path):
182            continue
183
184        _convert_mha_to_nifti(_find_image_path(path, case_id), target_image_path)
185        _convert_mha_to_nifti(label_path, target_label_path)
186
187
188def get_rumc_kidney_data(path: Union[os.PathLike, str], download: bool = False) -> str:
189    """Download the RUMC Kidney dataset.
190
191    The download is roughly 38 GB, and the volumes converted to nifti take up roughly 40 GB more.
192
193    Args:
194        path: Filepath to a folder where the data is downloaded for further processing.
195        download: Whether to download the data if it is not present.
196
197    Returns:
198        The folder where the dataset is downloaded and preprocessed.
199    """
200    data_dir = os.path.join(path, "preprocessed")
201    if os.path.exists(data_dir):
202        return data_dir
203
204    os.makedirs(path, exist_ok=True)
205
206    for fname, checksum in CHECKSUMS.items():
207        zip_path = os.path.join(path, fname)
208        util.download_source(path=zip_path, url=URL.format(fname), download=download, checksum=checksum)
209        util.unzip(zip_path=zip_path, dst=path)
210
211    _preprocess_inputs(path)
212
213    return data_dir
214
215
216def _get_split_map(path, data_dir):
217    split_path = os.path.join(path, "splits_rumc_kidney.json")
218    if os.path.exists(split_path):
219        with open(split_path) as f:
220            return json.load(f)
221
222    case_ids = [os.path.basename(p).split(".")[0] for p in glob(os.path.join(data_dir, "images", f"*{SUFFIX}"))]
223    train_ids, test_ids = train_test_split(natsorted(case_ids), test_size=0.25, random_state=42)
224    train_ids, val_ids = train_test_split(train_ids, test_size=0.1, random_state=42)
225
226    split_map = {"train": train_ids, "val": val_ids, "test": test_ids}
227    with open(split_path, "w") as f:
228        json.dump(split_map, f, indent=2)
229
230    return split_map
231
232
233def get_rumc_kidney_paths(
234    path: Union[os.PathLike, str],
235    split: Literal["train", "val", "test"],
236    download: bool = False,
237) -> Tuple[List[str], List[str]]:
238    """Get paths to the RUMC Kidney data.
239
240    Args:
241        path: Filepath to a folder where the data is downloaded for further processing.
242        split: Which data split to use.
243        download: Whether to download the data if it is not present.
244
245    Returns:
246        List of filepaths for the image data.
247        List of filepaths for the label data.
248    """
249    if split not in VALID_SPLITS:
250        raise ValueError(f"Invalid split '{split}'. Must be one of {VALID_SPLITS}.")
251
252    data_dir = get_rumc_kidney_data(path, download)
253    split_map = _get_split_map(path, data_dir)
254
255    raw_paths = [os.path.join(data_dir, "images", f"{case_id}{SUFFIX}") for case_id in split_map[split]]
256    label_paths = [os.path.join(data_dir, "labels", f"{case_id}{SUFFIX}") for case_id in split_map[split]]
257
258    missing = [p for p in raw_paths + label_paths if not os.path.exists(p)]
259    if missing:
260        raise RuntimeError(f"Could not find {len(missing)} files, e.g. '{missing[0]}'.")
261
262    return raw_paths, label_paths
263
264
265def get_rumc_kidney_dataset(
266    path: Union[os.PathLike, str],
267    patch_shape: Tuple[int, ...],
268    split: Literal["train", "val", "test"],
269    label_choice: Literal["all", "kidney", "abnormality"] = "all",
270    resize_inputs: bool = False,
271    download: bool = False,
272    **kwargs
273) -> Dataset:
274    """Get the RUMC Kidney dataset for kidney and kidney abnormality segmentation.
275
276    Args:
277        path: Filepath to a folder where the data is downloaded for further processing.
278        patch_shape: The patch shape to use for training.
279        split: Which data split to use.
280        label_choice: Which labels to use. 'all' keeps both classes, the other choices return a
281            binary mask for that class alone.
282        resize_inputs: Whether to resize inputs to the desired patch shape.
283        download: Whether to download the data if it is not present.
284        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
285
286    Returns:
287        The segmentation dataset.
288    """
289    raw_paths, label_paths = get_rumc_kidney_paths(path, split, download)
290
291    if label_choice != "all":
292        if label_choice not in LABEL_IDS:
293            raise ValueError(f"Invalid label choice '{label_choice}'. Must be 'all' or one of {tuple(LABEL_IDS)}.")
294
295        kwargs = util.update_kwargs(kwargs, "label_transform", _SelectLabel(LABEL_IDS[label_choice]))
296
297    if resize_inputs:
298        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
299        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
300            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
301        )
302
303    return torch_em.default_segmentation_dataset(
304        raw_paths=raw_paths, raw_key="data", label_paths=label_paths, label_key="data",
305        patch_shape=patch_shape, **kwargs
306    )
307
308
309def get_rumc_kidney_loader(
310    path: Union[os.PathLike, str],
311    batch_size: int,
312    patch_shape: Tuple[int, ...],
313    split: Literal["train", "val", "test"],
314    label_choice: Literal["all", "kidney", "abnormality"] = "all",
315    resize_inputs: bool = False,
316    download: bool = False,
317    **kwargs
318) -> DataLoader:
319    """Get the RUMC Kidney dataloader for kidney and kidney abnormality segmentation.
320
321    Args:
322        path: Filepath to a folder where the data is downloaded for further processing.
323        batch_size: The batch size for training.
324        patch_shape: The patch shape to use for training.
325        split: Which data split to use.
326        label_choice: Which labels to use. 'all' keeps both classes, the other choices return a
327            binary mask for that class alone.
328        resize_inputs: Whether to resize inputs to the desired patch shape.
329        download: Whether to download the data if it is not present.
330        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
331
332    Returns:
333        The DataLoader.
334    """
335    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
336    dataset = get_rumc_kidney_dataset(
337        path, patch_shape, split, label_choice, resize_inputs, download, **ds_kwargs
338    )
339    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)