torch_em.data.datasets.medical.ski10

  1"""The SKI10 dataset contains annotations for knee bone and cartilage segmentation in MRI.
  2
  3The data was curated for the SKI10 challenge (Segmentation of Knee Images 2010, https://ski10.grand-challenge.org),
  4which was held at MICCAI 2010. It consists of the 100 annotated training MRI of the challenge, which come from the
  5surgical planning program of Biomet Inc. The four annotated structures are the femoral bone and cartilage and the
  6tibial bone and cartilage, see `LABEL_IDS`. The 50 test MRI of the challenge were distributed without annotations
  7and are not included here.
  8
  9NOTE: The official challenge is closed and its data cannot be downloaded from the challenge website anymore.
 10This module uses the redistribution at https://huggingface.co/datasets/YongchengYAO/SKI10 (CC BY-NC-SA 4.0),
 11which converted the original 'mhd' / 'raw' volumes to nifti without changing the image or mask values.
 12
 13The volumes are stored with the slice axis last (the in-plane resolution is around 0.4 mm and the slice
 14thickness is 1 mm), so they are converted to hdf5 volumes with the slice axis first (the keys are 'raw' and
 15'labels') by this module.
 16
 17This dataset is from the publication 'Segmentation of Knee Images: A Grand Challenge' by Heimann et al.,
 18MICCAI Workshop on Medical Image Analysis for the Clinic (2010), which has no DOI. The challenge papers are
 19collected at https://doi.org/10.5281/zenodo.4781231.
 20Please cite it if you use this dataset in your research.
 21"""
 22
 23import os
 24from glob import glob
 25from tqdm import tqdm
 26from natsort import natsorted
 27from typing import Union, Tuple, List
 28
 29import numpy as np
 30
 31from torch.utils.data import Dataset, DataLoader
 32
 33import torch_em
 34
 35from .. import util
 36
 37
 38URL = "https://huggingface.co/datasets/YongchengYAO/SKI10/resolve/main/SKI10.zip"
 39CHECKSUM = "d367c1c68143f450e4cad92111afc064f116372442a4816718e20ea556507bc8"
 40
 41LABEL_IDS = {
 42    "background": 0, "femur_bone": 1, "femur_cartilage": 2, "tibia_bone": 3, "tibia_cartilage": 4,
 43}
 44
 45N_VOLUMES = 100
 46
 47
 48def _preprocess_inputs(data_dir, preprocessed_dir):
 49    import h5py
 50    import nibabel as nib
 51
 52    image_paths = natsorted(glob(os.path.join(data_dir, "image", "image-*.nii")))
 53    os.makedirs(preprocessed_dir, exist_ok=True)
 54
 55    for image_path in tqdm(image_paths, desc="Preprocessing the SKI10 volumes"):
 56        case_id = os.path.basename(image_path)[len("image-"):-len(".nii")]
 57        volume_path = os.path.join(preprocessed_dir, f"ski10_{case_id}.h5")
 58        if os.path.exists(volume_path):
 59            continue
 60
 61        label_path = os.path.join(data_dir, "label", f"labels-{case_id}.nii")
 62
 63        # The transpose maps the nifti axis order (X, Y, Z) to the (Z, Y, X) order used for the volumes,
 64        # so that the first axis is the slice axis of the acquisition.
 65        raw = np.asarray(nib.load(image_path).dataobj).T
 66        labels = np.asarray(nib.load(label_path).dataobj).T.astype("uint8")
 67
 68        # The file is written to a temporary path first, so that an interrupted run leaves no corrupt file.
 69        with h5py.File(f"{volume_path}.tmp", "w") as f:
 70            f.create_dataset("raw", data=raw, compression="gzip")
 71            f.create_dataset("labels", data=labels, compression="gzip")
 72
 73        os.rename(f"{volume_path}.tmp", volume_path)
 74
 75
 76def get_ski10_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 77    """Download the SKI10 dataset.
 78
 79    Args:
 80        path: Filepath to a folder where the data is downloaded for further processing.
 81        download: Whether to download the data if it is not present.
 82
 83    Returns:
 84        Filepath where the preprocessed data is stored.
 85    """
 86    preprocessed_dir = os.path.join(path, "preprocessed")
 87    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == N_VOLUMES:
 88        return preprocessed_dir
 89
 90    os.makedirs(path, exist_ok=True)
 91
 92    data_dir = os.path.join(path, "SKI10")
 93    if not os.path.exists(data_dir):
 94        zip_path = os.path.join(path, "SKI10.zip")
 95        util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 96        util.unzip(zip_path=zip_path, dst=path)
 97
 98    _preprocess_inputs(data_dir, preprocessed_dir)
 99    return preprocessed_dir
100
101
102def get_ski10_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
103    """Get paths to the SKI10 data.
104
105    Args:
106        path: Filepath to a folder where the data is downloaded for further processing.
107        download: Whether to download the data if it is not present.
108
109    Returns:
110        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
111    """
112    data_dir = get_ski10_data(path, download)
113    volume_paths = natsorted(glob(os.path.join(data_dir, "ski10_*.h5")))
114    assert len(volume_paths) > 0
115
116    return volume_paths
117
118
119def get_ski10_dataset(
120    path: Union[os.PathLike, str],
121    patch_shape: Tuple[int, ...],
122    resize_inputs: bool = False,
123    download: bool = False,
124    **kwargs
125) -> Dataset:
126    """Get the SKI10 dataset for knee bone and cartilage segmentation.
127
128    Args:
129        path: Filepath to a folder where the data is downloaded for further processing.
130        patch_shape: The patch shape to use for training.
131        resize_inputs: Whether to resize inputs to the desired patch shape.
132        download: Whether to download the data if it is not present.
133        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
134
135    Returns:
136        The segmentation dataset.
137    """
138    volume_paths = get_ski10_paths(path, download)
139
140    if resize_inputs:
141        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
142        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
143            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
144        )
145
146    return torch_em.default_segmentation_dataset(
147        raw_paths=volume_paths,
148        raw_key="raw",
149        label_paths=volume_paths,
150        label_key="labels",
151        patch_shape=patch_shape,
152        is_seg_dataset=True,
153        **kwargs
154    )
155
156
157def get_ski10_loader(
158    path: Union[os.PathLike, str],
159    batch_size: int,
160    patch_shape: Tuple[int, ...],
161    resize_inputs: bool = False,
162    download: bool = False,
163    **kwargs
164) -> DataLoader:
165    """Get the SKI10 dataloader for knee bone and cartilage segmentation.
166
167    Args:
168        path: Filepath to a folder where the data is downloaded for further processing.
169        batch_size: The batch size for training.
170        patch_shape: The patch shape to use for training.
171        resize_inputs: Whether to resize inputs to the desired patch shape.
172        download: Whether to download the data if it is not present.
173        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
174
175    Returns:
176        The DataLoader.
177    """
178    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
179    dataset = get_ski10_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
180    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)