torch_em.data.datasets.medical.cyst_x

  1"""The Cyst-X dataset contains annotations for whole-pancreas segmentation in T1-weighted and T2-weighted
  2abdominal MRI.
  3
  4The dataset consists of 1,461 full-volume MRI scans of 764 patients from 7 centers (723 T1-weighted and 738
  5T2-weighted scans), stored as nifti files. Each scan is paired with a binary mask of the whole pancreas
  6(0 = background, 1 = pancreas). NOTE: The masks are pancreas masks, they do not contain the pancreatic cysts,
  7ducts, mural nodules or tumors that the dataset was built to study (IPMN risk stratification).
  8
  9The data is located at https://huggingface.co/datasets/phy710/Cyst-X (also mirrored at https://osf.io/74vfs/),
 10released under a CC-BY-NC-4.0 license (non-commercial use only). The full collection is large, use `n_cases` to
 11only download a subset of it.
 12
 13The dataset and its code are described at https://github.com/NUBagciLab/Cyst-X.
 14Please cite the Cyst-X project if you use this dataset for your research.
 15"""
 16
 17import os
 18import json
 19from glob import glob
 20from concurrent import futures
 21from typing import Union, Tuple, List, Literal, Optional
 22
 23from tqdm import tqdm
 24
 25from torch.utils.data import Dataset, DataLoader
 26
 27import torch_em
 28
 29from .. import util
 30
 31
 32API_URL = "https://huggingface.co/api/datasets/phy710/Cyst-X"
 33FILE_URL = "https://huggingface.co/datasets/phy710/Cyst-X/resolve/main/{name}"
 34
 35SEQUENCES = ["t1", "t2"]
 36
 37
 38def _list_files(path):
 39    cache_path = os.path.join(path, "file_list.json")
 40    if os.path.exists(cache_path):
 41        with open(cache_path) as f:
 42            return json.load(f)
 43
 44    import requests
 45
 46    response = requests.get(API_URL, params={"blobs": "true"}, headers={"User-Agent": "Mozilla/5.0"})
 47    response.raise_for_status()
 48    files = {
 49        sibling["rfilename"]: sibling["lfs"]["sha256"]
 50        for sibling in response.json()["siblings"]
 51        if sibling["rfilename"].startswith("IPMN_images_masks/") and "lfs" in sibling
 52    }
 53
 54    os.makedirs(path, exist_ok=True)
 55    with open(cache_path, "w") as f:
 56        json.dump(files, f)
 57
 58    return files
 59
 60
 61def _download_case(path, sequence, case_id, files):
 62    for kind in ("images", "masks"):
 63        name = f"IPMN_images_masks/{sequence}/{kind}/{case_id}.nii.gz"
 64        os.makedirs(os.path.join(path, sequence, kind), exist_ok=True)
 65        util.download_source(
 66            path=os.path.join(path, sequence, kind, f"{case_id}.nii.gz"), url=FILE_URL.format(name=name),
 67            download=True, checksum=files[name],
 68        )
 69
 70
 71def get_cyst_x_data(
 72    path: Union[os.PathLike, str],
 73    sequence: Literal["t1", "t2"],
 74    n_cases: Optional[int] = None,
 75    n_workers: int = 8,
 76    download: bool = False,
 77) -> str:
 78    """Download the Cyst-X dataset.
 79
 80    NOTE: The full collection contains 1,461 scans. Use `n_cases` to only download a subset for a quick start.
 81
 82    Args:
 83        path: Filepath to a folder where the data is downloaded for further processing.
 84        sequence: The choice of MRI sequence. Either 't1' or 't2'.
 85        n_cases: The number of cases to download, sorted by case id. By default all cases are downloaded.
 86        n_workers: The number of parallel download workers.
 87        download: Whether to download the data if it is not present.
 88
 89    Returns:
 90        Filepath where the data is downloaded.
 91    """
 92    if sequence not in SEQUENCES:
 93        raise ValueError(f"'{sequence}' is not a valid sequence. Choose one of {SEQUENCES}.")
 94
 95    os.makedirs(path, exist_ok=True)
 96    files = _list_files(path)
 97
 98    prefix = f"IPMN_images_masks/{sequence}/images/"
 99    case_ids = sorted(name[len(prefix):-len(".nii.gz")] for name in files if name.startswith(prefix))
100    if n_cases is not None:
101        case_ids = case_ids[:n_cases]
102
103    missing = [
104        case_id for case_id in case_ids
105        if not all(
106            os.path.exists(os.path.join(path, sequence, kind, f"{case_id}.nii.gz")) for kind in ("images", "masks")
107        )
108    ]
109    if missing and not download:
110        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
111
112    with futures.ThreadPoolExecutor(n_workers) as pool:
113        tasks = [pool.submit(_download_case, path, sequence, case_id, files) for case_id in missing]
114        for task in tqdm(futures.as_completed(tasks), total=len(tasks), desc="Download Cyst-X cases"):
115            task.result()
116
117    return os.path.join(path, sequence)
118
119
120def get_cyst_x_paths(
121    path: Union[os.PathLike, str],
122    sequence: Literal["t1", "t2"],
123    n_cases: Optional[int] = None,
124    download: bool = False,
125) -> Tuple[List[str], List[str]]:
126    """Get paths to the Cyst-X data.
127
128    Args:
129        path: Filepath to a folder where the data is downloaded for further processing.
130        sequence: The choice of MRI sequence. Either 't1' or 't2'.
131        n_cases: The number of cases to use, sorted by case id. By default all cases are used.
132        download: Whether to download the data if it is not present.
133
134    Returns:
135        List of filepaths for the image data.
136        List of filepaths for the label data.
137    """
138    data_dir = get_cyst_x_data(path, sequence, n_cases, download=download)
139
140    raw_paths = sorted(glob(os.path.join(data_dir, "images", "*.nii.gz")))
141    if n_cases is not None:
142        raw_paths = raw_paths[:n_cases]
143    label_paths = [os.path.join(data_dir, "masks", os.path.basename(p)) for p in raw_paths]
144
145    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
146    assert all(os.path.exists(p) for p in label_paths)
147
148    return raw_paths, label_paths
149
150
151def get_cyst_x_dataset(
152    path: Union[os.PathLike, str],
153    patch_shape: Tuple[int, int, int],
154    sequence: Literal["t1", "t2"],
155    n_cases: Optional[int] = None,
156    resize_inputs: bool = False,
157    download: bool = False,
158    **kwargs
159) -> Dataset:
160    """Get the Cyst-X dataset for whole-pancreas segmentation in MRI.
161
162    Args:
163        path: Filepath to a folder where the data is downloaded for further processing.
164        patch_shape: The patch shape to use for training.
165        sequence: The choice of MRI sequence. Either 't1' or 't2'.
166        n_cases: The number of cases to use, sorted by case id. By default all cases are used.
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`.
170
171    Returns:
172        The segmentation dataset.
173    """
174    raw_paths, label_paths = get_cyst_x_paths(path, sequence, n_cases, download)
175
176    if resize_inputs:
177        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
178        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
179            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
180        )
181
182    return torch_em.default_segmentation_dataset(
183        raw_paths=raw_paths,
184        raw_key="data",
185        label_paths=label_paths,
186        label_key="data",
187        is_seg_dataset=True,
188        patch_shape=patch_shape,
189        ndim=3,
190        **kwargs
191    )
192
193
194def get_cyst_x_loader(
195    path: Union[os.PathLike, str],
196    batch_size: int,
197    patch_shape: Tuple[int, int, int],
198    sequence: Literal["t1", "t2"],
199    n_cases: Optional[int] = None,
200    resize_inputs: bool = False,
201    download: bool = False,
202    **kwargs
203) -> DataLoader:
204    """Get the Cyst-X dataloader for whole-pancreas segmentation in MRI.
205
206    Args:
207        path: Filepath to a folder where the data is downloaded for further processing.
208        batch_size: The batch size for training.
209        patch_shape: The patch shape to use for training.
210        sequence: The choice of MRI sequence. Either 't1' or 't2'.
211        n_cases: The number of cases to use, sorted by case id. By default all cases are used.
212        resize_inputs: Whether to resize the inputs to the patch shape.
213        download: Whether to download the data if it is not present.
214        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
215
216    Returns:
217        The DataLoader.
218    """
219    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
220    dataset = get_cyst_x_dataset(path, patch_shape, sequence, n_cases, resize_inputs, download, **ds_kwargs)
221    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)