torch_em.data.datasets.medical.scd

  1"""The SCD dataset contains annotations for left ventricular endocardium and epicardium
  2segmentation in short-axis cardiac cine-MRI.
  3
  4The Sunnybrook Cardiac Data (SCD), also known as the 2009 Cardiac MR Left Ventricle Segmentation
  5Challenge data, consists of 45 cine-MRI studies from a mix of patients and pathologies: healthy,
  6hypertrophy, heart failure with infarction and heart failure without infarction. Manual contours
  7of the left ventricular endocardium and epicardium are provided for the end-diastolic and
  8end-systolic frames (and, for a subset of slices, further frames), drawn by an expert
  9cardiologist. The dataset is located at https://www.cardiacatlas.org/sunnybrook-cardiac-data/
 10and is distributed under the CC0 1.0 Universal license.
 11
 12The labels are multiclass, see `LABEL_IDS`: 1 = left ventricular cavity (endocardium), 2 = left
 13ventricular myocardium (the region between the epicardial and endocardial contour). Slices without
 14an epicardial contour only have the cavity label.
 15
 16This dataset is from the publication http://hdl.handle.net/10380/3070.
 17Please cite it if you use this dataset for your research.
 18"""
 19
 20import os
 21import re
 22from glob import glob
 23from tqdm import tqdm
 24from natsort import natsorted
 25from typing import Union, Tuple, List
 26
 27import numpy as np
 28from skimage.draw import polygon
 29
 30from torch.utils.data import Dataset, DataLoader
 31
 32import torch_em
 33
 34from .. import util
 35
 36
 37URLS = {
 38    "images_1": "https://www.cardiacatlas.org/share/download.php?id=98&token=kUVtCTMzfjlHwunzFYST7yNGFnlUtV7W&download",  # noqa
 39    "images_2": "https://www.cardiacatlas.org/share/download.php?id=99&token=2euh4WH03DynHi9Y5elb2HSWJaJIeWER&download",  # noqa
 40    "images_3": "https://www.cardiacatlas.org/share/download.php?id=100&token=eT37AbXslu1JQp1GSZy2BmmWzWmXDdfx&download",  # noqa
 41    "images_4": "https://www.cardiacatlas.org/share/download.php?id=101&token=Qn0dY2lHTyTIBF5qNHRkbHdUSnWnS4Yy&download",  # noqa
 42    "images_5": "https://www.cardiacatlas.org/share/download.php?id=102&token=nYALJl6kS6cj5jgK9R8CKDpgQwZn1pm3&download",  # noqa
 43    "contours": "https://www.cardiacatlas.org/share/download.php?id=61&token=gFjv8psvCQu2vDYldUW9TYUIDic11fDt&download",  # noqa
 44    "patient_data": "https://www.cardiacatlas.org/share/download.php?id=66&token=Y1D66ieUdbXFlqmZ4icHJheJy44MXUPY&download",  # noqa
 45}
 46
 47CHECKSUMS = {
 48    "images_1": "708ce04db1ac33948a00b9052d44e9548c6807121a4841f4c35080d6db127b72",
 49    "images_2": "ea97bd11dfc1154bdcf02b6466c00637bc7e0dc04b5ddc3d94d68a7cd7a3e6a5",
 50    "images_3": "5b065fadef1b523bbabfec0d85327f375f743f2dd72810a5927b21b834b45deb",
 51    "images_4": "2a8d7054d9b89b5af4c9764b4b04eecee57b61b71b69f5ac69ec12629e9a615c",
 52    "images_5": "e65a818fe89c665c344317222a052360e210936c81f89ceb9d36e8e3c6e2d6e9",
 53    "contours": "93d8e06dfa53fc384e78810aeac86c336b5bdecbff1c9c175ec9254851cc2162",
 54    "patient_data": "c39eb11924d021e8ec51bb984daa57d1adadcc8eb7e342640181ce151dae5ded",
 55}
 56
 57LABEL_IDS = {"background": 0, "cavity": 1, "myocardium": 2}
 58
 59
 60def _normalize_original_id(original_id):
 61    """Zero-pad the trailing case number of an 'OriginalID' to match the contour directory naming,
 62    e.g. 'SC-HF-I-1' -> 'SC-HF-I-01'."""
 63    return re.sub(r"-(\d+)$", lambda m: f"-{int(m.group(1)):02d}", original_id)
 64
 65
 66def _patient_id_mapping(patient_csv_path):
 67    import csv
 68
 69    mapping = {}
 70    with open(patient_csv_path, newline="", encoding="utf-8-sig") as f:
 71        for row in csv.DictReader(f):
 72            mapping[_normalize_original_id(row["OriginalID"])] = row["PatientID"]
 73    return mapping
 74
 75
 76def _parse_contour(contour_path):
 77    points = np.loadtxt(contour_path)
 78    return points[:, 0], points[:, 1]
 79
 80
 81INSTANCE_PATTERN = re.compile(r"IM-\d+-(\d+)-[io]contour-manual\.txt$")
 82
 83
 84def _index_cinesax_instances(patient_dir):
 85    """Index the DICOM instances of the short-axis cine series ('CINESAX') of a patient folder by their
 86    instance number. The Cardiac Atlas Project redistribution stores several series per patient (long-axis
 87    cines, scouts, perfusion, ...) under arbitrarily numbered subfolders, and only the DICOM header's
 88    'SeriesDescription' reliably identifies the short-axis cine series that the manual contours refer to;
 89    the numeric prefix of the contour filenames (e.g. 'IM-0001-0048') does not correspond to the subfolder
 90    or series naming of this redistribution."""
 91    import pydicom
 92
 93    instances = {}
 94    for series_dir in sorted(p for p in glob(os.path.join(patient_dir, "*")) if os.path.isdir(p)):
 95        dicom_paths = natsorted(glob(os.path.join(series_dir, "*.dcm")))
 96        if not dicom_paths:
 97            continue
 98
 99        header = pydicom.dcmread(dicom_paths[0], stop_before_pixels=True)
100        if "CINESAX" not in getattr(header, "SeriesDescription", "").upper():
101            continue
102
103        for dicom_path in dicom_paths:
104            match = re.search(r"-(\d+)\.dcm$", os.path.basename(dicom_path))
105            if match is None:
106                continue
107            instances.setdefault(int(match.group(1)), dicom_path)
108
109    return instances
110
111
112def _rasterize_case(instances, contour_dir, preprocessed_dir, patient_id):
113    import pydicom
114    import imageio.v3 as imageio
115
116    icontours = natsorted(glob(os.path.join(contour_dir, "*-icontour-manual.txt")))
117
118    for icontour_path in icontours:
119        match = INSTANCE_PATTERN.search(os.path.basename(icontour_path))
120        if match is None or int(match.group(1)) not in instances:
121            continue
122
123        img_no = int(match.group(1))
124        stem = f"{img_no:04}"
125        gt_path = os.path.join(preprocessed_dir, f"{patient_id}_{stem}.tif")
126        raw_path = os.path.join(preprocessed_dir, f"{patient_id}_{stem}_raw.tif")
127        if os.path.exists(gt_path) and os.path.exists(raw_path):
128            continue
129
130        dcm = pydicom.dcmread(instances[img_no])
131        image = np.asarray(dcm.pixel_array)
132        shape = image.shape
133
134        labels = np.zeros(shape, dtype="uint8")
135
136        ocontour_path = os.path.join(contour_dir, os.path.basename(icontour_path).replace("icontour", "ocontour"))
137        if os.path.exists(ocontour_path):
138            x, y = _parse_contour(ocontour_path)
139            r, c = polygon(y, x, shape=shape)
140            labels[r, c] = LABEL_IDS["myocardium"]
141
142        x, y = _parse_contour(icontour_path)
143        r, c = polygon(y, x, shape=shape)
144        labels[r, c] = LABEL_IDS["cavity"]
145
146        imageio.imwrite(raw_path, image)
147        imageio.imwrite(gt_path, labels)
148
149
150def _preprocess_inputs(path, preprocessed_dir):
151    os.makedirs(preprocessed_dir, exist_ok=True)
152
153    patient_csv_path = os.path.join(path, "scd_patientdata.csv")
154    mapping = _patient_id_mapping(patient_csv_path)
155
156    contours_root = os.path.join(path, "SCD_ManualContours")
157    case_dirs = natsorted(glob(os.path.join(contours_root, "SC-*")))
158
159    for case_dir in tqdm(case_dirs, desc="Preprocessing the SCD studies"):
160        original_id = os.path.basename(case_dir)
161        patient_id = mapping.get(original_id)
162        if patient_id is None:
163            continue
164
165        contour_dir = os.path.join(case_dir, "contours-manual", "IRCCI-expert")
166        if not os.path.exists(contour_dir):
167            continue
168
169        patient_dir = os.path.join(path, patient_id)
170        instances = _index_cinesax_instances(patient_dir)
171        if not instances:
172            continue
173
174        _rasterize_case(instances, contour_dir, preprocessed_dir, patient_id)
175
176
177def get_scd_data(path: Union[os.PathLike, str], download: bool = False) -> str:
178    """Download the SCD dataset.
179
180    Args:
181        path: Filepath to a folder where the data is downloaded for further processing.
182        download: Whether to download the data if it is not present.
183
184    Returns:
185        Filepath where the preprocessed data is stored.
186    """
187    preprocessed_dir = os.path.join(path, "preprocessed")
188    if glob(os.path.join(preprocessed_dir, "*_raw.tif")):
189        return preprocessed_dir
190
191    os.makedirs(path, exist_ok=True)
192
193    for key in ["images_1", "images_2", "images_3", "images_4", "images_5"]:
194        zip_path = os.path.join(path, f"{key}.zip")
195        util.download_source(
196            path=zip_path, url=URLS[key], download=download, checksum=CHECKSUMS[key], verify=False
197        )
198        util.unzip(zip_path=zip_path, dst=path)
199
200    contours_dir = os.path.join(path, "SCD_ManualContours")
201    if not os.path.exists(contours_dir):
202        zip_path = os.path.join(path, "contours.zip")
203        util.download_source(
204            path=zip_path, url=URLS["contours"], download=download, checksum=CHECKSUMS["contours"], verify=False
205        )
206        util.unzip(zip_path=zip_path, dst=path)
207
208    patient_csv_path = os.path.join(path, "scd_patientdata.csv")
209    util.download_source(
210        path=patient_csv_path,
211        url=URLS["patient_data"],
212        download=download,
213        checksum=CHECKSUMS["patient_data"],
214        verify=False,
215    )
216
217    _preprocess_inputs(path, preprocessed_dir)
218    return preprocessed_dir
219
220
221def get_scd_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
222    """Get paths to the SCD data.
223
224    Args:
225        path: Filepath to a folder where the data is downloaded for further processing.
226        download: Whether to download the data if it is not present.
227
228    Returns:
229        List of filepaths for the image data.
230        List of filepaths for the label data.
231    """
232    data_dir = get_scd_data(path, download)
233
234    gt_paths = natsorted(glob(os.path.join(data_dir, "*.tif")))
235    gt_paths = [p for p in gt_paths if not p.endswith("_raw.tif")]
236    image_paths = [p.replace(".tif", "_raw.tif") for p in gt_paths]
237
238    return image_paths, gt_paths
239
240
241def get_scd_dataset(
242    path: Union[os.PathLike, str],
243    patch_shape: Tuple[int, int],
244    resize_inputs: bool = False,
245    download: bool = False,
246    **kwargs
247) -> Dataset:
248    """Get the SCD dataset for left ventricular cavity and myocardium segmentation.
249
250    Args:
251        path: Filepath to a folder where the data is downloaded for further processing.
252        patch_shape: The patch shape to use for training.
253        resize_inputs: Whether to resize the inputs to the expected patch shape.
254        download: Whether to download the data if it is not present.
255        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
256
257    Returns:
258        The segmentation dataset.
259    """
260    image_paths, gt_paths = get_scd_paths(path, download)
261
262    if resize_inputs:
263        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
264        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
265            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
266        )
267
268    return torch_em.default_segmentation_dataset(
269        raw_paths=image_paths,
270        raw_key=None,
271        label_paths=gt_paths,
272        label_key=None,
273        patch_shape=patch_shape,
274        is_seg_dataset=False,
275        **kwargs
276    )
277
278
279def get_scd_loader(
280    path: Union[os.PathLike, str],
281    batch_size: int,
282    patch_shape: Tuple[int, int],
283    resize_inputs: bool = False,
284    download: bool = False,
285    **kwargs
286) -> DataLoader:
287    """Get the SCD dataloader for left ventricular cavity and myocardium segmentation.
288
289    Args:
290        path: Filepath to a folder where the data is downloaded for further processing.
291        batch_size: The batch size for training.
292        patch_shape: The patch shape to use for training.
293        resize_inputs: Whether to resize the inputs to the expected patch shape.
294        download: Whether to download the data if it is not present.
295        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
296
297    Returns:
298        The DataLoader.
299    """
300    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
301    dataset = get_scd_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
302    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)