torch_em.data.datasets.medical.hvsmr

  1"""The HVSMR-2.0 dataset contains annotations for whole-heart segmentation in 3D cardiovascular MRI
  2of patients with congenital heart disease.
  3
  4This module implements the HVSMR-2.0 release (2024), not the older HVSMR 2016 challenge data
  5(https://segchd.csail.mit.edu), which labels only the blood pool and the ventricular myocardium.
  6HVSMR-2.0 consists of 60 cardiovascular MR scans acquired at Boston Children's Hospital, with manual
  7segmentations of the four cardiac chambers and the four great vessels. It does not contain a myocardium
  8label. The label ids are described in `LABEL_IDS`: 1 = left ventricle, 2 = right ventricle, 3 = left atrium,
  94 = right atrium, 5 = aorta, 6 = pulmonary artery, 7 = superior vena cava, 8 = inferior vena cava.
 10Some chambers are missing in single-ventricle and common-atrium patients, which is anatomy and not an
 11annotation error.
 12
 13The release is distributed in three variants, which are selected with the 'variant' argument:
 14- 'orig': the images cropped at the chin, with the original spacing and without normalization.
 15- 'cropped': the images cropped around the heart (the default, and the smallest download).
 16- 'cropped_norm': the heart-cropped images with normalized intensities.
 17Each variant also ships 'pat<n>_*_seg_endpoints.nii.gz' files, which delineate the optional extents of the
 18great vessels for a fairer evaluation. They are not exposed as labels by this module.
 19
 20The data is located at https://doi.org/10.6084/m9.figshare.c.7074755 and is licensed under CC BY 4.0.
 21
 22This dataset is from the publication https://doi.org/10.1038/s41597-024-03469-9.
 23Please cite it if you use this dataset in your research.
 24"""
 25
 26import os
 27from glob import glob
 28from natsort import natsorted
 29from typing import Union, Tuple, List, Literal
 30
 31from torch.utils.data import Dataset, DataLoader
 32
 33import torch_em
 34
 35from .. import util
 36
 37
 38URLS = {
 39    "orig": "https://ndownloader.figshare.com/files/44561774",
 40    "cropped": "https://ndownloader.figshare.com/files/44561792",
 41    "cropped_norm": "https://ndownloader.figshare.com/files/44561783",
 42}
 43
 44CHECKSUMS = {
 45    "orig": "1fb39b1a9ad040f5860eb39cdb6ee673d998e7c5a3b1b3fa06ae5d81d8fa3bb4",
 46    "cropped": "62338737c1cb8cf690f4d2dad770caae9590c4e757b2ecc44ce6d1f33aa2d005",
 47    "cropped_norm": "c1c101117e195b143ec4b66a5a35501766efe907a62db2655ef6024fa655f0fc",
 48}
 49
 50LABEL_IDS = {
 51    "background": 0,
 52    "LV": 1,
 53    "RV": 2,
 54    "LA": 3,
 55    "RA": 4,
 56    "AO": 5,
 57    "PA": 6,
 58    "SVC": 7,
 59    "IVC": 8,
 60}
 61
 62VARIANTS = list(URLS.keys())
 63
 64
 65def get_hvsmr_data(
 66    path: Union[os.PathLike, str],
 67    variant: Literal["orig", "cropped", "cropped_norm"] = "cropped",
 68    download: bool = False,
 69) -> str:
 70    """Download the HVSMR-2.0 dataset.
 71
 72    Args:
 73        path: Filepath to a folder where the data is downloaded for further processing.
 74        variant: The choice of image variant. Either 'orig', 'cropped' or 'cropped_norm'.
 75        download: Whether to download the data if it is not present.
 76
 77    Returns:
 78        Filepath where the data is downloaded.
 79    """
 80    if variant not in VARIANTS:
 81        raise ValueError(f"'{variant}' is not a valid variant. Please choose one of {VARIANTS}.")
 82
 83    data_dir = os.path.join(path, variant)
 84    if os.path.exists(data_dir):
 85        return data_dir
 86
 87    os.makedirs(path, exist_ok=True)
 88
 89    zip_path = os.path.join(path, f"{variant}.zip")
 90    util.download_source(path=zip_path, url=URLS[variant], download=download, checksum=CHECKSUMS[variant])
 91    util.unzip(zip_path=zip_path, dst=path)
 92
 93    return data_dir
 94
 95
 96def get_hvsmr_paths(
 97    path: Union[os.PathLike, str],
 98    variant: Literal["orig", "cropped", "cropped_norm"] = "cropped",
 99    download: bool = False,
100) -> Tuple[List[str], List[str]]:
101    """Get paths to the HVSMR-2.0 data.
102
103    Args:
104        path: Filepath to a folder where the data is downloaded for further processing.
105        variant: The choice of image variant. Either 'orig', 'cropped' or 'cropped_norm'.
106        download: Whether to download the data if it is not present.
107
108    Returns:
109        List of filepaths for the image data.
110        List of filepaths for the label data.
111    """
112    data_dir = get_hvsmr_data(path, variant, download)
113
114    # The 'cropped_norm' images are named after the 'cropped' variant, but with a '_norm' suffix.
115    image_suffix = "cropped_norm" if variant == "cropped_norm" else variant
116    label_suffix = "cropped" if variant == "cropped_norm" else variant
117
118    label_paths = natsorted(glob(os.path.join(data_dir, f"pat*_{label_suffix}_seg.nii.gz")))
119    raw_paths = [p.replace(f"_{label_suffix}_seg.nii.gz", f"_{image_suffix}.nii.gz") for p in label_paths]
120    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in raw_paths)
121
122    return raw_paths, label_paths
123
124
125def get_hvsmr_dataset(
126    path: Union[os.PathLike, str],
127    patch_shape: Tuple[int, ...],
128    variant: Literal["orig", "cropped", "cropped_norm"] = "cropped",
129    resize_inputs: bool = False,
130    download: bool = False,
131    **kwargs
132) -> Dataset:
133    """Get the HVSMR-2.0 dataset for whole-heart segmentation.
134
135    Args:
136        path: Filepath to a folder where the data is downloaded for further processing.
137        patch_shape: The patch shape to use for training.
138        variant: The choice of image variant. Either 'orig', 'cropped' or 'cropped_norm'.
139        resize_inputs: Whether to resize inputs to the desired patch shape.
140        download: Whether to download the data if it is not present.
141        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
142
143    Returns:
144        The segmentation dataset.
145    """
146    raw_paths, label_paths = get_hvsmr_paths(path, variant, download)
147
148    if resize_inputs:
149        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
150        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
151            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
152        )
153
154    return torch_em.default_segmentation_dataset(
155        raw_paths=raw_paths,
156        raw_key="data",
157        label_paths=label_paths,
158        label_key="data",
159        patch_shape=patch_shape,
160        is_seg_dataset=True,
161        **kwargs
162    )
163
164
165def get_hvsmr_loader(
166    path: Union[os.PathLike, str],
167    batch_size: int,
168    patch_shape: Tuple[int, ...],
169    variant: Literal["orig", "cropped", "cropped_norm"] = "cropped",
170    resize_inputs: bool = False,
171    download: bool = False,
172    **kwargs
173) -> DataLoader:
174    """Get the HVSMR-2.0 dataloader for whole-heart segmentation.
175
176    Args:
177        path: Filepath to a folder where the data is downloaded for further processing.
178        batch_size: The batch size for training.
179        patch_shape: The patch shape to use for training.
180        variant: The choice of image variant. Either 'orig', 'cropped' or 'cropped_norm'.
181        resize_inputs: Whether to resize inputs to the desired patch shape.
182        download: Whether to download the data if it is not present.
183        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
184
185    Returns:
186        The DataLoader.
187    """
188    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
189    dataset = get_hvsmr_dataset(path, patch_shape, variant, resize_inputs, download, **ds_kwargs)
190    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)