torch_em.data.datasets.medical.orcascore

  1"""The orCaScore dataset contains annotations for coronary artery calcification in cardiac CT.
  2
  3The data was curated for the orCaScore challenge (https://orcascore.grand-challenge.org), which was held at
  4MICCAI 2014 and evaluates automatic coronary artery calcium scoring. The training set consists of 32 patients,
  58 from each of four CT scanner vendors, and the test set of 40 patients without a reference standard.
  6Each examination consists of a non-contrast enhanced CT image (the file suffix 'CTI') and a contrast enhanced
  7CTA image (the file suffix 'CTAI'). The reference standard (the file suffix 'R') was annotated on the
  8non-contrast CT images only, so this module pairs the labels with the CT images.
  9
 10The reference standard labels every lesion with intensities above 130 HU by the coronary artery it belongs to,
 11see `LABEL_IDS`: 1 = left anterior descending artery (calcifications of the left main coronary artery are
 12labeled as LAD as well), 2 = left circumflex artery, 3 = right coronary artery.
 13
 14NOTE: The data is only handed out to registered participants and the challenge has been closed for new
 15registrations since August 2025, so it cannot be downloaded automatically. There is no openly published
 16copy of it: the data was behind a signed confidentiality agreement from the start, and the original
 17challenge website 'orcascore.isi.uu.nl' no longer resolves. To obtain the data, please follow these steps:
 18- Ask the challenge organizers for access to the training data, either through
 19  https://orcascore.grand-challenge.org or by writing to 'j.m.wolterink@utwente.nl'.
 20- Place the four downloaded archives 'Train_V1.rar', 'Train_V2.rar', 'Train_V3.rar' and 'Train_V4.rar' in the
 21  folder passed as 'path'. This module extracts them, or picks up the extracted MetaImage files
 22  ('TRV<vendor>P<patient>CTI.mhd' with the matching '.raw' or '.zraw') if they are already extracted.
 23
 24NOTE: The filenames and the archive names are taken from the archived download page of the original
 25challenge website, http://web.archive.org/web/20180418103434/http://orcascore.isi.uu.nl:80/download/. It
 26documents 'TR' for the training set (the test set uses 'TE'), 'V' for the vendor, 'P' for the patient,
 27the suffix 'CTI' for the non-contrast CT, 'CTAI' for the CTA and 'R' for the reference standard, and the
 28archive sizes 390 MiB, 815 MiB, 500 MiB and 608 MiB for the four training archives.
 29
 30The MetaImage volumes are converted to hdf5 volumes (the keys are 'raw' and 'labels') by this module.
 31
 32The dataset is located at https://orcascore.grand-challenge.org.
 33
 34This dataset is from the publication https://doi.org/10.1118/1.4945696.
 35Please cite it if you use this dataset in your research.
 36"""
 37
 38import os
 39import zlib
 40from glob import glob
 41from tqdm import tqdm
 42from natsort import natsorted
 43from typing import Union, Tuple, List
 44
 45import numpy as np
 46
 47from torch.utils.data import Dataset, DataLoader
 48
 49import torch_em
 50
 51from .. import util
 52
 53
 54LABEL_IDS = {"background": 0, "LAD": 1, "LCX": 2, "RCA": 3}
 55
 56ARCHIVE_NAMES = ["Train_V1.rar", "Train_V2.rar", "Train_V3.rar", "Train_V4.rar"]
 57
 58N_VOLUMES = 32
 59
 60METAIMAGE_DTYPES = {
 61    "MET_CHAR": "int8",
 62    "MET_UCHAR": "uint8",
 63    "MET_SHORT": "int16",
 64    "MET_USHORT": "uint16",
 65    "MET_INT": "int32",
 66    "MET_UINT": "uint32",
 67    "MET_FLOAT": "float32",
 68    "MET_DOUBLE": "float64",
 69}
 70
 71
 72def _read_metaimage(mhd_path):
 73    """Read a MetaImage volume (a '.mhd' header with a '.raw' or a zlib compressed '.zraw' data file)."""
 74    header = {}
 75    with open(mhd_path, "r") as f:
 76        for line in f:
 77            if "=" not in line:
 78                continue
 79            key, value = line.split("=", 1)
 80            header[key.strip()] = value.strip()
 81
 82    shape = [int(v) for v in header["DimSize"].split()][::-1]  # The data is stored with the first axis last.
 83    dtype = np.dtype(METAIMAGE_DTYPES[header["ElementType"]])
 84    dtype = dtype.newbyteorder(">" if header.get("BinaryDataByteOrderMSB", "False") == "True" else "<")
 85
 86    with open(os.path.join(os.path.dirname(mhd_path), header["ElementDataFile"]), "rb") as f:
 87        data = f.read()
 88    if header.get("CompressedData", "False") == "True":
 89        data = zlib.decompress(data)
 90
 91    return np.frombuffer(data, dtype=dtype).reshape(shape)
 92
 93
 94def _preprocess_inputs(data_dir, preprocessed_dir):
 95    import h5py
 96
 97    label_paths = natsorted(glob(os.path.join(data_dir, "**", "TRV*P*R.mhd"), recursive=True))
 98    os.makedirs(preprocessed_dir, exist_ok=True)
 99
100    for label_path in tqdm(label_paths, desc="Preprocessing the orCaScore cases"):
101        case_id = os.path.basename(label_path)[:-len("R.mhd")]
102        volume_path = os.path.join(preprocessed_dir, f"{case_id}.h5")
103        if os.path.exists(volume_path):
104            continue
105
106        raw = _read_metaimage(os.path.join(os.path.dirname(label_path), f"{case_id}CTI.mhd"))
107        labels = _read_metaimage(label_path)
108
109        # The file is written to a temporary path first, so that an interrupted run leaves no corrupt file.
110        with h5py.File(f"{volume_path}.tmp", "w") as f:
111            f.create_dataset("raw", data=raw, compression="gzip")
112            f.create_dataset("labels", data=labels.astype("uint8"), compression="gzip")
113
114        os.rename(f"{volume_path}.tmp", volume_path)
115
116
117def get_orcascore_data(path: Union[os.PathLike, str], download: bool = False) -> str:
118    """Obtain the orCaScore dataset.
119
120    Args:
121        path: Filepath to a folder where the manually downloaded data is stored.
122        download: Whether to download the data if it is not present. The data cannot be downloaded
123            automatically, so this raises if the data has not been downloaded manually.
124
125    Returns:
126        Filepath where the preprocessed data is stored.
127    """
128    preprocessed_dir = os.path.join(path, "preprocessed")
129    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == N_VOLUMES:
130        return preprocessed_dir
131
132    if not glob(os.path.join(path, "**", "TRV*P*R.mhd"), recursive=True):
133        archive_paths = [os.path.join(path, name) for name in ARCHIVE_NAMES]
134        if not any(os.path.exists(p) for p in archive_paths):
135            msg = "'torch_em' cannot download this dataset, because the orCaScore data is only handed out to "
136            msg += "registered participants and the challenge has been closed for new registrations. Please ask "
137            msg += "the organizers at 'https://orcascore.grand-challenge.org' or at 'j.m.wolterink@utwente.nl' for "
138            msg += f"the training data and place {ARCHIVE_NAMES} in '{path}'."
139            raise NotImplementedError(msg)
140
141        for archive_path in archive_paths:
142            util.unzip_rarfile(rar_path=archive_path, dst=path, remove=False)
143
144    _preprocess_inputs(path, preprocessed_dir)
145    return preprocessed_dir
146
147
148def get_orcascore_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
149    """Get paths to the orCaScore data.
150
151    Args:
152        path: Filepath to a folder where the manually downloaded data is stored.
153        download: Whether to download the data if it is not present.
154
155    Returns:
156        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
157    """
158    data_dir = get_orcascore_data(path, download)
159    return natsorted(glob(os.path.join(data_dir, "*.h5")))
160
161
162def get_orcascore_dataset(
163    path: Union[os.PathLike, str],
164    patch_shape: Tuple[int, ...],
165    resize_inputs: bool = False,
166    download: bool = False,
167    **kwargs
168) -> Dataset:
169    """Get the orCaScore dataset for coronary artery calcification segmentation.
170
171    Args:
172        path: Filepath to a folder where the manually downloaded data is stored.
173        patch_shape: The patch shape to use for training.
174        resize_inputs: Whether to resize inputs to the desired patch shape.
175        download: Whether to download the data if it is not present.
176        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
177
178    Returns:
179        The segmentation dataset.
180    """
181    volume_paths = get_orcascore_paths(path, download)
182
183    if resize_inputs:
184        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
185        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
186            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
187        )
188
189    return torch_em.default_segmentation_dataset(
190        raw_paths=volume_paths,
191        raw_key="raw",
192        label_paths=volume_paths,
193        label_key="labels",
194        patch_shape=patch_shape,
195        is_seg_dataset=True,
196        **kwargs
197    )
198
199
200def get_orcascore_loader(
201    path: Union[os.PathLike, str],
202    batch_size: int,
203    patch_shape: Tuple[int, ...],
204    resize_inputs: bool = False,
205    download: bool = False,
206    **kwargs
207) -> DataLoader:
208    """Get the orCaScore dataloader for coronary artery calcification segmentation.
209
210    Args:
211        path: Filepath to a folder where the manually downloaded data is stored.
212        batch_size: The batch size for training.
213        patch_shape: The patch shape to use for training.
214        resize_inputs: Whether to resize inputs to the desired patch shape.
215        download: Whether to download the data if it is not present.
216        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
217
218    Returns:
219        The DataLoader.
220    """
221    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
222    dataset = get_orcascore_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
223    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)