torch_em.data.datasets.medical.mnms

  1"""The M&Ms dataset contains annotations for cardiac structure segmentation in
  2multi-centre, multi-vendor and multi-disease cine cardiac MRI.
  3
  4The data was curated for the M&Ms challenge (https://www.ub.edu/mnms), which was held at MICCAI 2020 and
  5targets models that generalize across clinical centres and scanner vendors. The public 'OpenDataset' release
  6consists of 345 of the 375 studies of the cohort: 175 training studies (150 of them annotated), 34 validation
  7studies and 136 testing studies, all acquired in six clinical centres with scanners of four vendors
  8('A' = Siemens, 'B' = Philips, 'C' = General Electric, 'D' = Canon). This module exposes the 320 annotated
  9studies via the official splits ('train', 'val' and 'test'), and the vendor can be selected with the
 10'vendor' argument.
 11
 12Each study is a 4D short-axis cine acquisition, of which the end diastole (ED) and the end systole (ES) phase
 13are annotated. The two phases are extracted into separate volumes (selected with the 'phase' argument), so
 14that this module provides 640 annotated volumes, 300 of them in the training split. The label ids are
 15described in `LABEL_IDS`: 1 = left ventricle cavity, 2 = left ventricle myocardium, 3 = right ventricle.
 16NOTE: This is not the label order of the ACDC dataset, which numbers the same structures the other way round.
 17
 18The cine volumes are stored as 4D nifti files with the slice axis third, so the annotated phases are
 19converted to hdf5 volumes with the slice axis first (the keys are 'raw' and 'labels') by this module.
 20
 21NOTE: The official data at https://www.ub.edu/mnms is only handed out after signing a data use agreement, so
 22this module downloads a public mirror of the release at https://huggingface.co/datasets/zhuyinheng/mnms.
 23If the official 'OpenDataset' folder is placed in the folder passed as 'path', it is used instead.
 24
 25This module covers the first edition of the challenge. M&Ms-2 (https://www.ub.edu/mnms-2), which focuses on
 26the right ventricle and adds long-axis views, is a different dataset and would need a separate module.
 27
 28This dataset is from the publication https://doi.org/10.1109/TMI.2021.3090082.
 29Please cite it if you use this dataset in your research.
 30"""
 31
 32import os
 33import csv
 34from glob import glob
 35from tqdm import tqdm
 36from natsort import natsorted
 37from typing import Union, Tuple, List, Optional, Literal
 38
 39import numpy as np
 40
 41from torch.utils.data import Dataset, DataLoader
 42
 43import torch_em
 44
 45from .. import util
 46
 47
 48URL_BASE = "https://huggingface.co/datasets/zhuyinheng/mnms/resolve/main"
 49
 50CSV_NAME = "211230_M&Ms_Dataset_information_diagnosis_opendataset.csv"
 51
 52LABEL_IDS = {"background": 0, "LV": 1, "MYO": 2, "RV": 3}
 53
 54SPLITS = {"train": os.path.join("Training", "Labeled"), "val": "Validation", "test": "Testing"}
 55
 56PHASES = ["ED", "ES"]
 57
 58VENDORS = {"A": "Siemens", "B": "Philips", "C": "General Electric", "D": "Canon"}
 59
 60SUBJECT_IDS = {
 61    "train": [
 62        "A0S9V9", "A1D0Q7", "A1D9Z7", "A1E9Q1", "A1O8Z3", "A2C0I1", "A2N8V0", "A3B7E5", "A3H1O5", "A4B5U4", "A4J4S4",
 63        "A4U9V5", "A5E0T8", "A6B5G9", "A6D5F9", "A6M1Q7", "A7D9L8", "A7G0P5", "A7M7P8", "A7O4T6", "A8C9U8", "A8E1F4",
 64        "A9C5P4", "A9E3G9", "A9J5Q7", "A9J8W7", "B0I2Z0", "B0N3W8", "B2C2Z7", "B2D9M2", "B2D9O2", "B2F4K5", "B2G5R2",
 65        "B3D0N1", "B3O1S0", "B3P3R1", "B4O3V3", "B6D0U7", "B8H5H6", "B8J7R4", "B9E0Q1", "B9O1Q0", "C0K1P0", "C0S7W0",
 66        "C1G5Q0", "C1K8P5", "C2J0K3", "C2L5P7", "C2M6P8", "C3I2K3", "C4R8T7", "C4S8W9", "C5M4S2", "C6J5P1", "C8P3S7",
 67        "D0H9I4", "D0R0R9", "D1J5P6", "D1L4Q9", "D1M1S6", "D3D4Y5", "D3F3O5", "D3F9H9", "D3O9U9", "D4M3Q2", "D4N6W6",
 68        "D6H6O2", "D8E4F4", "D9L1Z3", "E0M3U7", "E0O0S0", "E3T0Z2", "E4M2Q7", "E4W8Z7", "E5E6O8", "E5F5V7", "E9H1U4",
 69        "E9H2K7", "E9L1W5", "E9V4Z8", "F0J2R8", "F1F3I6", "F2H5S1", "F3G5K5", "F4K3S1", "F5I9Q2", "F8N2S1", "G0H4J3",
 70        "G0I6P3", "G1N6S7", "G2J1M5", "G2M7W4", "G2O2S6", "G4L8Z7", "G4S9U3", "G5P4U3", "G7I5V7", "G8N2U5", "G9L0O9",
 71        "H0K3Q4", "H1I3W0", "H1J5W8", "H1M5Y6", "H1W2Y1", "H3U1Y1", "H4I2T8", "H5N0P0", "H6I0I6", "H7I4J3", "H7N4V9",
 72        "I0J5U3", "I2K2Y8", "I6N3P3", "I7T3U1", "J1T9Y1", "J4J9W6", "J6K6P5", "J6P5T8", "J8R5W2", "J9L6N9", "K2S1U6",
 73        "K4T7Y0", "K5L2U3", "K5P0Y1", "L1Q1Z5", "L1Q9V8", "L4Q2U3", "L5Q6T7", "M0P8U8", "M1R4S1", "M2P1R1", "M4P7Q6",
 74        "N1P8Q9", "N5S7Y1", "N7V9W9", "N8N9U0", "O0S9V7", "O3R8Y5", "P0S5Y0", "P5R1Y4", "P6U0Y0", "P9S7W2", "Q0U0V5",
 75        "Q3R9W7", "Q7V1Y5", "R4Y1Z9", "S1S3Z7", "T2T9Z9", "T9U9W2", "W5Z4Z8"
 76    ],
 77    "val": [
 78        "A5C2D2", "A9F3T5", "B0H7V0", "C4E9I1", "C5L0R0", "C6E0F9", "C8I7P7", "C8J7L5", "D1H6U2", "D1R0Y5", "D1S5T8",
 79        "D2U0V0", "D5G3W8", "D6N7Q8", "D7M8P9", "D7T3V8", "D8O0W2", "E1L8Y4", "E3F2U7", "E4O8P3", "F6J9L9", "G7Q2W0",
 80        "H3R6S9", "H6P7T1", "I4L4V7", "I7W4Y8", "J6M5O2", "K3R0Y7", "K5K6N1", "K5M7V5", "K9N0W0", "N2O7U5", "O7Q7U3",
 81        "P8V0Y7"
 82    ],
 83    "test": [
 84        "A1K2P5", "A2H5K9", "A2L1N6", "A3H5R1", "A3P9V7", "A4A8V9", "A4B9O6", "A4K8R4", "A4R4T0", "A5D0G0", "A5H1Q2",
 85        "A5P5W0", "A5Q1W8", "A6A8H0", "A6B7Y4", "A6J0Y2", "A7E4J0", "A7F4G2", "A8C5E9", "A9L7Y7", "B0L3Y2", "B1G9J3",
 86        "B2L0L2", "B3E2W8", "B3F0V9", "B3S2Z4", "B4E1K1", "B4S1Y2", "B5F8L9", "B5L5Y4", "B5T6V0", "B6I0T4", "B7F5P0",
 87        "B8P5Q9", "B9G4U2", "B9H8N8", "C0L7V1", "C0N8P4", "C6U4W8", "C7L8Z8", "C7M6W0", "C8O0P2", "D1H2O9", "D1L6T4",
 88        "D3K5Q2", "D3Q0W9", "D6E9U8", "D9F5P1", "D9I8O7", "E0J2Z9", "E0J7L9", "E1L7M3", "E3F5U2", "E3L8U8", "E4H7L4",
 89        "E4I9O7", "E5J6L2", "E5S7W7", "E6H0V9", "E6J4N8", "E6M6P2", "E7L0N6", "E9V9Z2", "F0I6U8", "F0K4T6", "F1K2S9",
 90        "F5I1Z8", "F9M1R2", "G1J5K3", "G1K1V3", "G3M5S4", "G4I7V2", "G7N8R7", "G7S6V0", "G8K0M3", "G8L0Z0", "G8R0Z9",
 91        "G9N5V9", "H1N8S6", "H2M9S1", "H7K5U5", "H7L8R8", "H7P5Z4", "H8K2K7", "I0I2J8", "I2J6Z6", "I4R8V6", "I5L3S2",
 92        "I6P4R0", "I6T4W8", "I8N8Y1", "I8Z0Z6", "J4J8Q3", "J6K4V3", "J9L4S2", "K3P3Y6", "K5L4S1", "K6N4N7", "K7L2Y6",
 93        "K7N0R7", "K7O3Q0", "L2V5Z0", "L5U7Y4", "L6T2T5", "L7Y7Z2", "L8M2U8", "L8N7P0", "L8N7Z0", "M2P5T8", "M4T4V6",
 94        "M6M9N1", "M6V2Y0", "N7P3T8", "N7W6Z8", "N9P5Z0", "N9Q4T8", "O4O6U5", "O4T6Y7", "O5U2U7", "O9V8W5", "P3P9S5",
 95        "P3R6Y5", "P3T5U1", "P8W4Z0", "Q1Q3T1", "Q3Q6R8", "Q4W5Z8", "Q5V8W3", "R1R6Y8", "R2R7Z5", "R3V5W7", "R6V5W3",
 96        "R8V0Y4", "T2Z1Z9", "V4W8Z5", "Y6Y9Z2"
 97    ],
 98}
 99
100
101def _get_metadata(path, download):
102    """Read the official information sheet, which holds the vendor and the ED and ES frame index per subject."""
103    csv_paths = glob(os.path.join(path, "**", CSV_NAME), recursive=True)
104    csv_path = csv_paths[0] if csv_paths else os.path.join(path, CSV_NAME)
105    util.download_source(path=csv_path, url=f"{URL_BASE}/{CSV_NAME.replace('&', '%26')}", download=download)
106
107    with open(csv_path, "r") as f:
108        return {row["External code"]: row for row in csv.DictReader(f)}
109
110
111def _download_volumes(path, split, subject_ids):
112    split_dir = os.path.join(path, "OpenDataset", SPLITS[split])
113    for subject_id in tqdm(subject_ids, desc=f"Downloading the M&Ms studies of '{split}'"):
114        for suffix in ["sa", "sa_gt"]:
115            volume_path = os.path.join(split_dir, subject_id, f"{subject_id}_{suffix}.nii.gz")
116            os.makedirs(os.path.dirname(volume_path), exist_ok=True)
117            url = f"{URL_BASE}/{SPLITS[split].replace(os.sep, '/')}/{subject_id}/{subject_id}_{suffix}.nii.gz"
118            util.download_source(path=volume_path, url=url, download=True)
119
120    return split_dir
121
122
123def _preprocess_inputs(split_dir, metadata, subject_ids, preprocessed_dir):
124    import h5py
125    import nibabel as nib
126
127    os.makedirs(preprocessed_dir, exist_ok=True)
128    for subject_id in tqdm(subject_ids, desc="Preprocessing the M&Ms studies"):
129        raw_volume = nib.load(os.path.join(split_dir, subject_id, f"{subject_id}_sa.nii.gz"))
130        label_volume = nib.load(os.path.join(split_dir, subject_id, f"{subject_id}_sa_gt.nii.gz"))
131
132        for phase in PHASES:
133            volume_path = os.path.join(preprocessed_dir, f"{subject_id}_{phase}.h5")
134            if os.path.exists(volume_path):
135                continue
136
137            # The transpose maps the nifti axis order (X, Y, Z) of one phase to the (Z, Y, X) order of the volumes.
138            frame = int(metadata[subject_id][phase])
139            raw = np.asarray(raw_volume.dataobj[..., frame]).T
140            labels = np.asarray(label_volume.dataobj[..., frame]).T
141
142            # The file is written to a temporary path first, so that an interrupted run leaves no corrupt file.
143            with h5py.File(f"{volume_path}.tmp", "w") as f:
144                f.create_dataset("raw", data=raw, compression="gzip")
145                f.create_dataset("labels", data=labels.astype("uint8"), compression="gzip")
146
147            os.rename(f"{volume_path}.tmp", volume_path)
148
149
150def get_mnms_data(
151    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False
152) -> str:
153    """Download the M&Ms dataset.
154
155    Args:
156        path: Filepath to a folder where the data is downloaded for further processing.
157        split: The choice of data split. Either 'train', 'val' or 'test'.
158        download: Whether to download the data if it is not present.
159
160    Returns:
161        Filepath where the preprocessed data is stored.
162    """
163    if split not in SPLITS:
164        raise ValueError(f"'{split}' is not a valid split. Please choose one of {list(SPLITS.keys())}.")
165
166    subject_ids = SUBJECT_IDS[split]
167    preprocessed_dir = os.path.join(path, "preprocessed", split)
168    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == len(subject_ids) * len(PHASES):
169        return preprocessed_dir
170
171    os.makedirs(path, exist_ok=True)
172
173    # Use the official release if it was downloaded manually, otherwise fetch the studies from the mirror.
174    mirror_dir = os.path.join(path, "OpenDataset", SPLITS[split])
175    split_dirs = [p for p in glob(os.path.join(path, "**", SPLITS[split]), recursive=True) if p != mirror_dir]
176    if split_dirs:
177        split_dir = split_dirs[0]
178    elif download:
179        split_dir = _download_volumes(path, split, subject_ids)
180    else:
181        raise RuntimeError(f"Cannot find the data at '{path}', but download was set to False.")
182
183    _preprocess_inputs(split_dir, _get_metadata(path, download), subject_ids, preprocessed_dir)
184    return preprocessed_dir
185
186
187def get_mnms_paths(
188    path: Union[os.PathLike, str],
189    split: Literal["train", "val", "test"],
190    phase: Optional[Literal["ED", "ES"]] = None,
191    vendor: Optional[Literal["A", "B", "C", "D"]] = None,
192    download: bool = False,
193) -> List[str]:
194    """Get paths to the M&Ms data.
195
196    Args:
197        path: Filepath to a folder where the data is downloaded for further processing.
198        split: The choice of data split. Either 'train', 'val' or 'test'.
199        phase: The choice of cardiac phase. Either 'ED' or 'ES'. If None, both phases are returned.
200        vendor: The choice of scanner vendor. One of 'A', 'B', 'C' or 'D'. If None, all vendors are returned.
201        download: Whether to download the data if it is not present.
202
203    Returns:
204        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
205    """
206    data_dir = get_mnms_data(path, split, download)
207
208    if phase is not None and phase not in PHASES:
209        raise ValueError(f"'{phase}' is not a valid phase. Please choose one of {PHASES}.")
210
211    volume_paths = natsorted(glob(os.path.join(data_dir, f"*_{'*' if phase is None else phase}.h5")))
212
213    if vendor is not None:
214        if vendor not in VENDORS:
215            raise ValueError(f"'{vendor}' is not a valid vendor. Please choose one of {list(VENDORS.keys())}.")
216
217        metadata = _get_metadata(path, download)
218        volume_paths = [
219            p for p in volume_paths if metadata[os.path.basename(p).rsplit("_", 1)[0]]["Vendor"] == vendor
220        ]
221
222    assert len(volume_paths) > 0, f"Could not find any preprocessed volumes in '{data_dir}'."
223    return volume_paths
224
225
226def get_mnms_dataset(
227    path: Union[os.PathLike, str],
228    patch_shape: Tuple[int, ...],
229    split: Literal["train", "val", "test"],
230    phase: Optional[Literal["ED", "ES"]] = None,
231    vendor: Optional[Literal["A", "B", "C", "D"]] = None,
232    resize_inputs: bool = False,
233    download: bool = False,
234    **kwargs
235) -> Dataset:
236    """Get the M&Ms dataset for cardiac structure segmentation.
237
238    Args:
239        path: Filepath to a folder where the data is downloaded for further processing.
240        patch_shape: The patch shape to use for training.
241        split: The choice of data split. Either 'train', 'val' or 'test'.
242        phase: The choice of cardiac phase. Either 'ED' or 'ES'. If None, both phases are returned.
243        vendor: The choice of scanner vendor. One of 'A', 'B', 'C' or 'D'. If None, all vendors are returned.
244        resize_inputs: Whether to resize inputs to the desired patch shape.
245        download: Whether to download the data if it is not present.
246        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
247
248    Returns:
249        The segmentation dataset.
250    """
251    volume_paths = get_mnms_paths(path, split, phase, vendor, download)
252
253    if resize_inputs:
254        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
255        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
256            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
257        )
258
259    return torch_em.default_segmentation_dataset(
260        raw_paths=volume_paths,
261        raw_key="raw",
262        label_paths=volume_paths,
263        label_key="labels",
264        patch_shape=patch_shape,
265        is_seg_dataset=True,
266        **kwargs
267    )
268
269
270def get_mnms_loader(
271    path: Union[os.PathLike, str],
272    batch_size: int,
273    patch_shape: Tuple[int, ...],
274    split: Literal["train", "val", "test"],
275    phase: Optional[Literal["ED", "ES"]] = None,
276    vendor: Optional[Literal["A", "B", "C", "D"]] = None,
277    resize_inputs: bool = False,
278    download: bool = False,
279    **kwargs
280) -> DataLoader:
281    """Get the M&Ms dataloader for cardiac structure segmentation.
282
283    Args:
284        path: Filepath to a folder where the data is downloaded for further processing.
285        batch_size: The batch size for training.
286        patch_shape: The patch shape to use for training.
287        split: The choice of data split. Either 'train', 'val' or 'test'.
288        phase: The choice of cardiac phase. Either 'ED' or 'ES'. If None, both phases are returned.
289        vendor: The choice of scanner vendor. One of 'A', 'B', 'C' or 'D'. If None, all vendors are returned.
290        resize_inputs: Whether to resize inputs to the desired patch shape.
291        download: Whether to download the data if it is not present.
292        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
293
294    Returns:
295        The DataLoader.
296    """
297    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
298    dataset = get_mnms_dataset(path, patch_shape, split, phase, vendor, resize_inputs, download, **ds_kwargs)
299    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)