torch_em.data.datasets.medical.cirrmri600

  1"""The CirrMRI600+ dataset contains annotations for liver segmentation in abdominal MRI of patients
  2with liver cirrhosis.
  3
  4The dataset consists of 628 abdominal MRI volumes: 310 T1-weighted and 318 T2-weighted scans of
  5cirrhotic patients, acquired with multiple scanners and annotated by physicians with binary liver masks
  6(1 = liver). The scans of the two sequences are distributed as separate archives, each with an official
  7'train' / 'val' / 'test' split, and can be selected via the 'sequence' argument.
  8
  9NOTE: This loader uses the 3D archives of the cirrhotic patients ('Cirrhosis_T1_3D' and
 10'Cirrhosis_T2_3D'). The 2D slice version of the T2 scans and the archive of healthy subjects are not used.
 11
 12The data is located at https://doi.org/10.17605/OSF.IO/CUK24 and released under a CC-BY-NC-4.0 license.
 13
 14This dataset is from the publication https://doi.org/10.1038/s41597-025-05201-7.
 15Please cite it if you use this dataset for your research.
 16"""
 17
 18import os
 19from glob import glob
 20from natsort import natsorted
 21from typing import Union, Tuple, Literal, List
 22
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30URLS = {
 31    "T1": "https://osf.io/download/47rxy/",
 32    "T2": "https://osf.io/download/72df5/",
 33}
 34
 35CHECKSUMS = {
 36    "T1": "36f533f0496ffc6bfb5c243a25a5e3927c682a538660484540e2cf565ce434d1",
 37    "T2": "539f0bc3bd3cdb0bb85b465ace8002bcb15631e8b59b81ca04dec0a7c287fda9",
 38}
 39
 40SEQUENCES = list(URLS.keys())
 41SPLITS = {"train": "train", "val": "valid", "test": "test"}
 42
 43
 44def get_cirrmri600_data(
 45    path: Union[os.PathLike, str], sequence: Literal["T1", "T2"] = "T2", download: bool = False
 46) -> str:
 47    """Download the CirrMRI600+ dataset.
 48
 49    Args:
 50        path: Filepath to a folder where the data is downloaded for further processing.
 51        sequence: The choice of MRI sequence. Either 'T1' or 'T2'.
 52        download: Whether to download the data if it is not present.
 53
 54    Returns:
 55        Filepath where the data is downloaded.
 56    """
 57    if sequence not in SEQUENCES:
 58        raise ValueError(f"'{sequence}' is not a valid sequence. Choose one of {SEQUENCES}.")
 59
 60    data_dir = os.path.join(path, f"Cirrhosis_{sequence}_3D")
 61    if os.path.exists(data_dir):
 62        return data_dir
 63
 64    os.makedirs(path, exist_ok=True)
 65
 66    zip_path = os.path.join(path, f"Cirrhosis_{sequence}_3D.zip")
 67    util.download_source(path=zip_path, url=URLS[sequence], download=download, checksum=CHECKSUMS[sequence])
 68    util.unzip(zip_path=zip_path, dst=path, remove=False)
 69
 70    assert os.path.exists(data_dir), f"The extraction of the archive did not create the expected folder in '{path}'."
 71
 72    return data_dir
 73
 74
 75def get_cirrmri600_paths(
 76    path: Union[os.PathLike, str],
 77    split: Literal["train", "val", "test"],
 78    sequence: Literal["T1", "T2"] = "T2",
 79    download: bool = False,
 80) -> Tuple[List[str], List[str]]:
 81    """Get paths to the CirrMRI600+ data.
 82
 83    Args:
 84        path: Filepath to a folder where the data is downloaded for further processing.
 85        split: The choice of data split. One of 'train', 'val' or 'test'.
 86        sequence: The choice of MRI sequence. Either 'T1' or 'T2'.
 87        download: Whether to download the data if it is not present.
 88
 89    Returns:
 90        List of filepaths for the image data.
 91        List of filepaths for the label data.
 92    """
 93    if split not in SPLITS:
 94        raise ValueError(f"'{split}' is not a valid split. Choose one of {list(SPLITS)}.")
 95
 96    data_dir = get_cirrmri600_data(path, sequence, download)
 97
 98    raw_paths = natsorted(glob(os.path.join(data_dir, f"{SPLITS[split]}_images", "*.nii.gz")))
 99    label_paths = [p.replace(f"{SPLITS[split]}_images", f"{SPLITS[split]}_masks") for p in raw_paths]
100
101    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in label_paths)
102
103    return raw_paths, label_paths
104
105
106def get_cirrmri600_dataset(
107    path: Union[os.PathLike, str],
108    patch_shape: Tuple[int, int, int],
109    split: Literal["train", "val", "test"],
110    sequence: Literal["T1", "T2"] = "T2",
111    resize_inputs: bool = False,
112    download: bool = False,
113    **kwargs
114) -> Dataset:
115    """Get the CirrMRI600+ dataset for liver segmentation in cirrhotic MRI.
116
117    Args:
118        path: Filepath to a folder where the data is downloaded for further processing.
119        patch_shape: The patch shape to use for training.
120        split: The choice of data split. One of 'train', 'val' or 'test'.
121        sequence: The choice of MRI sequence. Either 'T1' or 'T2'.
122        resize_inputs: Whether to resize the inputs to the patch shape.
123        download: Whether to download the data if it is not present.
124        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
125
126    Returns:
127        The segmentation dataset.
128    """
129    raw_paths, label_paths = get_cirrmri600_paths(path, split, sequence, download)
130
131    if resize_inputs:
132        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
133        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
134            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
135        )
136
137    return torch_em.default_segmentation_dataset(
138        raw_paths=raw_paths,
139        raw_key="data",
140        label_paths=label_paths,
141        label_key="data",
142        patch_shape=patch_shape,
143        is_seg_dataset=True,
144        ndim=3,
145        **kwargs
146    )
147
148
149def get_cirrmri600_loader(
150    path: Union[os.PathLike, str],
151    batch_size: int,
152    patch_shape: Tuple[int, int, int],
153    split: Literal["train", "val", "test"],
154    sequence: Literal["T1", "T2"] = "T2",
155    resize_inputs: bool = False,
156    download: bool = False,
157    **kwargs
158) -> DataLoader:
159    """Get the CirrMRI600+ dataloader for liver segmentation in cirrhotic MRI.
160
161    Args:
162        path: Filepath to a folder where the data is downloaded for further processing.
163        batch_size: The batch size for training.
164        patch_shape: The patch shape to use for training.
165        split: The choice of data split. One of 'train', 'val' or 'test'.
166        sequence: The choice of MRI sequence. Either 'T1' or 'T2'.
167        resize_inputs: Whether to resize the inputs to the patch shape.
168        download: Whether to download the data if it is not present.
169        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
170
171    Returns:
172        The DataLoader.
173    """
174    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
175    dataset = get_cirrmri600_dataset(path, patch_shape, split, sequence, resize_inputs, download, **ds_kwargs)
176    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)