torch_em.data.datasets.medical.skm_tea

  1"""The SKM-TEA dataset contains annotations for knee tissue segmentation in quantitative double-echo
  2steady-state (qDESS) knee MRI.
  3
  4It comprises 155 scans with two echoes each and dense segmentations of six soft tissue classes. The scans
  5come with an official train / val / test split. The label ids are:
  6- 0: background
  7- 1: patellar cartilage
  8- 2: femoral cartilage
  9- 3: tibial cartilage (medial)
 10- 4: tibial cartilage (lateral)
 11- 5: meniscus (medial)
 12- 6: meniscus (lateral)
 13
 14NOTE: The label ids were verified against the official dataset documentation
 15(https://github.com/StanfordMIMI/skm-tea/blob/main/DATASET.md), which documents the order of the one-hot
 16encoded 'seg' volumes, and on the data: the six channels do not overlap and the by far largest one is the
 17femoral cartilage.
 18
 19NOTE: The full dataset with all 155 scans requires registration and cannot be downloaded automatically.
 20Only the official sample of 3 scans (one per split), which the authors published at
 21https://huggingface.co/datasets/arjundd/skm-tea-mini, is downloaded if `download` is set to True.
 22Note that this sample must not be used for reporting metrics. To use the full dataset, follow these steps:
 23- Visit https://aimi.stanford.edu/skm-tea-knee-mri and follow the link to the dataset on the Stanford AIMI
 24  shared datasets portal (https://stanfordaimi.azurewebsites.net/datasets/4aaeafb9-c6e6-4e3c-9188-3aaaf0e0a9e7).
 25- Log in or create an account and follow the download instructions (the full dataset is ~900 GB compressed,
 26  but only the 'DICOM' track ('image_files') and the annotations are required here).
 27- Extract the data such that '<path>/image_files/<scan_id>.h5' and
 28  '<path>/annotations/<version>/{train,val,test}.json' exist, i.e. '<path>' is the 'skm-tea' folder described
 29  in https://github.com/StanfordMIMI/skm-tea/blob/main/DATASET.md.
 30
 31The dataset is located at https://aimi.stanford.edu/skm-tea-knee-mri.
 32
 33This dataset is from the publication https://openreview.net/forum?id=YDMFgD_qJuA.
 34Please cite it if you use this dataset in your research.
 35"""
 36
 37import os
 38import json
 39import warnings
 40from glob import glob
 41from tqdm import tqdm
 42from natsort import natsorted
 43from typing import Union, Tuple, Literal, List
 44
 45import numpy as np
 46
 47from torch.utils.data import Dataset, DataLoader
 48
 49import torch_em
 50
 51from .. import util
 52
 53
 54SAMPLE_URL = "https://huggingface.co/datasets/arjundd/skm-tea-mini/resolve/main/v1-release"
 55
 56SAMPLE_CHECKSUMS = {
 57    "image_files.tar.gz": "18af509e31f4468dbb5b237bf1a124508d00ade4ace0b90fae62947ce4709812",
 58    "train.json": "7a2b80f6ecae2eabe2489235302ac0003ac39f1266c4d356abbf5279f0ccdb7a",
 59    "val.json": "5d64909b69f89a101755ace309c4769c5884758957efa35d4da27b3b4b455986",
 60    "test.json": "b057cf984d7fcdd52318642ea2ec3a9241464575c3a91dc29f64e4fef51dad74",
 61}
 62
 63SAMPLE_VERSION = "v1.0.0"
 64
 65LABEL_IDS = {
 66    "background": 0,
 67    "patellar_cartilage": 1,
 68    "femoral_cartilage": 2,
 69    "tibial_cartilage_medial": 3,
 70    "tibial_cartilage_lateral": 4,
 71    "meniscus_medial": 5,
 72    "meniscus_lateral": 6,
 73}
 74
 75
 76def _preprocess_inputs(path, image_dir):
 77    import h5py
 78
 79    preprocessed_dir = os.path.join(path, "preprocessed")
 80    os.makedirs(preprocessed_dir, exist_ok=True)
 81
 82    for scan_path in tqdm(natsorted(glob(os.path.join(image_dir, "*.h5"))), desc="Preprocessing the SKM-TEA scans"):
 83        volume_path = os.path.join(preprocessed_dir, os.path.basename(scan_path))
 84        if os.path.exists(volume_path):
 85            continue
 86
 87        with h5py.File(scan_path, "r") as f:
 88            echo1, echo2, seg = f["echo1"][:], f["echo2"][:], f["seg"][:]
 89
 90        # The one-hot encoded segmentation (X, Y, Z, 6) is converted to a label volume with the ids 1 to 6.
 91        labels = np.where(seg.any(axis=-1), np.argmax(seg, axis=-1) + 1, 0).astype(np.uint8)
 92
 93        # The volumes are stored as (X, Y, Z) with Z (LR) being the sagittal slice direction.
 94        # We move the slice axis to the front, so that 2d training samples sagittal slices.
 95        # The file is written to a temporary path first, so that an interrupted run does not leave a corrupt file.
 96        with h5py.File(f"{volume_path}.tmp", "w") as f:
 97            f.create_dataset("raw/echo1", data=echo1.transpose(2, 0, 1), compression="gzip")
 98            f.create_dataset("raw/echo2", data=echo2.transpose(2, 0, 1), compression="gzip")
 99            rss = np.sqrt(echo1.astype("float32") ** 2 + echo2.astype("float32") ** 2)
100            f.create_dataset("raw/rss", data=rss.transpose(2, 0, 1), compression="gzip")
101            f.create_dataset("labels", data=labels.transpose(2, 0, 1), compression="gzip")
102
103        os.rename(f"{volume_path}.tmp", volume_path)
104
105    return preprocessed_dir
106
107
108def _download_sample_data(path, download):
109    msg = "Only the official sample of 3 SKM-TEA scans is downloaded, not the full dataset with 155 scans. "
110    msg += "See 'torch_em.data.datasets.medical.skm_tea' for how to obtain the full dataset."
111    warnings.warn(msg)
112
113    os.makedirs(path, exist_ok=True)
114
115    tar_path = os.path.join(path, "image_files.tar.gz")
116    util.download_source(
117        path=tar_path,
118        url=f"{SAMPLE_URL}/tarball/image_files.tar.gz",
119        download=download,
120        checksum=SAMPLE_CHECKSUMS["image_files.tar.gz"],
121    )
122    util.unzip_tarfile(tar_path=tar_path, dst=path)
123
124    annotation_dir = os.path.join(path, "annotations", SAMPLE_VERSION)
125    os.makedirs(annotation_dir, exist_ok=True)
126    for split in ["train", "val", "test"]:
127        split_path = os.path.join(annotation_dir, f"{split}.json")
128        if os.path.exists(split_path):
129            continue
130
131        util.download_source(
132            path=split_path,
133            url=f"{SAMPLE_URL}/annotations/{SAMPLE_VERSION}/{split}.json",
134            download=download,
135            checksum=SAMPLE_CHECKSUMS[f"{split}.json"],
136        )
137
138
139def get_skm_tea_data(path: Union[os.PathLike, str], download: bool = False) -> str:
140    """Download the SKM-TEA dataset.
141
142    Args:
143        path: Filepath to a folder where the data is downloaded for further processing.
144        download: Whether to download the data if it is not present.
145
146    Returns:
147        Filepath where the data is preprocessed.
148    """
149    preprocessed_dir = os.path.join(path, "preprocessed")
150    if os.path.exists(preprocessed_dir) and len(glob(os.path.join(preprocessed_dir, "*.h5"))) > 0:
151        return preprocessed_dir
152
153    image_dir = os.path.join(path, "image_files")
154    if not os.path.exists(image_dir):
155        if not download:
156            raise FileNotFoundError(
157                f"It's expected to place the downloaded SKM-TEA 'image_files' folder at '{image_dir}'. "
158                "See 'torch_em.data.datasets.medical.skm_tea' for the manual download instructions."
159            )
160
161        _download_sample_data(path, download)
162
163    return _preprocess_inputs(path, image_dir)
164
165
166def get_skm_tea_paths(
167    path: Union[os.PathLike, str],
168    split: Literal["train", "val", "test"],
169    version: str = "v1.0.0",
170    download: bool = False,
171) -> Tuple[List[str], List[str]]:
172    """Get paths to the SKM-TEA data.
173
174    Args:
175        path: Filepath to a folder where the data is downloaded for further processing.
176        split: The choice of data split.
177        version: The version of the annotations that define the split.
178        download: Whether to download the data if it is not present.
179
180    Returns:
181        List of filepaths for the image data.
182        List of filepaths for the label data.
183    """
184    data_dir = get_skm_tea_data(path, download)
185
186    if split not in ["train", "val", "test"]:
187        raise ValueError(f"'{split}' is not a valid split.")
188
189    split_path = os.path.join(path, "annotations", version, f"{split}.json")
190    if not os.path.exists(split_path):
191        raise FileNotFoundError(f"Could not find the annotation file for the '{split}' split at '{split_path}'.")
192
193    with open(split_path, "r") as f:
194        scan_ids = [image["scan_id"] for image in json.load(f)["images"]]
195
196    volume_paths = natsorted([os.path.join(data_dir, f"{scan_id}.h5") for scan_id in scan_ids])
197    volume_paths = [p for p in volume_paths if os.path.exists(p)]
198    assert len(volume_paths) > 0, f"Could not find any preprocessed scans for the '{split}' split."
199
200    return volume_paths, volume_paths
201
202
203def get_skm_tea_dataset(
204    path: Union[os.PathLike, str],
205    patch_shape: Tuple[int, ...],
206    split: Literal["train", "val", "test"],
207    echo: Literal["echo1", "echo2", "rss"] = "echo1",
208    resize_inputs: bool = False,
209    download: bool = False,
210    **kwargs
211) -> Dataset:
212    """Get the SKM-TEA dataset for knee tissue segmentation.
213
214    Args:
215        path: Filepath to a folder where the data is downloaded for further processing.
216        patch_shape: The patch shape to use for training.
217        split: The choice of data split.
218        echo: The qDESS echo to use as input. Either 'echo1', 'echo2' or 'rss' (root-sum-of-squares of both).
219        resize_inputs: Whether to resize inputs to the desired patch shape.
220        download: Whether to download the data if it is not present.
221        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
222
223    Returns:
224        The segmentation dataset.
225    """
226    if echo not in ["echo1", "echo2", "rss"]:
227        raise ValueError(f"'{echo}' is not a valid echo. Choose one of 'echo1', 'echo2' or 'rss'.")
228
229    raw_paths, label_paths = get_skm_tea_paths(path, split, download=download)
230
231    if resize_inputs:
232        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
233        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
234            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
235        )
236
237    return torch_em.default_segmentation_dataset(
238        raw_paths=raw_paths,
239        raw_key=f"raw/{echo}",
240        label_paths=label_paths,
241        label_key="labels",
242        patch_shape=patch_shape,
243        is_seg_dataset=True,
244        **kwargs
245    )
246
247
248def get_skm_tea_loader(
249    path: Union[os.PathLike, str],
250    batch_size: int,
251    patch_shape: Tuple[int, ...],
252    split: Literal["train", "val", "test"],
253    echo: Literal["echo1", "echo2", "rss"] = "echo1",
254    resize_inputs: bool = False,
255    download: bool = False,
256    **kwargs
257) -> DataLoader:
258    """Get the SKM-TEA dataloader for knee tissue segmentation.
259
260    Args:
261        path: Filepath to a folder where the data is downloaded for further processing.
262        batch_size: The batch size for training.
263        patch_shape: The patch shape to use for training.
264        split: The choice of data split.
265        echo: The qDESS echo to use as input. Either 'echo1', 'echo2' or 'rss' (root-sum-of-squares of both).
266        resize_inputs: Whether to resize inputs to the desired patch shape.
267        download: Whether to download the data if it is not present.
268        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
269
270    Returns:
271        The DataLoader.
272    """
273    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
274    dataset = get_skm_tea_dataset(path, patch_shape, split, echo, resize_inputs, download, **ds_kwargs)
275    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)